Repository: pandas-dev/pandas
Pull request: ENH: A new GroupBy method to slice rows preserving index and order
Closes #42864

New indexable attribute GroupBy.iloc to perform most grouped slices analagous to DataFrame.iloc

Row indexes can be integers or slices with +ve step
Row slices preserve the DataFrame order and grouping and extend the existing GroupBy head() and tail() methods
All column indexes are handled as for DataFrame.iloc

-ve slice step is not currently handled, and would have to reverse the order.
There is no plan to implement Integer lists since these do not preserve the order at all and cannot be easily vectorized.

Note
----
Neither GroupBy.nth() nor GroupBy.take() take a slice argument and neither of them preserve the original DataFrame order and index. They are both slow for large integer lists and take() is very slow for large group counts.

Example use case
--------------------
Supose that we have a multi-indexed DataFrame with a large primary index and a secondary sorted
in a different order for each primary index. 
To reduce the DataFrame to a middle slice of each secondary, we can groupby the primary index and then use iloc. 
This preserves the original DataFrame's order and indexing.
(See tests/groupby/test_groupby_iloc)

Problem statement:
ENH: A new GroupBy method to slice rows preserving index and order
#### Is your feature request related to a problem?

Pandas provides DataFrameGroupBy.head() and tail(), which efficiently slice the beginning and end of each group while preserving the order and index. I would like to be able to do a general row slice with the same properties. DataFrame has head(), tail() and iloc that behave in a compatible way. There is no corresponding DataFrameGroupBy.iloc.

#### Describe the solution you'd like
Provide a new DataFrameGroupBy method to slice rows per group

#### API breaking implications
None

#### Describe alternatives you've considered
The following are existing ways to extract, say, the second and third entry of each group, assuming that there are a large number of rows in each group (~10000):

1. grouped.apply(lambda x: x.iloc[1:3, :]) - Extremely slow. Does not preserve the order or indexing.
2. grouped.take([1, 2]) - Extremely slow. Does not preserve the order or indexing.
3. grouped.nth([1, 2]) - Quite fast for a small list. Does not preserve the order or indexing.
4. grouped.head(3).groupby('...').tail(2) - Quite fast. Does preserve index and ordering.
5. grouped._selected_obj[mask] where mask is built from grouped.cumcount() - Very fast. Does preserve index and ordering. But uses private attribute of DataFrameGroupBy and takes several lines of code.


#### Additional context
There are three options:

1. Add an option to an existing method to force it to preserve index and order. But take() is very slow and nth() is quite slow. Neither accept a slice argument, so a range list has to be provided.
2. Easiest: Add a new method taking a slice as an argument and implementing it as in 5 above.
3. Most logical and complete: Add a new iloc attribute analogous to DataFrame.iloc

Hints from issue discussion:
.nth already does this
As I pointed out, grouped.nth() does not preserve the order or index of the original df, particularly if it has a multiindex. So it is not compaible with head() or tail() and it takes a list rather than a slice argument.

Also, the speed of grouped.nth(list_of_ints) grows with the length of the list. It is over 10 times slower than alternative 5 when working with a large slice.
I appreciate that this has not yet been triaged, but I can propose a solution for GroupBy.iloc that addresses this issue. So I would like to take this.
ok would be ok with .iloc as long as it's clear how this is different than nth head and tail - eg the usecases r clear in the docs and api
take
I have implemented https://github.com/pandas-dev/pandas/pull/42947 and submitted a pull request. I'm not sure what happens next (this is my first contribution...)
Update to my  https://github.com/pandas-dev/pandas/pull/42947. I decided that the syntax and behaviour of my index was too different from DataFrame.iloc to use the same name. I implemented it as GroupBy.rows. I do understand that we are trying to reduce attributes rather than add to them, but I believe that my code adds useful functionality that is not otherwise available. It also resolves multiple requests for GroupBy.head and tail to handle negative arguments.

Commit message: Delete Example section

Patch to review:
diff --git a/pandas/core/groupby/generic.py b/pandas/core/groupby/generic.py
index 46e42326d4191..c6eacfeea2a32 100644
--- a/pandas/core/groupby/generic.py
+++ b/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
diff --git a/pandas/core/groupby/groupby.py b/pandas/core/groupby/groupby.py
index 068c4191e44e2..38eb16586913b 100644
--- a/pandas/core/groupby/groupby.py
+++ b/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,11 +2395,12 @@ 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)
@@ -2405,8 +2408,12 @@ def nth(
 
         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)
