fix(impute): keep large nullable integers exact; fill_missing dtype handling

Imputing a nullable integer column cast it to float64 whenever the fill value
was fractional, and computed mean/median in float64. For values beyond 2**53
that silently changed *present* values (2**53 + 1 became 2**53) with only a
"column cast to float64" note.

- New _util helpers: exceeds_float64_exact() detects such columns,
  exact_int_stat() computes mean/median in exact integer arithmetic (rounded
  half-to-even), and fill_na_exact() fills without a float64 cast for them.
  Columns within +-2**53 keep the existing behaviour (float64 cast note).
- steps/missing.py (impute=...) and engine/missing.py (_fill, KNN
  _assign_filled) use them; the report says the dtype was kept.
- fd.fill_missing crashed on nullable Int columns with a fractional
  mean/median (including the default method="auto"), on nullable boolean
  columns (numeric median into boolean), and on duplicate column labels. It now
  picks the fill value like fd.clean(impute=...) (booleans take the mode),
  fills through fill_na_exact, and rejects duplicate labels with a clear
  ValueError.

Closes #210
Refs #208

Author: kevincostner17

src/freshdata/_util.py

@@ -3,7 +3,9 @@
 from __future__ import annotations
 
 import hashlib
+import math
 import warnings
+from fractions import Fraction
 from typing import Any
 
 import pandas as pd
@@ -23,6 +25,64 @@ def safe_median(s: pd.Series) -> Any:
     return s.median()
 
 
