fix(polars): NaN as missing, finite outlier fences, LazyFrame inputs (#224)

- #200: the Polars engine only counted nulls, so float NaN from a native Polars
  DataFrame/LazyFrame, an Arrow table or Polars-written Parquet was treated as a
  value: empty rows/columns were kept and missing_before was wrong. Native
  sources now read float NaN as null at ingestion (pl.from_pandas already did
  this for pandas input). Outlier fences are computed over finite values only,
  matching pandas drop_infinite, so +-inf no longer blows up the quantiles.
- #203: a pl.LazyFrame source crashed whenever the run needed the pandas
  fallback ("cannot materialize source of type LazyFrame") and was rejected on
  the default path ("expected a pandas or polars DataFrame"). The fallback now
  collects it, and fd.clean(lazyframe) cleans it like a DataFrame and returns a
  LazyFrame. fallback_policy="error" still blocks before any collect.

Closes #200
Closes #203

Author: kevincostner17

src/freshdata/adapters/polars.py

@@ -31,9 +31,21 @@ def is_polars_frame(obj: object) -> bool:
     return isinstance(obj, pl.DataFrame)
 
 
+def is_polars_lazy(obj: object) -> bool:
+    """True for an uncollected ``pl.LazyFrame``."""
+    try:
+        pl = _polars_module()
+    except ImportError:
+        return False
+    return isinstance(obj, pl.LazyFrame)
+
+
 def to_pandas(df: object) -> pd.DataFrame:
     if isinstance(df, pd.DataFrame):
         return df
+    if is_polars_lazy(df):
+        lazy: Any = df
+        df = lazy.collect()  # the pandas pipeline needs the rows in memory
     if is_polars_frame(df):
         pl_df: Any = df
         try:
@@ -53,4 +65,7 @@ def from_pandas(df: pd.DataFrame, original: object | None = None) -> object:
     if is_polars_frame(original):
         pl = _polars_module()
         return pl.from_pandas(df)
+    if is_polars_lazy(original):
+        pl = _polars_module()
+        return pl.from_pandas(df).lazy()  # LazyFrame in, LazyFrame out
     return df

src/freshdata/cleaner.py

@@ -9,7 +9,7 @@
 import pandas as pd
 
 from ._util import memory_bytes
-from .adapters.polars import is_polars_frame, to_pandas
+from .adapters.polars import is_polars_frame, is_polars_lazy, to_pandas
 from .config import CleanConfig, merge_options
 from .engine import auto_missing, auto_outliers
 from .engine.cache import build_engine_cache
@@ -31,8 +31,10 @@ def _validate_input(df: object, config: CleanConfig) -> pd.DataFrame:
         raise TypeError(
             "freshdata works on DataFrames; got a Series. Convert it first with s.to_frame()."
         )
-    if not isinstance(df, pd.DataFrame) and not is_polars_frame(df):
-        raise TypeError(f"expected a pandas or polars DataFrame, got {type(df).__name__}")
+    if not (isinstance(df, pd.DataFrame) or is_polars_frame(df) or is_polars_lazy(df)):
+        raise TypeError(
+            f"expected a pandas or polars DataFrame/LazyFrame, got {type(df).__name__}"
+        )
     frame = to_pandas(df)
     if frame.columns.duplicated().any() and not config.column_names:
         dupes = sorted({str(c) for c in frame.columns[frame.columns.duplicated()]})

src/freshdata/execution/backends/_pandas.py

@@ -20,6 +20,9 @@
     from .._config import EngineConfig
 
 
+from .._lazy import has_polars
+
+
 def materialize_to_pandas(source: Any) -> pd.DataFrame:
     """Load *source* into a pandas DataFrame, reading file paths if needed."""
     import pandas as pd
@@ -35,6 +38,11 @@ def materialize_to_pandas(source: Any) -> pd.DataFrame:
         if low.endswith((".ipc", ".feather", ".arrow")):
             return pd.read_feather(source)
         raise ValueError(f"unsupported file type for path {source!r}")
+    if has_polars():
+        import polars as pl
+
+        if isinstance(source, pl.LazyFrame):
+            return source.collect().to_pandas()
     # polars frame
     to_pandas = getattr(source, "to_pandas", None)
     if callable(to_pandas):

src/freshdata/execution/backends/_polars.py

@@ -52,7 +52,21 @@ class PolarsEngine(ExecutionEngine):
     # -- source ingestion ---------------------------------------------------
 
     def _to_lazy(self, source: Any, pl: Any) -> tuple[Any, int]:
-        """Return ``(LazyFrame, memory_before_bytes)`` for *source*."""
+        """Return ``(LazyFrame, memory_before_bytes)`` for *source*.
+
+        Float ``NaN`` is read as null, because the pandas reference treats
+        ``NaN`` as missing. ``pl.from_pandas`` already does this; native Polars,
+        Arrow and file sources keep real ``NaN`` otherwise.
+        """
+        import pandas as pd
+
+        if isinstance(source, pd.DataFrame):
+            return pl.from_pandas(source).lazy(), int(source.memory_usage(deep=True).sum())
+        lf, memory_before = self._native_to_lazy(source, pl)
+        return lf.with_columns(pl.col(pl.Float32, pl.Float64).fill_nan(None)), memory_before
+
+    @staticmethod
+    def _native_to_lazy(source: Any, pl: Any) -> tuple[Any, int]:
         if isinstance(source, pl.LazyFrame):
             return source, 0
         if isinstance(source, pl.DataFrame):
@@ -76,10 +90,6 @@ def _to_lazy(self, source: Any, pl: Any) -> tuple[Any, int]:
                 return pl.from_arrow(table).lazy(), int(table.nbytes)
         except ImportError:
             pass
-        import pandas as pd
-
-        if isinstance(source, pd.DataFrame):
-            return pl.from_pandas(source).lazy(), int(source.memory_usage(deep=True).sum())
         raise TypeError(f"PolarsEngine: unsupported source type {type(source).__name__}")
 
     def _pandas_index_forces_fallback(self, source: Any) -> bool:
@@ -258,18 +268,24 @@ def _stage_outliers(  # noqa: PLR0915 — mirrors the pandas step's branchy poli
         clip = config.outliers == "clip"
         stat_aggs: list[Any] = []
         for n in numeric:
+            # Fences come from finite values only, like the pandas reference
+            # (steps.outliers.drop_infinite); inf is still tested against them.
+            values = (
+                pl.col(n).filter(pl.col(n).is_finite())
+                if schema[n] in (pl.Float32, pl.Float64) else pl.col(n)
+            )
             if method == "iqr":
-                stat_aggs.append(pl.col(n).quantile(0.25, "linear").alias(f"q1_{n}"))
-                stat_aggs.append(pl.col(n).quantile(0.75, "linear").alias(f"q3_{n}"))
+                stat_aggs.append(values.quantile(0.25, "linear").alias(f"q1_{n}"))
+                stat_aggs.append(values.quantile(0.75, "linear").alias(f"q3_{n}"))
             else:
-                stat_aggs.append(pl.col(n).mean().alias(f"m_{n}"))
-                stat_aggs.append(pl.col(n).std().alias(f"s_{n}"))
+                stat_aggs.append(values.mean().alias(f"m_{n}"))
+                stat_aggs.append(values.std().alias(f"s_{n}"))
             if clip:
                 # Skew-aware capping aggregates (see log_widened_bounds).
-                logs = pl.col(n).filter(pl.col(n) > 0).log()
-                stat_aggs.append(pl.col(n).skew(bias=False).alias(f"sk_{n}"))
-                stat_aggs.append(pl.col(n).min().alias(f"mn_{n}"))
-                stat_aggs.append(pl.col(n).count().alias(f"cnt_{n}"))
+                logs = values.filter(values > 0).log()
+                stat_aggs.append(values.skew(bias=False).alias(f"sk_{n}"))
+                stat_aggs.append(values.min().alias(f"mn_{n}"))
+                stat_aggs.append(values.count().alias(f"cnt_{n}"))
                 stat_aggs.append(logs.quantile(0.25, "linear").alias(f"lq1_{n}"))
                 stat_aggs.append(logs.quantile(0.75, "linear").alias(f"lq3_{n}"))
         stats = lf.select(stat_aggs).collect().row(0, named=True)

tests/test_execution/test_polars_engine.py

@@ -6,7 +6,9 @@
 import pytest
 
 import freshdata as fd
+from freshdata.config import CleanConfig
 from freshdata.execution import EngineConfig
+from freshdata.execution._config import FallbackError
 
 pl = pytest.importorskip("polars")
 
@@ -73,3 +75,71 @@ def test_projection_pushdown_drops_empty_before_collect(native_config):
     df = pd.DataFrame({"keep": [1, 2], "gone": [None, None]})
     out = fd.clean(df, config=native_config, engine="polars", output_format="polars")
     assert list(out.columns) == ["keep"]
+
+
+def _native_nan_source():
+    return pl.DataFrame({
+        "a": [1.0, float("nan"), 3.0],
+        "b": [float("nan")] * 3,
+        "s": ["x", None, "y"],
+    })
+
+
+def _steps(report):
+    return [(a.step, a.count) for a in report.actions]
+
+
+@pytest.mark.parametrize(
+    "make_source",
+    [lambda df: df, lambda df: df.lazy(), lambda df: df.to_arrow()],
+    ids=["dataframe", "lazyframe", "arrow"],
+)
+def test_native_nan_counts_as_missing_like_pandas(native_config, make_source):
+    src = _native_nan_source()
+    ref_out, ref = fd.clean(src.to_pandas(), config=native_config, engine="pandas",
+                            return_report=True)
+    out, report = fd.clean(make_source(src), config=native_config, engine="polars",
+                           return_report=True)
+    assert list(out.columns) == list(ref_out.columns) == ["a", "s"]
+    assert report.rows_after == ref.rows_after == 2
+    assert report.missing_before == ref.missing_before == 5
+    assert _steps(report) == _steps(ref)
+
+
+def test_polars_written_parquet_nan_counts_as_missing(tmp_path, native_config):
+    path = str(tmp_path / "nan.parquet")
+    _native_nan_source().write_parquet(path)
+    out, report = fd.clean(path, config=native_config, engine="polars", return_report=True)
+    assert list(out.columns) == ["a", "s"]
+    assert report.missing_before == 5
+
+
+def test_infinite_values_are_excluded_from_outlier_fences():
+    config = CleanConfig(strategy="conservative", fix_dtypes=False, verbose=False,
+                         outliers="clip", outlier_method="iqr")
+    df = pd.DataFrame({"x": [1.0, 2.0, 3.0, 4.0] + [float("inf")] * 4})
+    ref, ref_report = fd.clean(df, config=config, engine="pandas", return_report=True)
+    out, report = fd.clean(pl.from_pandas(df), config=config, engine="polars",
+                           return_report=True)
+    assert out["x"].tolist() == ref["x"].tolist()
+    assert _steps(report) == _steps(ref_report)
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [{}, {"engine": "auto"}, {"engine": "polars", "strategy": "balanced"},
+     {"engine": "polars", "strategy": "conservative"}],
+    ids=["default", "auto", "polars-balanced", "polars-conservative"],
+)
+def test_lazy_frame_source_survives_pandas_fallback(kwargs):
+    lf = pl.DataFrame({"a": [1.0, None, 3.0], "b": ["x", "y", "z"]}).lazy()
+    out = fd.clean(lf, verbose=False, **kwargs)
+    frame = out.collect() if isinstance(out, pl.LazyFrame) else out
+    assert len(frame) == 3
+
+
+def test_lazy_frame_fallback_still_honours_error_policy():
+    lf = pl.DataFrame({"a": [1.0, None, 3.0]}).lazy()
+    with pytest.raises(FallbackError):
+        fd.clean(lf, engine="polars", strategy="balanced", fallback_policy="error",
+                 verbose=False)

tests/test_polars_adapter.py

@@ -49,3 +49,10 @@ def boom(self, *a, **k):
     monkeypatch.setattr(pl.DataFrame, "to_pandas", boom)
     with pytest.raises(ModuleNotFoundError, match=r"freshdata-cleaner\[polars\]"):
         to_pandas(pl.DataFrame({"a": [1]}))
+
+
+def test_lazy_frame_round_trips_through_default_clean():
+    lf = pl.DataFrame({"a": [1.0, None, 3.0], "b": ["x", " y", "z"]}).lazy()
+    out = fd.clean(lf, verbose=False)
+    assert isinstance(out, pl.LazyFrame)  # LazyFrame in, LazyFrame out
+    assert out.collect()["b"].to_list() == ["x", "y", "z"]