From f2e65d75c29d68906d424aab67fed5b7ad1cfb41 Mon Sep 17 00:00:00 2001 From: Kevin Costner <120246174+kevincostner17@users.noreply.github.com> Date: Tue, 15 Sep 2026 14:50:32 +0530 Subject: [PATCH] fix(missforest): keep integer dtypes by rounding imputed values MissForest fills from a float64 working frame, so regressor predictions for a nullable Int64 column are fractional floats. Series.where refuses them on the masked integer array and _assign_success fell back to astype(object), leaving a column that mixed Python ints with fractional floats and had no report note. _fallback_fill had the same object fallback for a fractional median. For numpy and nullable integer columns both paths now round the imputed values half-to-even to Python ints and store them in the column's own dtype. The fallback median uses exact_int_stat for columns holding values beyond 2**53, so no digits are lost. The action rationale notes the rounding and the metadata gains a rounded_to_integer flag, set only when regressor predictions or a median were actually rounded. Numeric columns that still cannot hold the fill are cast to float64 with a rationale note instead of object. Float, bool, categorical and string targets keep their existing behaviour. Closes #263 --- src/freshdata/imputation/missforest.py | 112 ++++++++++++++--- tests/test_missforest_nullable_int.py | 163 +++++++++++++++++++++++++ 2 files changed, 261 insertions(+), 14 deletions(-) create mode 100644 tests/test_missforest_nullable_int.py diff --git a/src/freshdata/imputation/missforest.py b/src/freshdata/imputation/missforest.py index d6ebec3..ad519e8 100644 --- a/src/freshdata/imputation/missforest.py +++ b/src/freshdata/imputation/missforest.py @@ -13,9 +13,14 @@ import numpy as np import pandas as pd -from pandas.api.types import is_bool_dtype, is_datetime64_any_dtype, is_numeric_dtype - -from .._util import add_column, safe_median +from pandas.api.types import ( + is_bool_dtype, + is_datetime64_any_dtype, + is_integer_dtype, + is_numeric_dtype, +) + +from .._util import add_column, exact_int_stat, exceeds_float64_exact, safe_median from ..config import CleanConfig from ..engine.context import ColumnContext from ..report import CleanReport @@ -301,10 +306,12 @@ def _assign_success(self, df: pd.DataFrame, work: pd.DataFrame, plan: _ColumnPla missing_categories = [v for v in pd.unique(filled_values) if v not in s.cat.categories] if missing_categories: s = s.cat.add_categories(missing_categories) - try: - combined = s.where(s.notna(), filled_values) - except (TypeError, ValueError): - combined = s.astype(object).where(s.notna(), filled_values) + combined, dtype_note = _combine_filled(s, filled_values) + if dtype_note == "rounded" and plan.model_type != "regressor": + # Classifier predictions are observed labels, so nothing was + # rounded; the integer dtype is still kept. + dtype_note = None + rounded = dtype_note == "rounded" df[plan.column] = combined imputed = int(plan.missing_mask.sum()) indicator_added = self._maybe_indicator(df, plan) @@ -328,7 +335,8 @@ def _assign_success(self, df: pd.DataFrame, work: pd.DataFrame, plan: _ColumnPla column=str(plan.column), count=imputed, rationale="explicit MissForest imputation selected; random forests model " - "nonlinear relationships across mixed tabular predictors", + "nonlinear relationships across mixed tabular predictors" + + _dtype_rationale(dtype_note, s), risk=risk, confidence=confidence, model_id=f"missforest_{plan.model_type}", @@ -337,6 +345,7 @@ def _assign_success(self, df: pd.DataFrame, work: pd.DataFrame, plan: _ColumnPla imputed, fallback_reason=None, indicator_added=indicator_added, + rounded_to_integer=rounded, ), ) self.report.columns_imputed.append(str(plan.column)) @@ -349,7 +358,13 @@ def _fallback_fill( reason: str, ) -> None: s = df[col] - value = safe_median(s) if ctx.role == "numeric" and s.notna().any() else _mode_value(s) + use_median = ctx.role == "numeric" and s.notna().any() + if use_median and exceeds_float64_exact(s): + value: Any = exact_int_stat(s, "median") # a float median would lose digits + elif use_median: + value = safe_median(s) + else: + value = _mode_value(s) if value is None or pd.isna(value): value = "Missing" if ctx.role in ("categorical", "boolean") else None if value is None or pd.isna(value): @@ -357,10 +372,10 @@ def _fallback_fill( return if isinstance(s.dtype, pd.CategoricalDtype) and value not in s.cat.categories: s = s.cat.add_categories([value]) - try: - filled = s.fillna(value) - except (TypeError, ValueError): - filled = s.astype(object).fillna(value) + filled, dtype_note = _fillna_keeping_dtype(s, value) + if dtype_note == "rounded" and not use_median: + dtype_note = None # a mode is an observed value, nothing was rounded + rounded = dtype_note == "rounded" df[col] = filled imputed = ctx.n_missing confidence = 0.65 if ctx.n_rows >= self.config.missforest_min_rows_for_model else 0.55 @@ -369,7 +384,7 @@ def _fallback_fill( f"missforest fallback filled {imputed} missing value(s) with safe simple imputation", column=str(col), count=imputed, - rationale=f"MissForest not used: {reason}", + rationale=f"MissForest not used: {reason}" + _dtype_rationale(dtype_note, s), risk="medium", confidence=confidence, model_id="missforest_fallback", @@ -378,6 +393,7 @@ def _fallback_fill( imputed, fallback_reason=reason, indicator_added=False, + rounded_to_integer=rounded, ), ) self.report.columns_imputed.append(str(col)) @@ -407,6 +423,7 @@ def _preserve(self, df: pd.DataFrame, col: object, ctx: ColumnContext, reason: s "convergence_delta": None, "oob_score": None, "indicator_added": False, + "rounded_to_integer": False, }, ) self.report.columns_preserved.append(str(col)) @@ -443,6 +460,7 @@ def _metadata( *, fallback_reason: str | None, indicator_added: bool, + rounded_to_integer: bool = False, ) -> dict[str, Any]: selected = None if plan.model_type == "regressor": @@ -459,6 +477,7 @@ def _metadata( "convergence_delta": self._last_delta if fallback_reason is None else None, "oob_score": self._last_oob.get(plan.column) if fallback_reason is None else None, "indicator_added": indicator_added, + "rounded_to_integer": rounded_to_integer, } @staticmethod @@ -485,3 +504,68 @@ def _mode_value(s: pd.Series) -> Any | None: except TypeError: return None return counts.index[0] if len(counts) else None + + +def _is_integer_column(s: pd.Series) -> bool: + """True for numpy and nullable integer columns (never bool or categorical).""" + return is_integer_dtype(s.dtype) and not is_bool_dtype(s.dtype) + + +def _exact_integer(value: Any) -> int: + """Round *value* half-to-even to a Python int. + + Integers pass through untouched so values beyond 2**53 keep every digit. + Raises ``ValueError``/``OverflowError`` for non-finite values. + """ + if isinstance(value, (bool, np.bool_, int, np.integer)): + return int(value) + return int(round(float(value))) + + +def _rounded_to_dtype(values: pd.Series, dtype: Any) -> pd.Series: + """Round *values* to integers and store them in the integer *dtype*.""" + ints = [_exact_integer(v) for v in values] + return pd.Series(ints, index=values.index, dtype=object).astype(dtype) + + +def _combine_filled(s: pd.Series, filled_values: pd.Series) -> tuple[pd.Series, str | None]: + """Put *filled_values* into the missing cells of *s*, keeping numeric dtypes. + + Returns the combined series and a dtype note: ``"rounded"`` when an integer + column received rounded values in its own dtype, ``"float64"`` when a + numeric column had to be cast, otherwise ``None``. + """ + if _is_integer_column(s): + try: + return s.where(s.notna(), _rounded_to_dtype(filled_values, s.dtype)), "rounded" + except (TypeError, ValueError, OverflowError): + pass # e.g. a non-finite or out-of-range value; use the paths below + try: + return s.where(s.notna(), filled_values), None + except (TypeError, ValueError): + if is_numeric_dtype(s.dtype) and not is_bool_dtype(s.dtype): + return s.astype("float64").where(s.notna(), filled_values), "float64" + return s.astype(object).where(s.notna(), filled_values), None + + +def _fillna_keeping_dtype(s: pd.Series, value: Any) -> tuple[pd.Series, str | None]: + """``s.fillna(value)`` with the same dtype rules as :func:`_combine_filled`.""" + if _is_integer_column(s): + try: + return s.fillna(_exact_integer(value)), "rounded" + except (TypeError, ValueError, OverflowError): + pass + try: + return s.fillna(value), None + except (TypeError, ValueError): + if is_numeric_dtype(s.dtype) and not is_bool_dtype(s.dtype): + return s.astype("float64").fillna(value), "float64" + return s.astype(object).fillna(value), None + + +def _dtype_rationale(note: str | None, s: pd.Series) -> str: + if note == "rounded": + return f"; imputed values rounded to the nearest integer to keep dtype {s.dtype}" + if note == "float64": + return f"; column cast from {s.dtype} to float64 to hold imputed values" + return "" diff --git a/tests/test_missforest_nullable_int.py b/tests/test_missforest_nullable_int.py new file mode 100644 index 0000000..d6b8afd --- /dev/null +++ b/tests/test_missforest_nullable_int.py @@ -0,0 +1,163 @@ +"""MissForest keeps integer dtypes by rounding imputed values (issue #263).""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +import freshdata as fd + +pytest.importorskip("sklearn") + +ISOLATE = { + "drop_duplicates": False, + "drop_empty_rows": False, + "drop_empty_columns": False, + "fix_dtypes": False, + "verbose": False, +} + + +def _action(report: fd.CleanReport, column: str): + actions = [a for a in report if a.column == column and a.model_id.startswith("missforest")] + assert actions, f"no missforest action for {column!r}" + return actions[-1] + + +def _int_frame(n: int = 80, dtype: str = "Int64") -> pd.DataFrame: + x = np.random.default_rng(0).normal(size=n) + return pd.DataFrame({"a": x, "b": pd.array(np.round(x * 10).astype(int), dtype=dtype)}) + + +def _assert_python_ints(values: list) -> None: + for v in values: + assert isinstance(v, (int, np.integer)) and not isinstance(v, bool), v + + +@pytest.mark.parametrize("dtype", ["Int64", "Int32", "Int16"]) +def test_nullable_int_column_keeps_dtype_with_integer_predictions(dtype: str) -> None: + df = _int_frame(dtype=dtype) + df.loc[[1, 2, 5], "b"] = pd.NA + before = df["b"].copy() + + out, report = fd.clean(df, impute="missforest", return_report=True, **ISOLATE) + + assert str(out["b"].dtype) == dtype + assert out["b"].isna().sum() == 0 + _assert_python_ints(out.loc[[1, 2, 5], "b"].tolist()) + observed = before.notna() + pd.testing.assert_series_equal(out.loc[observed, "b"], before[observed]) + + action = _action(report, "b") + assert action.model_id == "missforest_regressor" + assert action.metadata["rounded_to_integer"] is True + assert "rounded to the nearest integer" in action.rationale + assert dtype in action.rationale + + +def test_nullable_int_predictions_are_rounded_regressor_output() -> None: + df = _int_frame() + df.loc[[1, 2, 5], "b"] = pd.NA + + out = fd.clean(df, impute="missforest", **ISOLATE) + as_float = fd.clean(df.astype({"b": "float64"}), impute="missforest", **ISOLATE) + + expected = [int(round(v)) for v in as_float.loc[[1, 2, 5], "b"]] + assert out.loc[[1, 2, 5], "b"].tolist() == expected + + +def test_float_column_keeps_fractional_predictions() -> None: + df = _int_frame().astype({"b": "float64"}) + df.loc[[1, 2, 5], "b"] = np.nan + + out, report = fd.clean(df, impute="missforest", return_report=True, **ISOLATE) + + assert out["b"].dtype == np.float64 + assert any(v != round(v) for v in out.loc[[1, 2, 5], "b"]) + action = _action(report, "b") + assert action.metadata["rounded_to_integer"] is False + assert "rounded" not in action.rationale + + +def test_numpy_int_column_has_nothing_to_impute() -> None: + # A numpy int64 column cannot hold NaN, so MissForest never touches it; + # it must come back unchanged while a sibling column is imputed. + df = _int_frame().astype({"b": "int64"}) + df.loc[[3, 4], "a"] = np.nan + + out, report = fd.clean(df, impute="missforest", return_report=True, **ISOLATE) + + assert out["b"].dtype == np.int64 + pd.testing.assert_series_equal(out["b"], df["b"]) + assert _action(report, "a").metadata["rounded_to_integer"] is False + + +def test_fallback_on_small_frame_keeps_int_dtype_and_reports_rounding() -> None: + # Fewer rows than missforest_min_rows_for_model routes to the simple + # median fallback, which reports the integer rounding. + df = pd.DataFrame( + { + "a": np.arange(6, dtype=float), + "b": pd.array([1, 2, None, 4, 7, None], dtype="Int64"), + } + ) + + out, report = fd.clean(df, impute="missforest", return_report=True, **ISOLATE) + + assert out["b"].dtype == "Int64" + # median of [1, 2, 4, 7] is 3.0 + assert out["b"].tolist() == [1, 2, 3, 4, 7, 3] + action = _action(report, "b") + assert action.model_id == "missforest_fallback" + assert action.metadata["fallback_reason"] + assert action.metadata["rounded_to_integer"] is True + assert "rounded to the nearest integer" in action.rationale + + +def test_fallback_rounds_half_median_to_even_integer() -> None: + df = pd.DataFrame( + { + "a": np.arange(5, dtype=float), + "b": pd.array([1, 2, None, 5, 8], dtype="Int64"), + } + ) + + out = fd.clean(df, impute="missforest", **ISOLATE) + + # median of [1, 2, 5, 8] is 3.5 -> half-to-even 4 + assert out["b"].dtype == "Int64" + assert out["b"].tolist() == [1, 2, 4, 5, 8] + _assert_python_ints(out["b"].tolist()) + + +def test_fallback_keeps_values_beyond_float64_exact_range() -> None: + big = 2**60 + df = pd.DataFrame( + { + "a": np.arange(5, dtype=float), + "b": pd.array([big + 1, big + 3, None, big + 6, big + 8], dtype="Int64"), + } + ) + + out = fd.clean(df, impute="missforest", **ISOLATE) + + assert out["b"].dtype == "Int64" + # exact median of the present values is big + 4.5 -> half-to-even big + 4; + # a float64 median would have lost the low digits entirely + assert out["b"].tolist() == [big + 1, big + 3, big + 4, big + 6, big + 8] + + +def test_regressor_keeps_present_values_beyond_float64_exact_range() -> None: + big = 2**60 + df = _int_frame() + df["b"] = pd.array([big + 2 * int(v) + 1 for v in df["b"]], dtype="Int64") + before = df["b"].copy() + df.loc[[1, 2, 5], "b"] = pd.NA + + out = fd.clean(df, impute="missforest", **ISOLATE) + + assert out["b"].dtype == "Int64" + observed = df["b"].notna() + assert out.loc[observed, "b"].tolist() == before[observed].tolist() + _assert_python_ints(out.loc[[1, 2, 5], "b"].tolist())