diff --git a/CHANGELOG.md b/CHANGELOG.md index 39308adf..6d1e0cda 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,13 @@ adheres to [Semantic Versioning](https://semver.org/). ## [Unreleased] ### Fixed +- Explicit imputation (`impute="mean"`, `"median"`, `"mode"`, `"auto"`, + `"missforest"` or an `impute_strategy` entry) no longer fills the declared + `id_columns` or `target_column`, as documented. Those columns keep their + missing values and the report records `skipped: identifier column` or + `skipped: target column`. An `impute_strategy` entry naming one of them is + ignored with a warning. Declared names resolve after column renaming, and the + columns can still serve as MissForest features for other columns. - `fd.clean` on a Spark DataFrame no longer raises `TypeError: cannot materialize source of type DataFrame`. Under the default `strategy="balanced"` the pandas fallback now materializes a Spark source diff --git a/src/freshdata/steps/missing.py b/src/freshdata/steps/missing.py index 6d709202..9b3d2322 100644 --- a/src/freshdata/steps/missing.py +++ b/src/freshdata/steps/missing.py @@ -7,6 +7,7 @@ from __future__ import annotations +from collections.abc import Iterable from typing import Any import pandas as pd @@ -48,14 +49,51 @@ def _strategy_for_column(col: object, config: CleanConfig) -> str | None: return config.impute +def _declared_roles(config: CleanConfig, columns: Iterable[object]) -> dict[str, str]: + """Declared identifier and target columns present in *columns*, by role. + + Names resolve like context-protected columns (exact, else snake case), so a + declared ``"Customer ID"`` still matches ``customer_id`` after renaming. A + column declared as both is reported as the target. + """ + from ..guard import _match_columns # noqa: PLC0415 — cycle-safe lazy import + + names = [str(c) for c in columns] + roles: dict[str, str] = {} + if config.target_column is not None: + for name in _match_columns([str(config.target_column)], names): + if name in names: + roles[name] = "target" + for name in _match_columns([str(c) for c in config.id_columns], names): + if name in names: + roles.setdefault(name, "identifier") + return roles + + def impute_missing(df: pd.DataFrame, config: CleanConfig, report: CleanReport) -> pd.DataFrame: - """Fill missing values per column according to explicit impute config.""" + """Fill missing values per column according to explicit impute config. + + Context-protected columns and the declared ``id_columns`` and + ``target_column`` are never filled, whatever ``impute`` or + ``impute_strategy`` says: imputing an identifier corrupts keys and imputing + the target leaks into it. They can still inform ``"missforest"`` as + features for other columns. + """ if config.impute is None and not config.impute_strategy: return df from ..guard import hard_protected_columns # noqa: PLC0415 — cycle-safe lazy import protected = hard_protected_columns(config, df.columns) + roles = _declared_roles(config, df.columns) + for name, declared_role in roles.items(): + if config.impute_strategy and name in config.impute_strategy: + report.add_warning( + f"impute_strategy for '{name}' ignored: it is the declared " + f"{declared_role} column") + # MissForest applies its own role gates (target and identifier columns are + # preserved with an audited fallback action), so declared roles stay in its + # column list and are reported there. missforest_columns = [ col for col in df.columns if str(col) not in protected @@ -78,6 +116,11 @@ def impute_missing(df: pd.DataFrame, config: CleanConfig, strategy = _strategy_for_column(col, config) if strategy is None or strategy == "missforest": continue + role = roles.get(str(col)) + if role is not None: + if int(df[col].isna().sum()) and df[col].notna().any(): + report.add("impute", f"skipped: {role} column", column=str(col)) + continue s = df[col] n_missing = int(s.isna().sum()) if n_missing == 0 or s.notna().sum() == 0: diff --git a/tests/test_missing.py b/tests/test_missing.py index b7c015b8..dece9e9a 100644 --- a/tests/test_missing.py +++ b/tests/test_missing.py @@ -1,4 +1,6 @@ +import numpy as np import pandas as pd +import pytest import freshdata as fd @@ -87,3 +89,84 @@ def test_boolean_mode_imputation_via_auto(): out = fd.clean(df, impute="auto", **KEEP_ROWS) assert out["b"].isna().sum() == 0 assert bool(out["b"].iloc[2]) is True + + +def _roles_frame() -> pd.DataFrame: + return pd.DataFrame( + { + "customer_id": [101.0, None, 103.0, 104.0, 105.0], + "churn": [0.0, 1.0, None, 1.0, 0.0], + "spend": [10.0, None, 30.0, 40.0, 50.0], + } + ) + + +@pytest.mark.parametrize("impute", ["mean", "median", "auto", "mode"]) +def test_declared_id_and_target_are_never_imputed(impute): + df = _roles_frame() + out, report = fd.clean( + df, + impute=impute, + id_columns=("customer_id",), + target_column="churn", + return_report=True, + **KEEP_ROWS, + ) + assert out["customer_id"].isna().sum() == 1 + assert out["churn"].isna().sum() == 1 + assert out["spend"].isna().sum() == 0 + notes = {(a.column, a.description) for a in report if a.step == "impute"} + assert ("customer_id", "skipped: identifier column") in notes + assert ("churn", "skipped: target column") in notes + + +def test_explicit_impute_strategy_for_target_is_ignored_with_warning(): + df = _roles_frame() + out, report = fd.clean( + df, + impute_strategy={"churn": "mean", "spend": "mean"}, + target_column="churn", + return_report=True, + **KEEP_ROWS, + ) + assert out["churn"].isna().sum() == 1 + assert out["spend"].isna().sum() == 0 + assert any( + "impute_strategy for 'churn' ignored: it is the declared target column" in w + for w in report.warnings + ) + + +def test_declared_id_matches_renamed_column(): + df = pd.DataFrame({"Customer ID": [1.0, None, 3.0, 4.0], "Spend": [1.0, None, 3.0, 4.0]}) + out = fd.clean(df, impute="mean", id_columns=("Customer ID",), **KEEP_ROWS) + id_col = next(c for c in out.columns if "customer" in str(c).lower()) + spend_col = next(c for c in out.columns if "spend" in str(c).lower()) + assert out[id_col].isna().sum() == 1 + assert out[spend_col].isna().sum() == 0 + + +def test_no_skip_note_when_declared_column_has_nothing_to_fill(): + df = _roles_frame().fillna({"customer_id": 102.0, "churn": 1.0}) + _, report = fd.clean( + df, + impute="mean", + id_columns=("customer_id",), + target_column="churn", + return_report=True, + **KEEP_ROWS, + ) + assert not [a for a in report if a.step == "impute" and "skipped:" in a.description] + + +def test_missforest_never_imputes_declared_target(): + pytest.importorskip("sklearn") + rng = np.random.default_rng(0) + n = 60 + x = rng.normal(size=n) + df = pd.DataFrame({"x": x, "y": x * 2.0 + rng.normal(scale=0.1, size=n)}) + df.loc[[3, 7, 11], "y"] = np.nan + df.loc[[5, 9], "x"] = np.nan + out = fd.clean(df, impute="missforest", target_column="y", **KEEP_ROWS) + assert out["y"].isna().sum() == 3 + assert out["x"].isna().sum() == 0