fix(cross-fitting): pass DataFrames to the wrapped transformer unchanged

docs/core_concepts/target_awareness.md

@@ -120,6 +120,13 @@ need cross-fitting when you build the training features once and reuse them dire
 The fitted spec records `cross_fitted=True` and the number of folds, so the choice is
 visible and serializable.
 
+`CrossFittedTransformer` can also wrap a `Preprocessor`, `ColumnTransformer`, or
+`Pipeline`. Pandas DataFrames keep their column names, dtypes, and index labels when
+passed to the wrapped estimator; each fold is selected by row position. Polars
+DataFrames are converted once to pandas before selecting folds, so nullable column
+dtypes remain consistent. Other array-like inputs retain the NumPy conversion, with
+one-dimensional input treated as a single column.
+
 ```{note}
 Cross-fitting matters most for strongly supervised encodings such as PLE, where the target
 directly determines the bins. For unsupervised methods it is unnecessary.

pretab/core/supervised.py

@@ -32,6 +32,7 @@
 from ._typing import TransformerLike
 from .parameters import validate_task
 from .representation import RepresentationSpecMixin
+from .validation import is_polars_frame, polars_to_pandas
 
 __all__ = ["CrossFittedTransformer", "in_controlled_context", "warn_target_leakage"]
 
@@ -88,6 +89,29 @@ def warn_target_leakage(estimator, y) -> None:
     )
 
 
+def _as_2d(X):
+    """Return ``X`` as 2D input for the wrapped transformer.
+
+    A pandas DataFrame is passed through unchanged, so column names and per-column
+    dtypes reach the wrapped estimator (e.g. a Preprocessor detecting numerical
+    vs categorical columns, or a name-based ColumnTransformer). A polars frame
+    is converted once to pandas, before selecting fold rows, so nullable column
+    dtypes remain consistent across folds. Other inputs are converted to NumPy,
+    with 1D input reshaped to a single column.
+    """
+    if is_polars_frame(X):
+        return polars_to_pandas(X)
+    if hasattr(X, "iloc") and getattr(X, "ndim", None) == 2:
+        return X
+    X = np.asarray(X)
+    return X.reshape(-1, 1) if X.ndim == 1 else X
+
+
+def _take_rows(X, indices):
+    """Select rows of a DataFrame or array by position."""
+    return X.iloc[indices] if hasattr(X, "iloc") else X[indices]
+
+
 class CrossFittedTransformer(RepresentationSpecMixin, TransformerMixin, BaseEstimator):
     """Cross-fit a supervised transformer to remove target leakage on training data.
 
@@ -99,7 +123,11 @@ class CrossFittedTransformer(RepresentationSpecMixin, TransformerMixin, BaseEsti
     Parameters
     ----------
     transformer : estimator
-        A supervised (target-aware) PreTab transformer to cross-fit.
+        A supervised (target-aware) PreTab transformer to cross-fit, or an
+        estimator such as a :class:`~pretab.Preprocessor` or ``Pipeline`` built from
+        them. A pandas DataFrame ``X`` is passed to it unchanged (each fold is a row
+        subset), so column names, dtypes and index labels are preserved. Polars
+        DataFrames are converted once to pandas before selecting fold rows.
     n_folds : int, default=5
         Number of cross-fitting folds. Must be at least 2.
     task : {"regression", "classification"}, default="regression"
@@ -135,15 +163,13 @@ def _make_splitter(self):
         return KFold(n_splits=self.n_folds, shuffle=self.shuffle, random_state=seed)
 
     def _fit_full(self, X, y):
-        """Validate inputs and fit ``estimator_`` on all data; return arrays."""
+        """Validate inputs and fit ``estimator_``; return 2D input and target array."""
         if y is None:
             raise IncompatibleParamsError("CrossFittedTransformer requires y at fit time; got y=None.")
         if not isinstance(self.n_folds, (int, np.integer)) or self.n_folds < 2:
             raise InvalidParamError(f"n_folds must be an integer >= 2; got {self.n_folds!r}.")
         validate_task(self.task, type(self).__name__)
-        X_arr = np.asarray(X)
-        if X_arr.ndim == 1:
-            X_arr = X_arr.reshape(-1, 1)
+        X_arr = _as_2d(X)
         y_arr = np.asarray(y).ravel()
         if len(X_arr) != len(y_arr):
             raise PretabDataError(f"X and y must have same length. Got {len(X_arr)} and {len(y_arr)}")
@@ -165,10 +191,7 @@ def fit(self, X, y=None):
     def transform(self, X):
         """Transform ``X`` using the transformer fit on all training data."""
         check_is_fitted(self, "estimator_")
-        X_arr = np.asarray(X)
-        if X_arr.ndim == 1:
-            X_arr = X_arr.reshape(-1, 1)
-        return self.estimator_.transform(X_arr)
+        return self.estimator_.transform(_as_2d(X))
 
     def fit_transform(self, X, y=None):
         """Fit and return leakage-free out-of-fold features for the training data."""
@@ -180,8 +203,8 @@ def fit_transform(self, X, y=None):
         try:
             for train_idx, test_idx in splitter.split(X_arr, y_arr):
                 fold = cast(TransformerLike, clone(self.transformer))
-                fold.fit(X_arr[train_idx], y_arr[train_idx])
-                fold_out = np.asarray(fold.transform(X_arr[test_idx]))
+                fold.fit(_take_rows(X_arr, train_idx), y_arr[train_idx])
+                fold_out = np.asarray(fold.transform(_take_rows(X_arr, test_idx)))
                 if fold_out.shape[1] != width:
                     raise IncompatibleParamsError(
                         "Cross-fitting requires a fixed output width across folds; expected "
@@ -201,6 +224,8 @@ def get_feature_names_out(self, input_features=None):
     def get_representation_spec(self, input_features=None):
         """Return the wrapped spec, flagged as cross-fitted."""
         check_is_fitted(self, "estimator_")
+        if input_features is None:
+            input_features = getattr(self.estimator_, "feature_names_in_", None)
         spec_fn = getattr(self.estimator_, "get_representation_spec", None)
         if spec_fn is not None:
             base = spec_fn(input_features)

tests/core/test_cross_fitted.py

@@ -6,6 +6,7 @@
 """
 
 import warnings
+from typing import cast
 
 import numpy as np
 import pytest
@@ -106,6 +107,70 @@ def test_invalid_task_rejected(data):
         CrossFittedTransformer(PLETransformer(), task="classificaton").fit_transform(X, y)
 
 
+# --- DataFrame input is passed through (issue #58) --------------------------------
+
+
+@pytest.fixture
+def mixed_frame():
+    import pandas as pd
+
+    rng = np.random.default_rng(0)
+    X = pd.DataFrame(
+        {"num": rng.normal(size=300), "city": rng.choice(["a", "b", "c"], size=300)},
+        index=rng.permutation(1000)[:300],
+    )
+    return X, X["num"].to_numpy() ** 2 + rng.normal(scale=0.1, size=300)
+
+
+def test_wrapped_preprocessor_keeps_dataframe_column_types(mixed_frame):
+    """Regression guard for issue #58: np.asarray turned the frame into an object
+    array, so the wrapped Preprocessor treated every column as categorical."""
+    from pretab import Preprocessor
+
+    X, y = mixed_frame
+    direct = Preprocessor(output_dim=6, random_state=0).fit(X, y)
+    cross_fitted = CrossFittedTransformer(Preprocessor(output_dim=6, random_state=0), n_folds=3, random_state=0)
+    out = cross_fitted.fit_transform(X, y)
+
+    wrapped = cast(Preprocessor, cross_fitted.estimator_)
+    assert wrapped.numerical_features_ == ["num"]
+    assert wrapped.categorical_features_ == ["city"]
+    assert out.shape == (300, direct.total_output_dim_)
+    assert list(cross_fitted.get_feature_names_out()) == list(direct.get_feature_names_out())
+    assert len(np.unique(out[:, 0])) > 1  # a supervised encoding, not an "unseen category" code
+    assert cross_fitted.transform(X).shape == out.shape
+
+
+def test_wrapped_column_transformer_selects_by_column_name(mixed_frame):
+    from sklearn.compose import ColumnTransformer
+
+    X, y = mixed_frame
+    wrapped = ColumnTransformer([("ple", PLETransformer(output_dim=4), ["num"])])
+    out = CrossFittedTransformer(wrapped, n_folds=3, random_state=0).fit_transform(X, y)
+    assert out.shape == (300, 4)
+
+
+def test_wrapped_preprocessor_accepts_feature_preprocessing_by_column_name(mixed_frame):
+    from pretab import Preprocessor
+
+    X, y = mixed_frame
+    wrapped = Preprocessor(
+        feature_preprocessing={"num": "ple"}, numerical_method="minmax", output_dim=4, random_state=0
+    )
+    cross_fitted = CrossFittedTransformer(wrapped, n_folds=3, random_state=0)
+    out = cross_fitted.fit_transform(X, y)
+
+    assert list(cross_fitted.get_feature_names_out()) == [f"num_num_ple{i}" for i in range(4)] + ["cat_city"]
+    assert out.shape == (300, 5)
+
+
+def test_one_dimensional_input_is_still_a_single_column():
+    rng = np.random.default_rng(0)
+    x = rng.normal(size=200)
+    out = CrossFittedTransformer(PLETransformer(output_dim=4), n_folds=3, random_state=0).fit_transform(x, x**2)
+    assert out.shape == (200, 4)
+
+
 @pytest.mark.parametrize("seed", range(5))
 def test_cross_fitting_a_discrete_feature_keeps_a_fixed_width(seed):
     """Regression guard for issue #57: tied features gave fold-dependent widths."""
@@ -118,3 +183,74 @@ def test_cross_fitting_a_discrete_feature_keeps_a_fixed_width(seed):
         RBFExpansionTransformer(output_dim=10, target_aware=True), n_folds=5, random_state=0
     ).fit_transform(x, y)
     assert out.shape == (200, 10)
+
+
+@pytest.mark.parametrize("wrapped_kind", ["preprocessor", "column_transformer", "pipeline"])
+def test_dataframe_representation_spec_defaults_to_fitted_names(mixed_frame, wrapped_kind):
+    from sklearn.compose import ColumnTransformer
+    from sklearn.pipeline import Pipeline
+
+    from pretab import Preprocessor
+
+    X, y = mixed_frame
+    if wrapped_kind == "preprocessor":
+        wrapped = Preprocessor(output_dim=4, random_state=0)
+    else:
+        wrapped = ColumnTransformer([("ple", PLETransformer(output_dim=4), ["num"])])
+        if wrapped_kind == "pipeline":
+            wrapped = Pipeline([("features", wrapped)])
+    cross_fitted = CrossFittedTransformer(wrapped, n_folds=3, random_state=0).fit(X, y)
+
+    spec = cross_fitted.get_representation_spec()
+    assert spec.input_features == ("num", "city")
+    assert spec.output_features == tuple(cross_fitted.get_feature_names_out())
+    assert spec.cross_fitted is True
+    assert spec.n_folds == 3
+    assert spec.uses_target is True
+    assert spec == type(spec).from_dict(spec.to_dict())
+    assert cross_fitted.get_representation_spec(["num", "city"]) == spec
+
+
+def test_dataframe_fold_rows_keep_index_dtypes_and_target_alignment(mixed_frame, monkeypatch):
+    import pandas as pd
+
+    from pretab import Preprocessor
+
+    X, y = mixed_frame
+    original = X.copy(deep=True)
+    fits, transforms = [], []
+    original_fit = Preprocessor.fit
+    original_transform = Preprocessor.transform
+
+    def record_fit(self, X, y=None):
+        assert isinstance(X, pd.DataFrame)
+        fits.append((X.copy(deep=True), np.asarray(y).copy()))
+        return original_fit(self, X, y)
+
+    def record_transform(self, X):
+        assert isinstance(X, pd.DataFrame)
+        transforms.append(X.copy(deep=True))
+        return original_transform(self, X)
+
+    monkeypatch.setattr(Preprocessor, "fit", record_fit)
+    monkeypatch.setattr(Preprocessor, "transform", record_transform)
+    cross_fitted = CrossFittedTransformer(
+        Preprocessor(output_dim=4, output_format="dense", random_state=0), n_folds=3, random_state=0
+    )
+    cross_fitted.fit_transform(X, y)
+    splits = list(KFold(n_splits=3, shuffle=True, random_state=0).split(X))
+
+    assert len(fits) == 4
+    pd.testing.assert_frame_equal(fits[0][0], X)
+    np.testing.assert_array_equal(fits[0][1], y)
+    for (actual_X, actual_y), (train_idx, _) in zip(fits[1:], splits, strict=True):
+        pd.testing.assert_frame_equal(actual_X, X.iloc[train_idx])
+        np.testing.assert_array_equal(actual_y, y[train_idx])
+    assert len(transforms) == 3
+    for actual_X, (_, test_idx) in zip(transforms, splits, strict=True):
+        pd.testing.assert_frame_equal(actual_X, X.iloc[test_idx])
+
+    unseen = X.iloc[:5][["city", "num"]]
+    cross_fitted.transform(unseen)
+    pd.testing.assert_frame_equal(transforms[-1], unseen)
+    pd.testing.assert_frame_equal(X, original)

tests/integration/test_polars_input.py

@@ -207,3 +207,47 @@ def search(X):
         return RepresentationSearchCV(Ridge(), ["minmax"], **params).fit(X, y)
 
     assert search(polars_frame).cv_results_ == pytest.approx(search(pandas_frame).cv_results_)
+
+
+@pytest.mark.parametrize("method", ["minmax", "ple", "bspline"])
+def test_cross_fitting_polars_input_matches_pandas_input(frames, method):
+    from pretab import CrossFittedTransformer
+
+    pandas_frame, polars_frame, y = frames
+    columns = ["income", "plan", "city", "flag"]
+    pandas_frame = pandas_frame[columns]
+    polars_frame = polars_frame.select(columns)
+
+    def cross_fit(X):
+        return CrossFittedTransformer(
+            Preprocessor(
+                numerical_method=method,
+                categorical_method="one-hot",
+                output_dim=6,
+                output_format="dense",
+                random_state=0,
+            ),
+            n_folds=3,
+            random_state=0,
+        ).fit_transform(X, y)
+
+    np.testing.assert_allclose(cross_fit(polars_frame), cross_fit(pandas_frame))
+
+
+def test_cross_fitting_converts_a_polars_frame_once(frames_with_one_null_integer):
+    from pretab import CrossFittedTransformer
+
+    pandas_frame, polars_frame, y = frames_with_one_null_integer
+
+    def cross_fit(X):
+        wrapped = CrossFittedTransformer(
+            Preprocessor(numerical_method="minmax", output_format="dense"), n_folds=5, random_state=0
+        )
+        return wrapped, wrapped.fit_transform(X, y)
+
+    from_pandas, expected = cross_fit(pandas_frame)
+    from_polars, actual = cross_fit(polars_frame)
+    np.testing.assert_allclose(actual, expected)
+    np.testing.assert_allclose(from_polars.transform(polars_frame), from_pandas.transform(pandas_frame))
+    np.testing.assert_array_equal(from_polars.get_feature_names_out(), from_pandas.get_feature_names_out())
+    assert from_polars.get_representation_spec().input_features == ("x", "kids")