fix(impute): never impute declared id_columns or target_column (#423)

* fix(impute): never impute declared id_columns or target_column

Explicit imputation (impute=/impute_strategy=) filled the declared
id_columns and target_column, contradicting the documented safe default
('never imputes an identifier, modifies a target column'). Only
context-protected columns were skipped.

Skip the declared identifier and target columns in the simple mean/median/
mode strategies too, resolving names after renaming like the protected-column
guard. Report 'skipped: identifier column' / 'skipped: target column' when a
value would have changed, and warn when an impute_strategy entry names one of
them. MissForest already gates these roles with its own audited action, so its
candidate list is unchanged. Defaults are unaffected (impute is off).

* fix(impute): rename loop var to satisfy mypy (str vs str|None)

Author: kevincostner17

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

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:

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