feat(polars): add any and all reductions

ibis-project · May 23, 2023 · 0bd3c01 · 0bd3c01
1 parent fc66380
commit 0bd3c01
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 6 deletions.
diff --git a/ibis/backends/polars/compiler.py b/ibis/backends/polars/compiler.py
@@ -646,6 +646,8 @@ def struct_column(op):
     ops.Median: 'median',
     ops.First: 'first',
     ops.Last: 'last',
+    ops.All: 'all',
+    ops.Any: 'any',
 }
 
 for reduction in _reductions.keys():
@@ -1050,3 +1052,21 @@ def execute_union(op, **kwargs):
 @translate.register(ops.Hash)
 def execute_hash(op, **kwargs):
     return translate(op.arg).hash()
+
+
+@translate.register(ops.NotAll)
+def execute_not_all(op, **kwargs):
+    arg = op.arg
+    if (op_where := op.where) is not None:
+        arg = ops.Where(op_where, arg, None)
+
+    return translate(arg).all().is_not()
+
+
+@translate.register(ops.NotAny)
+def execute_not_any(op, **kwargs):
+    arg = op.arg
+    if (op_where := op.where) is not None:
+        arg = ops.Where(op_where, arg, None)
+
+    return translate(arg).any().is_not()
diff --git a/ibis/backends/tests/test_aggregation.py b/ibis/backends/tests/test_aggregation.py
@@ -279,7 +279,7 @@ def mean_and_std(v):
             id='any',
             marks=[
                 pytest.mark.notimpl(
-                    ["polars", "datafusion"],
+                    ["datafusion"],
                     raises=com.OperationNotDefinedError,
                 ),
                 pytest.mark.broken(
@@ -295,7 +295,7 @@ def mean_and_std(v):
             id='notany',
             marks=[
                 pytest.mark.notimpl(
-                    ["polars", "datafusion"],
+                    ["datafusion"],
                     raises=com.OperationNotDefinedError,
                 ),
                 pytest.mark.broken(
@@ -316,7 +316,7 @@ def mean_and_std(v):
             id='any_negate',
             marks=[
                 pytest.mark.notimpl(
-                    ["polars", "datafusion"],
+                    ["datafusion"],
                     raises=com.OperationNotDefinedError,
                 ),
                 pytest.mark.broken(
@@ -337,7 +337,7 @@ def mean_and_std(v):
             id='all',
             marks=[
                 pytest.mark.notimpl(
-                    ["polars", "datafusion"],
+                    ["datafusion"],
                     raises=com.OperationNotDefinedError,
                 ),
                 pytest.mark.broken(
@@ -353,7 +353,7 @@ def mean_and_std(v):
             id='notall',
             marks=[
                 pytest.mark.notimpl(
-                    ["polars", "datafusion"],
+                    ["datafusion"],
                     raises=com.OperationNotDefinedError,
                 ),
                 pytest.mark.broken(
@@ -374,7 +374,7 @@ def mean_and_std(v):
             id='all_negate',
             marks=[
                 pytest.mark.notimpl(
-                    ["polars", "datafusion"],
+                    ["datafusion"],
                     raises=com.OperationNotDefinedError,
                 ),
                 pytest.mark.broken(