fix: Return correct value for when().then().else() on structs when …

…using `first()`\`last()` (#18969)
pola-rs · Sep 27, 2024 · 2dbb444 · 2dbb444
1 parent 653a4cd
commit 2dbb444
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 5 deletions.
diff --git a/crates/polars-core/src/chunked_array/ops/zip.rs b/crates/polars-core/src/chunked_array/ops/zip.rs
@@ -375,10 +375,10 @@ impl ChunkZip<StructType> for StructChunked {
                         .all(|(r, m)| r == m));
 
                     let combine = if l.null_count() == 0 {
-                        |r: Option<&Bitmap>, m: &Bitmap| r.map(|r| arrow::bitmap::or_not(r, m))
+                        |r: Option<&Bitmap>, m: &Bitmap| r.map(|r| arrow::bitmap::or(r, m))
                     } else {
                         |r: Option<&Bitmap>, m: &Bitmap| {
-                            Some(r.map_or_else(|| m.clone(), |r| arrow::bitmap::and(r, m)))
+                            Some(r.map_or_else(|| m.clone(), |r| arrow::bitmap::and_not(r, m)))
                         }
                     };
 
@@ -411,10 +411,10 @@ impl ChunkZip<StructType> for StructChunked {
                         .all(|(l, m)| l == m));
 
                     let combine = if r.null_count() == 0 {
-                        |r: Option<&Bitmap>, m: &Bitmap| r.map(|r| arrow::bitmap::or(r, m))
+                        |l: Option<&Bitmap>, m: &Bitmap| l.map(|l| arrow::bitmap::or_not(l, m))
                     } else {
-                        |r: Option<&Bitmap>, m: &Bitmap| {
-                            Some(r.map_or_else(|| m.clone(), |r| arrow::bitmap::and_not(r, m)))
+                        |l: Option<&Bitmap>, m: &Bitmap| {
+                            Some(l.map_or_else(|| m.clone(), |l| arrow::bitmap::and(l, m)))
                         }
                     };
 

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