Fix stream ordering bug in join of expression-based keys by wence- · Pull Request #23024 · NVIDIA/cudf

wence- · 2026-06-29T10:42:33Z

Description

When join keys are expressions, and not just simple column references, then evaluation of them might launch new kernels. Since Expr.evaluate does not take a stream (it gets the stream from the input DataFrame) we must launch the evaluation of input keys before the stream_ordered_after context manager. Otherwise, although accessing the to-be-join tables is safe on the join stream, accessing the keys is not: they are on the left and right streams respectively with (potentially) new work queued up.

To fix this, just compute the key columns on their respective table's stream before obtaining a joined stream for the join.

Closes Flaky test_join_on_expression_conditions in cudf-polars-polars tests #22967

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-29T10:47:34Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 64c2cb81-5c7d-4ba0-bcd4-d1ae88952a54

📥 Commits

Reviewing files that changed from the base of the PR and between 78052be and 13f576b.

📒 Files selected for processing (2)

python/cudf_polars/cudf_polars/dsl/ir.py
python/cudf_polars/cudf_polars/streaming/actor_graph/over.py

🚧 Files skipped from review as they are similar to previous changes (2)

python/cudf_polars/cudf_polars/streaming/actor_graph/over.py
python/cudf_polars/cudf_polars/dsl/ir.py

📝 Walkthrough

Summary by CodeRabbit

Bug Fixes
- Improved join handling for decimal columns by applying consistent type casting to both inputs before conditional joins.
- Refined join execution and slicing behavior for cross, semi, anti, and standard joins, including more direct handling for empty inputs.
- Updated streaming synchronization for scalar over() broadcasts to better coordinate work across CUDA streams.

Walkthrough

dsl/ir.py updates decimal cast handling in ConditionalJoin.do_evaluate and restructures Join.do_evaluate branching and result slicing. streaming/actor_graph/over.py updates scalar broadcast synchronization to use stream_ordered_after directly between DataFrame streams.

Changes

Join evaluation fixes

Layer / File(s)	Summary
ConditionalJoin decimal cast fix `python/cudf_polars/cudf_polars/dsl/ir.py`	Decimal cast targets are applied to `left` and `right` before `conditional_inner_join` is called.
Join branching and slicing `python/cudf_polars/cudf_polars/dsl/ir.py`	`Join.do_evaluate` unpacks options early, handles `"Cross"` in its own stream-ordered branch, evaluates non-Cross keys outside that block, returns early for semi/anti joins, and slices results directly in each path.

Broadcast stream synchronization

Layer / File(s)	Summary
Broadcast sync helper update `python/cudf_polars/cudf_polars/streaming/actor_graph/over.py`	`_evaluate_ir_broadcast_sync` uses `stream_ordered_after` directly for chunk and aggregate streams, and `_evaluate_broadcast_chunk` matches the updated helper signature.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Suggested reviewers

mroeschke

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Out of Scope Changes check	⚠️ Warning	The `streaming/actor_graph/over.py` change is unrelated to the join-key flakiness issue in `#22967`.	Move the `over()` stream-ordering adjustment to a separate PR or add linked-issue context showing why it belongs here.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly describes the main change: fixing a stream-ordering bug for expression-based join keys.
Description check	✅ Passed	The description is directly related to the join stream-ordering fix and explains the problem being addressed.
Linked Issues check	✅ Passed	The join changes evaluate expression-based keys before entering `stream_ordered_after`, matching `#22967`'s requirement.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

pentschev

LGTM, thanks Lawrence!

wence- · 2026-06-29T11:01:08Z

+        left = _apply_casts(left, left_casts)
+        right = _apply_casts(right, right_casts)


This potentially delivers new columns on left's (resp. right) stream, so can't be inside the context manager.

wence- · 2026-06-29T11:02:28Z

+            left_on = DataFrame(
+                broadcast(
+                    *(e.evaluate(left) for e in left_on_exprs), stream=left.stream
+                ),
+                stream=left.stream,
+            )
+            right_on = DataFrame(
+                broadcast(
+                    *(e.evaluate(right) for e in right_on_exprs), stream=right.stream
+                ),
+                stream=right.stream,
+            )


expr.evaluate runs on the input DataFrame's stream, so we now run it on the respective side's stream before getting a joined stream to actually do the join on.

wence- · 2026-06-29T11:41:57Z

