feat(duckdb): support array zip operation

ibis-project · May 22, 2023 · 2d14ccc · 2d14ccc
1 parent 4f3a4cd
commit 2d14ccc
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 1 deletion.
diff --git a/ibis/backends/duckdb/registry.py b/ibis/backends/duckdb/registry.py
@@ -302,6 +302,24 @@ def _map_merge(t, op):
     )
 
 
+def _array_zip(t, op):
+    args = tuple(map(t.translate, op.arg))
+
+    i = sa.literal_column("i", type_=sa.INTEGER)
+    dtype = op.output_dtype
+    return array_map(
+        sa.func.range(1, sa.func.greatest(*map(sa.func.array_length, args)) + 1),
+        i,
+        struct_pack(
+            {
+                name: sa.func.list_extract(arg, i)
+                for name, arg in zip(dtype.value_type.names, args)
+            },
+            type=t.get_sqla_type(dtype),
+        ),
+    )
+
+
 operation_registry.update(
     {
         ops.ArrayColumn: (
@@ -342,6 +360,7 @@ def _map_merge(t, op):
         ops.ArrayUnion: fixed_arity(
             lambda left, right: sa.func.list_distinct(sa.func.list_cat(left, right)), 2
         ),
+        ops.ArrayZip: _array_zip,
         ops.DayOfWeekName: unary(sa.func.dayname),
         ops.Literal: _literal,
         ops.Log2: unary(sa.func.log2),

diff --git a/ibis/backends/tests/test_array.py b/ibis/backends/tests/test_array.py
@@ -715,7 +715,6 @@ def test_unnest_struct(con):
         "dask",
         "datafusion",
         "druid",
-        "duckdb",
         "oracle",
         "pandas",
         "polars",

diff --git a/ibis/expr/types/arrays.py b/ibis/expr/types/arrays.py
@@ -774,6 +774,43 @@ def zip(self, other: ir.Array, *others: ir.Array) -> ir.Array:
         Array
             Array of structs where each struct field is an element of each input
             array.
+
+        Examples
+        --------
+        >>> import ibis
+        >>> ibis.options.interactive = True
+        >>> t = ibis.memtable({"numbers": [[3, 2], [], None], "strings": [["a", "c"], None, ["e"]]})
+        >>> t
+        ┏━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
+        ┃ numbers              ┃ strings              ┃
+        ┡━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
+        │ array<int64>         │ array<string>        │
+        ├──────────────────────┼──────────────────────┤
+        │ [3, 2]               │ ['a', 'c']           │
+        │ []                   │ NULL                 │
+        │ NULL                 │ ['e']                │
+        └──────────────────────┴──────────────────────┘
+        >>> expr = t.numbers.zip(t.strings)
+        >>> expr
+        ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
+        ┃ ArrayZip()                           ┃
+        ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
+        │ array<struct<f1: int64, f2: string>> │
+        ├──────────────────────────────────────┤
+        │ [{...}, {...}]                       │
+        │ []                                   │
+        │ [{...}]                              │
+        └──────────────────────────────────────┘
+        >>> expr.unnest()
+        ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
+        ┃ ArrayZip()                    ┃
+        ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
+        │ struct<f1: int64, f2: string> │
+        ├───────────────────────────────┤
+        │ {'f1': 3, 'f2': 'a'}          │
+        │ {'f1': 2, 'f2': 'c'}          │
+        │ {'f1': None, 'f2': 'e'}       │
+        └───────────────────────────────┘
         """
 
         return ops.ArrayZip((self, other, *others)).to_expr()