Implement streaming window functions in cudf-polars by Matt711 · Pull Request #22191 · NVIDIA/cudf

Matt711 · 2026-04-17T14:10:49Z

Description

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.

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 [FEA] Support .over() in streaming cuDF-Polars #22047
Closes [BUG] Q67 OOMs #22235
Depends on Add LocalRepartitioner utility #22439
Contributes to [Story] Eliminate fallbacks (single partition) for PDS-DS in the rapidsmpf streaming executor #21749 and [Story][FEA] Streaming ordered-window execution in cuDF-Polars #22032

Checklist

I am familiar with the Contributing Guidelines.
New or existing tests cover these changes.
The documentation is up to date with these changes.

copy-pr-bot · 2026-04-17T14:10:53Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Matt711

Review Guide: I recommend reading the PR description first, then looking over the tests. Then looking at the three execution strategies (and the corresponding tests). Finally, the full-order preservation logic.

If you think we should split up the PR, that's okay. Additionally if you think logic should be shared (especially in the scalar aggs - groupby case), we can discuss what specifically in your review. I abstracted some logic into a helper function like _make_hash_shuffle_metadata but in general I avoided it (in the groupby case) because it made it more difficult to understand IMO.

Matt711 · 2026-04-17T14:18:24Z

+    [False, True],
+    ids=["same_rank", "cross_rank"],
+)
+def test_over_multirank(


I tested this using rrun

Details

(rapids) coder ➜ ~/cudf $ rrun -n 2 python -m pytest python/cudf_polars/tests/experimental/test_spmd.py::test_over_multirank -x -v [rrun] All ranks launched. Waiting for completion... ============================= test session starts ============================== platform linux -- Python 3.14.4, pytest-9.0.3, pluggy-1.6.0 -- /home/coder/.conda/envs/rapids/bin/python ============================= test session starts ============================== platform linux -- Python 3.14.4, pytest-9.0.3, pluggy-1.6.0 -- /home/coder/.conda/envs/rapids/bin/python cachedir: .pytest_cache hypothesis profile 'default' benchmark: 5.2.3 (defaults: timer=time.perf_counter disable_gc=False min_rounds=5 min_time=0.000005 max_time=1.0 calibration_precision=10 warmup=False warmup_iterations=100000) rootdir: /home/coder/cudf/python/cudf_polars configfile: pyproject.toml plugins: cases-3.10.1, anyio-4.13.0, hypothesis-6.151.13, cov-7.1.0, xdist-3.8.0, benchmark-5.2.3, pytest_httpserver-1.1.5, rerunfailures-16.1 cachedir: .pytest_cache hypothesis profile 'default' benchmark: 5.2.3 (defaults: timer=time.perf_counter disable_gc=False min_rounds=5 min_time=0.000005 max_time=1.0 calibration_precision=10 warmup=False warmup_iterations=100000) rootdir: /home/coder/cudf/python/cudf_polars configfile: pyproject.toml plugins: cases-3.10.1, anyio-4.13.0, hypothesis-6.151.13, cov-7.1.0, xdist-3.8.0, benchmark-5.2.3, pytest_httpserver-1.1.5, rerunfailures-16.1 collecting ... collected 4 items collecting ... collected 4 items python/cudf_polars/tests/experimental/test_spmd.py::test_over_multirank[same_rank-scalar_sum] PASSED [ 25%] python/cudf_polars/tests/experimental/test_spmd.py::test_over_multirank[same_rank-scalar_sum] PASSED [ 25%] python/cudf_polars/tests/experimental/test_spmd.py::test_over_multirank[same_rank-nonscalar_rank] PASSED [ 50%] python/cudf_polars/tests/experimental/test_spmd.py::test_over_multirank[same_rank-nonscalar_rank] PASSED [ 50%] python/cudf_polars/tests/experimental/test_spmd.py::test_over_multirank[cross_rank-scalar_sum] PASSED [ 75%] python/cudf_polars/tests/experimental/test_spmd.py::test_over_multirank[cross_rank-scalar_sum] PASSED [ 75%] python/cudf_polars/tests/experimental/test_spmd.py::test_over_multirank[cross_rank-nonscalar_rank] XFAIL [100%] python/cudf_polars/tests/experimental/test_spmd.py::test_over_multirank[cross_rank-nonscalar_rank] XFAIL [100%] ========================= 3 passed, 1 xfailed in 3.45s ========================= ========================= 3 passed, 1 xfailed in 3.46s =========================

The streaming_engine fixture runs on multiple GPUs I think, so you can just use streaming_engine as a fixture here. Or do you need to explicitly run on two ranks with spmd so you can check the local "bit"?

question: Does this run multi-rank in CI?

Or do you need to explicitly run on two ranks with spmd so you can check the local "bit"?

Yes

question: Does this run multi-rank in CI?

No. I checked for rrun in CI jobs and didn't see anything. I'll file an issue about multi-rank in CI. But I'm hesitant because we have enough problems as-is. Some of which was improved by #22493

Matt711 · 2026-04-17T15:10:19Z

Moved this to WIP while I work throught the rapidsmpf test failures: https://github.com/rapidsai/cudf/actions/runs/24570156132/job/71845610400?pr=22191

Matt711 · 2026-04-17T19:43:11Z

/ok to test 0d6d6b1

rjzamora

This is really coming along @Matt711 . Thank you for all the great work here! I only have a few thoughts/suggestion this time around.

Just a note that I looked carefully at the changes outside of the new over.py files. As far as I'm concerned, the changes outside of these files look good, and so the PR is relatively "safe" to merge at this point (I don't expect collateral damage outside the over() code path(s)).

