Skip to content

Commit

Permalink
fix: incorrect Expr.replace when old is single NULL
Browse files Browse the repository at this point in the history
  • Loading branch information
nameexhaustion committed Jan 9, 2024
1 parent 67507ff commit 58197ab
Show file tree
Hide file tree
Showing 3 changed files with 24 additions and 4 deletions.
3 changes: 2 additions & 1 deletion crates/polars-core/src/series/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)?;
Expand Down
17 changes: 14 additions & 3 deletions crates/polars-ops/src/series/ops/replace.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
use std::ops::BitAnd;

use polars_core::prelude::*;
use polars_core::utils::try_get_supertype;
use polars_error::{polars_bail, polars_ensure, PolarsResult};
Expand Down Expand Up @@ -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
Expand Down
8 changes: 8 additions & 0 deletions py-polars/tests/unit/operations/test_replace.py
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down

0 comments on commit 58197ab

Please sign in to comment.