Repository: koaning/scikit-lego
Pull request #742 (2025-04-30 07:31:37)

RFC, patch: fix `GroupTransformer` return dtype
# Description

Allow object/strings dtype in `GroupTransformer` return object

Fixes #740 


## Type of change

- [x] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)


## Checklist:

- [x] My code follows the style guidelines (ruff)
- [ ] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation (also to the readme.md)
- [x] I have added tests that prove my fix is effective or that my feature works
- [x] I have added tests to check whether the new feature adheres to the sklearn convention
- [ ] New and existing unit tests pass locally with my changes


koaning: @FBruzzesi any idea why so many unit tests fail? There's no new pandas version, but it's odd to see so many issues pop up and only on specific Python/OS versions. 
FBruzzesi: > @FBruzzesi any idea why so many unit tests fail? There's no new pandas version, but it's odd to see so many issues pop up and only on specific Python/OS versions.

Hey sorry! I noticed the failure but didn't have the chance to dive deeper - I locally run ubuntu with py3.12 and have no issue. Let me take a closer look as soon as I can
FBruzzesi: @koaning it should be fixed now, issue was using `set_config` de-attached from the individual transformer
kamilm: @FBruzzesi Thanks so much for the speedy reply and patch, I will test and report back if it's still an issue. 

BTW this is probably my favorite transformer ;)

Patch:
diff --git a/sklego/meta/grouped_transformer.py b/sklego/meta/grouped_transformer.py
index d62aa64..683f509 100644
--- a/sklego/meta/grouped_transformer.py
+++ b/sklego/meta/grouped_transformer.py
@@ -215,7 +215,7 @@ class GroupedTransformer(TransformerMixin, MetaEstimatorMixin, BaseEstimator):
             for group_name, X_grp in frame.group_by(self.groups_)
         ]
 
-        output = np.zeros(shape=(n_samples, results[0][1].shape[1]))
+        output = np.empty(shape=(n_samples, results[0][1].shape[1]), dtype=results[0][1].dtype)
         for grp_index, grp_result in results:
             output[grp_index, :] = grp_result
 


Test patch:
diff --git a/tests/test_meta/test_grouped_transformer.py b/tests/test_meta/test_grouped_transformer.py
index c6bdaa0..d3237fc 100644
--- a/tests/test_meta/test_grouped_transformer.py
+++ b/tests/test_meta/test_grouped_transformer.py
@@ -6,7 +6,9 @@ import polars as pl
 import pyarrow as pa
 import pytest
 import sklearn
+from pandas.testing import assert_frame_equal
 from sklearn import clone
+from sklearn.impute import SimpleImputer
 from sklearn.linear_model import LinearRegression
 from sklearn.pipeline import Pipeline
 from sklearn.preprocessing import LabelEncoder, MinMaxScaler, StandardScaler, TargetEncoder
@@ -408,3 +410,35 @@ def test_set_output(penguins_df, frame_func, transform_output):
     meta = GroupedTransformer(StandardScaler(), groups="island").set_output(transform=transform_output)
     transformed = meta.fit_transform(X, y)
     assert isinstance(transformed, frame_func)
+
+
+def test_with_object_dtype():
+    # Example from https://github.com/koaning/scikit-lego/issues/740
+
+    data = {
+        "big": ["A", "A", "A", "A", "A", "B", "B", "B", "C", "C"],
+        "small": ["a", "a", None, "a", "a", "b", "b", None, "C", "C"],
+        "other": [0.1, 0.2, 0.3, 0.6, 0.5, 0.1, 0.3, 0.5, 0.6, 0.6],
+        "y": [1, 1, 0, 1, 0, 1, 1, 0, 0, 0],
+    }
+    df = pd.DataFrame(data=data)
+    X, y = df.drop(columns=["y"]), df["y"]
+
+    result = (
+        GroupedTransformer(
+            transformer=SimpleImputer(strategy="most_frequent", missing_values=None),
+            groups=["big"],
+            check_X=False,
+        )
+        .set_output(transform="pandas")
+        .fit_transform(X, y)
+    )
+
+    expected = pd.DataFrame(
+        {
+            "small": ["a", "a", "a", "a", "a", "b", "b", "b", "C", "C"],
+            "other": [0.1, 0.2, 0.3, 0.6, 0.5, 0.1, 0.3, 0.5, 0.6, 0.6],
+        }
+    ).astype("object")
+
+    assert_frame_equal(result, expected)