rjzamora

I think these will be my last suggestions. I'll approve once these changes are made (or I am corrected).

rjzamora

I think I'm happy with this once CI is green. Thanks @Matt711 !

Matt711 · 2026-05-14T17:27:20Z

/ok to test d2a0a60

wence-

I have only small further suggestions, so approving on the assumption they are easy fixes.

wence- · 2026-05-15T10:17:12Z

+    [False, True],
+    ids=["same_rank", "cross_rank"],
+)
+def test_over_multirank(


The streaming_engine fixture runs on multiple GPUs I think, so you can just use streaming_engine as a fixture here. Or do you need to explicitly run on two ranks with spmd so you can check the local "bit"?

question: Does this run multi-rank in CI?

wence- · 2026-05-15T10:18:56Z

+    # Small max_rows_per_partition forces many chunks, exercising the AllGather
+    # (scalar broadcast) and sort-and-split (non-scalar) paths across many
+    # partitions. Two values cover both single-row chunks and multi-row chunks
+    # so the within-chunk position sort is also exercised.
+    df = pl.LazyFrame(


question: How is the statement of this comment actually enforced by the test?

It actually wasn't. I switched to use streaming_engine_factory and parametrized it on max_rows_per_partition. So now we actually get two streaming engines.

wence- · 2026-05-15T10:31:53Z

+        max(comm.nranks, total_bytes // max(1, target_partition_size)),
+        max(1, total_count),
+    )
+    return sample, modulus


OK, let's leave it as an issue/todo.

wence- · 2026-05-15T10:45:36Z

+    # TODO: thread ir_context through repartition_by_index so each
+    # PackedData piece moves on its own pool stream rather than sharing one.
+    local = LocalRepartitioner(return_shuffle, local_count=n_chunks)
+    await local.repartition_by_index(
+        partition_col=chunk_index_column, stream=ir_context.get_cuda_stream()
+    )


This todo seems out of date?

No LocalRepartitioner.repartition_by_index itself still takes a single stream. I was using a single stream for each iteration in the loop before, and that is no longer the case.

wence- · 2026-05-15T10:57:17Z

Thanks for all the work here @Matt711!

rjzamora · 2026-05-15T12:38:17Z

Just a note that we (probably) want to retarget this to release/26.06

…e in tests

Matt711 · 2026-05-15T14:28:10Z

/ok to test 8162194

Matt711 · 2026-05-15T18:43:59Z

/merge

Matt711 · 2026-05-18T01:19:28Z

/ok to test bb0a75f

`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

Matt711 added 6 commits April 17, 2026 01:58

Support streaming over expression

5861da9

Add fast path for scalar aggs case, support preserving the full oreder

66be8e9

add more tests, protect against OOMs when preserving the full order

3b15601

check style

6b91f1b

add pre-shuffle test, and clarifying comments

ac73ede

add multi-rank test

0c8e3de

Matt711 requested a review from a team as a code owner April 17, 2026 14:10

Matt711 requested review from galipremsagar and vyasr April 17, 2026 14:10

Matt711 added feature request New feature or request non-breaking Non-breaking change labels Apr 17, 2026

github-actions Bot assigned Matt711 Apr 17, 2026

github-actions Bot added Python Affects Python cuDF API. cudf-polars Issues specific to cudf-polars labels Apr 17, 2026

github-project-automation Bot added this to cuDF Python Apr 17, 2026

GPUtester moved this to In Progress in cuDF Python Apr 17, 2026

Matt711 commented Apr 17, 2026

View reviewed changes

style

696bbc9

Matt711 force-pushed the fea/polars/streaming-over branch from 594810d to 696bbc9 Compare April 17, 2026 14:23

Merge branch 'main' into fea/polars/streaming-over

fdc40c9

Matt711 changed the title ~~Implement streaming window functions in cudf-polars~~ [WIP] Implement streaming window functions in cudf-polars Apr 17, 2026

Matt711 added 2 commits April 17, 2026 16:39

reject non decomposable aggregations

e477439

fix xfail condition for mg test

4a6eb6d

Matt711 requested a review from a team as a code owner April 17, 2026 17:03

Merge branch 'main' into fea/polars/streaming-over

0c48e1d

Matt711 removed the request for review from galipremsagar April 17, 2026 17:04

Merge branch 'main' into fea/polars/streaming-over

0d6d6b1

rjzamora reviewed May 13, 2026

View reviewed changes

Matt711 and others added 3 commits May 13, 2026 19:12

zip local partition ids with sequence numbers in the over emit loop

9491a9e

add a TODO for per stream work

92501aa

Merge branch 'main' into fea/polars/streaming-over

cf5b099

Matt711 requested a review from rjzamora May 13, 2026 21:22

distinct stream per chunk

96f24e0

rjzamora reviewed May 14, 2026

View reviewed changes

Comment thread python/cudf_polars/cudf_polars/experimental/rapidsmpf/over.py

Comment thread python/cudf_polars/cudf_polars/experimental/rapidsmpf/over.py Outdated

Comment thread python/cudf_polars/tests/experimental/test_rolling.py Outdated

Matt711 added 3 commits May 14, 2026 15:55

check is strictly partitioned

520badf

exercise both single- and multi-row chunks in test_over_many_partitions

0483abd

merge conflict

d2a0a60

rjzamora approved these changes May 14, 2026

View reviewed changes

wence- approved these changes May 15, 2026

View reviewed changes

Matt711 changed the base branch from main to release/26.06 May 15, 2026 12:44

Matt711 and others added 3 commits May 15, 2026 13:41

address small nits

5b75ae3

prune gw_nodes; tighten bound for reserve_extra; use streaming fixtur…

590317f

…e in tests

Merge branch 'release/26.06' into fea/polars/streaming-over

8162194

Merge branch 'release/26.06' into fea/polars/streaming-over

bb0a75f

rapids-bot Bot merged commit 1a10328 into NVIDIA:release/26.06 May 18, 2026
87 checks passed

github-project-automation Bot moved this from In Progress to Done in cuDF Python May 18, 2026

Matt711 deleted the fea/polars/streaming-over branch May 18, 2026 02:26

This was referenced May 18, 2026

Implement single-rank streaming rolling aggregation #22308

Open

Preserve filter fallback row order #22552

Closed

coderabbitai Bot mentioned this pull request May 19, 2026

Move collectives module #22578

Merged

3 tasks

Uh oh!

Conversation

Matt711 commented Apr 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

The over_actor: three strategies

Preserving full order

Checklist

Uh oh!

copy-pr-bot Bot commented Apr 17, 2026

Uh oh!

Matt711 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Matt711 commented Apr 17, 2026

Uh oh!

Matt711 commented Apr 17, 2026

Uh oh!

rjzamora left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

rjzamora left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

rjzamora left a comment

Choose a reason for hiding this comment

Uh oh!

Matt711 commented May 14, 2026

Uh oh!

wence- left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wence- commented May 15, 2026

Uh oh!

rjzamora commented May 15, 2026

Uh oh!

Matt711 commented May 15, 2026

Uh oh!

Matt711 commented May 15, 2026

Uh oh!

Matt711 commented May 18, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Matt711 commented Apr 17, 2026 •

edited

Loading

The `over_actor`: three strategies

rjzamora left a comment •

edited

Loading