Add LocalRepartitioner utility - #22439
Merged
Merged
Conversation
rjzamora
added
2 - In Progress
github-actions
Bot
added
Python
rjzamora
commented
May 9, 2026
| ) | ||
| ) | ||
|
|
||
| def insert_index(self, chunk: TableChunk, partition_map: TableChunk) -> None: |
Contributor
Author
There was a problem hiding this comment.
@Matt711 - I think we need something like this for the reverse shuffle? We can record the local chunk index ([0,N)) in a column, and then use that column as the index for the reverse shuffle? (see also repartition_by_index)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThis PR extends the RapidsMPF shuffle implementation by adding index-based partition insertion, refactoring extraction APIs, and introducing LocalRepartitioner for local repartitioning workflows. The implementation uses precomputed partition-maps and pylibcudf partitioning to reorder rows, stores communicator state on ShuffleManager, and integrates the new APIs into global shuffle emission. Two SPMD-distributed tests validate hash and index-based local repartitioning. ChangesLocalRepartitioner and Shuffle Extraction Refactoring
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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.
🧹 Nitpick comments (2)
python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/shuffle.py (1)
304-322: 💤 Low valueConsider adding bounds validation for
partition_col.If
partition_colis out of range (>= number of columns),cols[partition_col]will raise an IndexError. While this might be acceptable as a programming error, explicit validation with a clearer error message could improve debuggability.💡 Optional: Add bounds check
async with self._local_shuffle.inserting() as inserter: for table in self._iter_chunks(stream): cols = table.columns() + if partition_col >= len(cols): + raise ValueError( + f"partition_col {partition_col} out of range for table with {len(cols)} columns" + ) payload = plc.Table(🤖 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/experimental/rapidsmpf/collectives/shuffle.py` around lines 304 - 322, Validate partition_col before indexing into cols to avoid an IndexError and provide a clearer error message: check that partition_col is an int and 0 <= partition_col < len(cols) (where cols comes from each table returned by self._iter_chunks(stream)) and raise a ValueError with a descriptive message if it is out of range; perform this check once per table before building partition_map and payload (i.e., before using cols[partition_col] in the inserter.insert_index call that uses TableChunk.from_pylibcudf_table) so failures indicate the invalid partition_col and the table column count.python/cudf_polars/tests/experimental/test_shuffler.py (1)
145-272: 💤 Low valueConsider adding edge case tests for empty and single-row inputs.
The current tests validate the happy path with 12 rows. Per coding guidelines, consider adding tests for:
- Empty DataFrame (0 rows)
- Single-row input
local_countgreater than row countThese can be added as a follow-up if time permits.
🤖 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/experimental/test_shuffler.py` around lines 145 - 272, Add three edge-case variants for both test_local_repartitioner_hash and test_local_repartitioner_index: run each test with an empty DataFrame (0 rows), a single-row DataFrame (1 row), and with local_count greater than the number of rows to ensure partitioning logic handles small inputs; reuse the existing test bodies but change the input pl.DataFrame construction and parametrize or add new tests that call the same logic (referencing test_local_repartitioner_hash, test_local_repartitioner_index, LocalRepartitioner, repartition_by_hash, repartition_by_index, and the allgather_polars_dataframe validation) so you assert preserved routing and global row counts for these edge cases.
🤖 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/experimental/rapidsmpf/collectives/shuffle.py`:
- Around line 304-322: Validate partition_col before indexing into cols to avoid
an IndexError and provide a clearer error message: check that partition_col is
an int and 0 <= partition_col < len(cols) (where cols comes from each table
returned by self._iter_chunks(stream)) and raise a ValueError with a descriptive
message if it is out of range; perform this check once per table before building
partition_map and payload (i.e., before using cols[partition_col] in the
inserter.insert_index call that uses TableChunk.from_pylibcudf_table) so
failures indicate the invalid partition_col and the table column count.
In `@python/cudf_polars/tests/experimental/test_shuffler.py`:
- Around line 145-272: Add three edge-case variants for both
test_local_repartitioner_hash and test_local_repartitioner_index: run each test
with an empty DataFrame (0 rows), a single-row DataFrame (1 row), and with
local_count greater than the number of rows to ensure partitioning logic handles
small inputs; reuse the existing test bodies but change the input pl.DataFrame
construction and parametrize or add new tests that call the same logic
(referencing test_local_repartitioner_hash, test_local_repartitioner_index,
LocalRepartitioner, repartition_by_hash, repartition_by_index, and the
allgather_polars_dataframe validation) so you assert preserved routing and
global row counts for these edge cases.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: b8b4b9b4-1533-4925-b0c8-e08de2463de8
📒 Files selected for processing (2)
python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/shuffle.pypython/cudf_polars/tests/experimental/test_shuffler.py
3 tasks
Matt711
reviewed
May 13, 2026
There was a problem hiding this comment.
🧹 Nitpick comments (1)
python/cudf_polars/cudf_polars/utils/cuda_stream.py (1)
67-109: ⚡ Quick winDocument or enforce owner lifetimes in this helper.
stream_ordered_afteronly retainsCudaStreamLikeobjects, so it does not itself guarantee the deallocation behavior described in the docstring. A caller can still drop the owning dataframe/buffer beforefinallyruns, enqueueing its stream-ordered free before the reverse join. Either keep the upstream owners alive in this utility, or narrow the contract so callers must keep those owners alive until the context exits.As per coding guidelines, `python/**/*.{py,pyx}`: Use-after-free scenarios in device memory handling - Prevent access to freed device memory resources.💡 Minimal fix
- Get a joined CUDA stream with safe stream ordering for deallocation of inputs. + Get a joined CUDA stream ordered after the given upstream streams. @@ - upstreams - The streams being provided to stream-ordered operations. + upstreams + The streams being provided to stream-ordered operations. Callers must + keep the owning objects alive until the context exits; this helper only + establishes stream dependencies. @@ - This context manager provides two useful guarantees when working with - objects holding references to stream-ordered objects: + This context manager provides stream-ordering guarantees. It does not keep + the upstream-owning Python objects alive.🤖 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/utils/cuda_stream.py` around lines 67 - 109, stream_ordered_after currently only holds CudaStreamLike objects so callers can drop the actual owner objects (e.g., DataFrame/buffer) before the context exits, causing their stream-ordered free to be enqueued and leading to use-after-free; fix by adding an explicit mechanism to retain owner lifetimes: modify stream_ordered_after to accept an optional upstream_owners: Sequence[Any] (or rename keep_alive), store that sequence in a local variable referenced for the duration of the context so Python keeps those owners alive, and then proceed to yield the downstream stream and call join_cuda_streams(get_joined_cuda_stream(...)) in the finally block as before; update the docstring to state that callers must pass owners or that the helper will retain them until exit and reference the function names stream_ordered_after, get_joined_cuda_stream, and join_cuda_streams.
🤖 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/utils/cuda_stream.py`:
- Around line 67-109: stream_ordered_after currently only holds CudaStreamLike
objects so callers can drop the actual owner objects (e.g., DataFrame/buffer)
before the context exits, causing their stream-ordered free to be enqueued and
leading to use-after-free; fix by adding an explicit mechanism to retain owner
lifetimes: modify stream_ordered_after to accept an optional upstream_owners:
Sequence[Any] (or rename keep_alive), store that sequence in a local variable
referenced for the duration of the context so Python keeps those owners alive,
and then proceed to yield the downstream stream and call
join_cuda_streams(get_joined_cuda_stream(...)) in the finally block as before;
update the docstring to state that callers must pass owners or that the helper
will retain them until exit and reference the function names
stream_ordered_after, get_joined_cuda_stream, and join_cuda_streams.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 61aa97ae-5c5a-40d6-80a7-14e4b4cc0105
📒 Files selected for processing (3)
python/cudf_polars/cudf_polars/dsl/ir.pypython/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/shuffle.pypython/cudf_polars/cudf_polars/utils/cuda_stream.py
🚧 Files skipped from review as they are similar to previous changes (1)
- python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/shuffle.py
Member
|
/merge |
rjzamora
added
5 - Ready to Merge
rapids-bot Bot
pushed a commit
that referenced
this pull request
May 18, 2026
`over()` is, at its heart, a grouped aggregation followed by a broadcast back to the shape of the input. For each group `g` defined by the partition-by keys, evaluate the expression, then map the result back to every row that belongs to `g`.
```python
import polars as pl
df = pl.LazyFrame(
{
"x": [1, 2, 3, 4, 5, 6],
"g": [1, None, 1, None, 2, 1],
}
)
print(df.select(pl.col("x").sum().over("g")).collect())
```
```
shape: (6, 1)
┌─────┐
│ x │
│ --- │
│ i64 │
╞═════╡
│ 10 │
│ 6 │
│ 10 │
│ 6 │
│ 5 │
│ 10 │
└─────┘
```
Polars represents this with a `WindowMapping` enum. This PR adds support for the `group_to_rows` mapping in the RapidsMPF streaming executor (the variant where the output has the same number of rows as the input and each row receives the value computed for its group). The entry point is a new `over_actor` that selects one of three execution strategies at runtime based on the incoming channel metadata and expression shape.
### The `over_actor`: three strategies
**1. Chunkwise (already partitioned)**
If the incoming channel metadata shows the data is already hash-partitioned on the `over()` keys (or any prefix of them; being partitioned on `('a',)` is sufficient for `over('a', 'b')`, since every group is contained within one rank), the window function is trivially correct on each chunk in isolation. We evaluate chunkwise with no coordination at all.
**2. Scalar aggregations: AllGather + broadcast**
When every `GroupedWindow` in the expression is a scalar aggregation (`sum`, `mean`, `count`, etc.), we exploit the fact that these are decomposable: each worker computes partial aggregates chunkwise, an AllGather collects all workers' partial results, a single reduction produces the global aggregate per group, and then each original chunk has those results broadcast back into its row positions via a hash join on the partition keys.
**3. Non-scalar aggregations: forward-shuffle + return-shuffle**
Functions like `rank` are not decomposable; they require every row in the group to be visible at once. We hash-shuffle by the partition keys so that all rows belonging to group `g` land in the same rank for evaluation. The challenge is then twofold: putting rows back in the right order, *and* getting them back to the rank that owns the corresponding output chunk in the first place. Output channels are rank-local, so only the rank that received an input chunk is wired up to emit it, and the hash shuffle scatters rows by group with no regard for where they originated. We need an explicit return trip.
### Preserving full order
A lot of the implementation exists purely to put output chunks back in the same sequence-number order as the input. Getting this right across both strategies is where most of the complexity lives.
**Scalar aggregation path.** We can't produce any output until the global aggregate is known, so we buffer incoming chunks while simultaneously computing partial aggregates over them. Once the AllGather + final reduction completes, we iterate over the buffer and evaluate each chunk against the global aggregate, emitting results with their original sequence numbers. Order preservation falls out naturally: the buffer is in receive order and we never reorder it.
**Non-scalar shuffle path.** Each row is stamped with three pieces of origin metadata before it enters the forward shuffle: an `origin_rank` (which rank ingested it), a `chunk_index` (a rank-local 0-based counter, *not* the upstream message sequence number, which can collide when the input is the output of a prior shuffle), and a `position` within that input chunk. After the forward shuffle, each rank holds a mix of rows from every origin, but each row knows where it came from. We evaluate the window function on each local forward partition (so `rank` sees every row in the group), then route the results through a *return* shuffle keyed on `origin_rank`. The return shuffle uses `num_partitions = nranks` and `PartitionAssignment.CONTIGUOUS`, so partition `i` lives on rank `i`, and every row goes back to the rank that originally received it. Each rank then sorts the returned rows by `(chunk_index, position)`, splits at chunk-index transitions, drops the stamp columns, and emits one output chunk per input chunk in input order.
To avoid buffering every input chunk just to size the forward shuffle, the actor samples a small number of chunks up front (`_choose_modulus`), AllGathers a size estimate, picks the modulus, and then replays the sampled chunks back through a fresh channel via `replay_buffered_channel`. The forward-insert phase reads from that replay channel and streams rows into the shuffle as they arrive, never holding more than the shuffle's own internal buffering.
- Closes #22047
- Closes #22235
- Depends on #22439
- Contributes to #21749 and #22032
Authors:
- Matthew Murray (https://github.com/Matt711)
- Richard (Rick) Zamora (https://github.com/rjzamora)
Approvers:
- Richard (Rick) Zamora (https://github.com/rjzamora)
- Lawrence Mitchell (https://github.com/wence-)
URL: #22191
madsbk
pushed a commit
to madsbk/cudf
that referenced
this pull request
May 19, 2026
`over()` is, at its heart, a grouped aggregation followed by a broadcast back to the shape of the input. For each group `g` defined by the partition-by keys, evaluate the expression, then map the result back to every row that belongs to `g`.
```python
import polars as pl
df = pl.LazyFrame(
{
"x": [1, 2, 3, 4, 5, 6],
"g": [1, None, 1, None, 2, 1],
}
)
print(df.select(pl.col("x").sum().over("g")).collect())
```
```
shape: (6, 1)
┌─────┐
│ x │
│ --- │
│ i64 │
╞═════╡
│ 10 │
│ 6 │
│ 10 │
│ 6 │
│ 5 │
│ 10 │
└─────┘
```
Polars represents this with a `WindowMapping` enum. This PR adds support for the `group_to_rows` mapping in the RapidsMPF streaming executor (the variant where the output has the same number of rows as the input and each row receives the value computed for its group). The entry point is a new `over_actor` that selects one of three execution strategies at runtime based on the incoming channel metadata and expression shape.
### The `over_actor`: three strategies
**1. Chunkwise (already partitioned)**
If the incoming channel metadata shows the data is already hash-partitioned on the `over()` keys (or any prefix of them; being partitioned on `('a',)` is sufficient for `over('a', 'b')`, since every group is contained within one rank), the window function is trivially correct on each chunk in isolation. We evaluate chunkwise with no coordination at all.
**2. Scalar aggregations: AllGather + broadcast**
When every `GroupedWindow` in the expression is a scalar aggregation (`sum`, `mean`, `count`, etc.), we exploit the fact that these are decomposable: each worker computes partial aggregates chunkwise, an AllGather collects all workers' partial results, a single reduction produces the global aggregate per group, and then each original chunk has those results broadcast back into its row positions via a hash join on the partition keys.
**3. Non-scalar aggregations: forward-shuffle + return-shuffle**
Functions like `rank` are not decomposable; they require every row in the group to be visible at once. We hash-shuffle by the partition keys so that all rows belonging to group `g` land in the same rank for evaluation. The challenge is then twofold: putting rows back in the right order, *and* getting them back to the rank that owns the corresponding output chunk in the first place. Output channels are rank-local, so only the rank that received an input chunk is wired up to emit it, and the hash shuffle scatters rows by group with no regard for where they originated. We need an explicit return trip.
### Preserving full order
A lot of the implementation exists purely to put output chunks back in the same sequence-number order as the input. Getting this right across both strategies is where most of the complexity lives.
**Scalar aggregation path.** We can't produce any output until the global aggregate is known, so we buffer incoming chunks while simultaneously computing partial aggregates over them. Once the AllGather + final reduction completes, we iterate over the buffer and evaluate each chunk against the global aggregate, emitting results with their original sequence numbers. Order preservation falls out naturally: the buffer is in receive order and we never reorder it.
**Non-scalar shuffle path.** Each row is stamped with three pieces of origin metadata before it enters the forward shuffle: an `origin_rank` (which rank ingested it), a `chunk_index` (a rank-local 0-based counter, *not* the upstream message sequence number, which can collide when the input is the output of a prior shuffle), and a `position` within that input chunk. After the forward shuffle, each rank holds a mix of rows from every origin, but each row knows where it came from. We evaluate the window function on each local forward partition (so `rank` sees every row in the group), then route the results through a *return* shuffle keyed on `origin_rank`. The return shuffle uses `num_partitions = nranks` and `PartitionAssignment.CONTIGUOUS`, so partition `i` lives on rank `i`, and every row goes back to the rank that originally received it. Each rank then sorts the returned rows by `(chunk_index, position)`, splits at chunk-index transitions, drops the stamp columns, and emits one output chunk per input chunk in input order.
To avoid buffering every input chunk just to size the forward shuffle, the actor samples a small number of chunks up front (`_choose_modulus`), AllGathers a size estimate, picks the modulus, and then replays the sampled chunks back through a fresh channel via `replay_buffered_channel`. The forward-insert phase reads from that replay channel and streams rows into the shuffle as they arrive, never holding more than the shuffle's own internal buffering.
- Closes NVIDIA#22047
- Closes NVIDIA#22235
- Depends on NVIDIA#22439
- Contributes to NVIDIA#21749 and NVIDIA#22032
Authors:
- Matthew Murray (https://github.com/Matt711)
- Richard (Rick) Zamora (https://github.com/rjzamora)
Approvers:
- Richard (Rick) Zamora (https://github.com/rjzamora)
- Lawrence Mitchell (https://github.com/wence-)
URL: NVIDIA#22191
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
3 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.
Description
Typical Usge*
ShuffleManagerin aLocalRepartitionerafter insertion is finished (and before extracting chunks).LocalRepartitioner.repartition_by_hashorrepartition_by_indexto dictate how the local partition(s) should be re-partitioned locally.LocalRepartitioner.local_partitions()andLocalRepartitioner.extract_chunk()in the same way we do withShuffleManagerChecklist