+    # global_agg_df and chunk_df may live on different streams. Since we do
+    # an evaluation of values in chunk_df via Expr.evaluate, run the
+    # broadcast on chunk_dfs stream, making sure the global_agg stream
+    # waits.
+    with stream_ordered_after(
+        lambda: chunk_df.stream, upstreams=[global_agg_df.stream]
+    ) as stream:
        result_cols = [
            _broadcast_gw_sync(
                ne.value, chunk_df, global_agg_df, key_names, stream
            ).rename(ne.name)
            if isinstance(ne.value, GroupedWindow)
            else ne.evaluate(chunk_df, context=ExecutionContext.FRAME)


Both _broadcast_gw_sync and ne.evaluate return some expressions on chunk_df.stream, so we can't run those on a stream joined with chunk_df and global_agg_df. Instead, join global_agg_df against chunk's stream and use that.

When join keys are expressions, and not just simple column references, then evaluation of them might launch new kernels. Since Expr.evaluate does not take a stream (it gets the stream from the input DataFrame) we must launch the evaluation of input keys before the stream_ordered_after context manager. Otherwise, although accessing the to-be-join tables is safe on the join stream, accessing the keys is not: they are on the left and right streams respectively with (potentially) new work queued up. To fix this, just compute the key columns on their respective table's stream before obtaining a joined stream for the join. - Closes NVIDIA#22967

The same problem applied, Expr.evaluate always delivers expressions on its input DataFrame's stream.

wence- · 2026-06-29T13:23:19Z

/merge

wence- requested a review from a team as a code owner June 29, 2026 10:42

wence- requested a review from pentschev June 29, 2026 10:42

github-actions Bot assigned wence- Jun 29, 2026

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

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

wence- added bug Something isn't working and removed Python Affects Python cuDF API. cudf-polars Issues specific to cudf-polars labels Jun 29, 2026

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

wence- added the non-breaking Non-breaking change label Jun 29, 2026

wence- mentioned this pull request Jun 29, 2026

[BUG] Root cause and fix flaky failures in Polars CI runs due to the rapidsmpf transition #22307

Open

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

pentschev approved these changes Jun 29, 2026

View reviewed changes

TomAugspurger approved these changes Jun 29, 2026

View reviewed changes

wence- commented Jun 29, 2026

View reviewed changes

wence- added 2 commits June 29, 2026 13:09

Fix stream ordering issue in window execution

13f576b

The same problem applied, Expr.evaluate always delivers expressions on its input DataFrame's stream.

wence- force-pushed the wence/fix/join-expr-stream-ordering branch from 78052be to 13f576b Compare June 29, 2026 12:09

rapids-bot Bot merged commit 09b988f into NVIDIA:main Jun 29, 2026
108 checks passed

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

wence- deleted the wence/fix/join-expr-stream-ordering branch June 29, 2026 13:23

coderabbitai Bot mentioned this pull request Jul 1, 2026

Improve chunk sampling in dynamic actors #23051

Merged

3 tasks

TomAugspurger mentioned this pull request Jul 2, 2026

Flaky test_shuffle_group_by_reseed in cudf-polars-polars tests #22964

Open

Navigation Menu

Uh oh!

Fix stream ordering bug in join of expression-based keys - #23024

Fix stream ordering bug in join of expression-based keys#23024
rapids-bot[bot] merged 2 commits into
NVIDIA:mainfrom
wence-:wence/fix/join-expr-stream-ordering

wence- commented Jun 29, 2026

Uh oh!

coderabbitai Bot commented Jun 29, 2026 •

edited

Loading

Summary by CodeRabbit

Walkthrough

Changes

Estimated code review effort

Suggested reviewers

❌ Failed checks (1 warning)

Uh oh!

pentschev left a comment

Uh oh!

wence- Jun 29, 2026

Uh oh!

wence- Jun 29, 2026

Uh oh!

wence- Jun 29, 2026

Uh oh!

wence- commented Jun 29, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		left = _apply_casts(left, left_casts)
		right = _apply_casts(right, right_casts)

Uh oh!

Conversation

wence- commented Jun 29, 2026

Description

Checklist

Uh oh!

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

Uh oh!

Summary by CodeRabbit

Walkthrough

Changes

Estimated code review effort

Suggested reviewers

❌ Failed checks (1 warning)

Uh oh!

pentschev left a comment

Choose a reason for hiding this comment

Uh oh!

wence- Jun 29, 2026

Choose a reason for hiding this comment

Uh oh!

wence- Jun 29, 2026

Choose a reason for hiding this comment

Uh oh!

wence- Jun 29, 2026

Choose a reason for hiding this comment

Uh oh!

wence- commented Jun 29, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

coderabbitai Bot commented Jun 29, 2026 •

edited

Loading