Skip to content

address review feedback on dataframe_checks.py #989

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
solegalli merged 3 commits into from
Aug 24, 2026
Merged

address review feedback on dataframe_checks.py #989

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
192c74d
address review feedback on dataframe_checks.py
solegalli Aug 24, 2026
10d3eee
inline null/nan checks to match FBruzzesi's suggested one-liner
solegalli Aug 24, 2026
8f377c4
speed up numeric column detection in _check_contains_na
solegalli Aug 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 34 additions & 12 deletions feature_engine/dataframe_checks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -84,7 +91,9 @@ def check_y(

if nwd.is_into_series(y):
nw_y = nw.from_native(y, series_only=True)
if nw_y.is_null().any():
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():
Expand All @@ -95,7 +104,10 @@ 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():
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.")
Expand All @@ -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
Expand Down Expand Up @@ -213,7 +228,14 @@ 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():
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()
):
if error_msg == "simple":
raise ValueError(error_msg_simple)
else:
Expand Down
51 changes: 51 additions & 0 deletions tests/test_dataframe_checks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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],
Expand Down Expand Up @@ -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],
Expand Down Expand Up @@ -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
# --------------------------
Expand Down