fix(api): return NULL when NULL is passed to Array.zip (#8652)

Co-authored-by: Phillip Cloud <[email protected]>

ibis/backends/duckdb/compiler.py

@@ -187,14 +187,17 @@ def visit_ArrayZip(self, op, *, arg):
             ]
         )
         func = sge.Lambda(this=body, expressions=[i])
-        return self.f.list_apply(
+        zipped_arrays = self.f.list_apply(
             self.f.range(
                 1,
                 # DuckDB Range excludes upper bound
                 self.f.greatest(*map(self.f.len, arg)) + 1,
             ),
             func,
         )
+        # if any of the input arrays in arg are NULL, the result is NULL
+        any_arg_null = sg.or_(*(arr.is_(NULL) for arr in arg))
+        return self.if_(any_arg_null, NULL, zipped_arrays)
 
     def visit_MapGet(self, op, *, arg, key, default):
         return self.f.ifnull(

ibis/backends/tests/test_array.py

@@ -818,8 +818,7 @@ def test_unnest_struct(con):
     tm.assert_series_equal(result, expected)
 
 
-@builtin_array
-@pytest.mark.notimpl(
+array_zip_notimpl = pytest.mark.notimpl(
     [
         "dask",
         "datafusion",
@@ -833,6 +832,10 @@ def test_unnest_struct(con):
     ],
     raises=com.OperationNotDefinedError,
 )
+
+
+@builtin_array
+@array_zip_notimpl
 def test_zip(backend):
     t = backend.array_types
 
@@ -851,6 +854,22 @@ def test_zip(backend):
     assert len(x[0]) == len(s[0])
 
 
+@builtin_array
+@array_zip_notimpl
+@pytest.mark.notyet(
+    "clickhouse",
+    raises=ClickHouseDatabaseError,
+    reason="clickhouse nested types can't be null",
+)
+def test_zip_null(backend):
+    # the .map is workaround for https://github.com/ibis-project/ibis/issues/8641
+    a = ibis.literal([1, 2, 3], type="array<int64>").map(ibis._)
+    b = ibis.literal(None, type="array<int64>")
+    assert backend.connection.execute(a.zip(b)) is None
+    assert backend.connection.execute(b.zip(a)) is None
+    assert backend.connection.execute(b.zip(b)) is None
+
+
 @builtin_array
 @pytest.mark.notyet(
     ["clickhouse"],

ibis/expr/types/arrays.py

@@ -929,44 +929,41 @@ def zip(self, other: ArrayValue, *others: ArrayValue) -> ArrayValue:
         -------
         Array
             Array of structs where each struct field is an element of each input
-            array.
+            array. The fields are named `f1`, `f2`, `f3`, etc.
 
         Examples
         --------
         >>> import ibis
         >>> ibis.options.interactive = True
-        >>> t = ibis.memtable({"numbers": [[3, 2], [], None], "strings": [["a", "c"], None, ["e"]]})
+        >>> ibis.options.repr.interactive.max_depth = 2
+        >>> t = ibis.memtable(
+        ...     {
+        ...         "numbers": [[3, 2], [6, 7], [], None],
+        ...         "strings": [["a", "c"], ["d"], [], ["x", "y"]],
+        ...     }
+        ... )
         >>> t
         ┏━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
         ┃ numbers              ┃ strings              ┃
         ┡━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
         │ array<int64>         │ array<string>        │
         ├──────────────────────┼──────────────────────┤
         │ [3, 2]               │ ['a', 'c']           │
-        │ []                   │ NULL                 │
-        │ NULL                 │ ['e']                │
+        │ [6, 7]               │ ['d']                │
+        │ []                   │ []                   │
+        │ NULL                 │ ['x', 'y']           │
         └──────────────────────┴──────────────────────┘
-        >>> 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'}       │
-        └───────────────────────────────┘
+        >>> t.numbers.zip(t.strings)
+        ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
+        ┃ ArrayZip()                                    ┃
+        ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
+        │ array<struct<f1: int64, f2: string>>          │
+        ├───────────────────────────────────────────────┤
+        │ [{'f1': 3, 'f2': 'a'}, {'f1': 2, 'f2': 'c'}]  │
+        │ [{'f1': 6, 'f2': 'd'}, {'f1': 7, 'f2': None}] │
+        │ []                                            │
+        │ NULL                                          │
+        └───────────────────────────────────────────────┘
         """
 
         return ops.ArrayZip((self, other, *others)).to_expr()