fix when-then-chain (#3048)

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

@@ -17,7 +17,7 @@ fn prepare_mask(mask: &BooleanArray) -> BooleanArray {
 
     match mask.validity() {
         // nulls are set to true meaning we take from the right in the zip/ if_then_else kernel
-        Some(validity) if validity.null_count() == 0 => {
+        Some(validity) if validity.null_count() != 0 => {
             let mask = mask.values() & validity;
             BooleanArray::from_data_default(mask, None)
         }

polars/tests/it/lazy/expressions/arity.rs

@@ -75,5 +75,40 @@ fn includes_null_predicate_3038() -> Result<()> {
         "b" => [Some("good hit"), None, None],
     }?;
     assert!(res.frame_equal_missing(&exp_df));
+
+    let df = df! {
+        "a" => ["a1", "a2", "a3", "a4", "a2"],
+        "b" => [Some("tree"), None, None, None, None],
+    }?;
+    let res = df
+        .lazy()
+        .with_column(
+            when(col("b").map(
+                move |s| {
+                    s.utf8()?
+                        .to_lowercase()
+                        .contains("non-existent")
+                        .map(Into::into)
+                },
+                GetOutput::from_type(DataType::Boolean),
+            ))
+            .then(lit("weird-1"))
+            .when(col("a").eq(lit("a1".to_string())))
+            .then(lit("ok1"))
+            .when(col("a").eq(lit("a2".to_string())))
+            .then(lit("ok2"))
+            .when(lit(true))
+            .then(lit("ft"))
+            .otherwise(Expr::Literal(LiteralValue::Null))
+            .alias("c"),
+        )
+        .collect()?;
+    let exp_df = df! {
+        "a" => ["a1", "a2", "a3", "a4", "a2"],
+        "b" => [Some("tree"), None, None, None, None],
+        "c" => ["ok1", "ok2", "ft", "ft", "ok2"]
+    }?;
+    assert!(res.frame_equal_missing(&exp_df));
+
     Ok(())
 }