diff --git a/ibis/backends/sql/compilers/datafusion.py b/ibis/backends/sql/compilers/datafusion.py index 2bd7d2f0dc4a..c095753a9777 100644 --- 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 index c0f64fb8911c..52c056fdb172 100644 --- 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), ], ), ],