diff --git a/CHANGELOG.md b/CHANGELOG.md index 7631294..efb438d 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/src/freshdata/cleaner.py b/src/freshdata/cleaner.py index 94de877..6d03938 100644 --- a/src/freshdata/cleaner.py +++ b/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. diff --git a/src/freshdata/steps/missing.py b/src/freshdata/steps/missing.py index 9b3d232..1d46c2d 100644 --- a/src/freshdata/steps/missing.py +++ b/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: diff --git a/tests/test_missforest_declared_roles.py b/tests/test_missforest_declared_roles.py new file mode 100644 index 0000000..a209d8b --- /dev/null +++ b/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") + )