fix(datafusion): handle NULLs in array flatten

ibis-project · Sep 3, 2024 · ecc199f · ecc199f
1 parent 9ee910d
commit ecc199f
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/ibis/backends/sql/compilers/datafusion.py b/ibis/backends/sql/compilers/datafusion.py
@@ -499,5 +499,8 @@ def visit_GroupConcat(self, op, *, arg, sep, where, order_by):
             op, arg=arg, sep=sep, where=where, order_by=order_by
         )
 
+    def visit_ArrayFlatten(self, op, *, arg):
+        return self.if_(arg.is_(NULL), NULL, self.f.flatten(arg))
+
 
 compiler = DataFusionCompiler()
diff --git a/ibis/backends/tests/test_array.py b/ibis/backends/tests/test_array.py
@@ -1033,7 +1033,6 @@ def flatten_data():
                     reason="Arrays are never nullable",
                     raises=AssertionError,
                 ),
-                pytest.mark.notimpl(["datafusion"], raises=AssertionError),
             ],
         ),
         param(
@@ -1051,7 +1050,6 @@ def flatten_data():
                     raises=TypeError,
                     reason="comparison of nested arrays doesn't work in pandas testing module",
                 ),
-                pytest.mark.notimpl(["datafusion"], raises=AssertionError),
             ],
         ),
     ],