feat(datafusion): enable array flatten, group concat, and timestamp now

ibis-project · Jul 31, 2024 · 4d110a0 · 4d110a0
1 parent 09dae23
commit 4d110a0
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 9 deletions.
diff --git a/ibis/backends/sql/compilers/datafusion.py b/ibis/backends/sql/compilers/datafusion.py
@@ -37,15 +37,13 @@ class DataFusionCompiler(SQLGlotCompiler):
         ops.ArgMin,
         ops.ArrayDistinct,
         ops.ArrayFilter,
-        ops.ArrayFlatten,
         ops.ArrayMap,
         ops.ArrayZip,
         ops.BitwiseNot,
         ops.Clip,
         ops.CountDistinctStar,
         ops.DateDelta,
         ops.Greatest,
-        ops.GroupConcat,
         ops.IntervalFromInteger,
         ops.Least,
         ops.MultiQuantile,
@@ -55,7 +53,6 @@ class DataFusionCompiler(SQLGlotCompiler):
         ops.TimeDelta,
         ops.TimestampBucket,
         ops.TimestampDelta,
-        ops.TimestampNow,
         ops.TypeOf,
         ops.StringToDate,
         ops.StringToTimestamp,

diff --git a/ibis/backends/tests/test_aggregation.py b/ibis/backends/tests/test_aggregation.py
@@ -1233,7 +1233,7 @@ def test_date_quantile(alltypes):
         ),
     ],
 )
-@pytest.mark.notimpl(["datafusion", "polars"], raises=com.OperationNotDefinedError)
+@pytest.mark.notimpl(["polars"], raises=com.OperationNotDefinedError)
 @pytest.mark.notimpl(["exasol"], raises=ExaQueryError)
 @pytest.mark.notyet(["flink"], raises=Py4JJavaError)
 def test_group_concat(
@@ -1485,7 +1485,10 @@ def test_grouped_case(backend, con):
     backend.assert_frame_equal(result, expected)
 
 
-@pytest.mark.notimpl(["datafusion", "polars"], raises=com.OperationNotDefinedError)
+@pytest.mark.notimpl(["polars"], raises=com.OperationNotDefinedError)
+@pytest.mark.notimpl(
+    ["datafusion"], raises=Exception, reason="not supported in datafusion"
+)
 @pytest.mark.notimpl(["exasol"], raises=ExaQueryError)
 @pytest.mark.notyet(["flink"], raises=Py4JJavaError)
 @pytest.mark.notyet(["impala"], raises=ImpalaHiveServer2Error)

diff --git a/ibis/backends/tests/test_array.py b/ibis/backends/tests/test_array.py
@@ -987,7 +987,8 @@ def flatten_data():
                     ["clickhouse"],
                     reason="Arrays are never nullable",
                     raises=AssertionError,
-                )
+                ),
+                pytest.mark.notimpl(["datafusion"], raises=AssertionError),
             ],
         ),
         param(
@@ -1005,11 +1006,12 @@ def flatten_data():
                     raises=TypeError,
                     reason="comparison of nested arrays doesn't work in pandas testing module",
                 ),
+                pytest.mark.notimpl(["datafusion"], raises=AssertionError),
             ],
         ),
     ],
 )
-@pytest.mark.notimpl(["datafusion", "flink"], raises=com.OperationNotDefinedError)
+@pytest.mark.notimpl(["flink"], raises=com.OperationNotDefinedError)
 def test_array_flatten(backend, flatten_data, column, expected):
     data = flatten_data[column]
     t = ibis.memtable({column: data["data"]}, schema={column: data["type"]})

diff --git a/ibis/backends/tests/test_temporal.py b/ibis/backends/tests/test_temporal.py
@@ -1420,14 +1420,12 @@ def test_day_of_week_column_group_by(
     backend.assert_frame_equal(result, expected, check_dtype=False)
 
 
-@pytest.mark.notimpl(["datafusion"], raises=com.OperationNotDefinedError)
 def test_now(con):
     expr = ibis.now()
     result = con.execute(expr.name("tmp"))
     assert isinstance(result, datetime.datetime)
 
 
-@pytest.mark.notimpl(["datafusion"], raises=com.OperationNotDefinedError)
 def test_now_from_projection(alltypes):
     n = 2
     expr = alltypes.select(now=ibis.now()).limit(n)