feat(pyspark): implement count distinct

ibis-project · Mar 27, 2023 · db29e10 · db29e10
1 parent aea4ccd
commit db29e10
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/ibis/backends/pyspark/compiler.py b/ibis/backends/pyspark/compiler.py
@@ -579,6 +579,11 @@ def compile_count(t, op, **kwargs):
     return compile_aggregator(t, op, fn=F.count, **kwargs)
 
 
+@compiles(ops.CountDistinct)
+def compile_count_distinct(t, op, **kwargs):
+    return compile_aggregator(t, op, fn=F.count_distinct, **kwargs)
+
+
 @compiles(ops.CountStar)
 def compile_count_star(t, op, aggcontext=None, **kwargs):
     src_table = t.translate(op.arg, **kwargs)

diff --git a/ibis/backends/tests/test_aggregation.py b/ibis/backends/tests/test_aggregation.py
@@ -282,7 +282,7 @@ def mean_and_std(v):
             lambda t, where: t.bool_col[where].dropna().nunique(),
             id='nunique',
             marks=pytest.mark.notimpl(
-                ["pyspark", "datafusion"], raises=com.OperationNotDefinedError
+                ["datafusion"], raises=com.OperationNotDefinedError
             ),
         ),
         param(