diff --git a/README.md b/README.md index 99f60ef..96fdb29 100644 --- a/README.md +++ b/README.md @@ -34,8 +34,8 @@ public, discoverable protocol. ## Why PreTab? - **Familiar interface.** A scikit-learn `fit`/`transform`/`fit_transform` API that drops - into existing pipelines and works with both `pandas.DataFrame` and `numpy.ndarray` - inputs. + into existing pipelines and works with `pandas.DataFrame`, `polars.DataFrame` and + `numpy.ndarray` inputs. - **Automatic feature handling.** Feature-type detection and per-feature strategies let you describe intent once instead of wiring transformers by hand. - **Beyond scaling.** Spline bases, neural basis maps, piecewise-linear encoding, and @@ -80,8 +80,8 @@ print(X.shape) # (100, 22) ``` -> **Note:** PreTab accepts a `pandas.DataFrame` or a `numpy.ndarray` and infers numerical -> versus categorical columns either way. +> **Note:** PreTab accepts a `pandas.DataFrame`, a `polars.DataFrame` or a `numpy.ndarray` and +> infers numerical versus categorical columns either way. > **Tip:** Swap the global methods for a `feature_preprocessing` map, for example > `{"age": "ple", "income": "rbf", "city": "one-hot"}`, and PreTab fits each column with its diff --git a/docs/core_concepts/configuration.md b/docs/core_concepts/configuration.md index c5fa21d..8390e1a 100644 --- a/docs/core_concepts/configuration.md +++ b/docs/core_concepts/configuration.md @@ -51,6 +51,31 @@ example `{"feature_0": "rbf"}`. Check `numerical_features_` / `categorical_featu overrides, rather than guessing the column order. ``` +## Polars DataFrame input + +A `polars.DataFrame` is accepted as input to the `Preprocessor` and `RepresentationSearchCV` +wherever a pandas `DataFrame` is (temporal columns are not supported), so a scikit-learn +`Pipeline` configured with `set_output(transform="polars")` can hand its polars output straight +to the `Preprocessor`. PreTab reads the frame through pandas with the same column names, order +and dtypes, then detects column types as usual: `Float*` columns are numerical, `Int*` / +`UInt*` columns follow the integer cardinality rule (an integer column with nulls becomes +float, as in pandas), and `String`, `Categorical`, `Enum` and `Boolean` columns are +categorical. Polars nulls become missing values, so imputation and `missing_policy` treat them +like `NaN` in a pandas frame. The conversion does not need `pyarrow`. + +```python +import polars as pl +from sklearn.impute import SimpleImputer +from sklearn.pipeline import Pipeline + +X = pl.DataFrame({"age": [25.0, None, 47.0, 33.0], "income": [3.1e4, 5.2e4, None, 4.4e4]}) +y = [0.0, 1.0, 1.0, 0.0] + +pipe = Pipeline([("impute", SimpleImputer()), ("pretab", Preprocessor(numerical_method="minmax"))]) +pipe.set_output(transform="polars") +pipe.fit_transform(X, y) # a polars.DataFrame with columns num_age, num_income +``` + ## Per-feature overrides Columns rarely want identical treatment. The `feature_preprocessing` dict assigns a strategy diff --git a/docs/getting_started/installation.md b/docs/getting_started/installation.md index 924416f..e6b43ba 100644 --- a/docs/getting_started/installation.md +++ b/docs/getting_started/installation.md @@ -67,7 +67,9 @@ pip install "pretab[polars]" `polars` is only needed for the `set_output(transform="polars")` output path; every other output (`output_structure="matrix"`/`"blocks"`, `output_format="dense"`/`"sparse"`, `set_output(transform="pandas")`) works without it. Requesting `"polars"` output without the -extra installed raises a clear `OptionalDependencyError`. +extra installed raises a clear `OptionalDependencyError`. A `polars.DataFrame` is also accepted +as input to `fit` / `transform` (see +[Polars DataFrame input](../core_concepts/configuration.md#polars-dataframe-input)). ``` Use the convenience `all` extra to install every optional dependency at once: diff --git a/docs/homepage.md b/docs/homepage.md index 907ee97..4112bd8 100644 --- a/docs/homepage.md +++ b/docs/homepage.md @@ -34,8 +34,8 @@ compose standalone representations with scikit-learn transformers and estimators ::: :::{grid-item-card} 🧠 Smart defaults -Automatic feature-type detection (numerical vs. categorical) with support for both -`pandas.DataFrame` and `numpy.ndarray` inputs. +Automatic feature-type detection (numerical vs. categorical) with support for +`pandas.DataFrame`, `polars.DataFrame` and `numpy.ndarray` inputs. ::: :::: diff --git a/pretab/compose/feature_detection.py b/pretab/compose/feature_detection.py index ebdb1cc..188e5b6 100644 --- a/pretab/compose/feature_detection.py +++ b/pretab/compose/feature_detection.py @@ -1,4 +1,4 @@ -"""Coerce inputs to DataFrames and classify columns as numerical or categorical. +"""Coerce inputs to pandas DataFrames and classify columns as numerical or categorical. Feature-type detection decides which construction path each column takes. It is kept here, separate from orchestration, so the Preprocessor's ``fit`` reads as a @@ -8,27 +8,54 @@ import numpy as np import pandas as pd +from ..core.validation import is_polars_frame, polars_to_pandas from ..exceptions import PretabDataError, invalid_param_error -__all__ = ["detect_column_types", "to_dataframe"] +__all__ = ["detect_column_types", "has_column_labels", "to_dataframe"] + + +def has_column_labels(X) -> bool: + """Return True if ``X`` carries its own column labels. + + DataFrames (pandas or polars) and dicts of columns are matched by label; + anything else (a NumPy array, a list of rows, ...) is matched by position. + """ + return isinstance(X, dict) or hasattr(X, "columns") def to_dataframe(X, *, copy: bool = False) -> pd.DataFrame: - """Return ``X`` as a DataFrame, naming array columns ``feature_0``, ``feature_1`` .... + """Return ``X`` as a pandas DataFrame, naming array columns ``feature_0``, ``feature_1`` .... - Dicts and NumPy arrays are wrapped in a fresh DataFrame; an existing - DataFrame is returned as-is, or copied when ``copy`` is True. + Dicts are wrapped in a fresh DataFrame and a polars DataFrame is converted to + one with the same column names, order and (as far as possible) dtypes. Input + without column labels -- a NumPy array, a list of rows -- is read as a 2D + array. An existing pandas DataFrame is returned as-is, or copied when + ``copy`` is True. Raises ------ PretabDataError - If the resulting columns contain a duplicate label. A + If ``X`` is a pandas Series or array-like input that is not 2D, or if the + resulting columns contain a duplicate label. A :class:`~sklearn.compose.ColumnTransformer` keys its per-column steps by name, so a duplicate cannot be routed unambiguously. """ + if isinstance(X, pd.Series): + raise PretabDataError( + "Expected 2D input (a DataFrame or a 2D array), got a pandas Series.\n" + "Fix: pass X.to_frame() to use the Series as a single column." + ) if isinstance(X, dict): X = pd.DataFrame(X) - elif isinstance(X, np.ndarray): + elif is_polars_frame(X): + X = polars_to_pandas(X) + elif not has_column_labels(X): + X = np.asarray(X) + if X.ndim != 2: + raise PretabDataError( + f"Expected 2D input (a DataFrame or a 2D array), got an array with {X.ndim} dimension(s).\n" + "Fix: reshape a single feature to one column, e.g. X.reshape(-1, 1)." + ) X = pd.DataFrame(X, columns=pd.Index([f"feature_{i}" for i in range(X.shape[1])])) else: X = X.copy() if copy else X diff --git a/pretab/compose/search.py b/pretab/compose/search.py index 7ce1d63..e5a3428 100644 --- a/pretab/compose/search.py +++ b/pretab/compose/search.py @@ -18,6 +18,7 @@ from sklearn.utils.validation import check_is_fitted from ..core._typing import PredictorLike +from ..core.validation import is_polars_frame, polars_to_pandas from ..exceptions import InvalidParamError from ..preprocessor import Preprocessor @@ -80,13 +81,22 @@ def _make_preprocessor(self, method): return Preprocessor(numerical_method=method, **params) def fit(self, X, y=None): - """Search over ``methods`` and refit the best representation on all data.""" + """Search over ``methods`` and refit the best representation on all data. + + ``X`` may be a pandas or polars DataFrame or an array. A polars frame is + converted to pandas once, before the fold rows are taken, so every fold's + Preprocessor detects the same column types as the final refit. + """ methods = list(self.methods) if not methods: raise InvalidParamError("methods must be a non-empty sequence of numerical methods.") if y is None: raise InvalidParamError("RepresentationSearchCV requires y at fit time; got y=None.") y_arr = np.asarray(y).ravel() + if is_polars_frame(X): + # Converted once, so every fold's rows carry the dtypes of the whole frame, + # as a pandas frame's do (see polars_to_pandas). + X = polars_to_pandas(X) n_samples = X.shape[0] if hasattr(X, "shape") else len(X) cv = cast(BaseCrossValidator, check_cv(self.cv, y_arr, classifier=is_classifier(self.estimator))) # Reuse the same held-out rows for every candidate, including splitters diff --git a/pretab/core/validation.py b/pretab/core/validation.py index fbe7ea7..1cfc786 100644 --- a/pretab/core/validation.py +++ b/pretab/core/validation.py @@ -6,15 +6,51 @@ ``transform`` -- and so ``n_features_in_`` is recorded consistently. """ +import sys import warnings from typing import Literal import numpy as np +import pandas as pd from sklearn.utils.validation import check_array from ..exceptions import DataWarning, PretabDataError -__all__ = ["validate_2d_allow_nan"] +__all__ = ["is_polars_frame", "polars_to_pandas", "validate_2d_allow_nan"] + + +def is_polars_frame(X) -> bool: + """Return True if ``X`` is a polars DataFrame, without importing polars. + + A polars object can only exist once polars has been imported, so when the + module is absent from ``sys.modules`` ``X`` cannot be a polars frame. + """ + polars = sys.modules.get("polars") + return polars is not None and isinstance(X, polars.DataFrame) + + +def polars_to_pandas(X) -> pd.DataFrame: + """Convert a polars DataFrame to pandas, keeping column names, order and dtypes. + + ``polars.DataFrame.to_pandas`` requires pyarrow, which neither polars nor + PreTab depends on, so each column goes through ``Series.to_numpy`` instead: + integer, unsigned, float, boolean and temporal columns keep their NumPy + dtype (an integer column with nulls becomes float with NaN, as in pandas), + while string, categorical and enum columns -- and boolean columns with nulls + -- become ``object`` columns. Their nulls arrive as ``None`` and are mapped + to ``NaN``, the missing marker the imputers and missing indicators recognize. + + Convert a frame once and take row subsets of the result: converted on its + own, a subset can get another dtype (an integer column with a null is float + as a whole, but int in a subset without the null). + """ + columns = {} + for column in X.get_columns(): + values = column.to_numpy() + if values.dtype == object: + values = np.where(pd.isna(values), np.nan, values) + columns[column.name] = values + return pd.DataFrame(columns) def validate_2d_allow_nan(X, *, allow_nan: bool = True, reset: bool, estimator): diff --git a/pretab/preprocessor.py b/pretab/preprocessor.py index a8193b8..32c7709 100644 --- a/pretab/preprocessor.py +++ b/pretab/preprocessor.py @@ -121,7 +121,8 @@ class Preprocessor(TransformerMixin, BaseEstimator): -------- - Supports a wide range of preprocessing methods for numerical and categorical features. - Automatically detects feature types (numerical vs. categorical). - - Compatible with both pandas DataFrames and NumPy arrays. + - Compatible with pandas and polars DataFrames and NumPy arrays (a polars frame is read + through pandas, so it gets the same column-type detection). - Handles external embedding arrays for models that require learned representations. - Returns either a dictionary of transformed feature blocks or a single NumPy array. - Fully compatible with scikit-learn transformers and pipelines. @@ -467,8 +468,10 @@ def fit(self, X, y=None, embeddings=None): Parameters ---------- - X : pandas.DataFrame, numpy.ndarray, or dict - The input features. + X : pandas.DataFrame, polars.DataFrame, numpy.ndarray, or dict + The input features. A polars frame is converted to pandas with the same column + names, order and dtypes (no ``pyarrow`` needed); an array (or a list of rows) has + its columns named ``feature_0``, ``feature_1``, .... y : array-like, default=None Target values (used for decision tree-based methods). embeddings : np.ndarray or list of np.ndarray, optional @@ -612,7 +615,7 @@ def transform(self, X, embeddings=None, return_array=None): Parameters ---------- - X : pandas.DataFrame, numpy.ndarray, or dict + X : pandas.DataFrame, polars.DataFrame, numpy.ndarray, or dict Input features to transform. embeddings : np.ndarray or list of np.ndarray, optional External embeddings to attach to dictionary output. Required when @@ -675,7 +678,7 @@ def fit_transform(self, X, y=None, embeddings=None, return_array=None): Parameters ---------- - X : pandas.DataFrame, numpy.ndarray, or dict + X : pandas.DataFrame, polars.DataFrame, numpy.ndarray, or dict Input features. y : array-like, optional Target values. @@ -840,7 +843,7 @@ def estimate_output_shape(self, X) -> tuple: Parameters ---------- - X : pandas.DataFrame, numpy.ndarray, or dict + X : pandas.DataFrame, polars.DataFrame, numpy.ndarray, or dict Input whose row count drives the estimate; not transformed. Returns @@ -860,7 +863,7 @@ def estimate_memory(self, X) -> int: Parameters ---------- - X : pandas.DataFrame, numpy.ndarray, or dict + X : pandas.DataFrame, polars.DataFrame, numpy.ndarray, or dict Input whose row count drives the estimate; not transformed. Returns diff --git a/tests/compose/test_feature_detection.py b/tests/compose/test_feature_detection.py index 670e8c4..91c9efc 100644 --- a/tests/compose/test_feature_detection.py +++ b/tests/compose/test_feature_detection.py @@ -35,6 +35,83 @@ def test_to_dataframe_rejects_duplicate_columns(): to_dataframe(df) +def test_to_dataframe_converts_polars_frame_with_names_order_and_dtypes(): + """A polars frame used to be treated as pandas and fail on ``columns.duplicated``.""" + pl = pytest.importorskip("polars") + frame = pl.DataFrame( + { + "z_float": [0.5, None, 2.5], + "count": [3, 1, 2], + "count_with_null": [3, None, 2], + "flag": [True, False, True], + "flag_with_null": [True, None, False], + "city": ["b", "a", None], + } + ).with_columns(pl.col("city").cast(pl.Categorical).alias("city_cat")) + + df = to_dataframe(frame) + + assert isinstance(df, pd.DataFrame) + assert list(df.columns) == frame.columns + assert df.dtypes.to_dict() == { + "z_float": np.float64, + "count": np.int64, + "count_with_null": np.float64, + "flag": bool, + "flag_with_null": object, + "city": object, + "city_cat": object, + } + assert df["count"].tolist() == [3, 1, 2] + assert df["city"].tolist()[:2] == ["b", "a"] and df["city_cat"].tolist()[:2] == ["b", "a"] + # Polars nulls become NaN, the missing marker pandas and the imputers use. + assert df.isna().sum().to_dict() == { + "z_float": 1, + "count": 0, + "count_with_null": 1, + "flag": 0, + "flag_with_null": 1, + "city": 1, + "city_cat": 1, + } + assert np.isnan(df["city"][2]) and np.isnan(df["flag_with_null"][1]) + + +def test_detect_column_types_on_polars_frame(): + pl = pytest.importorskip("polars") + frame = pl.DataFrame( + { + "f": [0.1, 0.2, 0.3, 0.4, 0.5, 0.6], + "many": [1, 2, 3, 4, 5, 6], + "few": [1, 2, 1, 2, 1, 2], + "s": ["a", "b", "a", "b", "a", "b"], + "flag": [True, False, True, False, True, False], + } + ).with_columns(pl.col("s").cast(pl.Categorical).alias("cat")) + num, cat = detect_column_types(frame, cat_cutoff=0.5, treat_all_integers_as_numerical=False) + assert num == ["f", "many"] + assert cat == ["few", "s", "flag", "cat"] + + +def test_to_dataframe_reads_a_list_of_rows_as_an_array(): + df = to_dataframe([[1.0, 2.0], [3.0, 4.0]]) + assert list(df.columns) == ["feature_0", "feature_1"] + np.testing.assert_array_equal(df.to_numpy(), [[1.0, 2.0], [3.0, 4.0]]) + + +@pytest.mark.parametrize( + "X, match", + [ + (pd.Series([1.0, 2.0]), "got a pandas Series"), + (np.arange(3.0), "got an array with 1 dimension"), + ([1.0, 2.0], "got an array with 1 dimension"), + ], +) +def test_to_dataframe_rejects_input_that_is_not_2d(X, match): + with pytest.raises(PretabDataError, match=match): + to_dataframe(X) + + def test_float_cutoff_uses_unique_ratio(): df = pd.DataFrame({"x": [1, 2, 3, 1, 2, 3]}) # 3 unique of 6 -> ratio 0.5 num, cat = detect_column_types(df, cat_cutoff=0.6, treat_all_integers_as_numerical=False) diff --git a/tests/integration/test_polars_input.py b/tests/integration/test_polars_input.py new file mode 100644 index 0000000..d85b569 --- /dev/null +++ b/tests/integration/test_polars_input.py @@ -0,0 +1,142 @@ +"""Polars DataFrames as Preprocessor input. + +Regression guard: a polars frame used to be treated as a pandas frame and failed +with ``AttributeError: 'list' object has no attribute 'duplicated'``, so a +scikit-learn ``Pipeline`` with ``set_output(transform="polars")`` crashed when it +handed its polars output to the Preprocessor. +""" + +import numpy as np +import pandas as pd +import pytest + +from pretab import Preprocessor + +pl = pytest.importorskip("polars") + + +@pytest.fixture +def frames(): + """The same mixed-type data as a pandas and as a polars frame (Float64, Int64, + String and Categorical columns).""" + rng = np.random.default_rng(0) + n = 120 + pandas_frame = pd.DataFrame( + { + "income": rng.normal(5e4, 1e4, n), + "visits": rng.integers(0, 1000, n), + "rooms": rng.integers(1, 4, n), + "plan": rng.choice(["basic", "premium"], n), + "city": rng.choice(["a", "b", "c"], n), + } + ) + polars_frame = pl.DataFrame( + { + "income": pandas_frame["income"].to_numpy(), + "visits": pandas_frame["visits"].to_numpy(), + "rooms": pandas_frame["rooms"].to_numpy(), + "plan": pandas_frame["plan"].tolist(), + "city": pandas_frame["city"].tolist(), + } + ).with_columns(pl.col("city").cast(pl.Categorical)) + return pandas_frame, polars_frame, rng.normal(size=n) + + +@pytest.mark.parametrize( + "params", + [ + {"numerical_method": "minmax"}, + {"numerical_method": "bspline", "random_state": 0}, + {"numerical_method": "ple", "categorical_method": "one-hot", "random_state": 0}, + ], +) +def test_polars_input_matches_pandas_input(frames, params): + pandas_frame, polars_frame, y = frames + from_pandas = Preprocessor(**params).fit(pandas_frame, y) + from_polars = Preprocessor(**params).fit(polars_frame, y) + + assert from_polars.numerical_features_ == ["income", "visits"] + assert from_polars.categorical_features_ == ["rooms", "plan", "city"] + assert from_polars.categorical_features_ == from_pandas.categorical_features_ + np.testing.assert_array_equal(from_polars.get_feature_names_out(), from_pandas.get_feature_names_out()) + + expected = np.asarray(from_pandas.transform(pandas_frame)) + np.testing.assert_allclose(np.asarray(from_polars.transform(polars_frame)), expected) + np.testing.assert_allclose(np.asarray(Preprocessor(**params).fit_transform(polars_frame, y)), expected) + # A preprocessor fitted on pandas transforms the polars frame the same way. + np.testing.assert_allclose(np.asarray(from_pandas.transform(polars_frame)), expected) + # Columns are matched by name, so a reordered polars frame transforms the same. + reordered = polars_frame.select(polars_frame.columns[::-1]) + np.testing.assert_allclose(np.asarray(from_polars.transform(reordered)), expected) + + +def test_polars_input_with_polars_output(frames): + _, polars_frame, y = frames + pre = Preprocessor(numerical_method="minmax").set_output(transform="polars").fit(polars_frame, y) + out = pre.transform(polars_frame) + assert isinstance(out, pl.DataFrame) + assert out.columns == list(pre.get_feature_names_out()) + assert out.shape == (polars_frame.height, pre.total_output_dim_) + + +def test_polars_nulls_are_missing_values(): + frame = pl.DataFrame({"x": [1.0, None, 3.0, 4.0] * 10, "city": ["a", None, "b", "a"] * 10}) + pre = Preprocessor(numerical_method="minmax", categorical_method="one-hot", missing_policy="impute_with_indicator") + pre.fit(frame) + names = list(pre.get_feature_names_out()) + # The null is imputed and flagged, not encoded as a "None" category. + assert any("missingindicator_city" in name for name in names) + assert not any("None" in name for name in names) + + +def test_representation_search_on_polars_matches_pandas(frames): + """Each fold's Preprocessor gets a frame with the polars column types too, not + an object array that would turn every column categorical (and score a + different model than the final refit).""" + from sklearn.linear_model import Ridge + + from pretab import RepresentationSearchCV + + pandas_frame, polars_frame, y = frames + y = np.asarray(pandas_frame["income"]) / 1e4 + y + + def search(X): + return RepresentationSearchCV(Ridge(), ["minmax", "bspline"], cv=3, random_state=0).fit(X, y) + + from_pandas, from_polars = search(pandas_frame), search(polars_frame) + + assert from_polars.cv_results_ == pytest.approx(from_pandas.cv_results_) + assert from_polars.best_method_ == from_pandas.best_method_ + np.testing.assert_allclose(from_polars.predict(polars_frame), from_pandas.predict(pandas_frame)) + + +@pytest.fixture +def frames_with_one_null_integer(): + """An integer column with a single null, as polars and as the pandas frame it converts to.""" + rng = np.random.default_rng(0) + kids = rng.integers(0, 5, 300).astype(float) + kids[7] = np.nan + x = rng.normal(size=300) + polars_frame = pl.DataFrame( + {"x": x, "kids": pl.Series([None if np.isnan(v) else int(v) for v in kids], dtype=pl.Int64)} + ) + # As a whole the column converts to float (with NaN), so it is numerical. + pandas_frame = pd.DataFrame({"x": x, "kids": kids}) + return pandas_frame, polars_frame, x + np.nan_to_num(kids) + rng.normal(scale=0.1, size=300) + + +def test_search_converts_a_polars_frame_once(frames_with_one_null_integer): + """Rows of a fold without the null would convert to an int column and be + detected as categorical, unlike the refit on the whole frame.""" + from sklearn.linear_model import Ridge + from sklearn.model_selection import KFold + + from pretab import RepresentationSearchCV + + pandas_frame, polars_frame, y = frames_with_one_null_integer + + def search(X): + params = {"cv": KFold(5, shuffle=True, random_state=0), "random_state": 0} + return RepresentationSearchCV(Ridge(), ["minmax"], **params).fit(X, y) + + assert search(polars_frame).cv_results_ == pytest.approx(search(pandas_frame).cv_results_) diff --git a/tests/integration/test_preprocessor.py b/tests/integration/test_preprocessor.py index df9d224..b87f710 100644 --- a/tests/integration/test_preprocessor.py +++ b/tests/integration/test_preprocessor.py @@ -1,8 +1,12 @@ +import importlib.util + import numpy as np import pandas as pd import pytest from sklearn.base import clone from sklearn.exceptions import NotFittedError +from sklearn.impute import SimpleImputer +from sklearn.pipeline import Pipeline from sklearn.utils.validation import check_is_fitted from pretab.exceptions import IncompatibleParamsError, InvalidParamError, PretabDataError @@ -377,3 +381,32 @@ def test_unknown_feature_preprocessing_key_raises(sample_data): X, y = sample_data with pytest.raises(InvalidParamError, match="feature_preprocessing"): Preprocessor(feature_preprocessing={"nnum1": "minmax"}).fit(X, y) + + +# --- input containers (issue #94) ------------------------------------------------- + + +def test_a_list_of_rows_is_read_like_an_array(): + rows = [[1.0, 2.0], [3.0, 5.0], [4.0, 9.0]] + pre = Preprocessor(numerical_method="minmax").fit(rows) + assert pre.numerical_features_ == ["feature_0", "feature_1"] + np.testing.assert_allclose(np.asarray(pre.transform(rows)), np.asarray(pre.transform(np.asarray(rows)))) + + +_needs_polars = pytest.mark.skipif(importlib.util.find_spec("polars") is None, reason="polars is not installed") + + +@pytest.mark.parametrize("container", ["pandas", pytest.param("polars", marks=_needs_polars)]) +def test_pipeline_set_output_hands_frames_to_the_preprocessor(container): + """A Pipeline configured with set_output hands its frames (pandas or polars) + from one step to the next, so the Preprocessor must accept both.""" + rng = np.random.default_rng(0) + numeric = pd.DataFrame({"income": rng.normal(5e4, 1e4, 120), "visits": rng.integers(0, 1000, 120)}) + pipe = Pipeline([("impute", SimpleImputer()), ("pretab", Preprocessor(numerical_method="minmax"))]) + pipe.set_output(transform=container) + + out = pipe.fit_transform(numeric, rng.normal(size=120)) + + assert type(out).__module__.split(".")[0] == container + assert list(out.columns) == ["num_income", "num_visits"] + np.testing.assert_allclose(np.asarray(out), np.asarray(pipe.transform(numeric)))