Support ordered first/last groupby aggregations in cudf-polars

[rjzamora]

Closes #23592

Adds cudf-polars support for grouped sort_by(...).first() and sort_by(...).last() aggregations, including streaming decomposition for multi-partition execution.

This enables expressions like:

df.group_by("g").agg(
    pl.col("x").sort_by("t").first().alias("open"),
    pl.col("x").sort_by("t").last().alias("close"),
)

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[rjzamora]

/ok to test

[rjzamora]

/ok to test

[coderabbitai]

📝 Walkthrough

Summary by CodeRabbit

• New Features
  • Added support for ordered first and last aggregations using sort_by within grouped operations.
  • Extended support to streaming and distributed grouped processing.
  • Preserved sorted group-key order and stable ordering for ties where supported.
• Bug Fixes
  • Improved ordering consistency across grouped aggregation results.
  • Added validation and clear rejection for unsupported ordering configurations, nested aggregations, and incompatible options.

Walkthrough

The change adds SortedAgg for ordered first and last groupby aggregations. It supports in-memory, streaming, and SPMD execution, including stable ties, sorting options, validation, and unsupported repartition cases.

Changes

Sorted groupby aggregation support

Layer / File(s)	Summary
SortedAgg contract and translation python/cudf_polars/cudf_polars/dsl/expressions/aggregation.py, python/cudf_polars/cudf_polars/dsl/expr.py, python/cudf_polars/cudf_polars/dsl/utils/aggregations.py	Defines and exports SortedAgg for ordered first and last operations. Groupby translation builds it from sort_by expressions and validates unsupported nested aggregations.
In-memory sorted aggregation evaluation python/cudf_polars/cudf_polars/dsl/ir.py	Separates sorted aggregations from ordinary requests, sorts grouped values, reduces results, and restores request and group-key order.
Streaming sorted aggregation execution python/cudf_polars/cudf_polars/streaming/groupby.py, python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py	Carries payloads and sort keys through streaming reductions. Stable aggregations preserve order and reject unsupported repartition cases.
Validation and execution coverage python/cudf_polars/tests/test_groupby.py, python/cudf_polars/tests/streaming/test_groupby.py, python/cudf_polars/tests/streaming/test_spmd.py	Adds coverage for sorting options, stable ties, group order, streaming execution, SPMD parquet input, and invalid expressions.

Estimated code review effort: 4 (Complex) | ~60 minutes

Suggested reviewers: matt711, madsbk, vyasr

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 25.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes support for ordered first and last groupby aggregations in cudf-polars.
Description check	✅ Passed	The description directly explains the grouped sort_by first and last aggregation support and streaming execution changes.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (3)

python/cudf_polars/cudf_polars/dsl/ir.py (1)

2346-2349: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Fix the request_groups key annotation.

tuple[bool] describes a one-element tuple. SortedAgg.options stores variable-length tuples of nulls_last and descending flags, one per order key. Use tuple[bool, ...].

♻️ Proposed annotation fix

         request_groups: dict[
-            tuple[tuple[bool, tuple[bool], tuple[bool]], tuple[expr.Expr, ...]],
+            tuple[tuple[bool, tuple[bool, ...], tuple[bool, ...]], tuple[expr.Expr, ...]],
             list[expr.NamedExpr],
         ] = {}

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@python/cudf_polars/cudf_polars/dsl/ir.py` around lines 2346 - 2349, Update
the request_groups type annotation in the surrounding request-grouping logic so
the order-options component uses tuple[bool, ...] rather than tuple[bool],
matching the variable-length nulls_last and descending flags stored by
SortedAgg.options; leave the other key and value types unchanged.

python/cudf_polars/tests/test_groupby.py (1)

436-441: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add coverage for two sorted aggregations that use different order keys.

Every new test groups sorted aggregations under one (options, by_exprs) key. _evaluate_sorted_aggregations in python/cudf_polars/cudf_polars/dsl/ir.py builds one request_groups entry per distinct option and order-key pair, and it sorts and re-groups each entry separately. A query such as pl.col("val").sort_by("idx").first() together with pl.col("val").sort_by("seq", descending=True).first() exercises the multi-entry alignment path.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@python/cudf_polars/tests/test_groupby.py` around lines 436 - 441, Add test
coverage in the groupby aggregation tests for multiple sorted aggregations using
different order keys and sort directions, such as aggregations ordered by “idx”
and by descending “seq”. Build the query with both aggregations and validate it
through assert_gpu_result_equal with row-order checking, exercising
_evaluate_sorted_aggregations’ multi-entry alignment path.

python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py (1)

501-504: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Reuse the shared stable-sorted-agg predicate.

_has_stable_sorted_agg in python/cudf_polars/cudf_polars/streaming/groupby.py encodes the same rule, but it walks the expression tree with traversal. This check inspects only the top-level ne.value. The two predicates can diverge if a SortedAgg ever appears below the top level. Import and call the shared helper here.

♻️ Proposed change

 def _maintain_order(ir: GroupBy | Distinct) -> bool:
     if isinstance(ir, GroupBy):
-        return ir.maintain_order or any(
-            isinstance(ne.value, SortedAgg) and ne.value.options[0]
-            for ne in ir.agg_requests
-        )
+        return ir.maintain_order or _has_stable_sorted_agg(ir.agg_requests)

Import the helper from cudf_polars.streaming.groupby and drop the now-unused SortedAgg import if nothing else in this module uses it.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py` around lines
501 - 504, Update the stable-order check around the return expression to import
and call the shared _has_stable_sorted_agg helper from
cudf_polars.streaming.groupby, passing the aggregation requests instead of
inspecting only top-level values. Remove the SortedAgg import if it is unused
after this change.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@python/cudf_polars/cudf_polars/dsl/ir.py`:
- Around line 2346-2349: Update the request_groups type annotation in the
surrounding request-grouping logic so the order-options component uses
tuple[bool, ...] rather than tuple[bool], matching the variable-length
nulls_last and descending flags stored by SortedAgg.options; leave the other key
and value types unchanged.

In `@python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py`:
- Around line 501-504: Update the stable-order check around the return
expression to import and call the shared _has_stable_sorted_agg helper from
cudf_polars.streaming.groupby, passing the aggregation requests instead of
inspecting only top-level values. Remove the SortedAgg import if it is unused
after this change.

In `@python/cudf_polars/tests/test_groupby.py`:
- Around line 436-441: Add test coverage in the groupby aggregation tests for
multiple sorted aggregations using different order keys and sort directions,
such as aggregations ordered by “idx” and by descending “seq”. Build the query
with both aggregations and validate it through assert_gpu_result_equal with
row-order checking, exercising _evaluate_sorted_aggregations’ multi-entry
alignment path.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: dedc9819-027a-4a8f-85c9-22cb34453cf4

📥 Commits

Reviewing files that changed from the base of the PR and between 481e42a and 73f1744.

📒 Files selected for processing (9)

• python/cudf_polars/cudf_polars/dsl/expr.py
• python/cudf_polars/cudf_polars/dsl/expressions/aggregation.py
• python/cudf_polars/cudf_polars/dsl/ir.py
• python/cudf_polars/cudf_polars/dsl/utils/aggregations.py
• python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py
• python/cudf_polars/cudf_polars/streaming/groupby.py
• python/cudf_polars/tests/streaming/test_groupby.py
• python/cudf_polars/tests/streaming/test_spmd.py
• python/cudf_polars/tests/test_groupby.py