feat(dask): support Table.sample

ibis-project · Oct 17, 2023 · 09a7626 · 09a7626
1 parent 77215be
commit 09a7626
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/ibis/backends/dask/execution/generic.py b/ibis/backends/dask/execution/generic.py
@@ -552,3 +552,8 @@ def execute_table_array_view(op, _, **kwargs):
     # Need to compute dataframe in order to squeeze into a scalar
     ddf = execute(op.table)
     return ddf.compute().squeeze()
+
+
+@execute_node.register(ops.Sample, dd.DataFrame, object, object)
+def execute_sample(op, data, fraction, seed, **kwargs):
+    return data.sample(frac=fraction, random_state=seed)
diff --git a/ibis/backends/tests/test_generic.py b/ibis/backends/tests/test_generic.py
@@ -1532,7 +1532,6 @@ def test_dynamic_table_slice_with_computed_offset(backend):
 @pytest.mark.notimpl(
     [
         "bigquery",
-        "dask",
         "datafusion",
         "druid",
         "flink",
@@ -1560,7 +1559,6 @@ def test_sample(backend):
 @pytest.mark.notimpl(
     [
         "bigquery",
-        "dask",
         "datafusion",
         "druid",
         "flink",
@@ -1581,7 +1579,6 @@ def test_sample_memtable(con, backend):
     [
         "bigquery",
         "clickhouse",
-        "dask",
         "datafusion",
         "druid",
         "flink",