+#: float64 represents every integer up to this magnitude exactly.
+FLOAT64_EXACT_INT = 2**53
+
+
+def exceeds_float64_exact(s: pd.Series) -> bool:
+    """True for a nullable integer column holding a value float64 cannot represent.
+
+    Casting such a column to float64 silently changes *present* values
+    (``2**53 + 1`` becomes ``2**53``), so fill paths must keep its integer dtype.
+    """
+    if not isinstance(s.array, pd.arrays.IntegerArray):
+        return False
+    present = s.dropna()
+    if present.empty:
+        return False
+    return int(present.max()) > FLOAT64_EXACT_INT or int(present.min()) < -FLOAT64_EXACT_INT
+
+
+def exact_int_stat(s: pd.Series, strategy: str) -> int:
+    """Mean or median of an integer column in exact integer arithmetic.
+
+    Rounded half-to-even to the nearest integer so the result fits the column's
+    dtype. Only used for columns where :func:`exceeds_float64_exact` holds.
+    """
+    values = sorted(int(v) for v in s.dropna())
+    n = len(values)
+    if strategy == "mean":
+        return round(Fraction(sum(values), n))
+    mid = n // 2
+    if n % 2:
+        return values[mid]
+    return round(Fraction(values[mid - 1] + values[mid], 2))
+
+
+def fill_na_exact(s: pd.Series, value: Any) -> tuple[pd.Series, str]:
+    """``s.fillna(value)`` that never silently corrupts large nullable integers.
+
+    Returns the filled series and a note for the report. A column holding values
+    beyond 2**53 keeps its integer dtype and receives an integer fill value.
+    Any other numeric column that cannot hold a fractional *value* is cast to
+    float64, which is exact for it. Raises ``TypeError``/``ValueError`` when
+    *value* cannot be stored.
+    """
+    if exceeds_float64_exact(s):
+        if isinstance(value, float):
+            if not math.isfinite(value):
+                raise ValueError(f"cannot fill {s.dtype} with {value!r}")
+            value = int(round(value))
+        return s.fillna(value), f", kept {s.dtype} so values beyond 2**53 stay exact"
+    try:
+        return s.fillna(value), ""
+    except (TypeError, ValueError):
+        if pd.api.types.is_numeric_dtype(s) and isinstance(value, float):
+            # e.g. a fractional median into an integer column
+            return s.astype("float64").fillna(value), ", column cast to float64"
+        raise
+
+
 def add_column(df: pd.DataFrame, name: object, values: object) -> None:
     """Insert a new column in place, suppressing pandas' fragmentation notice.
 

src/freshdata/engine/missing.py

@@ -35,7 +35,13 @@
 import pandas as pd
 from pandas.api.types import is_bool_dtype, is_numeric_dtype
 
-from .._util import add_column, safe_median
+from .._util import (
+    add_column,
+    exact_int_stat,
+    exceeds_float64_exact,
+    fill_na_exact,
+    safe_median,
+)
 from ..config import CleanConfig
 from ..report import CleanReport
 from ..steps.missing import _mode_value
@@ -486,6 +492,8 @@ def _fill(df: pd.DataFrame, col: object, ctx: ColumnContext, report: CleanReport
             min_confidence=min_confidence, model_id=model_id, label=label,
         )
     s = df[col]
+    if label in ("mean", "median") and exceeds_float64_exact(s):
+        value = exact_int_stat(s, label)  # a float statistic would lose digits
     if value is None or pd.isna(value):
         _preserve(df, col, ctx, report,
                   rationale="no usable fill value could be derived "
@@ -494,18 +502,13 @@ def _fill(df: pd.DataFrame, col: object, ctx: ColumnContext, report: CleanReport
         return df
     if isinstance(s.dtype, pd.CategoricalDtype) and value not in s.cat.categories:
         s = s.cat.add_categories([value])
-    cast_note = ""
     try:
-        filled = s.fillna(value)
+        filled, cast_note = fill_na_exact(s, value)
     except (TypeError, ValueError):
-        if is_numeric_dtype(s) and isinstance(value, float):
-            filled = s.astype("float64").fillna(value)
-            cast_note = ", column cast to float64"
-        else:
-            _preserve(df, col, ctx, report,
-                      rationale=f"fill value not representable in dtype {s.dtype}",
-                      risk="medium", confidence=0.6, model_id="preserve")
-            return df
+        _preserve(df, col, ctx, report,
+                  rationale=f"fill value not representable in dtype {s.dtype}",
+                  risk="medium", confidence=0.6, model_id="preserve")
+        return df
     df[col] = filled
     shown = f"{value:.6g}" if isinstance(value, float) else repr(value)
     report.add(_STEP,
@@ -533,7 +536,11 @@ def _assign_filled(df: pd.DataFrame, col: object, ctx: ColumnContext,
     try:
         combined = s.where(s.notna(), filled_values)
     except (TypeError, ValueError):
-        combined = s.astype("float64").where(s.notna(), filled_values)
+        if exceeds_float64_exact(s):
+            # keep the integer dtype: a float64 cast would change present values
+            combined = s.where(s.notna(), filled_values.round().astype(s.dtype))
+        else:
+            combined = s.astype("float64").where(s.notna(), filled_values)
     df[col] = combined
     report.add(_STEP, f"filled {ctx.n_missing} missing value(s) with {label}",
                column=str(col), count=ctx.n_missing, rationale=rationale,

src/freshdata/simple.py

@@ -21,7 +21,8 @@
 import pandas as pd
 from pandas.api.types import is_bool_dtype, is_numeric_dtype
 
-from ._util import safe_median
+from ._util import fill_na_exact
+from .steps.missing import _fill_value
 
 __all__ = [
     "detect_outliers",
@@ -78,12 +79,6 @@ def _is_outlier_numeric(s: pd.Series) -> bool:
     return is_numeric_dtype(s) and not is_bool_dtype(s)
 
 
-def _mode_value(s: pd.Series) -> object | None:
-    """Most-frequent non-null value, or None when the column is all-null."""
-    modes = s.dropna().mode()
-    return modes.iloc[0] if not modes.empty else None
-
-
 def _outlier_mask(
     df: pd.DataFrame,
     cols: Sequence[str],
@@ -125,14 +120,22 @@ def fill_missing(
     """Fill missing values in ``columns`` (all columns by default).
 
     ``method`` is one of ``"auto"`` (median for numeric columns, mode for the
-    rest), ``"mean"``, ``"median"``, ``"mode"``, ``"constant"`` (fills with
-    ``value``), ``"ffill"`` or ``"bfill"``. Returns the filled frame -- a copy
-    unless ``inplace=True``.
+    rest, booleans included), ``"mean"``, ``"median"``, ``"mode"``,
+    ``"constant"`` (fills with ``value``), ``"ffill"`` or ``"bfill"``. A
+    fractional mean/median is filled into an integer column by casting it to
+    float64, except for nullable integers holding values beyond 2**53, which keep
+    their dtype and get a rounded, exactly-computed fill. Returns the filled
+    frame -- a copy unless ``inplace=True``.
     """
     if method not in _FILL_METHODS:
         raise ValueError(f"method must be one of {_FILL_METHODS}, got {method!r}")
     if method == "constant" and value is None:
         raise ValueError("method='constant' requires a value")
+    if not df.columns.is_unique:
+        duplicated = sorted({str(c) for c in df.columns[df.columns.duplicated()]})
+        raise ValueError(
+            f"fill_missing requires unique column labels; duplicated: {duplicated}"
+        )
     if not inplace:
         df = df.copy()
     cols = _resolve_columns(df, columns)
@@ -149,18 +152,16 @@ def fill_missing(
         elif method == "constant":
             df[col] = s.fillna(value)
         else:
-            strategy = method
-            if strategy == "auto":
-                strategy = "median" if is_numeric_dtype(s) else "mode"
-            if strategy == "mode":
-                fill = _mode_value(s)
-            elif is_numeric_dtype(s):
-                fill = s.mean() if strategy == "mean" else safe_median(s)
-            else:
-                fill = None  # mean/median are undefined for non-numeric columns
-            if fill is None:
+            # Same value choice as fd.clean(impute=...): booleans take the mode,
+            # mean/median are undefined (None) for non-numeric columns.
+            fill = _fill_value(s, method)
+            if fill is None or pd.isna(fill):
                 continue
-            df[col] = s.fillna(fill)
+            try:
+                df[col], _ = fill_na_exact(s, fill)
+            except (TypeError, ValueError):
+                continue  # the value cannot be stored in this dtype
+
         filled += before - int(df[col].isna().sum())
     if verbose:
         print(

src/freshdata/steps/missing.py

@@ -12,7 +12,7 @@
 import pandas as pd
 from pandas.api.types import is_bool_dtype, is_numeric_dtype
 
-from .._util import safe_median
+from .._util import exact_int_stat, exceeds_float64_exact, fill_na_exact, safe_median
 from ..config import CleanConfig
 from ..report import CleanReport
 
@@ -36,6 +36,8 @@ def _fill_value(s: pd.Series, strategy: str) -> Any | None:
     if strategy in ("mean", "median"):
         if not numeric:
             return None  # not defined for this dtype; caller reports the skip
+        if exceeds_float64_exact(s):
+            return exact_int_stat(s, strategy)  # a float statistic would lose digits
         return s.mean() if strategy == "mean" else safe_median(s)
     return _mode_value(s)
 
@@ -87,19 +89,13 @@ def impute_missing(df: pd.DataFrame, config: CleanConfig,
                            f"skipped ({strategy} is not defined for dtype {s.dtype})",
                            column=str(col))
             continue
-        cast_note = ""
         try:
-            filled = s.fillna(value)
+            filled, cast_note = fill_na_exact(s, value)
         except (TypeError, ValueError):
-            if is_numeric_dtype(s) and isinstance(value, float):
-                # e.g. fractional median into an integer column
-                filled = s.astype("float64").fillna(value)
-                cast_note = ", column cast to float64"
-            else:
-                # e.g. value not representable in this dtype
-                report.add("impute", f"skipped (could not fill dtype {s.dtype})",
-                           column=str(col))
-                continue
+            # e.g. value not representable in this dtype
+            report.add("impute", f"skipped (could not fill dtype {s.dtype})",
+                       column=str(col))
+            continue
         df[col] = filled
         shown = f"{value:.6g}" if isinstance(value, float) else repr(value)
         report.add("impute",

tests/test_nullable_int_median.py

@@ -6,12 +6,14 @@
 imputation path must survive those columns and keep Int64/float behaviour.
 """
 
