feat(clickhouse): implement bit aggs

ibis-project · May 23, 2022 · f94a5d2 · f94a5d2
1 parent d5fd215
commit f94a5d2
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 9 deletions.
diff --git a/ibis/backends/clickhouse/registry.py b/ibis/backends/clickhouse/registry.py
@@ -3,6 +3,7 @@
 
 import ibis
 import ibis.common.exceptions as com
+import ibis.expr.datatypes as dt
 import ibis.expr.operations as ops
 import ibis.expr.types as ir
 import ibis.util as util
@@ -591,6 +592,23 @@ def _cotangent(translator, expr):
     return f"cos({arg}) / sin({arg})"
 
 
+def _bit_agg(func):
+    def compile(translator, expr):
+        op = expr.op()
+        raw_arg = op.arg
+        arg = translator.translate(raw_arg)
+        if not isinstance((type := raw_arg.type()), dt.UnsignedInteger):
+            nbits = type._nbytes * 8
+            arg = f"reinterpretAsUInt{nbits}({arg})"
+
+        if (where := op.where) is not None:
+            return f"{func}If({arg}, {translator.translate(where)})"
+        else:
+            return f"{func}({arg})"
+
+    return compile
+
+
 # TODO: clickhouse uses different string functions
 #       for ascii and utf-8 encodings,
 
@@ -738,6 +756,9 @@ def _cotangent(translator, expr):
     ops.ArrayRepeat: _array_repeat_op,
     ops.ArraySlice: _array_slice_op,
     ops.Unnest: _unary("arrayJoin"),
+    ops.BitAnd: _bit_agg("groupBitAnd"),
+    ops.BitOr: _bit_agg("groupBitOr"),
+    ops.BitXor: _bit_agg("groupBitXor"),
 }
 
 

diff --git a/ibis/backends/tests/test_aggregation.py b/ibis/backends/tests/test_aggregation.py
@@ -254,25 +254,19 @@ def test_aggregate_grouped(
             lambda t, where: t.bigint_col.bit_and(where=where),
             lambda t, where: np.bitwise_and.reduce(t.bigint_col[where].values),
             id='bit_and',
-            marks=pytest.mark.notimpl(
-                ["clickhouse", "dask", "pandas", "sqlite"]
-            ),
+            marks=pytest.mark.notimpl(["dask", "pandas", "sqlite"]),
         ),
         param(
             lambda t, where: t.bigint_col.bit_or(where=where),
             lambda t, where: np.bitwise_or.reduce(t.bigint_col[where].values),
             id='bit_or',
-            marks=pytest.mark.notimpl(
-                ["clickhouse", "dask", "pandas", "sqlite"]
-            ),
+            marks=pytest.mark.notimpl(["dask", "pandas", "sqlite"]),
         ),
         param(
             lambda t, where: t.bigint_col.bit_xor(where=where),
             lambda t, where: np.bitwise_xor.reduce(t.bigint_col[where].values),
             id='bit_xor',
-            marks=pytest.mark.notimpl(
-                ["clickhouse", "dask", "pandas", "sqlite"]
-            ),
+            marks=pytest.mark.notimpl(["dask", "pandas", "sqlite"]),
         ),
     ],
 )