fix(datafusion): handle NULLs in array flatten

Author: cpcloud

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()

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),
             ],
         ),
     ],