|
| 1 | +"""MissForest honors resolved declarations after column-name normalization.""" |
| 2 | + |
| 3 | +from __future__ import annotations |
| 4 | + |
| 5 | +import pandas as pd |
| 6 | +import pytest |
| 7 | + |
| 8 | +import freshdata as fd |
| 9 | +from freshdata.imputation.missforest import MissForestImputer |
| 10 | + |
| 11 | +ISOLATE = { |
| 12 | + "strategy": "conservative", "fix_dtypes": False, "verbose": False, |
| 13 | + "drop_empty_rows": False, "drop_empty_columns": False, |
| 14 | + "missforest_add_indicators": False, |
| 15 | +} |
| 16 | + |
| 17 | + |
| 18 | +def _frame() -> pd.DataFrame: |
| 19 | + frame = pd.DataFrame({ |
| 20 | + "Observed Hours": [float(i % 12) for i in range(80)], |
| 21 | + "External Measure": [10 + 2 * (i % 12) + 0.5 * (i % 5) for i in range(80)], |
| 22 | + }) |
| 23 | + frame.loc[frame.index[::11], "External Measure"] = float("nan") |
| 24 | + return frame |
| 25 | + |
| 26 | + |
| 27 | +@pytest.mark.parametrize("role", ["target", "identifier", "both"]) |
| 28 | +@pytest.mark.parametrize("choice", ["global", "per-column"]) |
| 29 | +@pytest.mark.parametrize("declaration", ["External Measure", "external_measure"]) |
| 30 | +def test_normalized_declared_roles_are_preserved_without_loading_models( |
| 31 | + role, choice, declaration, monkeypatch |
| 32 | +): |
| 33 | + def forbidden(): |
| 34 | + raise AssertionError("a protected column must not train a model") |
| 35 | + |
| 36 | + monkeypatch.setattr(MissForestImputer, "_sklearn_forests", staticmethod(forbidden)) |
| 37 | + frame = _frame() |
| 38 | + original = frame.copy(deep=True) |
| 39 | + options = {} |
| 40 | + if role in ("target", "both"): |
| 41 | + options["target_column"] = declaration |
| 42 | + if role in ("identifier", "both"): |
| 43 | + options["id_columns"] = (declaration,) |
| 44 | + if choice == "global": |
| 45 | + options["impute"] = "missforest" |
| 46 | + else: |
| 47 | + options["impute_strategy"] = {"external_measure": "missforest"} |
| 48 | + |
| 49 | + out, report = fd.clean(frame, return_report=True, **ISOLATE, **options) |
| 50 | + |
| 51 | + pd.testing.assert_series_equal( |
| 52 | + out["external_measure"], original["External Measure"].rename("external_measure") |
| 53 | + ) |
| 54 | + pd.testing.assert_frame_equal(frame, original) |
| 55 | + (action,) = [a for a in report.actions if a.step == "impute"] |
| 56 | + assert action.count == 0 |
| 57 | + assert action.metadata["imputed_count"] == 0 |
| 58 | + assert action.metadata["selected_model_type"] is None |
| 59 | + expected = "identifier" if role == "identifier" else "target" |
| 60 | + assert action.metadata["fallback_reason"] == f"{expected} column is never imputed" |
| 61 | + |
| 62 | + |
| 63 | +@pytest.mark.parametrize("role", ["target", "identifier"]) |
| 64 | +@pytest.mark.parametrize("choice", ["global", "per-column"]) |
| 65 | +def test_original_declarations_still_work_when_renaming_is_disabled(role, choice): |
| 66 | + options = {"target_column": "External Measure"} if role == "target" else { |
| 67 | + "id_columns": ("External Measure",) |
| 68 | + } |
| 69 | + options.update({"impute": "missforest"} if choice == "global" else { |
| 70 | + "impute_strategy": {"External Measure": "missforest"} |
| 71 | + }) |
| 72 | + frame = _frame() |
| 73 | + out = fd.clean(frame, column_names=False, **ISOLATE, **options) |
| 74 | + pd.testing.assert_frame_equal(out, frame) |
| 75 | + |
| 76 | + |
| 77 | +def test_eligible_feature_still_trains_while_the_declared_target_is_preserved(): |
| 78 | + pytest.importorskip("sklearn") |
| 79 | + frame = _frame() |
| 80 | + frame["Other Measure"] = [20 + 3 * (i % 12) + 0.5 * (i % 5) for i in range(80)] |
| 81 | + frame.loc[frame.index[::13], "Other Measure"] = float("nan") |
| 82 | + out, report = fd.clean( |
| 83 | + frame, impute="missforest", target_column="External Measure", |
| 84 | + missforest_n_estimators=10, missforest_max_iter=2, return_report=True, **ISOLATE, |
| 85 | + ) |
| 86 | + pd.testing.assert_series_equal( |
| 87 | + out["external_measure"], frame["External Measure"].rename("external_measure") |
| 88 | + ) |
| 89 | + assert out["other_measure"].notna().all() |
| 90 | + (trained,) = [a for a in report.actions if a.column == "other_measure" and a.step == "impute"] |
| 91 | + assert trained.metadata["selected_model_type"] == "regressor" |
| 92 | + assert trained.count == 7 |
| 93 | + |
| 94 | + |
| 95 | +@pytest.mark.parametrize("role", ["target", "identifier", "both"]) |
| 96 | +@pytest.mark.parametrize("choice", ["global", "per-column"]) |
| 97 | +def test_raw_declaration_tracks_the_correct_deduplicated_header(role, choice): |
| 98 | + pytest.importorskip("sklearn") |
| 99 | + frame = _frame().rename(columns={"External Measure": "A B"}) |
| 100 | + frame["A-B"] = [20 + 3 * (i % 12) + 0.5 * (i % 5) for i in range(80)] |
| 101 | + frame.loc[frame.index[::13], "A-B"] = float("nan") |
| 102 | + options = {} |
| 103 | + if role in ("target", "both"): |
| 104 | + options["target_column"] = "A-B" |
| 105 | + if role in ("identifier", "both"): |
| 106 | + options["id_columns"] = ("A-B",) |
| 107 | + options.update({"impute": "missforest"} if choice == "global" else { |
| 108 | + "impute_strategy": {"a_b": "missforest", "a_b_2": "missforest"} |
| 109 | + }) |
| 110 | + out, report = fd.clean( |
| 111 | + frame, missforest_n_estimators=10, missforest_max_iter=2, |
| 112 | + return_report=True, **ISOLATE, **options, |
| 113 | + ) |
| 114 | + pd.testing.assert_series_equal(out["a_b_2"], frame["A-B"].rename("a_b_2")) |
| 115 | + assert out["a_b"].notna().all() |
| 116 | + (preserved,) = [a for a in report.actions if a.column == "a_b_2" and a.step == "impute"] |
| 117 | + expected = "identifier" if role == "identifier" else "target" |
| 118 | + assert preserved.metadata["fallback_reason"] == f"{expected} column is never imputed" |
| 119 | + |
| 120 | + |
| 121 | +def test_resolving_roles_does_not_mutate_the_callers_config(): |
| 122 | + frame = _frame() |
| 123 | + config = fd.CleanConfig( |
| 124 | + target_column="External Measure", impute="missforest", **ISOLATE |
| 125 | + ) |
| 126 | + out = fd.clean(frame, config=config) |
| 127 | + assert config.target_column == "External Measure" |
| 128 | + assert config.id_columns == () |
| 129 | + pd.testing.assert_series_equal( |
| 130 | + out["external_measure"], frame["External Measure"].rename("external_measure") |
| 131 | + ) |
0 commit comments