Add adjust_ordering utility - #22628
Conversation
rjzamora
added
feature request
|
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. |
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:
📝 WalkthroughWalkthroughAdds async operator ChangesOrderScheme Adjustment Collective
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Suggested labels
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: 1
🧹 Nitpick comments (4)
python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py (2)
137-139: 💤 Low valueUnnecessary single-element concatenation.
plc.concatenate.concatenateis called with a single-element list. Constructing the table directly would be more efficient.♻️ Suggested simplification
- payload = plc.concatenate.concatenate( - [plc.Table(payload_cols)], stream=stream, mr=br.device_mr - ) + payload = plc.Table(payload_cols)🤖 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/collectives/orderscheme.py` around lines 137 - 139, The call to plc.concatenate.concatenate with a single element is unnecessary; replace the concatenation with a direct table construction using plc.Table(payload_cols) instead of plc.concatenate.concatenate([plc.Table(payload_cols)], stream=stream, mr=br.device_mr). If plc.Table supports the stream and mr arguments, pass stream=stream and mr=br.device_mr to plc.Table; otherwise construct plc.Table(payload_cols) and set any required memory/stream settings via the appropriate Table APIs rather than using concatenate.
180-209: 💤 Low valueConsider adding a
Returnssection to the docstring.The function returns
None, but for completeness and API documentation consistency, consider documenting the return value.🤖 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/collectives/orderscheme.py` around lines 180 - 209, The function whose docstring begins "Adjust flat OrderScheme boundaries using contiguous partition ownership" currently omits a Returns section; update that docstring to add a Returns entry documenting that the function returns None (e.g., "Returns\n--------\nNone\n This function performs in-place adjustments and does not return a value.") so callers reading the docstring for parameters like input_scheme, output_scheme, and collective_id have a clear API contract.python/cudf_polars/tests/streaming/test_adjust_orderscheme.py (2)
318-320: 💤 Low valueVariable
keysshadows outer scope.The loop variable
keyson line 318 shadows the test inputkeysdefined on line 296. Consider renaming toexpected_keysfor clarity.Same pattern occurs on line 353.
✨ Suggested rename
- for pid, keys in expected.items(): - assert output[pid]["key"].to_list() == keys - assert output[pid]["val"].to_list() == keys + for pid, expected_keys in expected.items(): + assert output[pid]["key"].to_list() == expected_keys + assert output[pid]["val"].to_list() == expected_keys🤖 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_adjust_orderscheme.py` around lines 318 - 320, The loop variable name keys in the test loop (for pid, keys in expected.items()) shadows the outer test fixture/variable keys; rename the loop variable to expected_keys (or another unambiguous name) and update the two assertions to compare against expected_keys (output[pid]["key"].to_list() == expected_keys and output[pid]["val"].to_list() == expected_keys); apply the same rename for the identical pattern later in the file (the second loop around line 353) so there is no shadowing.
289-355: ⚖️ Poor tradeoffConsider adding edge case tests for all-null and single-element inputs.
Per coding guidelines, test files should cover edge cases including all-null columns and single-element Series. The empty partition case is well covered, but these additional cases would strengthen coverage of the
adjust_orderschemeoperator's handling of degenerate inputs.🤖 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_adjust_orderscheme.py` around lines 289 - 355, Add two new pytest cases that exercise _adjust_and_collect with degenerate inputs: one where the input frame produced by _frame contains columns that are all null (e.g., "key" or "val" series all-null) and one where the input has a single-element Series; reuse the existing patterns from test_adjust_orderscheme_emits_empty_owned_partitions and test_adjust_orderscheme_single_rank_no_collective to create context, comm, stream, input_scheme via _make_scheme and call _adjust_and_collect, then assert outputs by pid as in the existing tests; ensure you parametrize/skip for nranks like the other tests and reference the helper functions _frame, _make_scheme and _adjust_and_collect so the new tests integrate with the current test scaffolding.
🤖 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/orderscheme.py`:
- Around line 228-289: The receive/split/send loop in adjust_orderscheme doesn't
guard ch_in/ch_out and exchange finalization on exceptions; wrap the main async
work in shutdown_on_error(context, ch_in, ch_out) (or add a try/finally) so that
on any exception you always call await exchange.insert_finished(context) if
exchange is not None, drain/shutdown ch_out (await ch_out.drain(context) and
ch_out.shutdown/context-equivalent), and ensure ch_in is closed; locate this
around the loop using symbols adjust_orderscheme, ch_in, ch_out, exchange,
exchange.insert_finished, and ch_out.drain to implement the try/finally or async
context manager to guarantee cleanup.
---
Nitpick comments:
In
`@python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py`:
- Around line 137-139: The call to plc.concatenate.concatenate with a single
element is unnecessary; replace the concatenation with a direct table
construction using plc.Table(payload_cols) instead of
plc.concatenate.concatenate([plc.Table(payload_cols)], stream=stream,
mr=br.device_mr). If plc.Table supports the stream and mr arguments, pass
stream=stream and mr=br.device_mr to plc.Table; otherwise construct
plc.Table(payload_cols) and set any required memory/stream settings via the
appropriate Table APIs rather than using concatenate.
- Around line 180-209: The function whose docstring begins "Adjust flat
OrderScheme boundaries using contiguous partition ownership" currently omits a
Returns section; update that docstring to add a Returns entry documenting that
the function returns None (e.g., "Returns\n--------\nNone\n This function
performs in-place adjustments and does not return a value.") so callers reading
the docstring for parameters like input_scheme, output_scheme, and collective_id
have a clear API contract.
In `@python/cudf_polars/tests/streaming/test_adjust_orderscheme.py`:
- Around line 318-320: The loop variable name keys in the test loop (for pid,
keys in expected.items()) shadows the outer test fixture/variable keys; rename
the loop variable to expected_keys (or another unambiguous name) and update the
two assertions to compare against expected_keys (output[pid]["key"].to_list() ==
expected_keys and output[pid]["val"].to_list() == expected_keys); apply the same
rename for the identical pattern later in the file (the second loop around line
353) so there is no shadowing.
- Around line 289-355: Add two new pytest cases that exercise
_adjust_and_collect with degenerate inputs: one where the input frame produced
by _frame contains columns that are all null (e.g., "key" or "val" series
all-null) and one where the input has a single-element Series; reuse the
existing patterns from test_adjust_orderscheme_emits_empty_owned_partitions and
test_adjust_orderscheme_single_rank_no_collective to create context, comm,
stream, input_scheme via _make_scheme and call _adjust_and_collect, then assert
outputs by pid as in the existing tests; ensure you parametrize/skip for nranks
like the other tests and reference the helper functions _frame, _make_scheme and
_adjust_and_collect so the new tests integrate with the current test
scaffolding.
🪄 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: 0fee68f4-d44e-47a4-90b1-08a10f9beff4
📒 Files selected for processing (2)
python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.pypython/cudf_polars/tests/streaming/test_adjust_orderscheme.py
rjzamora
added
3 - Ready for Review
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
python/cudf_polars/tests/streaming/test_adjust_orderscheme.py (1)
305-341: ⚡ Quick winAdd null and single-row coverage before merge.
The new suite exercises empty inputs and chunk accumulation, but it never hits all-null or single-element partitions. Since
_make_schemefixesNullOrder.BEFORE, a regression around null boundary placement or degenerate one-row inputs would slip through here.As per coding guidelines,
python/**/test_*.py: "Ensure test files provide comprehensive edge case coverage (empty, all-null, single-element, mixed types)".Also applies to: 376-396
🤖 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_adjust_orderscheme.py` around lines 305 - 341, The test test_adjust_orderscheme_all_empty_input only covers empty inputs; add cases that exercise all-null partitions and single-row partitions to catch regressions in null ordering and degenerate partitions: create additional calls (or parametrize) to _adjust_and_collect using _frame with rows that are all null for the key column(s) and with single-row frames, using the same input_scheme/output_scheme setup (which uses _make_scheme and fixes NullOrder.BEFORE) and run both the single-rank branch and the collective branch (reserve_op_id / collective_id) as done currently; assert results with _assert_partition_output to verify behavior for all-null and single-element partitions.
🤖 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/tests/streaming/test_adjust_orderscheme.py`:
- Around line 157-164: The helper _assert_partition_output currently only checks
that "val" equals "key" which won't catch payload misalignment; change
tests/fixtures to produce distinct non-key payloads (make val different from
key) and update _assert_partition_output (and its callers) to assert keys and
vals independently (e.g., expected should include per-partition key lists and
per-partition val lists, and _assert_partition_output should compare
output[pid]["key"].to_list() to expected_keys and output[pid]["val"].to_list()
to expected_vals). Locate the helper by name (_assert_partition_output) and
update any fixtures or expected dicts used by tests to provide independent val
data so row/column misalignment is detected.
---
Nitpick comments:
In `@python/cudf_polars/tests/streaming/test_adjust_orderscheme.py`:
- Around line 305-341: The test test_adjust_orderscheme_all_empty_input only
covers empty inputs; add cases that exercise all-null partitions and single-row
partitions to catch regressions in null ordering and degenerate partitions:
create additional calls (or parametrize) to _adjust_and_collect using _frame
with rows that are all null for the key column(s) and with single-row frames,
using the same input_scheme/output_scheme setup (which uses _make_scheme and
fixes NullOrder.BEFORE) and run both the single-rank branch and the collective
branch (reserve_op_id / collective_id) as done currently; assert results with
_assert_partition_output to verify behavior for all-null and single-element
partitions.
🪄 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: 76ca8654-f462-49d1-aaaf-ea2fcf7e62c6
📒 Files selected for processing (2)
python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.pypython/cudf_polars/tests/streaming/test_adjust_orderscheme.py
🚧 Files skipped from review as they are similar to previous changes (1)
- python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py
rjzamora
added
3 - Ready for Review
rjzamora
changed the title
adjust_ordering utilityadjust_ordering utility
|
Could you give a bit more context (either here, or preferably in the |
|
@TomAugspurger - Added more info to the PR description, but I can revise and/or add info to the code/comments if necessary. |
TomAugspurger
left a comment
TomAugspurger
left a comment
There was a problem hiding this comment.
About halfway through reviewing, moving onto the collective part next.
| if comm.nranks > 1 and collective_id is None: | ||
| raise ValueError("collective_id is required when comm.nranks > 1.") | ||
|
|
||
| try: |
There was a problem hiding this comment.
Nitpick: this try/except stuff could be put in an async context manager, similar to shutdown_on_error
@asynccontextmanager
async def cleanup_on_error(...):
try:
yield
except Exception:
await gather_in_task_group(...)
raise
and _adjust_ordering_impl could be inlined here.
There was a problem hiding this comment.
Okay, I think I understand. I'm updating utils.py a bit so we have re-usable primitives for cleaning up channels when we are doing something like this at a finer granularity than an actor.
| [_PID_DTYPE], | ||
| stream, | ||
| ) | ||
| .to_polars()["split"] |
There was a problem hiding this comment.
Just confirming: this effectively introduces a stream synchronization here, right? So that we can get the list[int] split points back on the host?
Same with def _boundary_search_positions. Might be worth documenting that.
TomAugspurger
left a comment
TomAugspurger
left a comment
There was a problem hiding this comment.
One high-level question I have: why doesn't this return anything?
IIUC, the way it operates is to takes some data that's flowing through some Channels. And we'll somehow transform the messages sent on those Channels to honor the desired output_ordering
Will the caller need to know that the result is now ordered by output_ordering?
| while not self.input_done and not any( | ||
| pending_pid >= stop for pending_pid in self.pending | ||
| ): |
There was a problem hiding this comment.
I was a little worried about an infinite loop here since the loop body doesn't modify any of the loop conditions in some cases, but I believe that _store_chunk mutates self.pending, which eventually makes the any(...) False.
|
Thanks for the review @TomAugspurger ! I think I've addressed most comments/suggestions at this point.
The intention is for this utility to be used within an actor (e.g. We don't "return" anything, because we are requiring the properly-ordered data to be emitted from the output channel. If the ordering cannot be adjusted, there is a bug in the code, and we should have never called this utility to begin with.
The caller is guaranteed that the data flowing out of I suppose we could call this |
|
Makes sense, thanks. |
rjzamora
added
5 - Ready to Merge
|
/merge |
3 participants
Description
OrderSchememetadata, strict sort boundaries, and sort-aware execution #22128Depends on AddOrderScheme.get_boundariesAPI rapidsmpf#1039SparseAlltoAllto enforce a newOrderingfor a channel that is already partitioned with a compatibleOrdering- The utility will typically be used to adjust boundaries or convert from non-strict to strict boundaries.This utility is intended for operators that can exploit existing ordered/range-partitioned input, but require a different concrete boundary layout before they can do so safely. For example, a downstream operator may need to adjust from one strict
Orderingto another with different boundaries, or convert metadata that is ordered but non-strict into strict output partitions before using chunkwise execution.The immediate motivation is ordered join/groupby/sort planning in cudf-polars: once a stream is known to be ordered, we want to repartition only the boundary-overlap regions needed to align with the target operator, rather than falling back to a full shuffle or sort.
NOTE: "Strict" partitioning means that a unique value may only exist in one chunk. It is possible for the data to be ordered without "strict" partitioning, but we need to enforce strictness before doing a sort-based join or groupby.
Simple example: aligning ordered join inputs
Suppose two input streams are both ordered on the join keys, but their partition boundaries do not line up. A chunkwise join can only be used safely when corresponding output partitions cover the same key ranges.
adjust_orderingprovides the data-movement primitive for reshaping one ordered stream to match the other stream's strict boundaries, moving only the boundary-overlap pieces that need to change ranks.This lets a downstream join operate partition-by-partition without requiring a full hash shuffle or global sort.
The same primitive can also be used by future groupby/sort optimizations that need to turn ordered-but-misaligned or non-strict partitioning into "strict" operator-ready partitions.
Checklist