diff --git a/pandas/core/groupby/indexing.py b/pandas/core/groupby/indexing.py
new file mode 100644
index 0000000000000..0f9ccf344a079
--- /dev/null
+++ b/pandas/core/groupby/indexing.py
@@ -0,0 +1,211 @@
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Iterable,
+    cast,
+)
+
+import numpy as np
+
+from pandas._typing import (
+    FrameOrSeries,
+    PositionalIndexer,
+)
+from pandas.util._decorators import (
+    cache_readonly,
+    doc,
+)
+
+from pandas.core.dtypes.common import (
+    is_integer,
+    is_list_like,
+)
+
+if TYPE_CHECKING:
+    from pandas.core.groupby import groupby
+
+from pandas.core.indexes.api import CategoricalIndex
+
+
+class GroupByIndexingMixin:
+    """
+    Mixin for adding ._rows to GroupBy.
+    """
+
+    @property
+    def _rows(self) -> RowsGroupByIndexer:
+        if TYPE_CHECKING:
+            groupby_object = cast(groupby.GroupBy, self)
+        else:
+            groupby_object = self
+
+        return RowsGroupByIndexer(groupby_object)
+
+
+@doc(GroupByIndexingMixin._rows)
+class RowsGroupByIndexer:
+    def __init__(self, groupby_object: groupby.GroupBy):
+        self.groupby_object = groupby_object
+
+    def __getitem__(self, arg: PositionalIndexer | tuple) -> FrameOrSeries:
+        """
+        Positional index for selection by integer location per group.
+
+        Used to implement GroupBy._rows which is used to implement GroupBy.nth
+        in the case when the keyword dropna is None or absent.
+
+        Parameters
+        ----------
+        arg : PositionalIndexer | tuple
+            Allowed values are:
+            - int
+            - int valued iterable such as list or range
+            - slice with step either None or positive
+            - tuple of integers and slices
+
+        Returns
+        -------
+        Series
+            The filtered subset of the original groupby Series.
+        DataFrame
+            The filtered subset of the original groupby DataFrame.
+
+        See Also
+        --------
+        DataFrame.iloc : Purely integer-location based indexing for selection by
+            position.
+        GroupBy.head : Return first n rows of each group.
+        GroupBy.tail : Return last n rows of each group.
+        GroupBy.nth : Take the nth row from each group if n is an int, or a
+            subset of rows, if n is a list of ints.
+        """
+        with self.groupby_object._group_selection_context():
+            if isinstance(arg, tuple):
+                if all(is_integer(i) for i in arg):
+                    mask = self._handle_list(arg)
+                else:
+                    mask = self._handle_tuple(arg)
+
+            elif isinstance(arg, slice):
+                mask = self._handle_slice(arg)
+            elif is_integer(arg):
+                mask = self._handle_int(cast(int, arg))
+            elif is_list_like(arg):
+                mask = self._handle_list(cast(Iterable[int], arg))
+            else:
+                raise TypeError(
+                    f"Invalid index {type(arg)}. "
+                    "Must be integer, list-like, slice or a tuple of "
+                    "integers and slices"
+                )
+
+            ids, _, _ = self.groupby_object.grouper.group_info
+
+            # Drop NA values in grouping
+            mask &= ids != -1
+
+            if mask is None or mask is True:
+                result = self.groupby_object._selected_obj[:]
+            else:
+                result = self.groupby_object._selected_obj[mask]
+
+            if self.groupby_object.as_index:
+                result_index = self.groupby_object.grouper.result_index
+                result.index = result_index[ids[mask]]
+
+                if not self.groupby_object.observed and isinstance(
+                    result_index, CategoricalIndex
+                ):
+                    result = result.reindex(result_index)
+
+                result = self.groupby_object._reindex_output(result)
+                if self.groupby_object.sort:
+                    result = result.sort_index()
+
+            return result
+
+    def _handle_int(self, arg: int) -> bool | np.ndarray:
+        if arg >= 0:
+            return self._ascending_count == arg
+        else:
+            return self._descending_count == (-arg - 1)
+
+    def _handle_list(self, args: Iterable[int]) -> bool | np.ndarray:
+        positive = [arg for arg in args if arg >= 0]
+        negative = [-arg - 1 for arg in args if arg < 0]
+
+        mask: bool | np.ndarray = False
+
+        if positive:
+            mask |= np.isin(self._ascending_count, positive)
+
+        if negative:
+            mask |= np.isin(self._descending_count, negative)
+
+        return mask
+
+    def _handle_tuple(self, args: tuple) -> bool | np.ndarray:
+        mask: bool | np.ndarray = False
+
+        for arg in args:
+            if is_integer(arg):
+                mask |= self._handle_int(cast(int, arg))
+            elif isinstance(arg, slice):
+                mask |= self._handle_slice(arg)
+            else:
+                raise ValueError(
+                    f"Invalid argument {type(arg)}. Should be int or slice."
+                )
+
+        return mask
+
+    def _handle_slice(self, arg: slice) -> bool | np.ndarray:
+        start = arg.start
+        stop = arg.stop
+        step = arg.step
+
+        if step is not None and step < 0:
+            raise ValueError(f"Invalid step {step}. Must be non-negative")
+
+        mask: bool | np.ndarray = True
+
+        if step is None:
+            step = 1
+
+        if start is None:
+            if step > 1:
+                mask &= self._ascending_count % step == 0
+
+        elif start >= 0:
+            mask &= self._ascending_count >= start
+
+            if step > 1:
+                mask &= (self._ascending_count - start) % step == 0
+
+        else:
+            mask &= self._descending_count < -start
+
+            offset_array = self._descending_count + start + 1
+            limit_array = (
+                self._ascending_count + self._descending_count + (start + 1)
+            ) < 0
+            offset_array = np.where(limit_array, self._ascending_count, offset_array)
+
+            mask &= offset_array % step == 0
+
+        if stop is not None:
+            if stop >= 0:
+                mask &= self._ascending_count < stop
+            else:
+                mask &= self._descending_count >= -stop
+
+        return mask
+
+    @cache_readonly
+    def _ascending_count(self) -> np.ndarray:
+        return self.groupby_object._cumcount_array()
+
+    @cache_readonly
+    def _descending_count(self) -> np.ndarray:
+        return self.groupby_object._cumcount_array(ascending=False)
diff --git a/pandas/tests/groupby/conftest.py b/pandas/tests/groupby/conftest.py
index d699d05963b46..622c56d707ead 100644
--- a/pandas/tests/groupby/conftest.py
+++ b/pandas/tests/groupby/conftest.py
@@ -116,6 +116,27 @@ def three_group():
     )
 
 
