From 2b1c74ab309f9a1c9a45cc610eb860b013b0854f Mon Sep 17 00:00:00 2001 From: Kevin Costner <120246174+kevincostner17@users.noreply.github.com> Date: Tue, 15 Sep 2026 01:21:05 +0530 Subject: [PATCH] fix(polars): NaN as missing, finite outlier fences, LazyFrame inputs - #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 --- src/freshdata/adapters/polars.py | 15 +++++ src/freshdata/cleaner.py | 8 ++- src/freshdata/execution/backends/_pandas.py | 8 +++ src/freshdata/execution/backends/_polars.py | 42 +++++++++---- tests/test_execution/test_polars_engine.py | 70 +++++++++++++++++++++ tests/test_polars_adapter.py | 7 +++ 6 files changed, 134 insertions(+), 16 deletions(-) diff --git a/src/freshdata/adapters/polars.py b/src/freshdata/adapters/polars.py index c237d200..e571daa3 100644 --- a/src/freshdata/adapters/polars.py +++ b/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 diff --git a/src/freshdata/cleaner.py b/src/freshdata/cleaner.py index f4944253..5012aa24 100644 --- a/src/freshdata/cleaner.py +++ b/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()]}) diff --git a/src/freshdata/execution/backends/_pandas.py b/src/freshdata/execution/backends/_pandas.py index 26078522..f9af0562 100644 --- a/src/freshdata/execution/backends/_pandas.py +++ b/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): diff --git a/src/freshdata/execution/backends/_polars.py b/src/freshdata/execution/backends/_polars.py index e14043eb..7e59155d 100644 --- a/src/freshdata/execution/backends/_polars.py +++ b/src/freshdata/execution/backends/_polars.py @@ -51,7 +51,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): @@ -75,10 +89,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: @@ -253,18 +263,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) diff --git a/tests/test_execution/test_polars_engine.py b/tests/test_execution/test_polars_engine.py index b10f950f..adaef4a3 100644 --- a/tests/test_execution/test_polars_engine.py +++ b/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) diff --git a/tests/test_polars_adapter.py b/tests/test_polars_adapter.py index 5a387d69..9945d057 100644 --- a/tests/test_polars_adapter.py +++ b/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"]