+from fractions import Fraction
+
 import numpy as np
 import pandas as pd
 import pytest
 
 import freshdata as fd
-from freshdata._util import safe_median
+from freshdata._util import exact_int_stat, exceeds_float64_exact, fill_na_exact, safe_median
 from freshdata.engine import missing, model_select
 from freshdata.engine.utils import _has_outliers
 
@@ -92,6 +94,61 @@ def test_seasonal_imputation_global_median_on_nullable_int(dtype):
     assert out["v"].isna().sum() == 0
 
 
+BIG = 2**53 + 1  # float64 rounds this to 2**53
+
+
+def test_exact_int_stat_and_detection():
+    s = pd.Series(pd.array([BIG, 0, 1, None], dtype="Int64"))
+    assert exceeds_float64_exact(s)
+    assert not exceeds_float64_exact(pd.Series(pd.array([2**53, None], dtype="Int64")))
+    assert not exceeds_float64_exact(pd.Series([float(BIG), np.nan]))
+    assert exact_int_stat(s, "mean") == round(Fraction(BIG + 1, 3))
+    assert exact_int_stat(s, "median") == 1
+    even = pd.Series(pd.array([BIG, BIG + 2, None, BIG + 5, BIG + 7], dtype="Int64"))
+    # (BIG+2 + BIG+5) / 2 == 2**53 + 4.5, which rounds half-to-even to 2**53 + 4
+    assert exact_int_stat(even, "median") == BIG + 3
+
+
+def test_fill_na_exact_keeps_small_int_behaviour():
+    s = pd.Series(pd.array([1, None, 2], dtype="Int64"))
+    filled, note = fill_na_exact(s, 1.5)
+    assert filled.tolist() == [1.0, 1.5, 2.0]
+    assert note == ", column cast to float64"
+    filled, note = fill_na_exact(s, 1)
+    assert str(filled.dtype) == "Int64" and note == ""
+
+
+@pytest.mark.parametrize("impute", ["mean", "median", "auto"])
+def test_explicit_impute_keeps_int64_beyond_2_53_exact(impute):
+    df = pd.DataFrame({"x": pd.array([BIG, 0, 1, None], dtype="Int64"), "y": [1.0, 2.0, 3.0, 4.0]})
+    out, report = fd.clean(
+        df, impute=impute, strategy="conservative", return_report=True, **KEEP_ROWS
+    )
+    assert str(out["x"].dtype) == "Int64"
+    assert out["x"].iloc[:3].tolist() == [BIG, 0, 1]  # present values untouched
+    expected = round(Fraction(BIG + 1, 3)) if impute == "mean" else 1
+    assert out["x"].iloc[3] == expected
+    notes = [a.description for a in report if a.step == "impute" and a.column == "x"]
+    assert notes and "2**53" in notes[0]
+
+
+def test_default_engine_keeps_int64_beyond_2_53_exact():
+    base = 2**60
+    values = pd.array([base + (i % 7) for i in range(60)], dtype="Int64")
+    s = pd.Series(values)
+    s.iloc[[3, 17]] = pd.NA
+    df = pd.DataFrame({"v": s, "x": np.random.default_rng(0).normal(0, 1, 60)})
+    out, report = fd.clean(df, return_report=True, **KEEP_ROWS)
+    assert str(out["v"].dtype) == "Int64"
+    present = s.notna()
+    assert out["v"][present].tolist() == s[present].tolist()
+    filled = [a for a in report if a.step == "missing" and a.column == "v"]
+    if filled and "filled" in filled[0].description:
+        assert out["v"].isna().sum() == 0
+        assert base <= int(out["v"].iloc[3]) <= base + 6
+        assert "2**53" in filled[0].description
+
+
 def test_has_outliers_single_definition():
     assert missing._has_outliers is _has_outliers
     assert model_select._has_outliers is _has_outliers