+@pytest.fixture()
+def slice_test_df():
+    data = [
+        [0, "a", "a0_at_0"],
+        [1, "b", "b0_at_1"],
+        [2, "a", "a1_at_2"],
+        [3, "b", "b1_at_3"],
+        [4, "c", "c0_at_4"],
+        [5, "a", "a2_at_5"],
+        [6, "a", "a3_at_6"],
+        [7, "a", "a4_at_7"],
+    ]
+    df = DataFrame(data, columns=["Index", "Group", "Value"])
+    return df.set_index("Index")
+
+
+@pytest.fixture()
+def slice_test_grouped(slice_test_df):
+    return slice_test_df.groupby("Group", as_index=False)
+
+
 @pytest.fixture(params=sorted(reduction_kernels))
 def reduction_func(request):
     """
diff --git a/pandas/tests/groupby/test_nth.py b/pandas/tests/groupby/test_nth.py
index f0eef550b39ac..ad9b16583def7 100644
--- a/pandas/tests/groupby/test_nth.py
+++ b/pandas/tests/groupby/test_nth.py
@@ -270,7 +270,7 @@ def test_nth():
     result = s.groupby(g, sort=False).nth(0, dropna="all")
     tm.assert_series_equal(result, expected)
 
-    with pytest.raises(ValueError, match="For a DataFrame groupby"):
+    with pytest.raises(ValueError, match="For a DataFrame"):
         s.groupby(g, sort=False).nth(0, dropna=True)
 
     # doc example
@@ -706,3 +706,44 @@ def test_groupby_last_first_nth_with_none(method, nulls_fixture):
         result = getattr(data, method)()
 
     tm.assert_series_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+    "arg, expected_rows",
+    [
+        [slice(None, 3, 2), [0, 1, 4, 5]],
+        [slice(None, -2), [0, 2, 5]],
+        [[slice(None, 2), slice(-2, None)], [0, 1, 2, 3, 4, 6, 7]],
+        [[0, 1, slice(-2, None)], [0, 1, 2, 3, 4, 6, 7]],
+    ],
+)
+def test_slice(slice_test_df, slice_test_grouped, arg, expected_rows):
+    # Test slices     GH #42947
+
+    result = slice_test_grouped.nth(arg)
+    expected = slice_test_df.iloc[expected_rows]
+
+    tm.assert_frame_equal(result, expected)
+
+
+def test_invalid_argument(slice_test_grouped):
+    # Test for error on invalid argument
+
+    with pytest.raises(TypeError, match="Invalid index"):
+        slice_test_grouped.nth(3.14)
+
+
+def test_negative_step(slice_test_grouped):
+    # Test for error on negative slice step
+
+    with pytest.raises(ValueError, match="Invalid step"):
+        slice_test_grouped.nth(slice(None, None, -1))
+
+
+def test_np_ints(slice_test_df, slice_test_grouped):
+    # Test np ints work
+
+    result = slice_test_grouped.nth(np.array([0, 1]))
+    expected = slice_test_df.iloc[[0, 1, 2, 3, 4]]
+
+    tm.assert_frame_equal(result, expected)
diff --git a/pandas/tests/groupby/test_rows.py b/pandas/tests/groupby/test_rows.py
new file mode 100644
index 0000000000000..a8db715e668d7
--- /dev/null
+++ b/pandas/tests/groupby/test_rows.py
@@ -0,0 +1,281 @@
+""" Test GroupBy.rows positional grouped indexing GH#42864"""
+
+import random
+
+import pytest
+
+import pandas as pd
+import pandas._testing as tm
+
+
+@pytest.mark.parametrize(
+    "arg, expected_rows",
+    [
+        [0, [0, 1, 4]],
+        [2, [5]],
+        [5, []],
+        [-1, [3, 4, 7]],
+        [-2, [1, 6]],
+        [-6, []],
+    ],
+)
+def test_int(slice_test_df, slice_test_grouped, arg, expected_rows):
+    """Test single integer"""
+
+    result = slice_test_grouped._rows[arg]
+    expected = slice_test_df.iloc[expected_rows]
+
+    tm.assert_frame_equal(result, expected)
+
+
+def test_slice(slice_test_df, slice_test_grouped):
+    """Test single slice"""
+
+    result = slice_test_grouped._rows[0:3:2]
+    expected = slice_test_df.iloc[[0, 1, 4, 5]]
+
+    tm.assert_frame_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+    "arg, expected_rows",
+    [
+        [[0, 2], [0, 1, 4, 5]],
+        [[0, 2, -1], [0, 1, 3, 4, 5, 7]],
+        [range(0, 3, 2), [0, 1, 4, 5]],
+        [{0, 2}, [0, 1, 4, 5]],
+    ],
+    ids=[
+        "list",
+        "negative",
+        "range",
+        "set",
+    ],
+)
+def test_list(slice_test_df, slice_test_grouped, arg, expected_rows):
+    """Test lists of integers and integer valued iterables"""
+
+    result = slice_test_grouped._rows[arg]
+    expected = slice_test_df.iloc[expected_rows]
+
+    tm.assert_frame_equal(result, expected)
+
+
+def test_ints(slice_test_df, slice_test_grouped):
+    """Test tuple of ints"""
+
+    result = slice_test_grouped._rows[0, 2, -1]
+    expected = slice_test_df.iloc[[0, 1, 3, 4, 5, 7]]
+
+    tm.assert_frame_equal(result, expected)
+
+
+def test_slices(slice_test_df, slice_test_grouped):
+    """Test tuple of slices"""
+
+    result = slice_test_grouped._rows[:2, -2:]
+    expected = slice_test_df.iloc[[0, 1, 2, 3, 4, 6, 7]]
+
+    tm.assert_frame_equal(result, expected)
+
+
+def test_mix(slice_test_df, slice_test_grouped):
+    """Test mixed tuple of ints and slices"""
+
+    result = slice_test_grouped._rows[0, 1, -2:]
+    expected = slice_test_df.iloc[[0, 1, 2, 3, 4, 6, 7]]
+
+    tm.assert_frame_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+    "arg, expected_rows",
+    [
+        [0, [0, 1, 4]],
+        [[0, 2, -1], [0, 1, 3, 4, 5, 7]],
+        [(slice(None, 2), slice(-2, None)), [0, 1, 2, 3, 4, 6, 7]],
+    ],
+)
+def test_as_index(slice_test_df, arg, expected_rows):
+    result = slice_test_df.groupby("Group", sort=False)._rows[arg]
+    expected = slice_test_df.iloc[expected_rows].set_index("Group")
+
+    tm.assert_frame_equal(result, expected)
+
+
+def test_doc_examples():
+    """Test the examples in the documentation"""
+
+    df = pd.DataFrame(
+        [["a", 1], ["a", 2], ["a", 3], ["b", 4], ["b", 5]], columns=["A", "B"]
+    )
+
+    grouped = df.groupby("A", as_index=False)
+
+    result = grouped._rows[1:2]
+    expected = pd.DataFrame([["a", 2], ["b", 5]], columns=["A", "B"], index=[1, 4])
+
+    tm.assert_frame_equal(result, expected)
+
+    result = grouped._rows[1, -1]
+    expected = pd.DataFrame(
+        [["a", 2], ["a", 3], ["b", 5]], columns=["A", "B"], index=[1, 2, 4]
+    )
+
+    tm.assert_frame_equal(result, expected)
+
+
+@pytest.fixture()
+def multiindex_data():
+    ndates = 100
+    nitems = 20
+    dates = pd.date_range("20130101", periods=ndates, freq="D")
+    items = [f"item {i}" for i in range(nitems)]
+
+    data = {}
+    for date in dates:
+        nitems_for_date = nitems - random.randint(0, 12)
+        levels = [
+            (item, random.randint(0, 10000) / 100, random.randint(0, 10000) / 100)
+            for item in items[:nitems_for_date]
+        ]
+        levels.sort(key=lambda x: x[1])
+        data[date] = levels
+
+    return data
+
+
+def _make_df_from_data(data):
+    rows = {}
+    for date in data:
+        for level in data[date]:
+            rows[(date, level[0])] = {"A": level[1], "B": level[2]}
+
+    df = pd.DataFrame.from_dict(rows, orient="index")
+    df.index.names = ("Date", "Item")
+    return df
+
+
+def test_multiindex(multiindex_data):
+    """Test the multiindex mentioned as the use-case in the documentation"""
+
+    df = _make_df_from_data(multiindex_data)
+    result = df.groupby("Date", as_index=False).nth(slice(3, -3))
+
+    sliced = {date: multiindex_data[date][3:-3] for date in multiindex_data}
+    expected = _make_df_from_data(sliced)
+
+    tm.assert_frame_equal(result, expected)
+
+
+@pytest.mark.parametrize("arg", [1, 5, 30, 1000])
+@pytest.mark.parametrize("method", ["head", "tail"])
+@pytest.mark.parametrize("simulated", [True, False])
+def test_against_head_and_tail(arg, method, simulated):
+    """Test gives the same results as grouped head and tail"""
+
+    n_groups = 100
+    n_rows_per_group = 30
+
+    data = {
+        "group": [
+            f"group {g}" for j in range(n_rows_per_group) for g in range(n_groups)
+        ],
+        "value": [
+            f"group {g} row {j}"
+            for j in range(n_rows_per_group)
+            for g in range(n_groups)
+        ],
+    }
+    df = pd.DataFrame(data)
+    grouped = df.groupby("group", as_index=False)
+
+    if method == "head":
+        result = grouped._rows[:arg]
+
+        if simulated:
+            indices = []
+            for j in range(arg):
+                for i in range(n_groups):
+                    if j * n_groups + i < n_groups * n_rows_per_group:
+                        indices.append(j * n_groups + i)
+
+            expected = df.iloc[indices]
+
+        else:
+            expected = grouped.head(arg)
+
+    else:
+        result = grouped._rows[-arg:]
+
+        if simulated:
+            indices = []
+            for j in range(arg):
+                for i in range(n_groups):
+                    if (n_rows_per_group + j - arg) * n_groups + i >= 0:
+                        indices.append((n_rows_per_group + j - arg) * n_groups + i)
+
+            expected = df.iloc[indices]
+
+        else:
+            expected = grouped.tail(arg)
+
+    tm.assert_frame_equal(result, expected)
+
+
+@pytest.mark.parametrize("start", [None, 0, 1, 10, -1, -10])
+@pytest.mark.parametrize("stop", [None, 0, 1, 10, -1, -10])
+@pytest.mark.parametrize("step", [None, 1, 5])
+def test_against_df_iloc(start, stop, step):
+    """Test that a single group gives the same results as DataFame.iloc"""
+
+    n_rows = 30
+
+    data = {
+        "group": ["group 0"] * n_rows,
+        "value": list(range(n_rows)),
+    }
+    df = pd.DataFrame(data)
+    grouped = df.groupby("group", as_index=False)
+
+    result = grouped._rows[start:stop:step]
+    expected = df.iloc[start:stop:step]
+
+    tm.assert_frame_equal(result, expected)
+
+
+def test_series():
+    """Test grouped Series"""
+
+    ser = pd.Series([1, 2, 3, 4, 5], index=["a", "a", "a", "b", "b"])
+    grouped = ser.groupby(level=0)
+    result = grouped._rows[1:2]
+    expected = pd.Series([2, 5], index=["a", "b"])
+
+    tm.assert_series_equal(result, expected)
+
+
+@pytest.mark.parametrize("step", [1, 2, 3, 4, 5])
+def test_step(step):
+    """Test slice with various step values"""
+
+    data = [["x", f"x{i}"] for i in range(5)]
+    data += [["y", f"y{i}"] for i in range(4)]
+    data += [["z", f"z{i}"] for i in range(3)]
+    df = pd.DataFrame(data, columns=["A", "B"])
+
+    grouped = df.groupby("A", as_index=False)
+
+    result = grouped._rows[::step]
+
+    data = [["x", f"x{i}"] for i in range(0, 5, step)]
+    data += [["y", f"y{i}"] for i in range(0, 4, step)]
+    data += [["z", f"z{i}"] for i in range(0, 3, step)]
+
+    index = [0 + i for i in range(0, 5, step)]
+    index += [5 + i for i in range(0, 4, step)]
+    index += [9 + i for i in range(0, 3, step)]
+
+    expected = pd.DataFrame(data, columns=["A", "B"], index=index)
+
+    tm.assert_frame_equal(result, expected)


