feat(postgres): implement array union

ibis-project · Aug 5, 2023 · 6d3d518 · 6d3d518
1 parent 8d42794
commit 6d3d518
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 6 deletions.
diff --git a/ibis/backends/postgres/registry.py b/ibis/backends/postgres/registry.py
@@ -382,7 +382,7 @@ def translate(t, op):
     return translate
 
 
-def _literal(t, op):
+def _literal(_, op):
     dtype = op.output_dtype
     value = op.value
 
@@ -784,12 +784,24 @@ def _array_filter(t, op):
             ),
             2,
         ),
-        ops.ArrayDistinct: fixed_arity(
-            lambda arg: sa.func.array(
-                sa.select(
-                    sa.distinct(sa.func.unnest(arg).column_valued())
+        ops.ArrayUnion: fixed_arity(
+            lambda left, right: sa.func.array(
+                sa.union(
+                    sa.select(sa.func.unnest(left).column_valued()),
+                    sa.select(sa.func.unnest(right).column_valued()),
                 ).scalar_subquery()
             ),
+            2,
+        ),
+        ops.ArrayDistinct: fixed_arity(
+            lambda arg: sa.case(
+                (arg.is_(sa.null()), sa.null()),
+                else_=sa.func.array(
+                    sa.select(
+                        sa.distinct(sa.func.unnest(arg).column_valued())
+                    ).scalar_subquery()
+                ),
+            ),
             1,
         ),
         ops.ArrayPosition: fixed_arity(_array_position, 2),

diff --git a/ibis/backends/tests/test_array.py b/ibis/backends/tests/test_array.py
@@ -648,7 +648,7 @@ def test_array_sort(backend, con):
 
 
 @pytest.mark.notimpl(
-    ["dask", "datafusion", "impala", "mssql", "pandas", "polars", "postgres"],
+    ["dask", "datafusion", "impala", "mssql", "pandas", "polars"],
     raises=com.OperationNotDefinedError,
 )
 @pytest.mark.notyet(