diff --git a/crates/polars-core/src/series/mod.rs b/crates/polars-core/src/series/mod.rs index 1e1039af4a31a..6cca321e0c219 100644 --- a/crates/polars-core/src/series/mod.rs +++ b/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 { let (lhs, rhs) = coerce_lhs_rhs(self, other)?; diff --git a/crates/polars-ops/src/series/ops/replace.rs b/crates/polars-ops/src/series/ops/replace.rs index af5e4ebf2e502..9c06b54037e15 100644 --- a/crates/polars-ops/src/series/ops/replace.rs +++ b/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 { - 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 diff --git a/py-polars/tests/unit/operations/test_replace.py b/py-polars/tests/unit/operations/test_replace.py index fa7375626cc3c..618384ae377ec 100644 --- a/py-polars/tests/unit/operations/test_replace.py +++ b/py-polars/tests/unit/operations/test_replace.py @@ -447,6 +447,14 @@ 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_to_one() -> None: lf = pl.LazyFrame({"a": [1, 2, 2, 3]}) result = lf.select(pl.col("a").replace([2, 3], 100))