Code review comments:
diff_hunk: @@ -0,0 +1,211 @@
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Iterable,
+    cast,
+)
+
+import numpy as np
+
+from pandas._typing import (
+    FrameOrSeries,
+    PositionalIndexer,
+)
+from pandas.util._decorators import (
+    cache_readonly,
+    doc,
+)
+
+from pandas.core.dtypes.common import (
+    is_integer,
+    is_list_like,
+)
+
+if TYPE_CHECKING:
+    from pandas.core.groupby import groupby
+
+from pandas.core.indexes.api import CategoricalIndex
+
+
+class GroupByIndexingMixin:
+    """
+    Mixin for adding ._rows to GroupBy.
+    """
+
+    @property
+    def _rows(self) -> RowsGroupByIndexer:
+        if TYPE_CHECKING:
+            groupby_object = cast(groupby.GroupBy, self)
+        else:
+            groupby_object = self
+
+        return RowsGroupByIndexer(groupby_object)
+
+
+@doc(GroupByIndexingMixin._rows)
+class RowsGroupByIndexer:
+    def __init__(self, groupby_object: groupby.GroupBy):
+        self.groupby_object = groupby_object
+
+    def __getitem__(self, arg: PositionalIndexer | tuple) -> FrameOrSeries:
+        """
+        Positional index for selection by integer location per group.
+
+        Used to implement GroupBy._rows which is used to implement GroupBy.nth
+        in the case when the keyword dropna is None or absent.
+
+        Parameters
+        ----------
+        arg : PositionalIndexer | tuple
+            Allowed values are:
+            - int
+            - int valued iterable such as list or range
+            - slice with step either None or positive
+            - tuple of integers and slices
+
+        Returns
+        -------
+        Series
+            The filtered subset of the original groupby Series.
+        DataFrame
+            The filtered subset of the original groupby DataFrame.
+
+        See Also
+        --------
+        DataFrame.iloc : Purely integer-location based indexing for selection by
+            position.
+        GroupBy.head : Return first n rows of each group.
+        GroupBy.tail : Return last n rows of each group.
+        GroupBy.nth : Take the nth row from each group if n is an int, or a
+            subset of rows, if n is a list of ints.
+        """
+        with self.groupby_object._group_selection_context():
+            if isinstance(arg, tuple):
+                if all(is_integer(i) for i in arg):
+                    mask = self._handle_list(arg)
+                else:
+                    mask = self._handle_tuple(arg)
+
+            elif isinstance(arg, slice):
+                mask = self._handle_slice(arg)
+            elif is_integer(arg):
+                mask = self._handle_int(cast(int, arg))
+            elif is_list_like(arg):
+                mask = self._handle_list(cast(Iterable[int], arg))
+            else:
+                raise TypeError(
+                    f"Invalid index {type(arg)}. "
+                    "Must be integer, list-like, slice or a tuple of "
+                    "integers and slices"
+                )
+
+            ids, _, _ = self.groupby_object.grouper.group_info
+
+            # Drop NA values in grouping
+            mask &= ids != -1
+
+            if mask is None or mask is True:
+                result = self.groupby_object._selected_obj[:]
+            else:
+                result = self.groupby_object._selected_obj[mask]
+
+            if self.groupby_object.as_index:
+                result_index = self.groupby_object.grouper.result_index
+                result.index = result_index[ids[mask]]
+
+                if not self.groupby_object.observed and isinstance(
+                    result_index, CategoricalIndex
+                ):
+                    result = result.reindex(result_index)
+
+                result = self.groupby_object._reindex_output(result)
+                if self.groupby_object.sort:
+                    result = result.sort_index()
+
+            return result
+
+    def _handle_int(self, arg: int) -> bool | np.ndarray:
+        if arg >= 0:
+            return self._ascending_count == arg
+        else:
+            return self._descending_count == (-arg - 1)
+
+    def _handle_list(self, args: Iterable[int]) -> bool | np.ndarray:
path: pandas/core/groupby/indexing.py
text: @user1:
Is the result Boolean only when args is empty?

