fix(datafusion): handle array concat like everything else

ibis-project · Nov 5, 2024 · f6b4ca4 · f6b4ca4
1 parent 6e51493
commit f6b4ca4
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/ibis/backends/sql/compilers/datafusion.py b/ibis/backends/sql/compilers/datafusion.py
@@ -2,7 +2,7 @@
 
 import calendar
 import math
-from functools import partial
+from functools import partial, reduce
 from itertools import starmap
 
 import sqlglot as sg
@@ -533,5 +533,13 @@ def visit_ArrayFlatten(self, op, *, arg):
     def visit_RandomUUID(self, op):
         return self.f.anon.uuid()
 
+    def visit_ArrayConcat(self, op, *, arg):
+        return reduce(
+            lambda x, y: self.if_(
+                x.is_(NULL).or_(y.is_(NULL)), NULL, self.f.array_cat(x, y)
+            ),
+            map(partial(self.cast, to=op.dtype), arg),
+        )
+
 
 compiler = DataFusionCompiler()