[PR-12515] Adding SeriesGroupBy.nlargest and nsmallest

dask/dask Issue #10632: https://github.com/dask/dask/issues/10632

Description

dask.dataframe.SeriesGroupBy did not implement nlargest or nsmallest, both of which exist on the pandas SeriesGroupBy. Any group-wise top-N query — "the three highest values within each group" — had no direct equivalent in dask, forcing users into a full sort or a groupby().apply() with a user-defined function, neither of which parallelizes well.

Background

Dask's dataframe groupby reductions are built on the ApplyConcatApply (ACA) pattern from the expression system: a chunk function runs on each partition, a combine function merges intermediates in batches down a reduction tree, and an aggregate function produces the final result. Most single-method aggregations (sum, min, idxmax, value_counts) are declared declaratively as SingleAggregation subclasses that only name their chunk and aggregate callables.

Problem:

SeriesGroupBy.nlargest and .nsmallest were simply absent from the dask implementation.
The issue was raised in the context of the duckdb db-benchmark suite, whose query 8 is a group-wise top-N.
Because GroupBy.__getattr__ falls back to column selection, calling .nlargest produced a confusing AttributeError rather than a clear "not implemented".

Impact:

Group-wise top-N queries had no scalable implementation in dask.
Benchmarks comparing dask against DuckDB and polars could not run the query at all.
Code written against the pandas API broke when the frame was swapped for a dask dataframe.

Objective

Add SeriesGroupBy.nlargest and SeriesGroupBy.nsmallest that:

Match the pandas result exactly, including the (group_key, original_index) MultiIndex
Reduce in a tree rather than materialising each group
Support the n and keep parameters where they are decomposable across partitions
Fail loudly and explicitly where they are not

Root Cause Analysis

The feature was never implemented rather than broken, so the analysis is about whether the operation decomposes across partitions at all.

SeriesGroupBy.nlargest keeps the group keys in the result index:

>>> df.groupby("a")["b"].nlargest(3).index.nlevels
2   # (a, original_index)

This is the property that makes it decomposable. Because the group keys survive into the intermediate, the per-partition results can be regrouped on those index levels and reduced again — the top 3 of each partition's top 3 is the global top 3.

Regrouping adds a duplicate set of levels that has to be stripped:

chunk output      index: (a, orig)
regroup on level 0 → nlargest  index: (a, a, orig)
droplevel([0])    index: (a, orig)   ✅ matches pandas

The same shape already exists in dask for GroupBy.head / .tail, whose _head_aggregate helper performs exactly this droplevel — so the feature had a working precedent to follow.
keep="last" is the exception and does not decompose. pandas returns tied values in reverse order of appearance in that mode:
```
>>> s = pd.Series([0.5, 0.5, 0.5, 0.9], index=[14, 18, 20, 25])
>>> s.nsmallest(3, keep="last")
20    0.5
18    0.5
14    0.5
>>> s.nsmallest(3, keep="first")
14    0.5
18    0.5
20    0.5
```
Since the reduction concatenates intermediates in partition order, a reversed intermediate causes the second pass to break ties in favour of the wrong row. This surfaced as a real test failure — a group whose correct answer was index 20 returned index 14.

Proposed Solution

Add aggregation helpers in dask/dataframe/groupby.py, alongside the existing _head_aggregate:

def _nlargest_aggregate(series_gb, **kwargs):
    levels = kwargs.pop("index_levels")
    return series_gb.nlargest(**kwargs).droplevel(list(range(levels)))

Declare the expressions in dask/dataframe/dask_expr/_groupby.py:

class NLargest(SingleAggregation):
    groupby_chunk = M.nlargest
    groupby_aggregate = staticmethod(_nlargest_aggregate)

Expose the public methods on SeriesGroupBy, sharing one implementation, with the number of index levels to drop derived from the grouping keys.

Reject keep="last" explicitly:

raise NotImplementedError(
    "keep='last' is not supported for a dask SeriesGroupBy, only "
    "keep='first' and keep='all'."
)

Implementation Details

dask/dataframe/groupby.py modifications:
- Added _nlargest_aggregate and _nsmallest_aggregate, placed next to the _head_aggregate / _tail_aggregate helpers they mirror
dask/dataframe/dask_expr/_groupby.py modifications:
- Added NLargest and NSmallest as SingleAggregation subclasses, using M.nlargest and M.nsmallest directly as the chunk step
- Added SeriesGroupBy.nlargest and .nsmallest, both delegating to a shared _n_largest_smallest helper, decorated with @derived_from(pd.core.groupby.SeriesGroupBy)
- Deliberately did not override combine, unlike Head, which forces a plain concatenation. Letting the default combine run means each level of the reduction tree actually shrinks the intermediate rather than accumulating it
- Methods were added only to SeriesGroupBy, matching pandas — DataFrameGroupBy has no nlargest
docs/source/dataframe-api.rst:
- Added SeriesGroupBy.nlargest and SeriesGroupBy.nsmallest to the API listing

Testing & Validation

Test fixtures pin 21 of 60 rows to an identical value, so that ties — the entire subtlety of the keep parameter — are actually exercised rather than assumed away.

ID	Scenario	Expected Result	Status
1	`nlargest` / `nsmallest` across `n` ∈ 10	Results match pandas, including when `n` exceeds the group size	✅
2	Single key (`"foo"`) and composite key (`["foo", "bar"]`)	Correct number of index levels dropped in both cases	✅
3	`keep="first"` and `keep="all"` with ties present	Tie-breaking matches pandas exactly	✅
4	`keep="last"`	Raises `NotImplementedError` with "keep='last'" in message	✅
5	`split_every=2` (forces a multi-stage reduction tree)	Identical result to the single-stage reduction	✅
6	Grouping by a Series rather than a column name	Results match pandas	✅
7	Full existing groupby suite	No regressions	✅

Validation command:

pixi run test dask/dataframe/tests/test_groupby.py -k "nlargest or nsmallest"

Output:

56 passed, 2234 deselected in 9.12s

Full groupby suite, checking for regressions:

pixi run test dask/dataframe/tests/test_groupby.py