Use OrderScheme metadata to select order-aware groupby execution - #23306
Merged
Conversation
rjzamora
added
2 - In Progress
github-actions
Bot
added
Python
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughSummary by CodeRabbit
WalkthroughChangesThe streaming actor graph now distinguishes local, inter-rank, and flat partitioning checks. Ordered groupby reduction adjusts strict ordering boundaries, sort and over fast paths select checks based on duplicated input metadata, and window evaluation can restore input order. Streaming ordering and reduction
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py (1)
1185-1193: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
is_orderedmust verify key identity, not only key count.
python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py#L1185-L1193: compare requested columns and, forOrderKey, direction/null ordering before returningTrue.python/cudf_polars/tests/streaming/test_metadata.py#L963-L970: add equal-length mismatched-key regression cases.🤖 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/utils.py` around lines 1185 - 1193, Update is_ordered in python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py:1185-1193 to compare each requested column with the scheme’s ordering keys, including OrderKey direction and null-ordering attributes, before returning True; retain strict-boundary handling for valid subset matches. Add equal-length mismatched-key regression cases in python/cudf_polars/tests/streaming/test_metadata.py:963-970 covering differing columns and OrderKey ordering options.Source: Coding guidelines
🤖 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.
Inline comments:
In `@python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py`:
- Around line 611-614: Update the duplicated-input fast paths in sort.py at
lines 611-614 and over.py at lines 723-725: ensure inherited inter-rank ordering
or partitioning does not qualify as the local/strict condition for
already_sorted or chunkwise evaluation, while explicit local ordering and strict
local partitioning do. Use the existing metadata and partitioning checks around
the sort collective and over evaluation paths, and add or update coverage for
both behaviors.
In `@python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py`:
- Around line 547-630: Add a Python unit benchmark covering the ordered
reduction path centered on _ordered_adjust_reduce, comparing sorted
maintain-order groupby with the hash-shuffle implementation across multiple
group cardinalities and partition counts. Measure successful local/tree
reduction performance and verify the ordered path avoids shuffle, following the
repository’s existing benchmark conventions.
In `@python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py`:
- Around line 1178-1183: Retain the public is_strictly_sorted method as a
deprecated compatibility alias for the new is_ordered method. Have it emit the
project-standard deprecation warning and delegate all arguments, including
level, to is_ordered while preserving the existing return behavior for one
deprecation cycle.
In `@python/cudf_polars/tests/streaming/test_groupby.py`:
- Around line 114-132: Expand test_dynamic_groupby_after_sort_on_group_keys
coverage with cases for empty, single-row, nullable/all-null, and mixed-type
keys, plus maintain_order unique aggregation. Parameterize exact full-key and
strict-prefix sort/group ordering where applicable, and keep each case
validating GPU results against Polars through the streaming engine.
In `@python/cudf_polars/tests/streaming/test_metadata.py`:
- Around line 963-970: Add negative assertions to the ordering tests around
NormalizedPartitioning.is_ordered, covering equal-length order keys with
differing column indices, sort directions, and null-order settings. Verify each
mismatch returns False for the relevant flat, local, and inter_rank checks as
applicable, including both partitioning and nested partitioning cases, to
exercise the predicate in utils.py.
---
Outside diff comments:
In `@python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py`:
- Around line 1185-1193: Update is_ordered in
python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py:1185-1193 to
compare each requested column with the scheme’s ordering keys, including
OrderKey direction and null-ordering attributes, before returning True; retain
strict-boundary handling for valid subset matches. Add equal-length
mismatched-key regression cases in
python/cudf_polars/tests/streaming/test_metadata.py:963-970 covering differing
columns and OrderKey ordering options.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d9af5262-96cf-4344-b54c-8ffc7cbb19d9
📒 Files selected for processing (8)
python/cudf_polars/cudf_polars/engine/core.pypython/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.pypython/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.pypython/cudf_polars/cudf_polars/streaming/actor_graph/groupby.pypython/cudf_polars/cudf_polars/streaming/actor_graph/over.pypython/cudf_polars/cudf_polars/streaming/actor_graph/utils.pypython/cudf_polars/tests/streaming/test_groupby.pypython/cudf_polars/tests/streaming/test_metadata.py
💤 Files with no reviewable changes (1)
- python/cudf_polars/cudf_polars/engine/core.py
Comment on lines
+547
to
+630
| async def _ordered_adjust_reduce( | ||
| context: Context, | ||
| comm: Communicator, | ||
| decomposed: DecomposedGroupBy, | ||
| ir_context: IRExecutionContext, | ||
| ch_out: Channel[TableChunk], | ||
| ch_in: Channel[TableChunk], | ||
| metadata_in: ChannelMetadata, | ||
| collective_id: int, | ||
| target_partition_size: int, | ||
| *, | ||
| aggregated: TableChunk, | ||
| input_drained: bool, | ||
| input_ordering: Ordering, | ||
| tracer: ActorTracer | None = None, | ||
| ) -> None: | ||
| """Adjust locally aggregated data to strict ordering boundaries.""" | ||
| partial_input_ordering = _remap_ordering_keys( | ||
| input_ordering, | ||
| decomposed.shuffle_indices[: len(input_ordering.keys)], | ||
| context.br(), | ||
| ) | ||
| partial_output_ordering = get_strict_ordering(partial_input_ordering, context.br()) | ||
| ch_local = context.create_channel() | ||
| ch_adjusted = context.create_channel() | ||
| adjusted_metadata = _adjusted_ordering_metadata( | ||
| comm, metadata_in, partial_output_ordering | ||
| ) | ||
| metadata_out = _groupby_output_metadata( | ||
| decomposed.ir, | ||
| decomposed, | ||
| adjusted_metadata.local_count, | ||
| adjusted_metadata.partitioning, | ||
| adjusted_metadata.duplicated, | ||
| context=context, | ||
| ) | ||
| if tracer is not None: | ||
| tracer.decision = "adjust_ordering" | ||
|
|
||
| await send_metadata(ch_out, context, metadata_out) | ||
| if tracer is not None and metadata_out.duplicated: | ||
| tracer.set_duplicated() | ||
|
|
||
| async def reduce_adjusted_chunks() -> None: | ||
| extract_irs = [decomposed.reduction_ir] + ( | ||
| [decomposed.select_ir] if decomposed.select_ir else [] | ||
| ) | ||
| partition_id = 0 | ||
| while (msg := await ch_adjusted.recv(context)) is not None: | ||
| chunk = await evaluate_chunk( | ||
| context, | ||
| TableChunk.from_message(msg, br=context.br()), | ||
| *extract_irs, | ||
| ir_context=ir_context, | ||
| ) | ||
| await send_chunk(context, ch_out, chunk, partition_id, tracer=tracer) | ||
| partition_id += 1 | ||
| await ch_out.drain(context) | ||
|
|
||
| async with shutdown_channels_on_error(context, ch_local, ch_adjusted): | ||
| await gather_in_task_group( | ||
| _send_locally_aggregated_chunks( | ||
| context, | ||
| decomposed, | ||
| ir_context, | ||
| ch_local, | ||
| ch_in, | ||
| target_partition_size, | ||
| aggregated=aggregated, | ||
| input_drained=input_drained, | ||
| ), | ||
| adjust_ordering( | ||
| context, | ||
| comm, | ||
| decomposed.reduction_ir, | ||
| ir_context, | ||
| ch_adjusted, | ||
| ch_local, | ||
| partial_input_ordering, | ||
| partial_output_ordering, | ||
| collective_id=collective_id, | ||
| ), | ||
| reduce_adjusted_chunks(), | ||
| ) |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy lift
Add a benchmark for the ordered reduction path.
Benchmark sorted, maintain-order groupby against the hash-shuffle path across several cardinalities and partition counts. This PR’s main benefit is avoiding shuffle without penalizing successful local/tree reduction, so that contract should be measured.
As per coding guidelines, Python feature contributions must add unit benchmarks.
🤖 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
547 - 630, Add a Python unit benchmark covering the ordered reduction path
centered on _ordered_adjust_reduce, comparing sorted maintain-order groupby with
the hash-shuffle implementation across multiple group cardinalities and
partition counts. Measure successful local/tree reduction performance and verify
the ordered path avoids shuffle, following the repository’s existing benchmark
conventions.
Source: Coding guidelines
Comment on lines
+1178
to
+1183
| def is_ordered( | ||
| self, | ||
| order_keys: Sequence[int | OrderKey], | ||
| *, | ||
| level: PartitioningLevel = "flat", | ||
| ) -> bool: |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Preserve is_strictly_sorted as a deprecated compatibility alias.
Replacing this public method outright will raise AttributeError for existing callers. Retain an alias that warns and delegates to is_ordered for one deprecation cycle.
As per coding guidelines, “Detect and flag API breaking changes to public methods/attributes without deprecation warnings.”
🤖 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/utils.py` around lines
1178 - 1183, Retain the public is_strictly_sorted method as a deprecated
compatibility alias for the new is_ordered method. Have it emit the
project-standard deprecation warning and delegate all arguments, including
level, to is_ordered while preserving the existing return behavior for one
deprecation cycle.
Source: Coding guidelines
Comment on lines
+114
to
+132
| @pytest.mark.parametrize("group_keys", [("key", "subkey"), ("key",)]) | ||
| def test_dynamic_groupby_after_sort_on_group_keys(spmd_engine_factory, group_keys): | ||
| """Group sorted data by the full sort key set or a sorted-key prefix.""" | ||
| streaming_engine = spmd_engine_factory( | ||
| StreamingOptions(target_partition_size=128), | ||
| ) | ||
| df = pl.LazyFrame( | ||
| { | ||
| "key": [0] * 16 + [1] * 16 + [2] * 16 + [3] * 16, | ||
| "subkey": ([0] * 8 + [1] * 8) * 4, | ||
| "value": range(64), | ||
| } | ||
| ) | ||
| q = ( | ||
| df.sort("key", "subkey") | ||
| .group_by(*group_keys, maintain_order=True) | ||
| .agg(pl.col("value").sum()) | ||
| ) | ||
| assert_gpu_result_equal(q, engine=streaming_engine) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Expand coverage for ordering-boundary edge cases.
Add empty, single-row, nullable/all-null key, mixed-type key, and maintain-order unique cases. Exercise both exact and strict-prefix ordering where applicable.
As per coding guidelines, streaming tests must cover empty, all-null, single-element, and mixed-type inputs while comparing GPU results with Polars.
🤖 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/streaming/test_groupby.py` around lines 114 - 132,
Expand test_dynamic_groupby_after_sort_on_group_keys coverage with cases for
empty, single-row, nullable/all-null, and mixed-type keys, plus maintain_order
unique aggregation. Parameterize exact full-key and strict-prefix sort/group
ordering where applicable, and keep each case validating GPU results against
Polars through the streaming engine.
Source: Coding guidelines
rjzamora
commented
Jul 16, 2026
Comment on lines
+47
to
+53
| def get_strict_ordering(ordering: Ordering, br: BufferResource) -> Ordering: | ||
| """Return an equivalent Ordering with strict boundaries.""" | ||
| return Ordering( | ||
| ordering.keys, | ||
| ordering.get_boundaries(br), | ||
| strict_boundaries=True, | ||
| ) |
Contributor
Author
There was a problem hiding this comment.
Note: This could also belong to Ordering (e.g. Ordering.as_strict(br: BufferResource)).
Contributor
There was a problem hiding this comment.
I think this should be implemented on Ordering:
diff --git a/python/cudf_streaming/cudf_streaming/channel_metadata.pyx b/python/cudf_streaming/cudf_streaming/channel_metadata.pyx
index 70d5ebfadc..99b9d83e95 100644
--- a/python/cudf_streaming/cudf_streaming/channel_metadata.pyx
+++ b/python/cudf_streaming/cudf_streaming/channel_metadata.pyx
@@ -150,6 +150,11 @@ cdef class Ordering:
ret._handle = move(ordering)
return ret
+ def as_strict(self) -> Ordering:
+ return Ordering.from_cpp(
+ cpp_Ordering(self._handle.keys, self._handle.boundaries, True)
+ )
+
@property
def keys(self) -> tuple:
"""Sort keys, one per sort column."""and
diff --git a/python/cudf_streaming/cudf_streaming/channel_metadata.pxd b/python/cudf_streaming/cudf_streaming/channel_metadata.pxd
index f2037164a9..f59e8c8c8b 100644
--- a/python/cudf_streaming/cudf_streaming/channel_metadata.pxd
+++ b/python/cudf_streaming/cudf_streaming/channel_metadata.pxd
@@ -41,6 +41,9 @@ cdef extern from "<cudf_streaming/channel_metadata.hpp>" \
cpp_Ordering(
vector[cpp_OrderKey], unique_ptr[cpp_TableChunk], bool_t
) except +ex_handler
+ cpp_Ordering(
+ vector[cpp_OrderKey], shared_ptr[cpp_TableChunk], bool_t
+ ) except +ex_handler
vector[cpp_OrderKey] keys
shared_ptr[cpp_TableChunk] boundaries
bool_t strict_boundariesI think.
3 tasks
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/cudf_polars/cudf_polars/streaming/actor_graph/over.py (1)
291-308: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winInclude
positionin the sort key.
partition()does not preserve row order within a partition, so sorting only by(rank, chunk_index)can leave rows from the same input chunk in arbitrary order. Addcolumns[n_child + 1]here soshift/cumulative/rolling windows withoutorder_bykeep input order; add a regression test with multi-row chunks.🤖 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/over.py` around lines 291 - 308, Update the stable sort key in the sort_by_input_order branch of the relevant actor-graph evaluation function to include columns[n_child + 1] between the existing rank and chunk_index keys, preserving input row order within each chunk. Add a regression test covering multi-row chunks and validating shift, cumulative, or rolling window behavior without order_by.Sources: Coding guidelines, MCP tools
🤖 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.
Outside diff comments:
In `@python/cudf_polars/cudf_polars/streaming/actor_graph/over.py`:
- Around line 291-308: Update the stable sort key in the sort_by_input_order
branch of the relevant actor-graph evaluation function to include
columns[n_child + 1] between the existing rank and chunk_index keys, preserving
input row order within each chunk. Add a regression test covering multi-row
chunks and validating shift, cumulative, or rolling window behavior without
order_by.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 67365a93-822b-4547-bf90-bbcd0b8a7932
📒 Files selected for processing (1)
python/cudf_polars/cudf_polars/streaming/actor_graph/over.py
3 tasks
16 tasks
wence-
requested changes
Aug 12, 2026
wence-
left a comment
wence-
left a comment
Contributor
There was a problem hiding this comment.
I have one concern about the assignment of strict boundaries to ranks.
Contributor
Author
|
@wence- - Do you still have concerns about this? |
wence-
approved these changes
Aug 20, 2026
wence-
left a comment
wence-
left a comment
Contributor
There was a problem hiding this comment.
Small suggestions to make the "new Ordering" helpers method of Ordering. That way you can also avoid some tablechunk shenanigans
Comment on lines
+47
to
+53
| def get_strict_ordering(ordering: Ordering, br: BufferResource) -> Ordering: | ||
| """Return an equivalent Ordering with strict boundaries.""" | ||
| return Ordering( | ||
| ordering.keys, | ||
| ordering.get_boundaries(br), | ||
| strict_boundaries=True, | ||
| ) |
Contributor
There was a problem hiding this comment.
I think this should be implemented on Ordering:
diff --git a/python/cudf_streaming/cudf_streaming/channel_metadata.pyx b/python/cudf_streaming/cudf_streaming/channel_metadata.pyx
index 70d5ebfadc..99b9d83e95 100644
--- a/python/cudf_streaming/cudf_streaming/channel_metadata.pyx
+++ b/python/cudf_streaming/cudf_streaming/channel_metadata.pyx
@@ -150,6 +150,11 @@ cdef class Ordering:
ret._handle = move(ordering)
return ret
+ def as_strict(self) -> Ordering:
+ return Ordering.from_cpp(
+ cpp_Ordering(self._handle.keys, self._handle.boundaries, True)
+ )
+
@property
def keys(self) -> tuple:
"""Sort keys, one per sort column."""and
diff --git a/python/cudf_streaming/cudf_streaming/channel_metadata.pxd b/python/cudf_streaming/cudf_streaming/channel_metadata.pxd
index f2037164a9..f59e8c8c8b 100644
--- a/python/cudf_streaming/cudf_streaming/channel_metadata.pxd
+++ b/python/cudf_streaming/cudf_streaming/channel_metadata.pxd
@@ -41,6 +41,9 @@ cdef extern from "<cudf_streaming/channel_metadata.hpp>" \
cpp_Ordering(
vector[cpp_OrderKey], unique_ptr[cpp_TableChunk], bool_t
) except +ex_handler
+ cpp_Ordering(
+ vector[cpp_OrderKey], shared_ptr[cpp_TableChunk], bool_t
+ ) except +ex_handler
vector[cpp_OrderKey] keys
shared_ptr[cpp_TableChunk] boundaries
bool_t strict_boundariesI think.
Comment on lines
+460
to
+473
| def _remap_ordering_keys( | ||
| ordering: Ordering, | ||
| column_indices: tuple[int, ...], | ||
| br: BufferResource, | ||
| ) -> Ordering: | ||
| """Return ``ordering`` with keys remapped to another schema.""" | ||
| return Ordering( | ||
| [ | ||
| OrderKey(index, key.order, key.null_order) | ||
| for key, index in zip(ordering.keys, column_indices, strict=True) | ||
| ], | ||
| ordering.get_boundaries(br), | ||
| strict_boundaries=ordering.strict_boundaries, | ||
| ) |
Contributor
There was a problem hiding this comment.
Again, I think this should be Ordering.remap(new_order_keys).
Contributor
Author
Good suggestions. I think I'll do this in #23371 or an isolated follow-up PR to get this out of the way. Hopefully that's okay. |
rjzamora
added
5 - Ready to Merge
Contributor
Author
|
/merge |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Part of #22128
Now that
OrderSchememetadata andadjust_orderingexist (#22628),groupby_actorcan use ordered input metadata to avoid falling back to a hash shuffle when the input is already ordered on the group keys.The existing
partitioning.is_strictly_partitioned()check ingroupby_actoralready addresses the case that the input data is strictly ordered. The proposed change in this PR is to useadjust_orderingwhen the input has a compatibleOrdering(strict_boundaries=False)scheme.Notes:
If the cardinality of the
groupby/uniqueoperation is low, a basic tree reduction is likely to outperform a naiveadjust_ordering+"chunkwise" do_evaluate. Therefore, we only introduce theadjust_orderingstep after we are sure the local-aggregation path is not available. We do this by introducing a new_ordered_adjust_reducecode path that starts "adjusting" chunks into the requiredOrderingschema after the local-aggregation step.