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