fix(preprocessor): accept polars DataFrames as input

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

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

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:

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.
 :::
 
 ::::

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

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

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):

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

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)