Return narwhals frame from check_X / check_X_y (#1019)

check_X now returns the validated narwhals DataFrame instead of converting
back to the native frame, and check_X_y propagates that. The pandas
index-consistency check in check_X_y is updated to reach the native frame
via X.to_native(), and the now-unused IntoDataFrameT import / type hints
are replaced with IntoDataFrame.

Tests in test_dataframe_checks.py are updated for the new return contract:
they assert the result is a narwhals.DataFrame and compare X.to_native()
against the original.

Note: downstream transformers still expect a native frame from these
helpers; they will be adapted on their own narwhals-* branches.

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

Author: solegalli

feature_engine/dataframe_checks.py

@@ -8,11 +8,11 @@
 import narwhals.dependencies as nwd
 import narwhals.selectors as nws
 import numpy as np
-from narwhals.typing import IntoDataFrame, IntoDataFrameT, IntoSeries
+from narwhals.typing import IntoDataFrame, IntoSeries
 from sklearn.utils.validation import _check_y, check_consistent_length, column_or_1d
 
 
-def check_X(X: IntoDataFrameT) -> IntoDataFrameT:
+def check_X(X: IntoDataFrame) -> IntoDataFrame:
     """
     Checks that X is a dataframe from any library supported by narwhals (for example
     pandas, polars, modin, cuDF, or PyArrow).
@@ -34,8 +34,8 @@ def check_X(X: IntoDataFrameT) -> IntoDataFrameT:
 
     Returns
     -------
-    X : dataframe.
-        The validated dataframe in its native format.
+    X : narwhals dataframe.
+        The validated dataframe in narwhals format.
     """
     if nwd.is_into_dataframe(X):
         # from_native() raises narwhals.exceptions.DuplicateError, a ValueError
@@ -53,7 +53,7 @@ def check_X(X: IntoDataFrameT) -> IntoDataFrameT:
             f"(e.g. pandas, polars, PyArrow). Got {type(X)} instead."
         )
 
-    return nw_X.to_native()
+    return nw_X
 
 
 def check_y(
@@ -121,10 +121,10 @@ def check_y(
 
 
 def check_X_y(
-    X: IntoDataFrameT,
+    X: IntoDataFrame,
     y: Union[IntoSeries, IntoDataFrame, np.generic, np.ndarray, List],
     y_numeric: bool = False,
-) -> Tuple[IntoDataFrameT, Union[IntoSeries, IntoDataFrame, np.ndarray]]:
+) -> Tuple[IntoDataFrame, 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.
@@ -156,16 +156,16 @@ def check_X_y(
 
     Returns
     -------
-    X: dataframe
+    X: narwhals dataframe
     y: Series, DataFrame, or numpy array
     """
     X = check_X(X)
     y = check_y(y, y_numeric=y_numeric)
     check_consistent_length(X, y)
 
-    if nwd.is_pandas_dataframe(X):
+    if X.implementation.is_pandas():
         if nwd.is_pandas_series(y) or nwd.is_pandas_dataframe(y):
-            if not X.index.equals(y.index):
+            if not X.to_native().index.equals(y.index):
                 raise ValueError("The indexes of X and y do not match.")
 
     return X, y

tests/test_dataframe_checks.py

@@ -1,3 +1,4 @@
+import narwhals as nw
 import numpy as np
 import pandas as pd
 import polars as pl
@@ -28,8 +29,8 @@
 def test_check_X_returns_df_unchanged(make_df, assert_equal_fn):
     df = make_df({"a": [1, 2, 3], "b": [4.0, 5.0, 6.0]})
     X = check_X(df)
-    assert isinstance(X, type(df))
-    assert_equal_fn(X, df)
+    assert isinstance(X, nw.DataFrame)
+    assert_equal_fn(X.to_native(), df)
 
 
 @pytest.mark.parametrize(
@@ -45,7 +46,9 @@ def test_check_X_returns_df_with_mixed_dtypes(make_df, assert_equal_fn):
         "dob": pd.date_range("2020-02-24", periods=4, freq="min"),
     }
     df = make_df(data)
-    assert_equal_fn(check_X(df), df)
+    X = check_X(df)
+    assert isinstance(X, nw.DataFrame)
+    assert_equal_fn(X.to_native(), df)
 
 
 @pytest.mark.parametrize(
@@ -265,8 +268,8 @@ def test_check_X_y_returns_df_and_series_unchanged(
     df = make_df({"a": [1, 2, 3], "b": [4, 5, 6]})
     s = make_series([0, 1, 2])
     X, y = check_X_y(df, s)
-    assert isinstance(X, type(df)) and isinstance(y, type(s))
-    assert_frame_fn(X, df)
+    assert isinstance(X, nw.DataFrame) and isinstance(y, type(s))
+    assert_frame_fn(X.to_native(), df)
     assert_series_fn(y, s)
 
 
@@ -278,7 +281,8 @@ def test_check_X_y_returns_df_and_multioutput_y_unchanged(make_df, assert_frame_
     df = make_df({"a": [1, 2, 3, 4], "b": [5, 6, 7, 8]})
     d = make_df({"t1": [1, 2, 3, 4], "t2": [5, 6, 7, 8]})
     X, y = check_X_y(df, d)
-    assert_frame_fn(X, df)
+    assert isinstance(X, nw.DataFrame)
+    assert_frame_fn(X.to_native(), df)
     assert_frame_fn(y, d)
 
 
@@ -299,7 +303,8 @@ def test_check_X_y_with_array_like_y_returns_check_y_output(
 ):
     df = make_df({"a": [1, 2, 3], "b": [4, 5, 6]})
     X, y_out = check_X_y(df, y)
-    assert_frame_fn(X, df)
+    assert isinstance(X, nw.DataFrame)
+    assert_frame_fn(X.to_native(), df)
     np.testing.assert_array_equal(y_out, check_y(y))
 
 
@@ -308,7 +313,8 @@ def test_check_X_y_returns_pandas_with_non_typical_index():
     df = pd.DataFrame({"0": [1, 2, 3, 4], "1": [5, 6, 7, 8]}, index=[22, 99, 101, 212])
     s = pd.Series([1, 2, 3, 4], index=[22, 99, 101, 212])
     x, y = check_X_y(df, s)
-    assert_frame_equal(df, x)
+    assert isinstance(x, nw.DataFrame)
+    assert_frame_equal(df, x.to_native())
     assert_series_equal(s, y)