ENH: A new GroupBy method to slice rows preserving index and order

pandas/core/groupby/generic.py

@@ -74,7 +74,6 @@
     _agg_template,
     _apply_docs,
     _transform_template,
-    group_selection_context,
 )
 from pandas.core.indexes.api import (
     Index,
@@ -236,7 +235,7 @@ def apply(self, func, *args, **kwargs):
     def aggregate(self, func=None, *args, engine=None, engine_kwargs=None, **kwargs):
 
         if maybe_use_numba(engine):
-            with group_selection_context(self):
+            with self._group_selection_context():
                 data = self._selected_obj
             result = self._aggregate_with_numba(
                 data.to_frame(), func, *args, engine_kwargs=engine_kwargs, **kwargs
@@ -878,7 +877,7 @@ class DataFrameGroupBy(GroupBy[DataFrame]):
     def aggregate(self, func=None, *args, engine=None, engine_kwargs=None, **kwargs):
 
         if maybe_use_numba(engine):
-            with group_selection_context(self):
+            with self._group_selection_context():
                 data = self._selected_obj
             result = self._aggregate_with_numba(
                 data, func, *args, engine_kwargs=engine_kwargs, **kwargs

pandas/core/groupby/groupby.py

@@ -47,6 +47,7 @@ class providing the base-class of operations.
     F,
     FrameOrSeries,
     IndexLabel,
+    PositionalIndexer,
     RandomState,
     Scalar,
     T,
@@ -66,6 +67,7 @@ class providing the base-class of operations.
     is_bool_dtype,
     is_datetime64_dtype,
     is_float_dtype,
+    is_integer,
     is_integer_dtype,
     is_numeric_dtype,
     is_object_dtype,
@@ -98,6 +100,7 @@ class providing the base-class of operations.
     numba_,
     ops,
 )
+from pandas.core.groupby.indexing import GroupByIndexingMixin
 from pandas.core.indexes.api import (
     CategoricalIndex,
     Index,
@@ -547,18 +550,6 @@ def f(self):
         return attr
 
 
-@contextmanager
-def group_selection_context(groupby: GroupBy) -> Iterator[GroupBy]:
-    """
-    Set / reset the group_selection_context.
-    """
-    groupby._set_group_selection()
-    try:
-        yield groupby
-    finally:
-        groupby._reset_group_selection()
-
-
 _KeysArgType = Union[
     Hashable,
     List[Hashable],
@@ -568,7 +559,7 @@ def group_selection_context(groupby: GroupBy) -> Iterator[GroupBy]:
 ]
 
 
-class BaseGroupBy(PandasObject, SelectionMixin[FrameOrSeries]):
+class BaseGroupBy(PandasObject, SelectionMixin[FrameOrSeries], GroupByIndexingMixin):
     _group_selection: IndexLabel | None = None
     _apply_allowlist: frozenset[str] = frozenset()
     _hidden_attrs = PandasObject._hidden_attrs | {
@@ -916,7 +907,7 @@ def __getattr__(self, attr: str):
     def _make_wrapper(self, name: str) -> Callable:
         assert name in self._apply_allowlist
 
-        with group_selection_context(self):
+        with self._group_selection_context():
             # need to setup the selection
             # as are not passed directly but in the grouper
             f = getattr(self._obj_with_exclusions, name)
@@ -993,6 +984,17 @@ def _reset_group_selection(self) -> None:
             self._group_selection = None
             self._reset_cache("_selected_obj")
 
+    @contextmanager
+    def _group_selection_context(self) -> Iterator[GroupBy]:
+        """
+        Set / reset the _group_selection_context.
+        """
+        self._set_group_selection()
+        try:
+            yield self
+        finally:
+            self._reset_group_selection()
+
     def _iterate_slices(self) -> Iterable[Series]:
         raise AbstractMethodError(self)
 
@@ -1365,7 +1367,7 @@ def f(g):
                 # fails on *some* columns, e.g. a numeric operation
                 # on a string grouper column
 
-                with group_selection_context(self):
+                with self._group_selection_context():
                     return self._python_apply_general(f, self._selected_obj)
 
         return result
@@ -1449,7 +1451,7 @@ def _agg_general(
         npfunc: Callable,
     ):
 
-        with group_selection_context(self):
+        with self._group_selection_context():
             # try a cython aggregation if we can
             result = self._cython_agg_general(
                 how=alias,
@@ -1514,7 +1516,7 @@ 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 group_selection_context(self):
+            with self._group_selection_context():
                 data = self._selected_obj
             df = data if data.ndim == 2 else data.to_frame()
             result = self._transform_with_numba(
@@ -1904,7 +1906,7 @@ def var(self, ddof: int = 1):
             )
         else:
             func = lambda x: x.var(ddof=ddof)
-            with group_selection_context(self):
+            with self._group_selection_context():
                 return self._python_agg_general(func)
 
     @final
@@ -2096,7 +2098,7 @@ def ohlc(self) -> DataFrame:
 
     @doc(DataFrame.describe)
     def describe(self, **kwargs):
-        with group_selection_context(self):
+        with self._group_selection_context():
             result = self.apply(lambda x: x.describe(**kwargs))
             if self.axis == 1:
                 return result.T
@@ -2393,20 +2395,25 @@ def backfill(self, limit=None):
     @Substitution(name="groupby")
     @Substitution(see_also=_common_see_also)
     def nth(
-        self, n: int | list[int], dropna: Literal["any", "all", None] = None
-    ) -> DataFrame:
+        self,
+        arg: PositionalIndexer | tuple,
+        dropna: Literal["any", "all", None] = None,
+    ) -> FrameOrSeries:
         """
-        Take the nth row from each group if n is an int, or a subset of rows
-        if n is a list of ints.
+        Take the nth row from each group if n is an int, otherwise a subset of rows.
 
         If dropna, will take the nth non-null row, dropna is either
         'all' or 'any'; this is equivalent to calling dropna(how=dropna)
         before the groupby.
 
         Parameters
         ----------
-        n : int or list of ints
-            A single nth value for the row or a list of nth values.
+        arg : int, slice or list of ints and slices
+            A single nth value for the row or a list of nth values or slices.
+
+            .. versionchanged:: 1.4.0
+                Added slice and lists containiing slices
+
         dropna : {'any', 'all', None}, default None
             Apply the specified dropna operation before counting which row is
             the nth row.
@@ -2444,6 +2451,12 @@ def nth(
         1  2.0
         2  3.0
         2  5.0
+        >>> g.nth(slice(None, -1))
+             B
+        A
+        1  NaN
+        1  2.0
+        2  3.0
 
         Specifying `dropna` allows count ignoring ``NaN``
 
@@ -2468,58 +2481,27 @@ def nth(
         1  1  2.0
         4  2  5.0
         """
-        valid_containers = (set, list, tuple)
-        if not isinstance(n, (valid_containers, int)):
-            raise TypeError("n needs to be an int or a list/set/tuple of ints")
-
         if not dropna:
+            if isinstance(arg, Iterable):
+                return self._rows[tuple(arg)]
 
-            if isinstance(n, int):
-                nth_values = [n]
-            elif isinstance(n, valid_containers):
-                nth_values = list(set(n))
-
-            nth_array = np.array(nth_values, dtype=np.intp)
-            with group_selection_context(self):
-
-                mask_left = np.in1d(self._cumcount_array(), nth_array)
-                mask_right = np.in1d(
-                    self._cumcount_array(ascending=False) + 1, -nth_array
-                )
-                mask = mask_left | mask_right
-
-                ids, _, _ = self.grouper.group_info
-
-                # Drop NA values in grouping
-                mask = mask & (ids != -1)
-
-                out = self._selected_obj[mask]
-                if not self.as_index:
-                    return out
-
-                result_index = self.grouper.result_index
-                out.index = result_index[ids[mask]]
-
-                if not self.observed and isinstance(result_index, CategoricalIndex):
-                    out = out.reindex(result_index)
-
-                out = self._reindex_output(out)
-                return out.sort_index() if self.sort else out
+            return self._rows[arg]
 
         # dropna is truthy
-        if isinstance(n, valid_containers):
-            raise ValueError("dropna option with a list of nth values is not supported")
+        if not is_integer(arg):
+            raise ValueError("dropna option only supported for an integer argument")
 
         if dropna not in ["any", "all"]:
             # Note: when agg-ing picker doesn't raise this, just returns NaN
             raise ValueError(
-                "For a DataFrame groupby, dropna must be "
+                "For a DataFrame or Series groupby.nth, dropna must be "
                 "either None, 'any' or 'all', "
                 f"(was passed {dropna})."
             )
 
         # old behaviour, but with all and any support for DataFrames.
         # modified in GH 7559 to have better perf
+        n = cast(int, arg)
         max_len = n if n >= 0 else -1 - n
         dropped = self.obj.dropna(how=dropna, axis=self.axis)
 
@@ -2784,7 +2766,7 @@ def ngroup(self, ascending: bool = True):
         5    0
         dtype: int64
         """
-        with group_selection_context(self):
+        with self._group_selection_context():
             index = self._selected_obj.index
             result = self._obj_1d_constructor(
                 self.grouper.group_info[0], index, dtype=np.int64
@@ -2848,7 +2830,7 @@ def cumcount(self, ascending: bool = True):
         5    0
         dtype: int64
         """
-        with group_selection_context(self):
+        with self._group_selection_context():
             index = self._selected_obj._get_axis(self.axis)
             cumcounts = self._cumcount_array(ascending=ascending)
             return self._obj_1d_constructor(cumcounts, index)