tests/test_simple.py

@@ -79,6 +79,34 @@ def test_fill_missing_verbose(capsys):
     assert "fill_missing" in capsys.readouterr().out
 
 
+def test_fill_missing_nullable_int_fractional_median_casts_to_float():
+    out = fd.fill_missing(pd.DataFrame({"x": pd.array([1, None, 2], dtype="Int64")}))
+    assert out["x"].tolist() == [1.0, 1.5, 2.0]
+
+
+def test_fill_missing_nullable_boolean_uses_mode():
+    df = pd.DataFrame({"b": pd.array([True, None, False, False], dtype="boolean")})
+    out = fd.fill_missing(df)
+    assert str(out["b"].dtype) == "boolean"
+    assert out["b"].tolist() == [True, False, False, False]
+    # mean/median are not defined for booleans: left as-is, no crash
+    assert fd.fill_missing(df, method="median")["b"].isna().sum() == 1
+
+
+def test_fill_missing_duplicate_labels_raise_clear_error():
+    df = pd.DataFrame([[1.0, None], [None, 2.0]], columns=["x", "x"])
+    with pytest.raises(ValueError, match="unique column labels"):
+        fd.fill_missing(df)
+
+
+def test_fill_missing_keeps_int64_beyond_2_53_exact():
+    big = 2**53 + 1
+    df = pd.DataFrame({"x": pd.array([big, None, big + 2], dtype="Int64")})
+    out = fd.fill_missing(df, method="mean")
+    assert str(out["x"].dtype) == "Int64"
+    assert out["x"].tolist() == [big, big + 1, big + 2]
+
+
 def test_resolve_columns_missing_raises():
     with pytest.raises(KeyError, match="columns not found"):
         fd.fill_missing(pd.DataFrame({"a": [1]}), columns="ghost")