@author:
True. Can happen eg.  `gp._rows[[]]`

diff_hunk: @@ -0,0 +1,211 @@
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Iterable,
+    cast,
+)
+
+import numpy as np
+
+from pandas._typing import (
+    FrameOrSeries,
+    PositionalIndexer,
+)
+from pandas.util._decorators import (
+    cache_readonly,
+    doc,
+)
+
+from pandas.core.dtypes.common import (
+    is_integer,
+    is_list_like,
+)
+
+if TYPE_CHECKING:
+    from pandas.core.groupby import groupby
+
+from pandas.core.indexes.api import CategoricalIndex
+
+
+class GroupByIndexingMixin:
+    """
+    Mixin for adding ._rows to GroupBy.
+    """
+
+    @property
+    def _rows(self) -> RowsGroupByIndexer:
+        if TYPE_CHECKING:
+            groupby_object = cast(groupby.GroupBy, self)
+        else:
+            groupby_object = self
+
+        return RowsGroupByIndexer(groupby_object)
+
+
+@doc(GroupByIndexingMixin._rows)
+class RowsGroupByIndexer:
+    def __init__(self, groupby_object: groupby.GroupBy):
+        self.groupby_object = groupby_object
+
+    def __getitem__(self, arg: PositionalIndexer | tuple) -> FrameOrSeries:
+        """
+        Positional index for selection by integer location per group.
+
+        Used to implement GroupBy._rows which is used to implement GroupBy.nth
+        in the case when the keyword dropna is None or absent.
+
+        Parameters
+        ----------
+        arg : PositionalIndexer | tuple
+            Allowed values are:
+            - int
+            - int valued iterable such as list or range
+            - slice with step either None or positive
+            - tuple of integers and slices
+
+        Returns
+        -------
+        Series
+            The filtered subset of the original groupby Series.
+        DataFrame
+            The filtered subset of the original groupby DataFrame.
+
+        See Also
+        --------
+        DataFrame.iloc : Purely integer-location based indexing for selection by
+            position.
+        GroupBy.head : Return first n rows of each group.
+        GroupBy.tail : Return last n rows of each group.
+        GroupBy.nth : Take the nth row from each group if n is an int, or a
+            subset of rows, if n is a list of ints.
+        """
+        with self.groupby_object._group_selection_context():
+            if isinstance(arg, tuple):
+                if all(is_integer(i) for i in arg):
+                    mask = self._handle_list(arg)
+                else:
+                    mask = self._handle_tuple(arg)
+
+            elif isinstance(arg, slice):
+                mask = self._handle_slice(arg)
+            elif is_integer(arg):
+                mask = self._handle_int(cast(int, arg))
+            elif is_list_like(arg):
+                mask = self._handle_list(cast(Iterable[int], arg))
+            else:
+                raise TypeError(
+                    f"Invalid index {type(arg)}. "
+                    "Must be integer, list-like, slice or a tuple of "
+                    "integers and slices"
+                )
+
+            ids, _, _ = self.groupby_object.grouper.group_info
+
+            # Drop NA values in grouping
+            mask &= ids != -1
+
+            if mask is None or mask is True:
+                result = self.groupby_object._selected_obj[:]
+            else:
+                result = self.groupby_object._selected_obj[mask]
+
+            if self.groupby_object.as_index:
+                result_index = self.groupby_object.grouper.result_index
+                result.index = result_index[ids[mask]]
+
+                if not self.groupby_object.observed and isinstance(
+                    result_index, CategoricalIndex
+                ):
+                    result = result.reindex(result_index)
+
+                result = self.groupby_object._reindex_output(result)
+                if self.groupby_object.sort:
+                    result = result.sort_index()
+
+            return result
+
+    def _handle_int(self, arg: int) -> bool | np.ndarray:
path: pandas/core/groupby/indexing.py
text: @user1:
In what case does this return a bool?

