Repository: koaning/scikit-lego

Issue:
[FEATURE] Enable GroupedTransformer to work with .set_output(transform="pandas")
In Scikit-learn the `set_output` API was recently introduced: https://scikit-learn.org/stable/auto_examples/miscellaneous/plot_set_output.html

However, unless I am mistaken, Scikit-lego is not yet compatible with this. It would be nice to have this compatibility.

I came across this incompatibility when creating:

```python
ColumnTransformer(
            transformers=[
                (
                    "cat",
                    OrdinalEncoder(
                        handle_unknown="use_encoded_value", unknown_value=-1
                    ),
                    categorical_features,
                ),
                (
                    "cont",
                    GroupedTransformer(
                        StandardScaler(), groups=["GROEP_A", "GROEP_B"]
                    ),
                    continuous_features
                ),
            ],
            verbose_feature_names_out=False,
        ).set_output(transform="pandas")

```

Which results in:

```
ValueError: Unable to configure output for GroupedTransformer(groups=['GROEP_A', 'GROEP_B'],
                   transformer=StandardScaler()) because `set_output` is not available.
```


Discussion:
FBruzzesi: @koaning tests failing seem unrelated

Patch:
diff --git a/sklego/meta/_grouped_utils.py b/sklego/meta/_grouped_utils.py
index 47e71df..32bec84 100644
--- a/sklego/meta/_grouped_utils.py
+++ b/sklego/meta/_grouped_utils.py
@@ -25,7 +25,7 @@ def parse_X_y(X, y, groups, check_X=True, **kwargs) -> nw.DataFrame:
         X = nw.from_native(pd.DataFrame(X))
 
     # Check groups and feaures values
-    if groups is not None:
+    if groups:
         _validate_groups_values(X, groups)
 
         if check_X:
diff --git a/sklego/meta/grouped_transformer.py b/sklego/meta/grouped_transformer.py
index fe737ec..b603bd1 100644
--- a/sklego/meta/grouped_transformer.py
+++ b/sklego/meta/grouped_transformer.py
@@ -1,4 +1,4 @@
-from typing import Union
+from typing import List, Union
 
 import narwhals.stable.v1 as nw
 import numpy as np
@@ -108,14 +108,19 @@ class GroupedTransformer(TransformerMixin, MetaEstimatorMixin, BaseEstimator):
         """
         self.__check_transformer()
         self.fallback_ = None
-        self.groups_ = as_list(self.groups) if self.groups is not None else None
+        self.groups_ = as_list(self.groups) if self.groups is not None else []
 
         X = nw.from_native(X, strict=False, eager_only=True)
-        if not isinstance(X, nw.DataFrame) and self.groups_ is not None:
+
+        if isinstance(X, nw.DataFrame):
+            self.feature_names_out_ = [c for c in X.columns if c not in self.groups_]
+
+        else:
             # Accounts for negative indices if X is an array
             self.groups_ = [
                 X.shape[1] + group if isinstance(group, int) and group < 0 else group for group in self.groups_
             ]
+            self.feature_names_out_ = [f"x{i}" for i in range(X.shape[1] - len(self.groups_))]
 
         frame = parse_X_y(X, y, self.groups_, check_X=self.check_X, **self._check_kwargs)
 
@@ -203,3 +208,7 @@ class GroupedTransformer(TransformerMixin, MetaEstimatorMixin, BaseEstimator):
 
     def _more_tags(self):
         return {"allow_nan": True}
+
+    def get_feature_names_out(self) -> List[str]:
+        "Alias for the `feature_names_out_` attribute defined during fit."
+        return self.feature_names_out_


Test patch:
diff --git a/tests/test_meta/test_grouped_transformer.py b/tests/test_meta/test_grouped_transformer.py
index fe1d3e7..875bf60 100644
--- a/tests/test_meta/test_grouped_transformer.py
+++ b/tests/test_meta/test_grouped_transformer.py
@@ -4,6 +4,7 @@ import numpy as np
 import pandas as pd
 import polars as pl
 import pytest
+import sklearn
 from sklearn import clone
 from sklearn.linear_model import LinearRegression
 from sklearn.pipeline import Pipeline
@@ -25,6 +26,7 @@ def test_sklearn_compatible_estimator(estimator, check):
         "check_dtype_object",  # custom message
         "check_estimators_empty_data_messages",  # custom message
         "check_estimators_pickle",  # Fails if input contains nan
+        "check_fit1d",
     }:
         pytest.skip()
 
@@ -389,3 +391,16 @@ def test_transform_with_y(transformer):
     )
 
     assert np.allclose(X_naive.to_numpy(), X_grouped)
+
+
+@pytest.mark.parametrize(("frame_func", "transform_output"), [(pd.DataFrame, "pandas"), (pl.DataFrame, "polars")])
+def test_set_output(penguins_df, frame_func, transform_output):
+    if transform_output == "polars" and sklearn.__version__ < "1.4.0":
+        pytest.skip()
+
+    X = frame_func(penguins_df.drop(columns=["sex"]))
+    y = penguins_df["sex"]
+
+    meta = GroupedTransformer(StandardScaler(), groups="island").set_output(transform=transform_output)
+    transformed = meta.fit_transform(X, y)
+    assert isinstance(transformed, frame_func)