feat(pyspark): support Table.sample

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

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/ibis/backends/pyspark/compiler.py b/ibis/backends/pyspark/compiler.py
@@ -300,6 +300,12 @@ def compile_limit(t, op, **kwargs):
         return df
 
 
+@compiles(ops.Sample)
+def compile_sample(t, op, **kwargs):
+    df = t.translate(op.table, **kwargs)
+    return df.sample(fraction=op.fraction, seed=op.seed)
+
+
 @compiles(ops.And)
 def compile_and(t, op, **kwargs):
     return t.translate(op.left, **kwargs) & t.translate(op.right, **kwargs)

diff --git a/ibis/backends/tests/test_generic.py b/ibis/backends/tests/test_generic.py
@@ -1537,7 +1537,6 @@ def test_dynamic_table_slice_with_computed_offset(backend):
         "flink",
         "impala",
         "polars",
-        "pyspark",
         "snowflake",
     ]
 )
@@ -1564,7 +1563,6 @@ def test_sample(backend):
         "flink",
         "impala",
         "polars",
-        "pyspark",
         "snowflake",
     ]
 )
@@ -1588,7 +1586,6 @@ def test_sample_memtable(con, backend):
         "oracle",
         "polars",
         "postgres",
-        "pyspark",
         "snowflake",
         "sqlite",
         "trino",