Improve chunk sampling in dynamic actors by rjzamora · Pull Request #23051 · NVIDIA/cudf

rjzamora · 2026-06-30T21:08:42Z

Description

We currently un-spill the chunks that we sample for dynamic-planning purposes in both join_actor and sort_actor. This contributes unnecessary memory pressure. When we sample a chunk, we only need to know the size and row count, which are both available on the TableChunk (regardless of its CPU/GPU residency). This PR updates the _sample_chunks utility to leverage the existing ChunkStore utility, and avoid manual un-spilling.
We currently have redundant sampling infrastructure for join_actor and sort_actor. This PR updates these actors to use the same _sample_chunks utility.
We currently use sample_chunk_count to limit the number of chunks we are willing to sample for dynamic planning purposes, even if the first 1-2 chunks are empty. This PR updates the meaning of sample_chunk_count so that we sample until the buffered byte size exceeds target_partition_size * sample_chunk_count (or the channel is drained). This improves the quality of our sampling statistics. (EDIT: After benchmarking, I decided to roll back this change - for now)

Together, these changes make dynamic sampling more stable and effective. In the case that we are dealing with skewed data, it also allows us to increase sample_chunk_count to improve the plan without adding unnecessary memory pressure.

Checklist

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

coderabbitai · 2026-06-30T21:11:34Z

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

📝 Walkthrough

Walkthrough

This PR switches streaming chunk buffering and sampling to ChunkStore, updates sort and join strategy code to pass that container through, adds a multi-partition join fallback for expression-based keys, and adjusts the related config doc and streaming join test.

Changes

ChunkStore Sampling and Replay Refactor

Layer / File(s)	Summary
ChunkStore container and sampling/replay core logic `python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py`	`TableSizeStats.chunks` changes to `ChunkStore`; `_sample_chunks` changes its inputs to `sample_chunk_count` and `target_partition_size`, uses a sequential `ChunkStore`-backed sampling loop with byte-budget stopping, extrapolates totals from sampled data, and `replay_buffered_channel` iterates `ChunkStore` messages directly.
Sort actor size estimation via ChunkStore `python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py`	`_sample_chunks_for_size_estimate` delegates to `_sample_chunks`, uses `sample.total_size` for the size estimate, returns `tuple[ChunkStore, int]`, and computes `num_partitions` with ceil division.
Join chunkwise strategy uses ChunkStore `python/cudf_polars/cudf_polars/streaming/actor_graph/join.py`	The chunkwise-aligned branch of `_choose_strategy` now initializes `left_sample.chunks` and `right_sample.chunks` with `ChunkStore(context)` and sets `total_chunks` from channel `local_count`.
Sample chunk count documentation update `python/cudf_polars/cudf_polars/utils/config.py`	`DynamicPlanningOptions.sample_chunk_count` documentation is rewritten to describe it as a multiplier deriving the sampling budget from target partition size.

Join fallback for expression keys

Layer / File(s)	Summary
Multi-partition join expression guard `python/cudf_polars/cudf_polars/streaming/join.py`	`Join` lowering now imports `Col` and falls back when either join-key side contains non-`Col` expressions in the multi-partition path.
Expression-key join warning test `python/cudf_polars/tests/streaming/test_join.py`	The computed-expression join test now uses warning fallback mode and asserts the multi-partition join warning around the GPU result check.

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

Possibly related PRs

rapidsai/cudf#23024: Also changes streaming join behavior for expression-based join keys, with related multi-partition join handling.
rapidsai/cudf#22318: Also changes streaming join strategy logic in python/cudf_polars/cudf_polars/streaming/actor_graph/join.py, covering a related multi-partition join path.

Suggested reviewers: vyasr, madsbk, Matt711

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately summarizes the main change: improving chunk sampling for dynamic actor planning.
Docstring Coverage	✅ Passed	Docstring coverage is 85.71% which is sufficient. The required threshold is 80.00%.
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.
Description check	✅ Passed	The description matches the changeset, covering shared chunk sampling, ChunkStore adoption, and dynamic-planning sampling behavior.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

coderabbitai

Actionable comments posted: 1

🤖 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`:
- Line 254: The partition count calculation in sort.py underestimates the number
of dynamic sort partitions because `global_size // target_partition_size` floors
non-exact multiples. Update the logic in the partition-sizing code that assigns
`num_partitions` to use ceiling division instead, so values like `1.9 *
target_partition_size` produce 2 partitions rather than 1. Keep the existing
`max(1, ...)` guard, but change the division behavior in the `num_partitions`
expression to preserve the target-size budget.

🪄 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: 552a10c8-caeb-423d-b16e-eb51ac442b4f

📥 Commits

Reviewing files that changed from the base of the PR and between aa8cfca and 199fe90.

📒 Files selected for processing (4)

python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py
python/cudf_polars/cudf_polars/streaming/actor_graph/join.py
python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py
python/cudf_polars/cudf_polars/utils/config.py

Matt711 · 2026-07-01T00:02:18Z

Nice! One potential real test failure to look at

FAILED tests/streaming/test_join.py::test_join_computed_expr_right_key[ray] - AssertionError: DataFrames are different (height (row count) mismatch)
[left]: 1024
[right]: 0

TomAugspurger

Looks like a nice cleanup.

rjzamora · 2026-07-01T15:15:07Z

One potential real test failure to look at

Thanks for calling that out @Matt711 - Looks like this change may have moved the test_join_computed_expr_right_key test from a broadcast join to a shuffle join (and thus exposed an existing bug). Looking into it now.

