REF: avoid unnecessary argument to groupby numba funcs (#43487)

Author: jbrockmendel

pandas/core/groupby/generic.py

@@ -228,9 +228,10 @@ def aggregate(self, func=None, *args, engine=None, engine_kwargs=None, **kwargs)
         if maybe_use_numba(engine):
             with group_selection_context(self):
                 data = self._selected_obj
-            result, index = self._aggregate_with_numba(
+            result = self._aggregate_with_numba(
                 data.to_frame(), func, *args, engine_kwargs=engine_kwargs, **kwargs
             )
+            index = self._group_keys_index
             return self.obj._constructor(result.ravel(), index=index, name=data.name)
 
         relabeling = func is None
@@ -924,9 +925,10 @@ def aggregate(self, func=None, *args, engine=None, engine_kwargs=None, **kwargs)
         if maybe_use_numba(engine):
             with group_selection_context(self):
                 data = self._selected_obj
-            result, index = self._aggregate_with_numba(
+            result = self._aggregate_with_numba(
                 data, func, *args, engine_kwargs=engine_kwargs, **kwargs
             )
+            index = self._group_keys_index
             return self.obj._constructor(result, index=index, columns=data.columns)
 
         relabeling, func, columns, order = reconstruct_func(func, **kwargs)

pandas/core/groupby/groupby.py

@@ -1269,7 +1269,6 @@ def _transform_with_numba(self, data, func, *args, engine_kwargs=None, **kwargs)
         data and indices into a Numba jitted function.
         """
         starts, ends, sorted_index, sorted_data = self._numba_prep(func, data)
-        group_keys = self.grouper.group_keys_seq
 
         numba_transform_func = numba_.generate_numba_transform_func(
             kwargs, func, engine_kwargs
@@ -1279,7 +1278,6 @@ def _transform_with_numba(self, data, func, *args, engine_kwargs=None, **kwargs)
             sorted_index,
             starts,
             ends,
-            len(group_keys),
             len(data.columns),
             *args,
         )
@@ -1302,15 +1300,13 @@ def _aggregate_with_numba(self, data, func, *args, engine_kwargs=None, **kwargs)
         data and indices into a Numba jitted function.
         """
         starts, ends, sorted_index, sorted_data = self._numba_prep(func, data)
-        index = self._group_keys_index
 
         numba_agg_func = numba_.generate_numba_agg_func(kwargs, func, engine_kwargs)
         result = numba_agg_func(
             sorted_data,
             sorted_index,
             starts,
             ends,
-            len(index),
             len(data.columns),
             *args,
         )
@@ -1319,7 +1315,7 @@ def _aggregate_with_numba(self, data, func, *args, engine_kwargs=None, **kwargs)
         if cache_key not in NUMBA_FUNC_CACHE:
             NUMBA_FUNC_CACHE[cache_key] = numba_agg_func
 
-        return result, index
+        return result
 
     # -----------------------------------------------------------------
     # apply/agg/transform

pandas/core/groupby/numba_.py

@@ -59,9 +59,7 @@ def generate_numba_agg_func(
     kwargs: dict[str, Any],
     func: Callable[..., Scalar],
     engine_kwargs: dict[str, bool] | None,
-) -> Callable[
-    [np.ndarray, np.ndarray, np.ndarray, np.ndarray, int, int, Any], np.ndarray
-]:
+) -> Callable[[np.ndarray, np.ndarray, np.ndarray, np.ndarray, int, Any], np.ndarray]:
     """
     Generate a numba jitted agg function specified by values from engine_kwargs.
 
@@ -100,10 +98,13 @@ def group_agg(
         index: np.ndarray,
         begin: np.ndarray,
         end: np.ndarray,
-        num_groups: int,
         num_columns: int,
         *args: Any,
     ) -> np.ndarray:
+
+        assert len(begin) == len(end)
+        num_groups = len(begin)
+
         result = np.empty((num_groups, num_columns))
         for i in numba.prange(num_groups):
             group_index = index[begin[i] : end[i]]
@@ -119,9 +120,7 @@ def generate_numba_transform_func(
     kwargs: dict[str, Any],
     func: Callable[..., np.ndarray],
     engine_kwargs: dict[str, bool] | None,
-) -> Callable[
-    [np.ndarray, np.ndarray, np.ndarray, np.ndarray, int, int, Any], np.ndarray
-]:
+) -> Callable[[np.ndarray, np.ndarray, np.ndarray, np.ndarray, int, Any], np.ndarray]:
     """
     Generate a numba jitted transform function specified by values from engine_kwargs.
 
@@ -160,10 +159,13 @@ def group_transform(
         index: np.ndarray,
         begin: np.ndarray,
         end: np.ndarray,
-        num_groups: int,
         num_columns: int,
         *args: Any,
     ) -> np.ndarray:
+
+        assert len(begin) == len(end)
+        num_groups = len(begin)
+
         result = np.empty((len(values), num_columns))
         for i in numba.prange(num_groups):
             group_index = index[begin[i] : end[i]]