feat(trino): implement bitwise scalar/column ops

ibis-project · Dec 21, 2022 · ac4876c · ac4876c
1 parent dfeb600
commit ac4876c
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/ibis/backends/tests/test_generic.py b/ibis/backends/tests/test_generic.py
@@ -703,7 +703,7 @@ def test_select_filter_select(backend, alltypes, df):
         param(lambda t: t.int_col, lambda _: 3, id="col_scalar"),
     ],
 )
-@pytest.mark.notimpl(["bigquery", "dask", "datafusion", "pandas", "snowflake", "trino"])
+@pytest.mark.notimpl(["bigquery", "dask", "datafusion", "pandas", "snowflake"])
 def test_bitwise_columns(backend, con, alltypes, df, op, left_fn, right_fn):
     expr = op(left_fn(alltypes), right_fn(alltypes)).name("tmp")
     result = con.execute(expr)
@@ -737,7 +737,7 @@ def test_bitwise_columns(backend, con, alltypes, df, op, left_fn, right_fn):
         param(rshift, lambda t: t.int_col, lambda _: 3, id="rshift_col_scalar"),
     ],
 )
-@pytest.mark.notimpl(["bigquery", "dask", "datafusion", "pandas", "trino"])
+@pytest.mark.notimpl(["bigquery", "dask", "datafusion", "pandas"])
 @pyspark_no_bitshift
 def test_bitwise_shift(backend, alltypes, df, op, left_fn, right_fn):
     expr = op(left_fn(alltypes), right_fn(alltypes)).name("tmp")
@@ -767,23 +767,23 @@ def test_bitwise_shift(backend, alltypes, df, op, left_fn, right_fn):
     ("left", "right"),
     [param(4, L(2), id="int_col"), param(L(4), 2, id="col_int")],
 )
-@pytest.mark.notimpl(["bigquery", "dask", "datafusion", "pandas", "trino"])
+@pytest.mark.notimpl(["bigquery", "dask", "datafusion", "pandas"])
 def test_bitwise_scalars(con, op, left, right):
     expr = op(left, right)
     result = con.execute(expr)
     expected = op(4, 2)
     assert result == expected
 
 
-@pytest.mark.notimpl(["bigquery", "dask", "datafusion", "pandas", "snowflake", "trino"])
+@pytest.mark.notimpl(["bigquery", "dask", "datafusion", "pandas", "snowflake"])
 def test_bitwise_not_scalar(con):
     expr = ~L(2)
     result = con.execute(expr)
     expected = -3
     assert result == expected
 
 
-@pytest.mark.notimpl(["bigquery", "dask", "datafusion", "pandas", "snowflake", "trino"])
+@pytest.mark.notimpl(["bigquery", "dask", "datafusion", "pandas", "snowflake"])
 def test_bitwise_not_col(backend, alltypes, df):
     expr = (~alltypes.int_col).name("tmp")
     result = expr.execute()

diff --git a/ibis/backends/trino/registry.py b/ibis/backends/trino/registry.py
@@ -11,7 +11,6 @@
 
 operation_registry = sqlalchemy_operation_registry.copy()
 
-# TODO: trino doesn't support `& |` for bitwise ops, it wants `bitwise_and` and `bitwise_or``
 
 def _arbitrary(t, op):
     if op.how == "heavy":
@@ -45,6 +44,12 @@ def _json_get_item(t, op):
         ops.Arbitrary: _arbitrary,
         ops.BitAnd: reduction(sa.func.bitwise_and_agg),
         ops.BitOr: reduction(sa.func.bitwise_or_agg),
+        ops.BitwiseAnd: fixed_arity(sa.func.bitwise_and, 2),
+        ops.BitwiseOr: fixed_arity(sa.func.bitwise_or, 2),
+        ops.BitwiseXor: fixed_arity(sa.func.bitwise_xor, 2),
+        ops.BitwiseLeftShift: fixed_arity(sa.func.bitwise_left_shift, 2),
+        ops.BitwiseRightShift: fixed_arity(sa.func.bitwise_right_shift, 2),
+        ops.BitwiseNot: unary(sa.func.bitwise_not),
         ops.ArrayCollect: reduction(sa.func.array_agg),
         ops.JSONGetItem: _json_get_item,
     }