fix: Struct outer nullabillity (#18156)

crates/polars-compute/src/if_then_else/mod.rs

@@ -100,7 +100,7 @@ impl<T: NotSimdPrimitive> IfThenElseKernel for PrimitiveArray<T> {
     }
 }
 
-fn if_then_else_validity(
+pub fn if_then_else_validity(
     mask: &Bitmap,
     if_true: Option<&Bitmap>,
     if_false: Option<&Bitmap>,

crates/polars-core/src/chunked_array/ops/full.rs

@@ -1,4 +1,4 @@
-use arrow::bitmap::MutableBitmap;
+use arrow::bitmap::{Bitmap, MutableBitmap};
 
 use crate::chunked_array::builder::get_list_builder;
 use crate::prelude::*;
@@ -189,7 +189,9 @@ impl ListChunked {
 impl ChunkFullNull for StructChunked {
     fn full_null(name: &str, length: usize) -> StructChunked {
         let s = vec![Series::new_null("", length)];
-        StructChunked::from_series(name, &s).unwrap()
+        StructChunked::from_series(name, &s)
+            .unwrap()
+            .with_outer_validity(Some(Bitmap::new_zeroed(length)))
     }
 }
 

crates/polars-core/src/chunked_array/ops/zip.rs

@@ -1,6 +1,6 @@
 use arrow::bitmap::Bitmap;
 use arrow::compute::utils::{combine_validities_and, combine_validities_and_not};
-use polars_compute::if_then_else::IfThenElseKernel;
+use polars_compute::if_then_else::{if_then_else_validity, IfThenElseKernel};
 
 #[cfg(feature = "object")]
 use crate::chunked_array::object::ObjectArray;
@@ -62,7 +62,7 @@ fn combine_validities_chunked<
 
 impl<T> ChunkZip<T> for ChunkedArray<T>
 where
-    T: PolarsDataType,
+    T: PolarsDataType<IsStruct = FalseT>,
     T::Array: for<'a> IfThenElseKernel<Scalar<'a> = T::Physical<'a>>,
     ChunkedArray<T>: ChunkExpandAtIndex<T>,
 {
@@ -206,3 +206,104 @@ impl<T: PolarsObject> IfThenElseKernel for ObjectArray<T> {
             .collect_arr()
     }
 }
+
+#[cfg(feature = "dtype-struct")]
+impl ChunkZip<StructType> for StructChunked {
+    fn zip_with(
+        &self,
+        mask: &BooleanChunked,
+        other: &ChunkedArray<StructType>,
+    ) -> PolarsResult<ChunkedArray<StructType>> {
+        let (l, r, mask) = align_chunks_ternary(self, other, mask);
+
+        // Prepare the boolean arrays such that Null maps to false.
+        // This prevents every field doing that.
+        // # SAFETY
+        // We don't modify the length and update the null count.
+        let mut mask = mask.into_owned();
+        unsafe {
+            for arr in mask.downcast_iter_mut() {
+                let bm = bool_null_to_false(arr);
+                *arr = BooleanArray::from_data_default(bm, None);
+            }
+            mask.set_null_count(0);
+        }
+
+        // Zip all the fields.
+        let fields = l
+            .fields_as_series()
+            .iter()
+            .zip(r.fields_as_series())
+            .map(|(lhs, rhs)| lhs.zip_with_same_type(&mask, &rhs))
+            .collect::<PolarsResult<Vec<_>>>()?;
+
+        let mut out = StructChunked::from_series(self.name(), &fields)?;
+
+        // Zip the validities.
+        if (l.null_count + r.null_count) > 0 {
+            let validities = l
+                .chunks()
+                .iter()
+                .zip(r.chunks())
+                .map(|(l, r)| (l.validity(), r.validity()));
+
+            fn broadcast(v: Option<&Bitmap>, arr: &ArrayRef) -> Bitmap {
+                if v.unwrap().get(0).unwrap() {
+                    Bitmap::new_with_value(true, arr.len())
+                } else {
+                    Bitmap::new_zeroed(arr.len())
+                }
+            }
+
+            // # SAFETY
+            // We don't modify the length and update the null count.
+            unsafe {
+                for ((arr, (lv, rv)), mask) in out
+                    .chunks_mut()
+                    .iter_mut()
+                    .zip(validities)
+                    .zip(mask.downcast_iter())
+                {
+                    // TODO! we can optimize this and use a kernel that is able to broadcast wo/ allocating.
+                    let (lv, rv) = match (lv.map(|b| b.len()), rv.map(|b| b.len())) {
+                        (Some(1), Some(1)) if arr.len() != 1 => {
+                            let lv = broadcast(lv, arr);
+                            let rv = broadcast(rv, arr);
+                            (Some(lv), Some(rv))
+                        },
+                        (Some(a), Some(b)) if a == b => (lv.cloned(), rv.cloned()),
+                        (Some(1), _) => {
+                            let lv = broadcast(lv, arr);
+                            (Some(lv), rv.cloned())
+                        },
+                        (_, Some(1)) => {
+                            let rv = broadcast(rv, arr);
+                            (lv.cloned(), Some(rv))
+                        },
+                        (None, Some(_)) | (Some(_), None) | (None, None) => {
+                            (lv.cloned(), rv.cloned())
+                        },
+                        (Some(a), Some(b)) => {
+                            polars_bail!(InvalidOperation: "got different sizes in 'zip' operation, got length: {a} and {b}")
+                        },
+                    };
+
+                    // broadcast mask
+                    let validity = if mask.len() != arr.len() && mask.len() == 1 {
+                        if mask.get(0).unwrap() {
+                            lv
+                        } else {
+                            rv
+                        }
+                    } else {
+                        if_then_else_validity(mask.values(), lv.as_ref(), rv.as_ref())
+                    };
+
+                    *arr = arr.with_validity(validity);
+                }
+            }
+            out.compute_len();
+        }
+        Ok(out)
+    }
+}

crates/polars-core/src/series/implementations/struct__.rs

@@ -56,15 +56,9 @@ impl PrivateSeries for SeriesWrap<StructChunked> {
 
     #[cfg(feature = "zip_with")]
     fn zip_with_same_type(&self, mask: &BooleanChunked, other: &Series) -> PolarsResult<Series> {
-        let other = other.struct_()?;
-        let fields = self
-            .0
-            .fields_as_series()
-            .iter()
-            .zip(other.fields_as_series())
-            .map(|(lhs, rhs)| lhs.zip_with_same_type(mask, &rhs))
-            .collect::<PolarsResult<Vec<_>>>()?;
-        StructChunked::from_series(self.0.name(), &fields).map(|ca| ca.into_series())
+        self.0
+            .zip_with(mask, other.struct_()?)
+            .map(|ca| ca.into_series())
     }
 
     #[cfg(feature = "algorithm_group_by")]

crates/polars-core/src/series/ops/null.rs

@@ -1,3 +1,5 @@
+use arrow::bitmap::Bitmap;
+
 #[cfg(feature = "object")]
 use crate::chunked_array::object::registry::get_object_builder;
 use crate::prelude::*;
@@ -53,9 +55,14 @@ impl Series {
                     .iter()
                     .map(|fld| Series::full_null(fld.name(), size, fld.data_type()))
                     .collect::<Vec<_>>();
-                StructChunked::from_series(name, &fields)
-                    .unwrap()
-                    .into_series()
+                let ca = StructChunked::from_series(name, &fields).unwrap();
+
+                if !fields.is_empty() {
+                    ca.with_outer_validity(Some(Bitmap::new_zeroed(size)))
+                        .into_series()
+                } else {
+                    ca.into_series()
+                }
             },
             DataType::Null => Series::new_null(name, size),
             DataType::Unknown(kind) => {

py-polars/tests/unit/datatypes/test_struct.py

@@ -623,15 +623,15 @@ def test_struct_categorical_5843() -> None:
 def test_empty_struct() -> None:
     # List<struct>
     df = pl.DataFrame({"a": [[{}]]})
-    assert df.to_dict(as_series=False) == {"a": [[{"": None}]]}
+    assert df.to_dict(as_series=False) == {"a": [[None]]}
 
     # Struct one not empty
     df = pl.DataFrame({"a": [[{}, {"a": 10}]]})
     assert df.to_dict(as_series=False) == {"a": [[{"a": None}, {"a": 10}]]}
 
     # Empty struct
     df = pl.DataFrame({"a": [{}]})
-    assert df.to_dict(as_series=False) == {"a": [{"": None}]}
+    assert df.to_dict(as_series=False) == {"a": [None]}
 
 
 @pytest.mark.parametrize(
@@ -710,7 +710,7 @@ def test_struct_null_cast() -> None:
         .lazy()
         .select([pl.lit(None, dtype=pl.Null).cast(dtype, strict=True)])
         .collect()
-    ).to_dict(as_series=False) == {"literal": [{"a": None, "b": None, "c": None}]}
+    ).to_dict(as_series=False) == {"literal": [None]}
 
 
 def test_nested_struct_in_lists_cast() -> None:
@@ -976,3 +976,23 @@ def test_named_exprs() -> None:
     res = df.select(pl.struct(schema=schema, b=pl.col("a")))
     assert res.to_dict(as_series=False) == {"b": [{"b": 1}]}
     assert res.schema["b"] == pl.Struct(schema)
+
+
+def test_struct_outer_nullability_zip_18119() -> None:
+    df = pl.Series("int", [0, 1, 2, 3], dtype=pl.Int64).to_frame()
+    assert df.lazy().with_columns(
+        result=pl.when(pl.col("int") >= 1).then(
+            pl.struct(
+                a=pl.when(pl.col("int") % 2 == 1).then(True),
+                b=pl.when(pl.col("int") >= 2).then(False),
+            )
+        )
+    ).collect().to_dict(as_series=False) == {
+        "int": [0, 1, 2, 3],
+        "result": [
+            None,
+            {"a": True, "b": None},
+            {"a": None, "b": False},
+            {"a": True, "b": False},
+        ],
+    }

py-polars/tests/unit/test_queries.py

@@ -241,8 +241,8 @@ def map_expr(name: str) -> pl.Expr:
     ).to_dict(as_series=False) == {
         "groups": [1, 2, 3, 4],
         "out": [
-            {"sum": None, "count": None},
-            {"sum": None, "count": None},
+            None,
+            None,
             {"sum": 1, "count": 1},
             {"sum": 2, "count": 1},
         ],