Matt711 · 2026-07-01T15:28:59Z

One potential real test failure to look at

Thanks for calling that out @Matt711 - Looks like this change may have moved the test_join_computed_expr_right_key test from a broadcast join to a shuffle join (and thus exposed an existing bug). Looking into it now.

Thanks, if it ends up being #21641. I'm planning on refreshing my PR soon that closes it. And so you could xfail it.

rjzamora · 2026-07-01T16:10:32Z

+                ir,
+                rec,
+                msg="Multi-partition Join not supported for keys with expressions.",
+            )


@Matt711 - It seems like the test_join_computed_expr_right_key failure was related to the fact that #21692 never landed, yet we have been failing to protect against computed keys in a shuffle join.

The test was previously "missing" the problem, because we were choosing a broadcast join.

The fix I added here is the "simplest" one. However, something like 21692 is definitely better.

OK will prioritize 21692

coderabbitai

Actionable comments posted: 1

🤖 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/join.py`:
- Around line 212-222: The expression-key guard in join lowering is firing too
early in the streaming join path. Update the logic in the join lowering function
around the existing key check on ir.left_on and ir.right_on so it runs only
after output_count and dynamic_planning are known, and only falls back when a
hash shuffle is actually required. Keep single-partition cases that reconstruct
the node directly from being sent to _lower_ir_fallback, and preserve the
existing multi-partition/dynamic-planning fallback behavior for computed keys.

🪄 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: 59892d1e-bdec-4698-8bea-83fa6331caf7

📥 Commits

Reviewing files that changed from the base of the PR and between 8f32b5c and 26ea4ae.

📒 Files selected for processing (2)

python/cudf_polars/cudf_polars/streaming/join.py
python/cudf_polars/tests/streaming/test_join.py

Matt711

Notes:

I checked TPC-DS queries to see if this would cause fallback. We should be good 👍🏾
The cpp-linter failing job will be fixed in #23066

rjzamora · 2026-07-02T19:22:16Z

/merge

rjzamora added 5 commits June 30, 2026 13:11

allow sampled chunks to be spilled

666ce23

require ChunkStore

6ad5b3a

unify sort logic

c6243ee

Merge remote-tracking branch 'upstream/main' into fix-join-sampling

41a43a8

use sample_chunk_count as a multiplier

199fe90

rjzamora self-assigned this Jun 30, 2026

rjzamora requested a review from a team as a code owner June 30, 2026 21:08

rjzamora added the 2 - In Progress Currently a work in progress label Jun 30, 2026

rjzamora requested a review from pentschev June 30, 2026 21:08

rjzamora added improvement Improvement / enhancement to an existing function non-breaking Non-breaking change labels Jun 30, 2026

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

github-project-automation Bot added this to cuDF Python Jun 30, 2026

GPUtester moved this to In Progress in cuDF Python Jun 30, 2026

coderabbitai Bot reviewed Jun 30, 2026

View reviewed changes

Comment thread python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py Outdated

accept coderabbit suggestion

04fbec2

TomAugspurger approved these changes Jul 1, 2026

View reviewed changes

Comment thread python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py

Comment thread python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py Outdated

rjzamora added 2 commits July 1, 2026 07:59

use Tom's suggestions

a3aaf69

Merge remote-tracking branch 'upstream/main' into fix-join-sampling

8f32b5c

fix existing bug exposed by this PR

26ea4ae

rjzamora commented Jul 1, 2026

View reviewed changes

coderabbitai Bot reviewed Jul 1, 2026

View reviewed changes

Comment thread python/cudf_polars/cudf_polars/streaming/join.py Outdated

take coderabbit suggestion

51ec0fc

Matt711 approved these changes Jul 1, 2026

View reviewed changes

partially roll-back sample_chunk_count meaning change

d92399c

rjzamora added 6 commits July 1, 2026 11:21

Merge remote-tracking branch 'upstream/main' into fix-join-sampling

c4f5e12

further revert some changes

98354b0

Merge remote-tracking branch 'upstream/main' into fix-join-sampling

c0742d3

update inject_gpu_engine.py

41ce201

Merge remote-tracking branch 'upstream/main' into fix-join-sampling

256fb5e

Merge branch 'main' into fix-join-sampling

32a0ee8

rjzamora added 5 - Ready to Merge Testing and reviews complete, ready to merge and removed 2 - In Progress Currently a work in progress labels Jul 2, 2026

rapids-bot Bot merged commit 6d1a24d into NVIDIA:main Jul 2, 2026
109 of 110 checks passed

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

rjzamora deleted the fix-join-sampling branch July 2, 2026 19:22

Uh oh!

Conversation

rjzamora commented Jun 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Checklist

Uh oh!

coderabbitai Bot commented Jun 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reviews paused

Walkthrough

Changes

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Matt711 commented Jul 1, 2026

Uh oh!

TomAugspurger left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

rjzamora commented Jul 1, 2026

Uh oh!

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

Uh oh!

Uh oh!

rjzamora Jul 1, 2026

Choose a reason for hiding this comment

Uh oh!

Matt711 Jul 1, 2026

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Matt711 left a comment

Choose a reason for hiding this comment

Uh oh!

rjzamora commented Jul 2, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

rjzamora commented Jun 30, 2026 •

edited

Loading

coderabbitai Bot commented Jun 30, 2026 •

edited

Loading

Matt711 commented Jul 1, 2026 •

edited

Loading