@author:
It doesn't. The bool was just for consistancy. I have taken it out.

diff_hunk: @@ -0,0 +1,211 @@
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Iterable,
+    cast,
+)
+
+import numpy as np
+
+from pandas._typing import (
+    FrameOrSeries,
+    PositionalIndexer,
+)
+from pandas.util._decorators import (
+    cache_readonly,
+    doc,
+)
+
+from pandas.core.dtypes.common import (
+    is_integer,
+    is_list_like,
+)
+
+if TYPE_CHECKING:
+    from pandas.core.groupby import groupby
+
+from pandas.core.indexes.api import CategoricalIndex
+
+
+class GroupByIndexingMixin:
+    """
+    Mixin for adding ._rows to GroupBy.
+    """
+
+    @property
+    def _rows(self) -> RowsGroupByIndexer:
+        if TYPE_CHECKING:
+            groupby_object = cast(groupby.GroupBy, self)
+        else:
+            groupby_object = self
+
+        return RowsGroupByIndexer(groupby_object)
+
+
+@doc(GroupByIndexingMixin._rows)
+class RowsGroupByIndexer:
+    def __init__(self, groupby_object: groupby.GroupBy):
+        self.groupby_object = groupby_object
+
+    def __getitem__(self, arg: PositionalIndexer | tuple) -> FrameOrSeries:
+        """
+        Positional index for selection by integer location per group.
+
+        Used to implement GroupBy._rows which is used to implement GroupBy.nth
+        in the case when the keyword dropna is None or absent.
+
+        Parameters
+        ----------
+        arg : PositionalIndexer | tuple
+            Allowed values are:
+            - int
+            - int valued iterable such as list or range
+            - slice with step either None or positive
+            - tuple of integers and slices
+
+        Returns
+        -------
+        Series
+            The filtered subset of the original groupby Series.
+        DataFrame
+            The filtered subset of the original groupby DataFrame.
+
+        See Also
+        --------
+        DataFrame.iloc : Purely integer-location based indexing for selection by
+            position.
+        GroupBy.head : Return first n rows of each group.
+        GroupBy.tail : Return last n rows of each group.
+        GroupBy.nth : Take the nth row from each group if n is an int, or a
+            subset of rows, if n is a list of ints.
+        """
+        with self.groupby_object._group_selection_context():
+            if isinstance(arg, tuple):
+                if all(is_integer(i) for i in arg):
+                    mask = self._handle_list(arg)
+                else:
+                    mask = self._handle_tuple(arg)
+
+            elif isinstance(arg, slice):
+                mask = self._handle_slice(arg)
+            elif is_integer(arg):
+                mask = self._handle_int(cast(int, arg))
+            elif is_list_like(arg):
+                mask = self._handle_list(cast(Iterable[int], arg))
+            else:
+                raise TypeError(
+                    f"Invalid index {type(arg)}. "
+                    "Must be integer, list-like, slice or a tuple of "
+                    "integers and slices"
+                )
+
+            ids, _, _ = self.groupby_object.grouper.group_info
+
+            # Drop NA values in grouping
+            mask &= ids != -1
path: pandas/core/groupby/indexing.py
text: @user1:
This is removing NA values after the determination of the mask - isn't this in conflict with the docstring?

> 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.

@author:
No, and this is confusing! There are:

1. The dropna df method mentioned here.
2. The dropna groupby parameter in e.g.  `df.groupby("A", dropna=False)`
3. The dropna nth parameter e.g. `gp.nth(3, dropna="any")`

2 refers to whether nans are dropped from the group column. That is the one being implemented here in code copied from nth (and maybe going back there. See Jeff's recent comment...)
3 refers to dropping rows based on nans in other columns and is the one in the nth docstring that you are referring to. That is separate code in nth that I have not touched (apart from tidying up the typing)


@user1:
Ah, right, thanks!