smaller-diff implementation

Author: jbrockmendel

pandas/core/groupby/generic.py

@@ -219,9 +219,16 @@ def apply(self, func, *args, **kwargs) -> Series:
     def aggregate(self, func=None, *args, engine=None, engine_kwargs=None, **kwargs):
 
         if maybe_use_numba(engine):
-            return self._aggregate_with_numba(
-                func, *args, engine_kwargs=engine_kwargs, **kwargs
+            data = self._obj_with_exclusions
+            result = self._aggregate_with_numba(
+                data.to_frame(), func, *args, engine_kwargs=engine_kwargs, **kwargs
             )
+            index = self.grouper.result_index
+            result = self.obj._constructor(result.ravel(), index=index, name=data.name)
+            if not self.as_index:
+                result = self._insert_inaxis_grouper(result)
+                result.index = default_index(len(result))
+            return result
 
         relabeling = func is None
         columns = None
@@ -1257,9 +1264,16 @@ class DataFrameGroupBy(GroupBy[DataFrame]):
     def aggregate(self, func=None, *args, engine=None, engine_kwargs=None, **kwargs):
 
         if maybe_use_numba(engine):
-            return self._aggregate_with_numba(
-                func, *args, engine_kwargs=engine_kwargs, **kwargs
+            data = self._obj_with_exclusions
+            result = self._aggregate_with_numba(
+                data, func, *args, engine_kwargs=engine_kwargs, **kwargs
             )
+            index = self.grouper.result_index
+            result = self.obj._constructor(result, index=index, columns=data.columns)
+            if not self.as_index:
+                result = self._insert_inaxis_grouper(result)
+                result.index = default_index(len(result))
+            return result
 
         relabeling, func, columns, order = reconstruct_func(func, **kwargs)
         func = maybe_mangle_lambdas(func)

pandas/core/groupby/groupby.py

@@ -1270,18 +1270,17 @@ def _transform_with_numba(
         return result.take(np.argsort(sorted_index), axis=0)
 
     @final
-    def _aggregate_with_numba(self, func, *args, engine_kwargs=None, **kwargs):
+    def _aggregate_with_numba(
+        self, data: DataFrame, func, *args, engine_kwargs=None, **kwargs
+    ):
         """
         Perform groupby aggregation routine with the numba engine.
 
         This routine mimics the data splitting routine of the DataSplitter class
         to generate the indices of each group in the sorted data and then passes the
         data and indices into a Numba jitted function.
         """
-        data = self._obj_with_exclusions
-        df = data if data.ndim == 2 else data.to_frame()
-
-        starts, ends, sorted_index, sorted_data = self._numba_prep(df)
+        starts, ends, sorted_index, sorted_data = self._numba_prep(data)
         numba_.validate_udf(func)
         numba_agg_func = numba_.generate_numba_agg_func(
             func, **get_jit_arguments(engine_kwargs, kwargs)
@@ -1291,18 +1290,10 @@ def _aggregate_with_numba(self, func, *args, engine_kwargs=None, **kwargs):
             sorted_index,
             starts,
             ends,
-            len(df.columns),
+            len(data.columns),
             *args,
         )
-
-        index = self.grouper.result_index
-        if data.ndim == 1:
-            result_kwargs = {"name": data.name}
-            result = result.ravel()
-        else:
-            result_kwargs = {"columns": data.columns}
-        result = data._constructor(result, index=index, **result_kwargs)
-        return self._wrap_aggregated_output(result)
+        return result
 
     # -----------------------------------------------------------------
     # apply/agg/transform