REF: avoid group_selection_context (#51096)

Author: jbrockmendel

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
             )

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