From 1644d07cf103c33b756f674008268f8bce862941 Mon Sep 17 00:00:00 2001
From: Brock <jbrockmendel@gmail.com>
Date: Sun, 12 Sep 2021 11:41:54 -0700
Subject: [PATCH] REF: share _wrap_transform_fast_result

---
 pandas/core/groupby/generic.py | 23 -----------------------
 pandas/core/groupby/groupby.py | 20 ++++++++++++++++++++
 2 files changed, 20 insertions(+), 23 deletions(-)

diff --git a/pandas/core/groupby/generic.py b/pandas/core/groupby/generic.py
index 23d1edeae200f..ed79ba6206dba 100644
--- a/pandas/core/groupby/generic.py
+++ b/pandas/core/groupby/generic.py
@@ -503,16 +503,6 @@ def _transform_general(self, func: Callable, *args, **kwargs) -> Series:
     def _can_use_transform_fast(self, result) -> bool:
         return True
 
-    def _wrap_transform_fast_result(self, result: Series) -> Series:
-        """
-        fast version of transform, only applicable to
-        builtin/cythonizable functions
-        """
-        ids, _, _ = self.grouper.group_info
-        result = result.reindex(self.grouper.result_index, copy=False)
-        out = algorithms.take_nd(result._values, ids)
-        return self.obj._constructor(out, index=self.obj.index, name=self.obj.name)
-
     def filter(self, func, dropna: bool = True, *args, **kwargs):
         """
         Return a copy of a Series excluding elements from groups that
@@ -1274,19 +1264,6 @@ def _can_use_transform_fast(self, result) -> bool:
             self._obj_with_exclusions.columns
         )
 
-    def _wrap_transform_fast_result(self, result: DataFrame) -> DataFrame:
-        """
-        Fast transform path for aggregations
-        """
-        obj = self._obj_with_exclusions
-
-        # for each col, reshape to size of original frame by take operation
-        ids, _, _ = self.grouper.group_info
-        result = result.reindex(self.grouper.result_index, copy=False)
-        output = result.take(ids, axis=0)
-        output.index = obj.index
-        return output
-
     def _define_paths(self, func, *args, **kwargs):
         if isinstance(func, str):
             fast_path = lambda group: getattr(group, func)(*args, **kwargs)
diff --git a/pandas/core/groupby/groupby.py b/pandas/core/groupby/groupby.py
index 11cd315553f41..c394b0a6bde66 100644
--- a/pandas/core/groupby/groupby.py
+++ b/pandas/core/groupby/groupby.py
@@ -1559,6 +1559,26 @@ def _transform(self, func, *args, engine=None, engine_kwargs=None, **kwargs):
             # only reached for DataFrameGroupBy
             return self._transform_general(func, *args, **kwargs)
 
+    @final
+    def _wrap_transform_fast_result(self, result: FrameOrSeries) -> FrameOrSeries:
+        """
+        Fast transform path for aggregations.
+        """
+        obj = self._obj_with_exclusions
+
+        # for each col, reshape to size of original frame by take operation
+        ids, _, _ = self.grouper.group_info
+        result = result.reindex(self.grouper.result_index, copy=False)
+
+        if self.obj.ndim == 1:
+            # i.e. SeriesGroupBy
+            out = algorithms.take_nd(result._values, ids)
+            output = obj._constructor(out, index=obj.index, name=obj.name)
+        else:
+            output = result.take(ids, axis=0)
+            output.index = obj.index
+        return output
+
     # -----------------------------------------------------------------
     # Utilities