fix: Expr.replace to single value did not replace NULLs (pola-rs#13551

)

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

@@ -64,8 +64,18 @@ impl PrivateSeries for NullChunked {
     }
 
     #[cfg(feature = "zip_with")]
-    fn zip_with_same_type(&self, _mask: &BooleanChunked, _other: &Series) -> PolarsResult<Series> {
-        Ok(self.clone().into_series())
+    fn zip_with_same_type(&self, mask: &BooleanChunked, other: &Series) -> PolarsResult<Series> {
+        let len = match (self.len(), mask.len(), other.len()) {
+            (a, b, c) if a == b && b == c => a,
+            (1, a, b) | (a, 1, b) | (a, b, 1) if a == b => a,
+            (a, 1, 1) | (1, a, 1) | (1, 1, a) => a,
+            (_, 0, _) => 0,
+            _ => {
+                polars_bail!(ShapeMismatch: "shapes of `self`, `mask` and `other` are not suitable for `zip_with` operation")
+            },
+        };
+
+        Ok(Self::new(self.name().into(), len).into_series())
     }
     fn explode_by_offsets(&self, offsets: &[i64]) -> Series {
         ExplodeByOffsets::explode_by_offsets(self, offsets)

crates/polars-core/src/series/mod.rs

@@ -421,7 +421,8 @@ impl Series {
     }
 
     /// Create a new ChunkedArray with values from self where the mask evaluates `true` and values
-    /// from `other` where the mask evaluates `false`
+    /// from `other` where the mask evaluates `false`. This function automatically broadcasts unit
+    /// length inputs.
     #[cfg(feature = "zip_with")]
     pub fn zip_with(&self, mask: &BooleanChunked, other: &Series) -> PolarsResult<Series> {
         let (lhs, rhs) = coerce_lhs_rhs(self, other)?;

crates/polars-ops/src/series/ops/replace.rs

@@ -1,3 +1,5 @@
+use std::ops::BitOr;
+
 use polars_core::prelude::*;
 use polars_core::utils::try_get_supertype;
 use polars_error::{polars_bail, polars_ensure, PolarsResult};
@@ -62,9 +64,18 @@ fn replace_by_single(
     new: &Series,
     default: &Series,
 ) -> PolarsResult<Series> {
-    let mask = is_in(s, old)?;
-    let new_broadcast = new.new_from_index(0, default.len());
-    new_broadcast.zip_with(&mask, default)
+    let mask = if old.null_count() == old.len() {
+        s.is_null()
+    } else {
+        let mask = is_in(s, old)?;
+
+        if old.null_count() == 0 {
+            mask
+        } else {
+            mask.bitor(s.is_null())
+        }
+    };
+    new.zip_with(&mask, default)
 }
 
 /// General case for replacing by multiple values

py-polars/tests/unit/functions/test_whenthen.py → py-polars/tests/unit/functions/test_when_then.py

@@ -514,3 +514,14 @@ def test_when_predicates_kwargs() -> None:
         ),
         pl.DataFrame({"misc": ["?", "z in (a|b), y<0", "?", "y=1"]}),
     )
+
+
+def test_when_then_null_broadcast() -> None:
+    assert (
+        pl.select(
+            pl.when(pl.repeat(True, 2, dtype=pl.Boolean)).then(
+                pl.repeat(None, 1, dtype=pl.Null)
+            )
+        ).height
+        == 2
+    )

py-polars/tests/unit/operations/test_replace.py

@@ -447,6 +447,21 @@ def test_replace_fast_path_one_to_one() -> None:
     assert_frame_equal(result, expected)
 
 
+def test_replace_fast_path_one_null_to_one() -> None:
+    # https://github.com/pola-rs/polars/issues/13391
+    lf = pl.LazyFrame({"a": [1, None]})
+    result = lf.select(pl.col("a").replace(None, 100))
+    expected = pl.LazyFrame({"a": [1, 100]})
+    assert_frame_equal(result, expected)
+
+
+def test_replace_fast_path_many_with_null_to_one() -> None:
+    lf = pl.LazyFrame({"a": [1, 2, None]})
+    result = lf.select(pl.col("a").replace([1, None], 100))
+    expected = pl.LazyFrame({"a": [100, 2, 100]})
+    assert_frame_equal(result, expected)
+
+
 def test_replace_fast_path_many_to_one() -> None:
     lf = pl.LazyFrame({"a": [1, 2, 2, 3]})
     result = lf.select(pl.col("a").replace([2, 3], 100))