feat(datafusion): add array and strings functions (#8895)

ibis-project · Apr 5, 2024 · 2f23223 · 2f23223
1 parent 4dfc19f
commit 2f23223
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 15 deletions.
diff --git a/ibis/backends/datafusion/compiler.py b/ibis/backends/datafusion/compiler.py
@@ -39,10 +39,7 @@ class DataFusionCompiler(SQLGlotCompiler):
             ops.ArrayDistinct,
             ops.ArrayFilter,
             ops.ArrayFlatten,
-            ops.ArrayIntersect,
             ops.ArrayMap,
-            ops.ArraySort,
-            ops.ArrayUnion,
             ops.ArrayZip,
             ops.BitwiseNot,
             ops.Clip,
@@ -62,9 +59,7 @@ class DataFusionCompiler(SQLGlotCompiler):
             ops.TimestampNow,
             ops.TypeOf,
             ops.Unnest,
-            ops.EndsWith,
             ops.StringToTimestamp,
-            ops.Levenshtein,
         )
     )
 
@@ -82,6 +77,9 @@ class DataFusionCompiler(SQLGlotCompiler):
         ops.StringLength: "character_length",
         ops.RandomUUID: "uuid",
         ops.RegexSplit: "regex_split",
+        ops.EndsWith: "ends_with",
+        ops.ArrayIntersect: "array_intersect",
+        ops.ArrayUnion: "array_union",
     }
 
     def _aggregate(self, funcname: str, *args, where):

diff --git a/ibis/backends/tests/test_array.py b/ibis/backends/tests/test_array.py
@@ -689,7 +689,7 @@ def test_array_unique(con, input, expected):
 
 @builtin_array
 @pytest.mark.notimpl(
-    ["datafusion", "flink", "polars"],
+    ["flink", "polars"],
     raises=com.OperationNotDefinedError,
 )
 @pytest.mark.broken(
@@ -709,7 +709,7 @@ def test_array_sort(con):
 
 
 @builtin_array
-@pytest.mark.notimpl(["datafusion", "polars"], raises=com.OperationNotDefinedError)
+@pytest.mark.notimpl(["polars"], raises=com.OperationNotDefinedError)
 @pytest.mark.parametrize(
     ("a", "b", "expected_array"),
     [
@@ -729,6 +729,11 @@ def test_array_sort(con):
                     raises=GoogleBadRequest,
                     reason="BigQuery doesn't support arrays with null elements",
                 ),
+                pytest.mark.notyet(
+                    ["datafusion"],
+                    raises=AssertionError,
+                    reason="DataFusion transforms null elements to NAN",
+                ),
             ],
         ),
         param(
@@ -752,7 +757,7 @@ def test_array_union(con, a, b, expected_array):
 
 @builtin_array
 @pytest.mark.notimpl(
-    ["dask", "datafusion", "pandas", "polars", "flink"],
+    ["dask", "pandas", "polars", "flink"],
     raises=com.OperationNotDefinedError,
 )
 @pytest.mark.notimpl(

diff --git a/ibis/backends/tests/test_string.py b/ibis/backends/tests/test_string.py
@@ -524,9 +524,7 @@ def uses_java_re(t):
             lambda t: t.int_col == 1,
             id="endswith",
             marks=[
-                pytest.mark.notimpl(
-                    ["datafusion", "mssql"], raises=com.OperationNotDefinedError
-                ),
+                pytest.mark.notimpl(["mssql"], raises=com.OperationNotDefinedError),
             ],
         ),
         param(
@@ -542,9 +540,7 @@ def uses_java_re(t):
             lambda t: t.date_string_col.str.endswith("/10"),
             id="endswith-simple",
             marks=[
-                pytest.mark.notimpl(
-                    ["datafusion", "mssql"], raises=com.OperationNotDefinedError
-                ),
+                pytest.mark.notimpl(["mssql"], raises=com.OperationNotDefinedError),
             ],
         ),
         param(
@@ -893,7 +889,6 @@ def test_multiple_subs(con):
     [
         "clickhouse",
         "dask",
-        "datafusion",
         "druid",
         "impala",
         "mssql",