feat(pandas): implement bitwise aggregations

ibis-project · May 23, 2022 · 37ff328 · 37ff328
1 parent e478479
commit 37ff328
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 3 deletions.
diff --git a/ibis/backends/pandas/execution/generic.py b/ibis/backends/pandas/execution/generic.py
@@ -717,6 +717,30 @@ def execute_count_frame(op, data, _, **kwargs):
     return len(data)
 
 
+@execute_node.register(ops.BitAnd, pd.Series, (pd.Series, type(None)))
+def execute_bit_and_series(_, data, mask, aggcontext=None, **kwargs):
+    return aggcontext.agg(
+        data[mask] if mask is not None else data,
+        np.bitwise_and.reduce,
+    )
+
+
+@execute_node.register(ops.BitOr, pd.Series, (pd.Series, type(None)))
+def execute_bit_or_series(_, data, mask, aggcontext=None, **kwargs):
+    return aggcontext.agg(
+        data[mask] if mask is not None else data,
+        np.bitwise_or.reduce,
+    )
+
+
+@execute_node.register(ops.BitXor, pd.Series, (pd.Series, type(None)))
+def execute_bit_xor_series(_, data, mask, aggcontext=None, **kwargs):
+    return aggcontext.agg(
+        data[mask] if mask is not None else data,
+        np.bitwise_xor.reduce,
+    )
+
+
 @execute_node.register((ops.Not, ops.Negate), (bool, np.bool_))
 def execute_not_bool(_, data, **kwargs):
     return not data

diff --git a/ibis/backends/tests/test_aggregation.py b/ibis/backends/tests/test_aggregation.py
@@ -255,7 +255,7 @@ def test_aggregate_grouped(
             lambda t, where: np.bitwise_and.reduce(t.bigint_col[where].values),
             id='bit_and',
             marks=[
-                pytest.mark.notimpl(["dask", "pandas"]),
+                pytest.mark.notimpl(["dask"]),
                 pytest.mark.notyet(["impala", "pyspark"]),
             ],
         ),
@@ -264,7 +264,7 @@ def test_aggregate_grouped(
             lambda t, where: np.bitwise_or.reduce(t.bigint_col[where].values),
             id='bit_or',
             marks=[
-                pytest.mark.notimpl(["dask", "pandas"]),
+                pytest.mark.notimpl(["dask"]),
                 pytest.mark.notyet(["impala", "pyspark"]),
             ],
         ),
@@ -273,7 +273,7 @@ def test_aggregate_grouped(
             lambda t, where: np.bitwise_xor.reduce(t.bigint_col[where].values),
             id='bit_xor',
             marks=[
-                pytest.mark.notimpl(["dask", "pandas"]),
+                pytest.mark.notimpl(["dask"]),
                 pytest.mark.notyet(["impala", "pyspark"]),
             ],
         ),