PERF: Try fast/slow paths only once in DataFrameGroupby.transform

doc/source/whatsnew/v1.4.0.rst

@@ -159,7 +159,7 @@ Deprecations
 Performance improvements
 ~~~~~~~~~~~~~~~~~~~~~~~~
 - Performance improvement in :meth:`.GroupBy.sample`, especially when ``weights`` argument provided (:issue:`34483`)
--
+- Performance improvement in :meth:`.GroupBy.transform` for user-defined functions (:issue:`41598`)
 
 .. ---------------------------------------------------------------------------
 

pandas/core/groupby/generic.py

@@ -1308,22 +1308,10 @@ def _transform_general(self, func, *args, **kwargs):
         gen = self.grouper.get_iterator(obj, axis=self.axis)
         fast_path, slow_path = self._define_paths(func, *args, **kwargs)
 
-        for name, group in gen:
-            if group.size == 0:
-                continue
-            object.__setattr__(group, "name", name)
-
-            # Try slow path and fast path.
-            try:
-                path, res = self._choose_path(fast_path, slow_path, group)
-            except TypeError:
-                return self._transform_item_by_item(obj, fast_path)
-            except ValueError as err:
-                msg = "transform must return a scalar value for each group"
-                raise ValueError(msg) from err
-
+        def process_result(
+            group: DataFrame | Series, res: DataFrame | Series
+        ) -> DataFrame | Series:
             if isinstance(res, Series):
-
                 # we need to broadcast across the
                 # other dimension; this will preserve dtypes
                 # GH14457
@@ -1339,10 +1327,32 @@ def _transform_general(self, func, *args, **kwargs):
                         columns=group.columns,
                         index=group.index,
                     )
-
-                applied.append(r)
+                return r
             else:
-                applied.append(res)
+                return res
+
+        try:
+            name, group = next(gen)
+        except StopIteration:
+            pass
+        else:
+            object.__setattr__(group, "name", name)
+            try:
+                path, res = self._choose_path(fast_path, slow_path, group)
+            except TypeError:
+                return self._transform_item_by_item(obj, fast_path)
+            except ValueError as err:
+                msg = "transform must return a scalar value for each group"
+                raise ValueError(msg) from err
+            if group.size > 0:
+                applied.append(process_result(group, res))
+
+        for name, group in gen:
+            if group.size == 0:
+                continue
+            object.__setattr__(group, "name", name)
+            res = path(group)
+            applied.append(process_result(group, res))
 
         concat_index = obj.columns if self.axis == 0 else obj.index
         other_axis = 1 if self.axis == 0 else 0  # switches between 0 & 1