fix(rust, python): adjust for null values in str.replace fast path (#…

…10132)
pola-rs · Jul 28, 2023 · f27c9b7 · f27c9b7
1 parent a207241
commit f27c9b7
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 2 deletions.
diff --git a/polars/polars-ops/src/chunked_array/strings/replace.rs b/polars/polars-ops/src/chunked_array/strings/replace.rs
@@ -61,8 +61,14 @@ pub(super) fn replace_lit_n_char(
 
             // set the end of this string region
             // safety: invariant of Utf8Array tells us that there is a next offset.
-            if let Some(next) = offsets_iter.next() {
-                end = *next as usize - 1;
+
+            // must loop to skip null values, as they have the same offsets
+            for next in offsets_iter.by_ref() {
+                let new_end = *next as usize - 1;
+                if new_end != end {
+                    end = new_end;
+                    break;
+                }
             }
         }
     }

diff --git a/py-polars/tests/unit/namespaces/test_string.py b/py-polars/tests/unit/namespaces/test_string.py
@@ -723,3 +723,17 @@ def test_titlecase() -> None:
             "And\tA\t Tab",
         ]
     }
+
+
+def test_string_replace_with_nulls_10124() -> None:
+    df = pl.DataFrame({"col1": ["S", "S", "S", None, "S", "S", "S", "S"]})
+
+    assert df.select(
+        pl.col("col1"),
+        pl.col("col1").str.replace("S", "O", n=1).alias("n_1"),
+        pl.col("col1").str.replace("S", "O", n=3).alias("n_3"),
+    ).to_dict(False) == {
+        "col1": ["S", "S", "S", None, "S", "S", "S", "S"],
+        "n_1": ["O", "O", "O", None, "O", "O", "O", "O"],
+        "n_3": ["O", "O", "O", None, "O", "O", "O", "O"],
+    }