Ignore "index"-columns  of transformations in `polars.DataFrame` objects

### Describe the bug

I would like to be able to use a decorator in order to wrap `transform` methods. The wrapped methods should ignore all columns in a `polars.DataFrame` starting with a certain prefix, i.e. "index".

I need to preserve the index of my data in all transformations, i.e. they should not be used in `fit` and `transform` calls. The simplest solution that I came up with was decorating the `transform` (and `fit`) methods of existing transformers.

The problem is that the `set_output` API seems to have an unwanted interaction with my decorator. I am not sure how I can avoid this without having to re-implement existing transformations, so that they can appropriately handle the "index"-columns .




As this issue lies somewhere between bug, feature request, and usage question I'd be open to (and very happy to receive) suggestions that help to solve the specific problem. Especially insights regarding the (unwanted) interaction of the decorator with the `set_output` API would be very helpful.

Thanks a lot for this great library and all the work that goes into maintaining and extending it!


### Steps/Code to Reproduce

```python

import functools
import polars as pl
import polars.selectors as cs
import sklearn.base
import sklearn.preprocessing


def ignore_index_columns_fit(fit_method):
    @functools.wraps(fit_method)
    def wrapper(estimator: sklearn.base.BaseEstimator, X: pl.DataFrame, y: pl.DataFrame | None = None):
        X = X.select(~cs.starts_with("index"))
        if y is not None:
            y = y.select(~cs.starts_with("index"))
        return fit_method(estimator, X=X, y=y)
    return wrapper


def ignore_index_columns_transform(transform_method):
    @functools.wraps(transform_method)
    def wrapper(transformer: sklearn.base.BaseEstimator, X: pl.DataFrame):
        index, data = X.select(cs.starts_with("index")), X.select(~cs.starts_with("index"))
        data_transformed = transform_method(transformer, data)
        assert isinstance(data_transformed, pl.DataFrame)
        result = pl.concat([index, data_transformed], how="horizontal")
        return result
    return wrapper


class StandardScaler(sklearn.preprocessing.StandardScaler):
    @ignore_index_columns_fit
    def fit(self, X, y=None):
        return super().fit(X=X, y=y)

    @ignore_index_columns_transform
    def transform(self, X):
        return super().transform(X)

if __name__ == "__main__":
    X = pl.from_dict({
        "index/level_1": ["a", "a", "b"],
        "index/level_2": [1, 2, 1],
        "feature1": [1.0, 2.0, 3.0],
    })
    transformation = StandardScaler().set_output(transform="polars")
```

### Expected Results

No error is thrown, the "index" columns are unaffected by the transformation.

### Actual Results

```
Xt = transformation.fit_transform(X)
    >>> Traceback (most recent call last):
  File ".../decorated_transform.py", line 46, in <module>
    Xt = transformation.fit_transform(X)
  File ".../sklearn/utils/_set_output.py", line 313, in wrapped
    data_to_wrap = f(self, X, *args, **kwargs)
  File ".../sklearn/base.py", line 1098, in fit_transform
    return self.fit(X, **fit_params).transform(X)
  File ".../sklearn/utils/_set_output.py", line 326, in wrapped
    return _wrap_data_with_container(method, data_to_wrap, X, self)
  File ".../sklearn/utils/_set_output.py", line 301, in _wrap_data_with_container
    return adapter.create_container(
  File ".../sklearn/utils/_set_output.py", line 165, in create_container
    return self.rename_columns(X_output, columns)
  File ".../sklearn/utils/_set_output.py", line 175, in rename_columns
    X.columns = columns
  File ".../polars/dataframe/frame.py", line 798, in columns
    self._df.set_column_names(names)
polars.exceptions.ShapeError: 1 column names provided for a DataFrame of width 3
```

### Versions

```shell
System:
    python: 3.10.14 | packaged by conda-forge | (main, Mar 20 2024, 12:51:49) [Clang 16.0.6 ]
executable: .../.miniforge3/envs/.../bin/python
   machine: macOS-14.5-arm64-arm-64bit
Python dependencies:
      sklearn: 1.5.1
          pip: 24.0
   setuptools: 70.1.1
        numpy: 1.26.4
        scipy: 1.14.0
       Cython: None
       pandas: 2.2.2
   matplotlib: 3.7.3
       joblib: 1.4.2
threadpoolctl: 3.5.0
Built with OpenMP: True
threadpoolctl info:
       user_api: blas
   internal_api: openblas
    num_threads: 10
         prefix: libopenblas
       filepath: .../.miniforge3/envs/selection/lib/libopenblas.0.dylib
        version: 0.3.27
threading_layer: openmp
   architecture: VORTEX
       user_api: openmp
   internal_api: openmp
    num_threads: 10
         prefix: libomp
       filepath: .../.miniforge3/envs/.../lib/libomp.dylib
        version: None
```

COMMENT:
@adrinjalali thanks for the quick reply. For completeness I ended up implementing a solution which uses (as suggested in your comment) a `ColumnTransformer` instead.

Labels: Bug, Needs Triage