fix(missforest): keep integer dtypes by rounding imputed values

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,18 +358,24 @@ 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):
             self._preserve(df, col, ctx, reason)
             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 ""

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())