fix(missforest): preserve declared roles after header normalization

CHANGELOG.md

@@ -8,6 +8,12 @@ adheres to [Semantic Versioning](https://semver.org/).
 
 ### Fixed
 
+- MissForest now preserves declared identifiers and targets whose headers were
+  normalized during cleaning, including headers that need deduplication. Exact
+  raw declarations follow the positional rename map, and MissForest uses the
+  same resolved roles as simple imputation instead of training a model for a
+  protected column (#525).
+
 - Model time budgets now stop at their deadline, including a zero-second
   budget when the clock has not advanced. This prevents optional model calls
   from starting after the time allowance has been consumed.

src/freshdata/cleaner.py

@@ -145,6 +145,22 @@ def run_pipeline(  # noqa: PLR0915 - fixed-order pipeline orchestration
     original_columns = list(out.columns)
     if config.column_names:
         out = normalize_column_names(out, report)
+        if config.target_column is not None or config.id_columns:
+            # Resolve exact raw declarations through the positional rename map
+            # before collisions lose their original spelling (A B / A-B).
+            role_names: dict[str, str] = {}
+            for original, renamed in zip(original_columns, out.columns):
+                role_names.setdefault(str(original), str(renamed))
+            config = dataclasses.replace(
+                config,
+                target_column=(
+                    role_names.get(config.target_column, config.target_column)
+                    if config.target_column is not None else None
+                ),
+                id_columns=tuple(
+                    role_names.get(str(name), str(name)) for name in config.id_columns
+                ),
+            )
         _emit_progress(progress_callback, "column_names", "after", out)
     # Full column evidence for compliance reports, in the report's (post-rename)
     # namespace, so untouched columns are still visible without the source frame.

src/freshdata/steps/missing.py

@@ -104,10 +104,17 @@ def impute_missing(df: pd.DataFrame, config: CleanConfig,
         from ..engine.context import build_contexts  # noqa: PLC0415
         from ..imputation.missforest import MissForestImputer  # noqa: PLC0415
 
+        contexts = build_contexts(df, config)
+        # The shared declaration matcher resolves names after header repairs.
+        # Carry those roles into MissForest's own preservation/audit gates.
+        for col in missforest_columns:
+            resolved_role = roles.get(str(col))
+            if resolved_role is not None:
+                contexts[col].role = "id" if resolved_role == "identifier" else "target"
         df = MissForestImputer(config, report).impute(
             df,
             missforest_columns,
-            build_contexts(df, config),
+            contexts,
         )
 
     for col in df.columns:

tests/test_missforest_declared_roles.py

@@ -0,0 +1,131 @@
+"""MissForest honors resolved declarations after column-name normalization."""
+
+from __future__ import annotations
+
+import pandas as pd
+import pytest
+
+import freshdata as fd
+from freshdata.imputation.missforest import MissForestImputer
+
+ISOLATE = {
+    "strategy": "conservative", "fix_dtypes": False, "verbose": False,
+    "drop_empty_rows": False, "drop_empty_columns": False,
+    "missforest_add_indicators": False,
+}
+
+
+def _frame() -> pd.DataFrame:
+    frame = pd.DataFrame({
+        "Observed Hours": [float(i % 12) for i in range(80)],
+        "External Measure": [10 + 2 * (i % 12) + 0.5 * (i % 5) for i in range(80)],
+    })
+    frame.loc[frame.index[::11], "External Measure"] = float("nan")
+    return frame
+
+
+@pytest.mark.parametrize("role", ["target", "identifier", "both"])
+@pytest.mark.parametrize("choice", ["global", "per-column"])
+@pytest.mark.parametrize("declaration", ["External Measure", "external_measure"])
+def test_normalized_declared_roles_are_preserved_without_loading_models(
+    role, choice, declaration, monkeypatch
+):
+    def forbidden():
+        raise AssertionError("a protected column must not train a model")
+
+    monkeypatch.setattr(MissForestImputer, "_sklearn_forests", staticmethod(forbidden))
+    frame = _frame()
+    original = frame.copy(deep=True)
+    options = {}
+    if role in ("target", "both"):
+        options["target_column"] = declaration
+    if role in ("identifier", "both"):
+        options["id_columns"] = (declaration,)
+    if choice == "global":
+        options["impute"] = "missforest"
+    else:
+        options["impute_strategy"] = {"external_measure": "missforest"}
+
+    out, report = fd.clean(frame, return_report=True, **ISOLATE, **options)
+
+    pd.testing.assert_series_equal(
+        out["external_measure"], original["External Measure"].rename("external_measure")
+    )
+    pd.testing.assert_frame_equal(frame, original)
+    (action,) = [a for a in report.actions if a.step == "impute"]
+    assert action.count == 0
+    assert action.metadata["imputed_count"] == 0
+    assert action.metadata["selected_model_type"] is None
+    expected = "identifier" if role == "identifier" else "target"
+    assert action.metadata["fallback_reason"] == f"{expected} column is never imputed"
+
+
+@pytest.mark.parametrize("role", ["target", "identifier"])
+@pytest.mark.parametrize("choice", ["global", "per-column"])
+def test_original_declarations_still_work_when_renaming_is_disabled(role, choice):
+    options = {"target_column": "External Measure"} if role == "target" else {
+        "id_columns": ("External Measure",)
+    }
+    options.update({"impute": "missforest"} if choice == "global" else {
+        "impute_strategy": {"External Measure": "missforest"}
+    })
+    frame = _frame()
+    out = fd.clean(frame, column_names=False, **ISOLATE, **options)
+    pd.testing.assert_frame_equal(out, frame)
+
+
+def test_eligible_feature_still_trains_while_the_declared_target_is_preserved():
+    pytest.importorskip("sklearn")
+    frame = _frame()
+    frame["Other Measure"] = [20 + 3 * (i % 12) + 0.5 * (i % 5) for i in range(80)]
+    frame.loc[frame.index[::13], "Other Measure"] = float("nan")
+    out, report = fd.clean(
+        frame, impute="missforest", target_column="External Measure",
+        missforest_n_estimators=10, missforest_max_iter=2, return_report=True, **ISOLATE,
+    )
+    pd.testing.assert_series_equal(
+        out["external_measure"], frame["External Measure"].rename("external_measure")
+    )
+    assert out["other_measure"].notna().all()
+    (trained,) = [a for a in report.actions if a.column == "other_measure" and a.step == "impute"]
+    assert trained.metadata["selected_model_type"] == "regressor"
+    assert trained.count == 7
+
+
+@pytest.mark.parametrize("role", ["target", "identifier", "both"])
+@pytest.mark.parametrize("choice", ["global", "per-column"])
+def test_raw_declaration_tracks_the_correct_deduplicated_header(role, choice):
+    pytest.importorskip("sklearn")
+    frame = _frame().rename(columns={"External Measure": "A B"})
+    frame["A-B"] = [20 + 3 * (i % 12) + 0.5 * (i % 5) for i in range(80)]
+    frame.loc[frame.index[::13], "A-B"] = float("nan")
+    options = {}
+    if role in ("target", "both"):
+        options["target_column"] = "A-B"
+    if role in ("identifier", "both"):
+        options["id_columns"] = ("A-B",)
+    options.update({"impute": "missforest"} if choice == "global" else {
+        "impute_strategy": {"a_b": "missforest", "a_b_2": "missforest"}
+    })
+    out, report = fd.clean(
+        frame, missforest_n_estimators=10, missforest_max_iter=2,
+        return_report=True, **ISOLATE, **options,
+    )
+    pd.testing.assert_series_equal(out["a_b_2"], frame["A-B"].rename("a_b_2"))
+    assert out["a_b"].notna().all()
+    (preserved,) = [a for a in report.actions if a.column == "a_b_2" and a.step == "impute"]
+    expected = "identifier" if role == "identifier" else "target"
+    assert preserved.metadata["fallback_reason"] == f"{expected} column is never imputed"
+
+
+def test_resolving_roles_does_not_mutate_the_callers_config():
+    frame = _frame()
+    config = fd.CleanConfig(
+        target_column="External Measure", impute="missforest", **ISOLATE
+    )
+    out = fd.clean(frame, config=config)
+    assert config.target_column == "External Measure"
+    assert config.id_columns == ()
+    pd.testing.assert_series_equal(
+        out["external_measure"], frame["External Measure"].rename("external_measure")
+    )