diff --git a/feature_engine/dataframe_checks.py b/feature_engine/dataframe_checks.py index ae3be3c7b..b3a07be12 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,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(): @@ -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.") @@ -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,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: 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 # --------------------------