Repository: pandas-dev/pandas
Pull request: BUG/PERF: groupby.transform with unobserved categories
- [x] closes #55326 
- [x] [Tests added and passed](https://pandas.pydata.org/pandas-docs/dev/development/contributing_codebase.html#writing-tests) if fixing a bug or adding a new feature
- [x] All [code checks passed](https://pandas.pydata.org/pandas-docs/dev/development/contributing_codebase.html#pre-commit).
- [ ] Added [type annotations](https://pandas.pydata.org/pandas-docs/dev/development/contributing_codebase.html#type-hints) to new arguments/methods/functions.
- [x] Added an entry in the latest `doc/source/whatsnew/vX.X.X.rst` file if fixing a bug or adding a new feature.


Problem statement:
BUG: groupby.transform with a reducer and unobserved categories coerces dtype
The implementation of groupby.transform with a reduction like `sum` or `min` calls the corresponding reduction and then reindexes the output to the input's index in `_wrap_transform_fast_result`. However, when there are unobserved categories, their presence can induce NA values in the reduction, coercing the dtype to float. This results in unnecessarily coercing the result to float as well.

```
df = pd.DataFrame({'a': pd.Categorical([1, 1, 2], categories=[1, 2, 3]), 'b': [3, 4, 5]})
gb = df.groupby('a', observed=False)
result = gb.transform('min')
print(result)
#      b
# 0  3.0
# 1  3.0
# 2  5.0
```

The expected output of column `b` here should be [3, 3, 5] rather than floats.

Hints from issue discussion:
take
@rhshadrach Did you mean, `b` should be [3, 3, 5] rather than floats? Instead of [3, 5, 5].
Yes - thanks. I've updated the OP.

Commit message: Move orig_obj arg in constructor to last param, to account for possible empty param

Patch to review:
diff --git a/pandas/core/groupby/generic.py b/pandas/core/groupby/generic.py
index 0a048d11d0b4d..bbfc35b8e30b7 100644
--- a/pandas/core/groupby/generic.py
+++ b/pandas/core/groupby/generic.py
@@ -2044,6 +2044,8 @@ def _gotitem(self, key, ndim: int, subset=None):
         elif ndim == 1:
             if subset is None:
                 subset = self.obj[key]
+
+            orig_obj = self.orig_obj if not self.observed else None
             return SeriesGroupBy(
                 subset,
                 self.keys,
@@ -2056,6 +2058,7 @@ def _gotitem(self, key, ndim: int, subset=None):
                 group_keys=self.group_keys,
                 observed=self.observed,
                 dropna=self.dropna,
+                orig_obj=orig_obj,
             )
 
         raise AssertionError("invalid ndim for _gotitem")
diff --git a/pandas/core/groupby/groupby.py b/pandas/core/groupby/groupby.py
index bc37405b25a16..7661902a65579 100644
--- a/pandas/core/groupby/groupby.py
+++ b/pandas/core/groupby/groupby.py
@@ -586,6 +586,7 @@ class BaseGroupBy(PandasObject, SelectionMixin[NDFrameT], GroupByIndexingMixin):
         "keys",
         "level",
         "obj",
+        "orig_obj",
         "observed",
         "sort",
     }
@@ -1106,6 +1107,7 @@ def __init__(
         group_keys: bool = True,
         observed: bool = False,
         dropna: bool = True,
+        orig_obj: NDFrameT | None = None,
     ) -> None:
         self._selection = selection
 
@@ -1117,6 +1119,7 @@ def __init__(
         self.sort = sort
         self.group_keys = group_keys
         self.dropna = dropna
+        self.orig_obj = obj if orig_obj is None else orig_obj
 
         if grouper is None:
             grouper, exclusions, obj = get_grouper(
@@ -1879,24 +1882,70 @@ def _transform(self, func, *args, engine=None, engine_kwargs=None, **kwargs):
 
         else:
             # i.e. func in base.reduction_kernels
+            if self.observed:
+                return self._reduction_kernel_transform(
+                    func, *args, engine=engine, engine_kwargs=engine_kwargs, **kwargs
+                )
 
-            # GH#30918 Use _transform_fast only when we know func is an aggregation
-            # If func is a reduction, we need to broadcast the
-            # result to the whole group. Compute func result
-            # and deal with possible broadcasting below.
-            with com.temp_setattr(self, "as_index", True):
-                # GH#49834 - result needs groups in the index for
-                # _wrap_transform_fast_result
-                if func in ["idxmin", "idxmax"]:
-                    func = cast(Literal["idxmin", "idxmax"], func)
-                    result = self._idxmax_idxmin(func, True, *args, **kwargs)
-                else:
-                    if engine is not None:
-                        kwargs["engine"] = engine
-                        kwargs["engine_kwargs"] = engine_kwargs
-                    result = getattr(self, func)(*args, **kwargs)
+            grouper, exclusions, obj = get_grouper(
+                self.orig_obj,
+                self.keys,
+                level=self.level,
+                sort=self.sort,
+                observed=True,
+                dropna=self.dropna,
+            )
+            exclusions = frozenset(exclusions) if exclusions else frozenset()
+            obj_has_not_changed = self.orig_obj.equals(self.obj)
+
+            with (
+                com.temp_setattr(self, "observed", True),
+                com.temp_setattr(self, "_grouper", grouper),
+                com.temp_setattr(self, "exclusions", exclusions),
+                com.temp_setattr(self, "obj", obj, condition=obj_has_not_changed),
+            ):
+                return self._reduction_kernel_transform(
+                    func, *args, engine=engine, engine_kwargs=engine_kwargs, **kwargs
+                )
+
+            # with com.temp_setattr(self, "as_index", True):
+            #     # GH#49834 - result needs groups in the index for
+            #     # _wrap_transform_fast_result
+            #     if func in ["idxmin", "idxmax"]:
+            #         func = cast(Literal["idxmin", "idxmax"], func)
+            #         result = self._idxmax_idxmin(func, True, *args, **kwargs)
+            #     else:
+            #         if engine is not None:
+            #             kwargs["engine"] = engine
+            #             kwargs["engine_kwargs"] = engine_kwargs
+            #         result = getattr(self, func)(*args, **kwargs)
+
+            #         print("result with observed = False\n", result.to_string())
+            #         r = self._wrap_transform_fast_result(result)
+            #         print("reindexed result", r.to_string())
+            #         return r
+
+    @final
+    def _reduction_kernel_transform(
+        self, func, *args, engine=None, engine_kwargs=None, **kwargs
+    ):
+        # GH#30918 Use _transform_fast only when we know func is an aggregation
+        # If func is a reduction, we need to broadcast the
+        # result to the whole group. Compute func result
+        # and deal with possible broadcasting below.
+        with com.temp_setattr(self, "as_index", True):
+            # GH#49834 - result needs groups in the index for
+            # _wrap_transform_fast_result
+            if func in ["idxmin", "idxmax"]:
+                func = cast(Literal["idxmin", "idxmax"], func)
+                result = self._idxmax_idxmin(func, True, *args, **kwargs)
+            else:
+                if engine is not None:
+                    kwargs["engine"] = engine
+                    kwargs["engine_kwargs"] = engine_kwargs
+                result = getattr(self, func)(*args, **kwargs)
 
-            return self._wrap_transform_fast_result(result)
+        return self._wrap_transform_fast_result(result)
 
     @final
     def _wrap_transform_fast_result(self, result: NDFrameT) -> NDFrameT:
diff --git a/pandas/tests/groupby/transform/test_transform.py b/pandas/tests/groupby/transform/test_transform.py
index 245fb9c7babd7..81fdd49492a66 100644
--- a/pandas/tests/groupby/transform/test_transform.py
+++ b/pandas/tests/groupby/transform/test_transform.py
@@ -1232,9 +1232,9 @@ def test_categorical_and_not_categorical_key(observed):
     tm.assert_frame_equal(result, expected_explicit)
 
     # Series case
-    result = df_with_categorical.groupby(["A", "C"], observed=observed)["B"].transform(
-        "sum"
-    )
+    gb = df_with_categorical.groupby(["A", "C"], observed=observed)
+    gbp = gb["B"]
+    result = gbp.transform("sum")
     expected = df_without_categorical.groupby(["A", "C"])["B"].transform("sum")
     tm.assert_series_equal(result, expected)
     expected_explicit = Series([4, 2, 4], name="B")
@@ -1535,3 +1535,151 @@ def test_transform_sum_one_column_with_matching_labels_and_missing_labels():
     result = df.groupby(series, as_index=False).transform("sum")
     expected = DataFrame({"X": [-93203.0, -93203.0, np.nan]})
     tm.assert_frame_equal(result, expected)
+
+
+# GH#58084
+def test_min_one_unobserved_category_no_type_coercion():
+    df = DataFrame({"A": Categorical([1, 1, 2], categories=[1, 2, 3]), "B": [3, 4, 5]})
+    df["B"] = df["B"].astype("int32")
+    gb = df.groupby("A", observed=False)
+    result = gb.transform("min")
+
+    expected = DataFrame({"B": [3, 3, 5]}, dtype="int32")
+    tm.assert_frame_equal(expected, result)
+
+
+# GH#58084
+def test_min_multiple_unobserved_categories_no_type_coercion():
+    df = DataFrame(
+        {
+            "X": Categorical(
+                ["432945", "randomcat", -4325466, "randomcat", -4325466, -4325466],
+                categories=[
+                    1,
+                    "randomcat",
+                    100,
+                    333,
+                    "cat43543",
+                    -4325466,
+                    54665,
+                    -546767,
+                    "432945",
+                    767076,
+                ],
+            ),
+            "Y": [0, 940645, np.iinfo(np.int64).min, 9449, 100044444, 40],
+        }
+    )
+    df["Y"] = df["Y"].astype("int64")
+
+    gb = df.groupby("X", observed=False)
+    result = gb.transform("min")
+
+    expected = DataFrame(
+        {
+            "Y": [
+                0,
+                9449,
+                np.iinfo(np.int64).min,
+                9449,
+                np.iinfo(np.int64).min,
+                np.iinfo(np.int64).min,
+            ]
+        },
+        dtype="int64",
+    )
+    tm.assert_frame_equal(expected, result)
+
+
+# GH#58084
+def test_min_float32_multiple_unobserved_categories_no_type_coercion():
+    df = DataFrame(
+        {
+            "X": Categorical(
+                ["cat43543", -4325466, 54665, "cat43543", -4325466, 54665],
+                categories=[
+                    1,
+                    "randomcat",
+                    100,
+                    333,
+                    "cat43543",
+                    -4325466,
+                    54665,
+                    -546767,
+                    "432945",
+                    767076,
+                ],
+            ),
+            "Y": [
+                0.3940429,
+                940645.49,
+                np.finfo(np.float32).min,
+                9449.03333,
+                100044444.403294,
+                40.3020909,
+            ],
+        }
+    )
+    df["Y"] = df["Y"].astype("float32")
+
+    gb = df.groupby("X", observed=False)
+    result = gb.transform("min")
+
+    expected = DataFrame(
+        {
+            "Y": [
+                0.3940429,
+                940645.49,
+                np.finfo(np.float32).min,
+                0.3940429,
+                940645.49,
+                np.finfo(np.float32).min,
+            ]
+        },
+        dtype="float32",
+    )
+    tm.assert_frame_equal(expected, result)
+
+
+# GH#58084
+def test_min_all_empty_data_no_type_coercion():
+    df = DataFrame(
+        {
+            "X": Categorical(
+                [],
+                categories=[
+                    1,
+                    "randomcat",
+                    100,
+                    333,
+                    "cat43543",
+                    -4325466,
+                    54665,
+                    -546767,
+                    "432945",
+                    767076,
+                ],
+            ),
+            "Y": [],
+        }
+    )
+    df["Y"] = df["Y"].astype("int32")
+
+    gb = df.groupby("X", observed=False)
+    result = gb.transform("min")
+
+    expected = DataFrame({"Y": []}, dtype="int32")
+    tm.assert_frame_equal(expected, result)
+
+
+# GH#58084
+def test_min_one_dim_no_type_coercion():
+    df = DataFrame({"Y": [9435, -5465765, 5055, 0, 954960]})
+    df["Y"] = df["Y"].astype("int32")
+    categories = Categorical([1, 2, 2, 5, 1], categories=[1, 2, 3, 4, 5])
+
+    gb = df.groupby(categories, observed=False)
+    result = gb.transform("min")
+
+    expected = DataFrame({"Y": [9435, -5465765, -5465765, 0, 9435]}, dtype="int32")
+    tm.assert_frame_equal(expected, result)


Code review comments:
diff_hunk: @@ -1879,24 +1882,70 @@ def _transform(self, func, *args, engine=None, engine_kwargs=None, **kwargs):
 
         else:
             # i.e. func in base.reduction_kernels
+            if self.observed:
+                return self._reduction_kernel_transform(
+                    func, *args, engine=engine, engine_kwargs=engine_kwargs, **kwargs
+                )
 
-            # GH#30918 Use _transform_fast only when we know func is an aggregation
-            # If func is a reduction, we need to broadcast the
-            # result to the whole group. Compute func result
-            # and deal with possible broadcasting below.
-            with com.temp_setattr(self, "as_index", True):
-                # GH#49834 - result needs groups in the index for
-                # _wrap_transform_fast_result
-                if func in ["idxmin", "idxmax"]:
-                    func = cast(Literal["idxmin", "idxmax"], func)
-                    result = self._idxmax_idxmin(func, True, *args, **kwargs)
-                else:
-                    if engine is not None:
-                        kwargs["engine"] = engine
-                        kwargs["engine_kwargs"] = engine_kwargs
-                    result = getattr(self, func)(*args, **kwargs)
+            grouper, exclusions, obj = get_grouper(
+                self.orig_obj,
+                self.keys,
+                level=self.level,
+                sort=self.sort,
+                observed=True,
+                dropna=self.dropna,
+            )
+            exclusions = frozenset(exclusions) if exclusions else frozenset()
+            obj_has_not_changed = self.orig_obj.equals(self.obj)
+
+            with (
+                com.temp_setattr(self, "observed", True),
+                com.temp_setattr(self, "_grouper", grouper),
+                com.temp_setattr(self, "exclusions", exclusions),
+                com.temp_setattr(self, "obj", obj, condition=obj_has_not_changed),
path: pandas/core/groupby/groupby.py
text: @user1:
Why can't we unconditionally set `obj` here?

diff_hunk: @@ -1879,24 +1882,70 @@ def _transform(self, func, *args, engine=None, engine_kwargs=None, **kwargs):
 
         else:
             # i.e. func in base.reduction_kernels
+            if self.observed:
+                return self._reduction_kernel_transform(
+                    func, *args, engine=engine, engine_kwargs=engine_kwargs, **kwargs
+                )
 
-            # GH#30918 Use _transform_fast only when we know func is an aggregation
-            # If func is a reduction, we need to broadcast the
-            # result to the whole group. Compute func result
-            # and deal with possible broadcasting below.
-            with com.temp_setattr(self, "as_index", True):
-                # GH#49834 - result needs groups in the index for
-                # _wrap_transform_fast_result
-                if func in ["idxmin", "idxmax"]:
-                    func = cast(Literal["idxmin", "idxmax"], func)
-                    result = self._idxmax_idxmin(func, True, *args, **kwargs)
-                else:
-                    if engine is not None:
-                        kwargs["engine"] = engine
-                        kwargs["engine_kwargs"] = engine_kwargs
-                    result = getattr(self, func)(*args, **kwargs)
+            grouper, exclusions, obj = get_grouper(
+                self.orig_obj,
+                self.keys,
+                level=self.level,
+                sort=self.sort,
+                observed=True,
+                dropna=self.dropna,
+            )
path: pandas/core/groupby/groupby.py
text: @user1:
I think we'll want to cache this on the groupby instance - we do not want to have to recompute it if the groupby is reused.