From 192c74d1e2ba3bb5f1b3fd4ff8882b245c2156b3 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Mon, 24 Aug 2026 09:53:42 +0200 Subject: [PATCH 1/3] address review feedback on dataframe_checks.py Follow-up to FBruzzesi's review on PR #965: - Clarify docstrings for check_X, check_y, check_X_y in terms of which dataframe libraries are safe to pass in (pandas, polars, PyArrow, modin, cuDF), instead of narwhals-specific "eager" terminology. - Use narwhals' IntoDataFrameT instead of IntoDataFrame for check_X and check_X_y, since both return the same concrete dataframe type they receive. - Fix a null/NaN detection bug: in polars, is_null() does not catch an explicit float("nan") value (only None counts as null), so check_y and _check_contains_na could silently miss NaNs in polars data. Now also check is_nan() for numeric columns/series, keeping numpy for the finite/inf checks since it benchmarks as fast or faster there. Co-Authored-By: Claude Sonnet 5 --- feature_engine/dataframe_checks.py | 44 +++++++++++++++++++------- tests/test_dataframe_checks.py | 51 ++++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+), 12 deletions(-) diff --git a/feature_engine/dataframe_checks.py b/feature_engine/dataframe_checks.py index ae3be3c7b..d1d73ca09 100644 --- a/feature_engine/dataframe_checks.py +++ b/feature_engine/dataframe_checks.py @@ -2,23 +2,27 @@ transform(). """ -from typing import List, Union +from typing import List, Tuple, Union import narwhals as nw import narwhals.dependencies as nwd +import narwhals.selectors as nws import numpy as np -from narwhals.typing import IntoDataFrame, IntoSeries +from narwhals.typing import IntoDataFrame, IntoDataFrameT, IntoSeries from sklearn.utils.validation import _check_y, check_consistent_length, column_or_1d -def check_X(X: IntoDataFrame): +def check_X(X: IntoDataFrameT) -> IntoDataFrameT: """ Checks that X is a dataframe from any library supported by narwhals (for example pandas, polars, modin, cuDF, or PyArrow). Parameters ---------- - X : dataframe (pandas, polars, or any other library supported by narwhals). + X : dataframe (pandas, polars, PyArrow, modin, or cuDF). Feature-engine does + not support libraries that build a deferred query plan (for example Dask, + DuckDB, PySpark, Ibis, or a polars LazyFrame). Convert those to an eager + dataframe (e.g. `LazyFrame.collect()`) before passing them in. The input to check and transform. Raises @@ -63,8 +67,11 @@ def check_y( Parameters ---------- - y : Series or DataFrame (pandas, polars, or any other library supported by - narwhals), np.array, list + y : Series or DataFrame (pandas, polars, PyArrow, modin, or cuDF), np.array, + list. Feature-engine does not support libraries that build a deferred + query plan (for example Dask, DuckDB, PySpark, Ibis, or a polars + LazyFrame). Convert those to an eager dataframe (e.g. `LazyFrame.collect()`) + before passing them in. The input to check. y_numeric : bool, default=False @@ -84,7 +91,10 @@ def check_y( if nwd.is_into_series(y): nw_y = nw.from_native(y, series_only=True) - if nw_y.is_null().any(): + has_na = nw_y.is_null().any() + if nw_y.dtype.is_numeric(): + has_na = has_na or nw_y.is_nan().any() + if has_na: raise ValueError("y contains NaN values.") if nw_y.dtype.is_numeric(): if not np.isfinite(nw_y.to_numpy()).all(): @@ -95,7 +105,9 @@ def check_y( if nwd.is_into_dataframe(y): nw_y = nw.from_native(y, eager_only=True) - if nw_y.select(nw.all().is_null().any()).to_numpy().any(): + has_null = nw_y.select(nw.all().is_null().any()).to_numpy().any() + has_nan = nw_y.select(nws.numeric().is_nan().any()).to_numpy().any() + if has_null or has_nan: raise ValueError("y contains NaN values.") if not np.isfinite(nw_y.to_numpy()).all(): raise ValueError("y contains infinity values.") @@ -109,17 +121,20 @@ def check_y( def check_X_y( - X: IntoDataFrame, + X: IntoDataFrameT, y: Union[IntoSeries, IntoDataFrame, np.generic, np.ndarray, List], y_numeric: bool = False, -): +) -> Tuple[IntoDataFrameT, Union[IntoSeries, IntoDataFrame, np.ndarray]]: """ Ensures X and y are compatible dataframe/array-like objects with a consistent number of rows. If both are pandas objects, checks that their indexes match. Parameters ---------- - X: dataframe (pandas, polars, or any other library supported by narwhals) + X: dataframe (pandas, polars, PyArrow, modin, or cuDF). Feature-engine does + not support libraries that build a deferred query plan (for example Dask, + DuckDB, PySpark, Ibis, or a polars LazyFrame). Convert those to an eager + dataframe (e.g. `LazyFrame.collect()`) before passing them in. The input to check. y: Series, DataFrame (pandas, polars, or any other library supported by @@ -213,7 +228,12 @@ def _check_contains_na( "`missing_values='ignore'` when initialising this transformer." ) nw_X = nw.from_native(X, eager_only=True) - if nw_X.select(nw.col(variables).is_null().any()).to_numpy().any(): + has_null = nw_X.select(nw.col(variables).is_null().any()).to_numpy().any() + numeric_vars = [v for v in variables if nw_X.schema[v].is_numeric()] + has_nan = False + if numeric_vars: + has_nan = nw_X.select(nw.col(numeric_vars).is_nan().any()).to_numpy().any() + if has_null or has_nan: if error_msg == "simple": raise ValueError(error_msg_simple) else: diff --git a/tests/test_dataframe_checks.py b/tests/test_dataframe_checks.py index 085a82fad..ccef69112 100644 --- a/tests/test_dataframe_checks.py +++ b/tests/test_dataframe_checks.py @@ -132,6 +132,18 @@ def test_check_y_series_raises_nan_error(make_series): check_y(s) +@pytest.mark.parametrize( + "make_series", + [pd.Series, pl.Series], +) +def test_check_y_series_raises_nan_error_for_explicit_nan(make_series): + # in polars, an explicit float("nan") is not a null value, so it is only + # caught if is_nan() is checked in addition to is_null() + s = make_series([0.0, float("nan"), 2.0]) + with pytest.raises(ValueError, match="y contains NaN values."): + check_y(s) + + @pytest.mark.parametrize( "make_series", [pd.Series, pl.Series], @@ -190,6 +202,18 @@ def test_check_y_dataframe_raises_nan_error(make_df): check_y(d) +@pytest.mark.parametrize( + "make_df", + [pd.DataFrame, pl.DataFrame], +) +def test_check_y_dataframe_raises_nan_error_for_explicit_nan(make_df): + # in polars, an explicit float("nan") is not a null value, so it is only + # caught if is_nan() is checked in addition to is_null() + d = make_df({"t1": [0.0, float("nan"), 2.0], "t2": [5.0, 6.0, 7.0]}) + with pytest.raises(ValueError, match="y contains NaN values."): + check_y(d) + + @pytest.mark.parametrize( "make_df", [pd.DataFrame, pl.DataFrame], @@ -372,6 +396,33 @@ def test_contains_na_ignores_columns_not_in_variables(make_df): assert _check_contains_na(df, ["City"]) is None +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_contains_na_raises_for_explicit_nan_in_numeric_column(make_df): + # in polars, an explicit float("nan") is not a null value, so it is only + # caught if is_nan() is checked in addition to is_null() + msg = ( + "Some of the variables in the dataset contain NaN. Check and " + "remove those before using this transformer." + ) + df = make_df({"Age": [20.0, float("nan"), 19.0], "City": ["a", "b", "c"]}) + with pytest.raises(ValueError, match=msg): + _check_contains_na(df, ["Age", "City"]) + + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_contains_na_raises_for_mix_of_null_and_nan_across_dtypes(make_df): + # a numeric column with a NaN and a string column with a null should both + # still be caught, and the numeric-only is_nan() scoping must not error out + # on the string column + msg = ( + "Some of the variables in the dataset contain NaN. Check and " + "remove those before using this transformer." + ) + df = make_df({"Age": [20.0, float("nan"), 19.0], "City": ["a", None, "c"]}) + with pytest.raises(ValueError, match=msg): + _check_contains_na(df, ["Age", "City"]) + + # -------------------------- # test _check_contains_inf # -------------------------- From 10d3eee4ff377cdf4c3f56e7246c794d3af0fe22 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Mon, 24 Aug 2026 10:02:09 +0200 Subject: [PATCH 2/3] inline null/nan checks to match FBruzzesi's suggested one-liner Collapses the has_na/has_null/has_nan accumulator variables into a single short-circuiting if-condition, as suggested in review. This also avoids an unnecessary is_nan() call when is_null() already found a null value. Co-Authored-By: Claude Sonnet 5 --- feature_engine/dataframe_checks.py | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/feature_engine/dataframe_checks.py b/feature_engine/dataframe_checks.py index d1d73ca09..9f855d0eb 100644 --- a/feature_engine/dataframe_checks.py +++ b/feature_engine/dataframe_checks.py @@ -91,10 +91,9 @@ def check_y( if nwd.is_into_series(y): nw_y = nw.from_native(y, series_only=True) - has_na = nw_y.is_null().any() - if nw_y.dtype.is_numeric(): - has_na = has_na or nw_y.is_nan().any() - if has_na: + if nw_y.is_null().any() or ( + nw_y.dtype.is_numeric() and nw_y.is_nan().any() + ): raise ValueError("y contains NaN values.") if nw_y.dtype.is_numeric(): if not np.isfinite(nw_y.to_numpy()).all(): @@ -105,9 +104,10 @@ def check_y( if nwd.is_into_dataframe(y): nw_y = nw.from_native(y, eager_only=True) - has_null = nw_y.select(nw.all().is_null().any()).to_numpy().any() - has_nan = nw_y.select(nws.numeric().is_nan().any()).to_numpy().any() - if has_null or has_nan: + if ( + nw_y.select(nw.all().is_null().any()).to_numpy().any() + or nw_y.select(nws.numeric().is_nan().any()).to_numpy().any() + ): raise ValueError("y contains NaN values.") if not np.isfinite(nw_y.to_numpy()).all(): raise ValueError("y contains infinity values.") @@ -228,12 +228,11 @@ def _check_contains_na( "`missing_values='ignore'` when initialising this transformer." ) nw_X = nw.from_native(X, eager_only=True) - has_null = nw_X.select(nw.col(variables).is_null().any()).to_numpy().any() numeric_vars = [v for v in variables if nw_X.schema[v].is_numeric()] - has_nan = False - if numeric_vars: - has_nan = nw_X.select(nw.col(numeric_vars).is_nan().any()).to_numpy().any() - if has_null or has_nan: + if nw_X.select(nw.col(variables).is_null().any()).to_numpy().any() or ( + numeric_vars + and nw_X.select(nw.col(numeric_vars).is_nan().any()).to_numpy().any() + ): if error_msg == "simple": raise ValueError(error_msg_simple) else: From 8f377c4f286145e6e9209088bfe3a9abc68c8ea1 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Mon, 24 Aug 2026 11:18:47 +0200 Subject: [PATCH 3/3] speed up numeric column detection in _check_contains_na The schema-based list comprehension rebuilt narwhals' full column schema on every single-column access, making it scale roughly quadratically with column count on pandas (benchmarked up to ~500x slower than necessary at 200 columns). Switch to the pandas fast-path / narwhals-selector pattern already used in variable_handling (find_numerical_variables, check_numerical_variables) for the same "which of these columns are numeric" problem. Co-Authored-By: Claude Sonnet 5 --- feature_engine/dataframe_checks.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/feature_engine/dataframe_checks.py b/feature_engine/dataframe_checks.py index 9f855d0eb..b3a07be12 100644 --- a/feature_engine/dataframe_checks.py +++ b/feature_engine/dataframe_checks.py @@ -228,7 +228,10 @@ def _check_contains_na( "`missing_values='ignore'` when initialising this transformer." ) nw_X = nw.from_native(X, eager_only=True) - numeric_vars = [v for v in variables if nw_X.schema[v].is_numeric()] + if nwd.is_pandas_dataframe(X): + numeric_vars = list(X[variables].select_dtypes(include="number").columns) + else: + numeric_vars = nw_X.select(variables).select(nw.selectors.numeric()).columns if nw_X.select(nw.col(variables).is_null().any()).to_numpy().any() or ( numeric_vars and nw_X.select(nw.col(numeric_vars).is_nan().any()).to_numpy().any()