diff --git a/tests/test_map.py b/tests/test_map.py index 3d1bda5..ebe57bc 100644 --- a/tests/test_map.py +++ b/tests/test_map.py @@ -49,6 +49,27 @@ def test_map_values_numeric(self) -> None: result, _ = plan.process(df) assert result["code"].to_list() == ["One", "Two", "Three", "One", "Two"] + def test_map_values_none_key_never_matches(self) -> None: + """Test that a None key does not match null values.""" + df = pl.DataFrame({"status": ["A", None]}) + plan = TransformPlan().map_values("status", {"A": "x", None: "was_null"}) + result, _ = plan.process(df) + assert result["status"].to_list() == ["x", None] + + def test_map_values_large_mapping_chunked(self) -> None: + """Test a large mapping on a multi-chunk frame (used to segfault).""" + n_keys = 2000 + df = pl.concat( + [pl.DataFrame({"id": [i]}) for i in range(2)], + rechunk=False, + ) + assert df["id"].n_chunks() == 2 + + mapping = {i: f"name {i}" for i in range(n_keys)} + plan = TransformPlan().map_values("id", mapping) + result, _ = plan.process(df, validate=True) + assert result["id"].to_list() == ["name 0", "name 1"] + class TestMapDiscretize: """Tests for map_discretize operation.""" @@ -167,6 +188,43 @@ def test_map_case_new_column(self, map_df: pl.DataFrame) -> None: assert "status_code" in result.columns assert "status" in result.columns # Original preserved + def test_map_case_large_case_list_chunked(self) -> None: + """Test a long case list on a multi-chunk frame (used to segfault).""" + df = pl.concat( + [pl.DataFrame({"id": [i]}) for i in range(2)], + rechunk=False, + ) + assert df["id"].n_chunks() == 2 + + cases = [(i, f"name {i}") for i in range(2000)] + plan = TransformPlan().map_case( + "id", cases=cases, default="other", new_column="name" + ) + result, _ = plan.process(df, validate=True) + assert result["name"].to_list() == ["name 0", "name 1"] + + def test_map_case_duplicate_condition_first_wins(self) -> None: + """Test that the first matching case wins.""" + df = pl.DataFrame({"status": ["A", "B"]}) + plan = TransformPlan().map_case( + "status", + cases=[("A", "first"), ("A", "second")], + default="other", + ) + result, _ = plan.process(df) + assert result["status"].to_list() == ["first", "other"] + + def test_map_case_none_condition_never_matches(self) -> None: + """Test that a None condition does not match null values.""" + df = pl.DataFrame({"status": ["A", None]}) + plan = TransformPlan().map_case( + "status", + cases=[("A", "x"), (None, "was_null")], + default="other", + ) + result, _ = plan.process(df) + assert result["status"].to_list() == ["x", "other"] + def test_map_case_empty_cases(self, map_df: pl.DataFrame) -> None: """Test case mapping with empty cases list.""" plan = TransformPlan().map_case("status", cases=[], default="default") diff --git a/transformplan/backends/polars.py b/transformplan/backends/polars.py index 96fe379..9b7e864 100644 --- a/transformplan/backends/polars.py +++ b/transformplan/backends/polars.py @@ -866,23 +866,24 @@ def map_values( default: Any, keep_unmapped: bool, # noqa: FBT001 ) -> pl.DataFrame: - expr = pl.col(column) - items = list(mapping.items()) - if not items: + # A when/then chain nests one expression per entry and overflows the + # stack for large mappings; replace_strict is flat and size-independent. + # A `None` key never matched under the chain (``col == None`` is null, + # not true), so it is dropped to keep nulls falling through to the + # fallback -- matching the DuckDB backend, where ``= NULL`` is never true. + if not mapping: return data - first_key, first_val = items[0] - chain = pl.when(expr == first_key).then(pl.lit(first_val)) - - for key, val in items[1:]: - chain = chain.when(expr == key).then(pl.lit(val)) - - if keep_unmapped: - chain = chain.otherwise(expr) - else: - chain = chain.otherwise(pl.lit(default)) + expr = pl.col(column) + fallback = expr if keep_unmapped else pl.lit(default) - return data.with_columns(chain.alias(column)) + pairs = {k: v for k, v in mapping.items() if k is not None} + if not pairs: + # Every key was None, so nothing can ever match: the fallback wins. + return data.with_columns(fallback.alias(column)) + return data.with_columns( + expr.replace_strict(pairs, default=fallback).alias(column) + ) def map_case( self, @@ -892,18 +893,24 @@ def map_case( default: Any, new_column: str, ) -> pl.DataFrame: - if not cases: + # A when/then chain nests one expression per case and overflows the + # stack for long case lists; replace_strict is flat and size-independent. + # `None` conditions never matched under the chain (``col == None`` is + # null, not true), and duplicate conditions resolved first-match-wins, + # so both are normalised away before handing the cases to polars. + mapping: dict[Any, Any] = {} + for cond_val, result_val in cases: + if cond_val is not None and cond_val not in mapping: + mapping[cond_val] = result_val + + if not mapping: return data.with_columns(pl.lit(default).alias(new_column)) - col = pl.col(column) - cond_val, result_val = cases[0] - chain = pl.when(col == cond_val).then(pl.lit(result_val)) - - for cond_val, result_val in cases[1:]: - chain = chain.when(col == cond_val).then(pl.lit(result_val)) - - chain = chain.otherwise(pl.lit(default)) - return data.with_columns(chain.alias(new_column)) + return data.with_columns( + pl.col(column) + .replace_strict(mapping, default=pl.lit(default)) + .alias(new_column) + ) def map_from_column( self,