diff --git a/pandas/core/groupby/generic.py b/pandas/core/groupby/generic.py index 68c7ba72729e6..72411f6ac4e96 100644 --- a/pandas/core/groupby/generic.py +++ b/pandas/core/groupby/generic.py @@ -242,8 +242,7 @@ def apply(self, func, *args, **kwargs) -> Series: def aggregate(self, func=None, *args, engine=None, engine_kwargs=None, **kwargs): if maybe_use_numba(engine): - with self._group_selection_context(): - data = self._selected_obj + data = self._obj_with_exclusions result = self._aggregate_with_numba( data.to_frame(), func, *args, engine_kwargs=engine_kwargs, **kwargs ) @@ -1234,8 +1233,7 @@ class DataFrameGroupBy(GroupBy[DataFrame]): def aggregate(self, func=None, *args, engine=None, engine_kwargs=None, **kwargs): if maybe_use_numba(engine): - with self._group_selection_context(): - data = self._selected_obj + data = self._obj_with_exclusions result = self._aggregate_with_numba( data, func, *args, engine_kwargs=engine_kwargs, **kwargs ) diff --git a/pandas/core/groupby/groupby.py b/pandas/core/groupby/groupby.py index eba97f41dc66a..2be73d21e3d88 100644 --- a/pandas/core/groupby/groupby.py +++ b/pandas/core/groupby/groupby.py @@ -1277,7 +1277,7 @@ def _wrap_applied_output( # numba @final - def _numba_prep(self, data): + def _numba_prep(self, data: DataFrame): ids, _, ngroups = self.grouper.group_info sorted_index = get_group_index_sorter(ids, ngroups) sorted_ids = algorithms.take_nd(ids, sorted_index, allow_fill=False) @@ -1337,7 +1337,9 @@ def _numba_agg_general( return data._constructor(result, index=index, **result_kwargs) @final - def _transform_with_numba(self, data, func, *args, engine_kwargs=None, **kwargs): + def _transform_with_numba( + self, data: DataFrame, func, *args, engine_kwargs=None, **kwargs + ): """ Perform groupby transform routine with the numba engine. @@ -1363,7 +1365,9 @@ def _transform_with_numba(self, data, func, *args, engine_kwargs=None, **kwargs) return result.take(np.argsort(sorted_index), axis=0) @final - def _aggregate_with_numba(self, data, 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. @@ -1529,15 +1533,13 @@ def _agg_general( npfunc: Callable, ): - with self._group_selection_context(): - # try a cython aggregation if we can - result = self._cython_agg_general( - how=alias, - alt=npfunc, - numeric_only=numeric_only, - min_count=min_count, - ) - return result.__finalize__(self.obj, method="groupby") + result = self._cython_agg_general( + how=alias, + alt=npfunc, + numeric_only=numeric_only, + min_count=min_count, + ) + return result.__finalize__(self.obj, method="groupby") def _agg_py_fallback( self, values: ArrayLike, ndim: int, alt: Callable @@ -1647,9 +1649,7 @@ def _cython_transform( def _transform(self, func, *args, engine=None, engine_kwargs=None, **kwargs): if maybe_use_numba(engine): - # TODO: tests with self._selected_obj.ndim == 1 on DataFrameGroupBy - with self._group_selection_context(): - data = self._selected_obj + data = self._obj_with_exclusions df = data if data.ndim == 2 else data.to_frame() result = self._transform_with_numba( df, func, *args, engine_kwargs=engine_kwargs, **kwargs