fix(streaming): exact large integers and non-string column labels (#221)

- StreamingCleaner cast nullable Int*/UInt* columns to float64 before filling
  them with the running mean/median, so present values beyond 2**53 changed
  (2**53 + 1 became 2**53). Such columns now keep their integer dtype and get
  the rounded running statistic via _util.fill_na_exact; the report says so.
  Columns within +-2**53 keep the float64 behaviour.
- StreamingCleaner.clean_batch and fd.clean_timeseries raised KeyError on
  non-string column labels (e.g. 0, 7): the imputer, drift detection and the
  time-series steps indexed the frame with str(label). The str name now only
  keys the running state and report entries; frame access uses the original
  label.

Closes #208
Closes #209

Author: kevincostner17

src/freshdata/streaming/_cleaner.py

@@ -20,6 +20,7 @@
 if TYPE_CHECKING:
     from ..context import ContextPolicy
 
+from .._util import exceeds_float64_exact, fill_na_exact
 from ..cleaner import run_pipeline
 from ..config import CleanConfig, merge_options
 from ..engine.context import infer_role
@@ -117,7 +118,7 @@ def clean_batch(self, batch: Any) -> tuple[pd.DataFrame, CleanReport]:
         # In time-series mode the TS processor owns missing-value policy for its
         # numeric columns (short-gap interpolation, long-gap preservation), so the
         # generic statistical imputer must not pre-fill them.
-        ts_skip = (set(self._ts.numeric_targets(cleaned, self._roles))
+        ts_skip = ({str(c) for c in self._ts.numeric_targets(cleaned, self._roles)}
                    if self._ts is not None else set())
         # Context-protected columns are never touched by the statistical imputer
         # either (the representation pass already byte-guards them); leave their
@@ -413,11 +414,13 @@ def _impute(self, df: pd.DataFrame, report: CleanReport,
             miss = int(df[col].isna().sum())
             if miss == 0:
                 continue
-            df = self._impute_column(df, name, miss, report)
+            df = self._impute_column(df, col, name, miss, report)
         return df
 
-    def _impute_column(self, df: pd.DataFrame, col: str, miss: int,
+    def _impute_column(self, df: pd.DataFrame, frame_col: object, col: str, miss: int,
                        report: CleanReport) -> pd.DataFrame:
+        """Fill one column. *frame_col* is the frame's own label (it may be an int);
+        *col* is its str name, which keys the running state and the report."""
         cs = self.state.columns[col]
         role, ratio = cs.role, cs.missing_ratio
         band = _band(ratio, self.config)
@@ -440,7 +443,7 @@ def _impute_column(self, df: pd.DataFrame, col: str, miss: int,
                 return self._preserve(df, col, miss, report,
                                       rationale="datetime without a usable order; fill would "
                                                 "invent timestamps", risk="medium")
-            df[col] = df[col].ffill().bfill()
+            df[frame_col] = df[frame_col].ffill().bfill()
             return self._record(df, report, col, miss, "forward/backward fill within batch",
                                 rationale="datetime column with monotonic order", confidence=0.8)
 
@@ -457,24 +460,30 @@ def _impute_column(self, df: pd.DataFrame, col: str, miss: int,
                 return self._preserve(df, col, miss, report,
                                       rationale="no running statistic available yet",
                                       risk="medium", confidence=0.5)
-            df[col] = (df[col].astype("float64") if df[col].dtype.kind in "iu"
-                       else df[col]).fillna(value)
-            return self._record(df, report, col, miss, f"{label} ({value:.6g})",
+            s = df[frame_col]
+            note = ""
+            if exceeds_float64_exact(s):
+                # A float64 cast would change present values beyond 2**53: keep the
+                # integer dtype and fill with the (rounded) running statistic.
+                df[frame_col], note = fill_na_exact(s, value)
+            else:
+                df[frame_col] = (s.astype("float64") if s.dtype.kind in "iu" else s).fillna(value)
+            return self._record(df, report, col, miss, f"{label} ({value:.6g}{note})",
                                 rationale=rationale, confidence=0.8 if band == "low" else 0.7)
 
         # categorical / boolean
         mode, mode_ratio = cs.mode(), cs.mode_ratio()
         threshold = 0.5 if band == "low" else 0.6
         if mode is not None and mode_ratio is not None and mode_ratio >= threshold:
-            df[col] = df[col].fillna(mode)
+            df[frame_col] = df[frame_col].fillna(mode)
             return self._record(df, report, col, miss, f"running mode ({mode!r})",
                                 rationale=f"dominant category ({100 * mode_ratio:.0f}% of seen)",
                                 confidence=0.8 if band == "low" else 0.7)
         sentinel = "Unknown" if band == "low" else "Missing"
-        s = df[col]
+        s = df[frame_col]
         if isinstance(s.dtype, pd.CategoricalDtype) and sentinel not in s.cat.categories:
             s = s.cat.add_categories([sentinel])
-        df[col] = s.fillna(sentinel)
+        df[frame_col] = s.fillna(sentinel)
         return self._record(df, report, col, miss, f'sentinel "{sentinel}"',
                             rationale="no dominant category; sentinel keeps the gap visible",
                             confidence=0.7, risk="low" if band == "low" else "medium")

src/freshdata/streaming/_drift.py

@@ -41,7 +41,10 @@ def detect_drift(
     if not state.schema_baseline:
         return findings
 
-    batch_cols = [str(c) for c in df.columns]
+    # Baseline and running state are keyed by str label; the frame keeps its own
+    # labels (e.g. ints), so index it through this map.
+    labels = {str(c): c for c in df.columns}
+    batch_cols = list(labels)
     baseline = set(state.schema_baseline)
     seen = set(batch_cols)
 
@@ -61,7 +64,7 @@ def detect_drift(
     for col in batch_cols:
         if col not in baseline:
             continue
-        s = df[col]
+        s = df[labels[col]]
         baseline_dtype = state.baseline_dtypes.get(col)
         if baseline_dtype is not None and str(s.dtype) != baseline_dtype:
             findings.append(DriftFinding(

src/freshdata/streaming/_timeseries.py

@@ -243,7 +243,7 @@ def process(self, df: pd.DataFrame, report: CleanReport, *,
 
         roles = roles or self._infer_roles(df)
         numeric_cols = self.numeric_targets(df, roles)
-        self.last_numeric_cols = list(numeric_cols)
+        self.last_numeric_cols = [str(c) for c in numeric_cols]
 
         # 3. Short-gap interpolation, then 4. seasonal imputation of what's left.
         df = self._interpolate(df, numeric_cols, report)
@@ -415,11 +415,11 @@ def fill_group(g: pd.DataFrame, col: str = col) -> pd.DataFrame:
                 report.add("timeseries_interpolation",
                            f"interpolated {filled_total} short-gap value(s) "
                            f"(<= {cfg.max_interpolation_gap} step gap, method={method})",
-                           column=col, count=filled_total, risk="low",
+                           column=str(col), count=filled_total, risk="low",
                            confidence=0.8 if method in ("time", "linear") else 0.7,
                            rationale="short consecutive gap in an ordered series",
                            model_id=f"interp_{method}")
-                report.columns_imputed.append(col)
+                report.columns_imputed.append(str(col))
         return df
 
     @staticmethod
@@ -486,10 +486,10 @@ def fill_group(g: pd.DataFrame, col: str = col) -> pd.DataFrame:
                 rationale = (f"{seasonal_filled} from matching {cfg.seasonal_period} season, "
                              f"{fallback_filled} from rolling/global median fallback")
                 report.add("seasonal_imputation",
-                           f"seasonally imputed {total} value(s)", column=col, count=total,
+                           f"seasonally imputed {total} value(s)", column=str(col), count=total,
                            risk="medium", confidence=conf, rationale=rationale,
                            model_id="seasonal_median")
-                report.columns_imputed.append(col)
+                report.columns_imputed.append(str(col))
         return df
 
     # -- step 6: windowed anomaly detection ------------------------------------
@@ -535,7 +535,7 @@ def score_group(g: pd.DataFrame, col: str = col, flags: pd.Series = flags,
                 action_note = "flagged and capped"
             report.add("windowed_anomaly",
                        f"{action_note} {n_flag} windowed anomaly(ies) (method={method}, "
-                       f"window={win})", column=col, count=n_flag, risk="medium",
+                       f"window={win})", column=str(col), count=n_flag, risk="medium",
                        confidence=0.7,
                        rationale=f"rolling {method} score beyond {thr}",
                        model_id=f"anomaly_{method}")
@@ -595,7 +595,7 @@ def numeric_targets(self, df: pd.DataFrame,
                 continue
             if not pd.api.types.is_numeric_dtype(df[c]):
                 continue
-            cols.append(name)
+            cols.append(c)  # the frame's own label (may be an int), not str(c)
         return cols
 
 

tests/test_streaming_cleaner.py

@@ -238,3 +238,27 @@ def test_global_duplicate_window_dedup():
     assert len(cleaned) == 0  # every row is a cross-batch duplicate
     assert report.duplicates_removed == 5
     assert any(a.step == "duplicates" for a in report)
+
+
+def test_integer_column_labels_are_imputed_without_keyerror():
+    cleaner = fd.StreamingCleaner(warmup_batches=0, verbose=False)
+    batch = pd.DataFrame({0: [1.0, None, 3.0, 4.0], 1: [1.0, 2.0, 3.0, 4.0]})
+    cleaned, report = cleaner.clean_batch(batch)
+    assert list(cleaned.columns) == [0, 1]
+    assert any(a.step == "missing" and a.column == "0" for a in report)
+    # The second batch runs drift detection against the str-keyed baseline.
+    cleaned2, report2 = cleaner.clean_batch(batch.copy())
+    assert list(cleaned2.columns) == [0, 1]
+    assert any(a.step == "missing" and a.column == "0" for a in report2)
+
+
+def test_numeric_fill_keeps_int64_beyond_2_53_exact():
+    big = 2**53 + 1
+    df = pd.DataFrame({"x": pd.array([big, None, 2**53 + 3, big], dtype="Int64"),
+                       "y": [1.0, 2.0, 3.0, 4.0]})
+    out, report = fd.StreamingCleaner(warmup_batches=0, verbose=False).clean_batch(df)
+    assert str(out["x"].dtype) == "Int64"
+    assert [int(out["x"][i]) for i in (0, 2, 3)] == [big, 2**53 + 3, big]
+    assert out["x"].isna().sum() == 0
+    notes = [a.description for a in report if a.step == "missing" and a.column == "x"]
+    assert notes and "2**53" in notes[0]

tests/test_streaming_timeseries.py

@@ -336,3 +336,14 @@ def test_no_anomalies_still_emits_flag_column():
         df, timestamp_column="t", anomaly_window_size=10, anomaly_method="rolling_zscore")
     assert "x_anomaly" in out.columns
     assert not out["x_anomaly"].any()
+
+
+def test_integer_column_labels_in_timeseries_mode():
+    df = pd.DataFrame({"t": pd.date_range("2024", periods=4, freq="h"),
+                       7: ["a", None, "a", "a"],
+                       8: [1.0, np.nan, 3.0, 4.0]})
+    out, report = fd.clean_timeseries(df, timestamp_column="t", max_interpolation_gap=1,
+                                      return_report=True)
+    assert 7 in out.columns and 8 in out.columns
+    assert out[8].tolist() == [1.0, 2.0, 3.0, 4.0]
+    assert any(a.step == "timeseries_interpolation" and a.column == "8" for a in report)