From d66c99333c3a13cb75026c88405e45b6c1a1be30 Mon Sep 17 00:00:00 2001
From: nameexhaustion <simonlin.rqmmw@slmail.me>
Date: Tue, 9 Jan 2024 11:35:59 +1100
Subject: [PATCH] fix: incorrect `Expr.replace` when old is single NULL

---
 crates/polars-core/src/series/mod.rs            |  3 ++-
 crates/polars-ops/src/series/ops/replace.rs     | 17 ++++++++++++++---
 py-polars/tests/unit/operations/test_replace.py |  8 ++++++++
 3 files changed, 24 insertions(+), 4 deletions(-)

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<Series> {
         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<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
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))