feat(datafusion): add BitwiseNot, Clip and IntervalFromInteger

Author: mesejo

ibis/backends/sql/compilers/datafusion.py

@@ -33,11 +33,8 @@ class DataFusionCompiler(SQLGlotCompiler):
         ops.ArrayFilter,
         ops.ArrayMap,
         ops.ArrayZip,
-        ops.BitwiseNot,
-        ops.Clip,
         ops.CountDistinctStar,
         ops.DateDelta,
-        ops.IntervalFromInteger,
         ops.RowID,
         ops.Strftime,
         ops.TimeDelta,
@@ -599,11 +596,13 @@ def visit_ArrayAny(self, op, *, arg):
             ),
         )
 
-    def visit_Greatest(self, op, *, arg):
-        return self.f.greatest(*arg)
+    def visit_BitwiseNot(self, op, *, arg):
+        # https://stackoverflow.com/q/69648488/4001592
+        return sge.BitwiseXor(this=arg, expression=sge.Literal.number(-1))
 
-    def visit_Least(self, op, *, arg):
-        return self.f.least(*arg)
+    def visit_IntervalFromInteger(self, op, *, arg, unit):
+        unit = unit.name.lower()
+        return sg.cast(self.f.concat(self.cast(arg, dt.string), f" {unit}"), "interval")
 
 
 compiler = DataFusionCompiler()

ibis/backends/tests/test_numeric.py

@@ -1362,7 +1362,6 @@ def test_random_different_per_row(alltypes):
         ),
     ],
 )
-@pytest.mark.notimpl(["datafusion"], raises=com.OperationNotDefinedError)
 def test_clip(backend, alltypes, df, ibis_func, pandas_func):
     result = ibis_func(alltypes.int_col).execute()
     expected = pandas_func(df.int_col).astype(result.dtype)
@@ -1371,7 +1370,7 @@ def test_clip(backend, alltypes, df, ibis_func, pandas_func):
     backend.assert_series_equal(result, expected, check_names=False)
 
 
-@pytest.mark.notimpl(["datafusion", "polars"], raises=com.OperationNotDefinedError)
+@pytest.mark.notimpl(["polars"], raises=com.OperationNotDefinedError)
 @pytest.mark.notyet(
     ["druid"],
     raises=PyDruidProgrammingError,
@@ -1509,7 +1508,7 @@ def test_bitwise_scalars(con, op, left, right):
     assert result == expected
 
 
-@pytest.mark.notimpl(["datafusion", "exasol"], raises=com.OperationNotDefinedError)
+@pytest.mark.notimpl(["exasol"], raises=com.OperationNotDefinedError)
 @pytest.mark.notimpl(["oracle"], raises=OracleDatabaseError)
 @flink_no_bitwise
 def test_bitwise_not_scalar(con):
@@ -1519,7 +1518,7 @@ def test_bitwise_not_scalar(con):
     assert result == expected
 
 
-@pytest.mark.notimpl(["datafusion", "exasol"], raises=com.OperationNotDefinedError)
+@pytest.mark.notimpl(["exasol"], raises=com.OperationNotDefinedError)
 @pytest.mark.notimpl(["oracle"], raises=OracleDatabaseError)
 @flink_no_bitwise
 def test_bitwise_not_col(backend, alltypes, df):

ibis/backends/tests/test_temporal.py

@@ -552,9 +552,7 @@ def test_date_truncate(backend, alltypes, df, unit):
         ),
     ],
 )
-@pytest.mark.notimpl(
-    ["datafusion", "druid", "exasol"], raises=com.OperationNotDefinedError
-)
+@pytest.mark.notimpl(["druid", "exasol"], raises=com.OperationNotDefinedError)
 def test_integer_to_interval_timestamp(
     backend, con, alltypes, df, unit, displacement_type
 ):
@@ -628,7 +626,7 @@ def convert_to_offset(offset, displacement_type=displacement_type):
         param("D", marks=sqlite_without_ymd_intervals),
     ],
 )
-@pytest.mark.notimpl(["datafusion", "druid"], raises=com.OperationNotDefinedError)
+@pytest.mark.notimpl(["druid"], raises=com.OperationNotDefinedError)
 @pytest.mark.notimpl(["exasol"], raises=com.OperationNotDefinedError)
 def test_integer_to_interval_date(backend, con, alltypes, df, unit):
     interval = alltypes.int_col.as_interval(unit)