fix: Wrong result on when().then().otherwise() on struct when both …

…result are broadcast (#19000) Co-authored-by: ritchie <[email protected]>
pola-rs · Sep 29, 2024 · 570e3a2 · 570e3a2
1 parent ebd5302
commit 570e3a2
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/crates/polars-core/src/chunked_array/ops/zip.rs b/crates/polars-core/src/chunked_array/ops/zip.rs
@@ -335,7 +335,7 @@ impl ChunkZip<StructType> for StructChunked {
             let rechunked_validity = match (l.len(), r.len()) {
                 (1, 1) if length != 1 => match (l.null_count() == 0, r.null_count() == 0) {
                     (true, true) => None,
-                    (true, false) => {
+                    (false, true) => {
                         if mask.chunks().len() == 1 {
                             let m = mask.chunks()[0]
                                 .as_any()
@@ -350,7 +350,7 @@ impl ChunkZip<StructType> for StructChunked {
                             )
                         }
                     },
-                    (false, true) => {
+                    (true, false) => {
                         if mask.chunks().len() == 1 {
                             let m = mask.chunks()[0]
                                 .as_any()

diff --git a/py-polars/tests/unit/functions/test_when_then.py b/py-polars/tests/unit/functions/test_when_then.py
@@ -625,6 +625,20 @@ def test_when_then_else_struct_18961() -> None:
     )
     assert expected == ans
 
+    df = pl.DataFrame({"left": v1, "right": v2, "mask": [True, False]})
+
+    expected2 = [None, {"foo": 0, "bar": "1"}]
+    ans = (
+        df.select(
+            pl.when(pl.col.mask)
+            .then(pl.col.left.first())
+            .otherwise(pl.col.right.first())
+        )
+        .get_column("left")
+        .to_list()
+    )
+    assert expected2 == ans
+
 
 def test_when_then_supertype_15975() -> None:
     df = pl.DataFrame({"a": [1, 2, 3]})