Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rapids-bot merged 2 commits into from
Jun 29, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
370eb77
Fix stream ordering bug in join of expression-based keys
wence- Jun 29, 2026
13f576b
Fix stream ordering issue in window execution
wence- Jun 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
67 changes: 32 additions & 35 deletions python/cudf_polars/cudf_polars/dsl/ir.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2176,17 +2176,16 @@ def do_evaluate(
context: IRExecutionContext,
) -> DataFrame:
"""Evaluate and return a dataframe."""
left_casts, right_casts = _collect_decimal_binop_casts(
predicate_wrapper.predicate
)
left_on = _apply_casts(left, left_casts)
right_on = _apply_casts(right, right_casts)
with context.stream_ordered_after(left, right) as stream:
left_casts, right_casts = _collect_decimal_binop_casts(
predicate_wrapper.predicate
)
_, _, zlice, suffix, _, _ = options

lg, rg = plc.join.conditional_inner_join(
_apply_casts(left, left_casts).table,
_apply_casts(right, right_casts).table,
predicate_wrapper.ast,
stream=stream,
left_on.table, right_on.table, predicate_wrapper.ast, stream=stream
)
left_result = DataFrame.from_table(
plc.copying.gather(
Expand Down Expand Up @@ -2469,9 +2468,9 @@ def do_evaluate(
context: IRExecutionContext,
) -> DataFrame:
"""Evaluate and return a dataframe."""
with context.stream_ordered_after(left, right) as stream:
how, nulls_equal, zlice, suffix, coalesce, maintain_order = options
if how == "Cross":
how, nulls_equal, zlice, suffix, coalesce, maintain_order = options
if how == "Cross":
with context.stream_ordered_after(left, right) as stream:
# Separate implementation, since cross_join returns the
# result, not the gather maps
if right.num_rows == 0:
Expand All @@ -2490,7 +2489,9 @@ def do_evaluate(
),
stream=stream,
)
result = DataFrame([*left_cols, *right_cols], stream=stream)
return DataFrame([*left_cols, *right_cols], stream=stream).slice(
zlice
)
else:
columns = plc.join.cross_join(
left.table, right.table, stream=stream
Expand All @@ -2509,25 +2510,25 @@ def do_evaluate(
left=False,
stream=stream,
)
result = DataFrame([*left_cols, *right_cols], stream=stream).slice(
return DataFrame([*left_cols, *right_cols], stream=stream).slice(
zlice
)

else:
# how != "Cross"
# TODO: Waiting on clarity based on https://github.com/pola-rs/polars/issues/17184
left_on = DataFrame(
broadcast(
*(e.evaluate(left) for e in left_on_exprs), stream=stream
),
stream=stream,
)
right_on = DataFrame(
broadcast(
*(e.evaluate(right) for e in right_on_exprs), stream=stream
),
stream=stream,
)
else:
# how != "Cross"
# TODO: Waiting on clarity based on https://github.com/pola-rs/polars/issues/17184
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,
)
Comment on lines +2519 to +2530

@wence- wence- Jun 29, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

with context.stream_ordered_after(left, right) as stream:
null_equality = (
plc.types.NullEquality.EQUAL
if nulls_equal
Expand All @@ -2540,16 +2541,15 @@ def do_evaluate(
table = plc.copying.gather(
left.table, lg, left_policy, stream=stream
)
result = DataFrame.from_table(
return DataFrame.from_table(
table, left.column_names, left.dtypes, stream=stream
)
).slice(zlice)
else:
if how == "Right":
# Right join is a left join with the tables swapped
left, right = right, left
left_on, right_on = right_on, left_on
maintain_order = Join.SWAPPED_ORDER[maintain_order]

lg, rg = join_fn(
left_on.table, right_on.table, null_equality, stream=stream
)
Expand Down Expand Up @@ -2627,10 +2627,7 @@ def do_evaluate(
if name in left.column_names_set
}
)
result = left.with_columns(right.columns, stream=stream)
result = result.slice(zlice)

return result
return left.with_columns(right.columns, stream=stream).slice(zlice)


class HStack(IR):
Expand Down
15 changes: 8 additions & 7 deletions python/cudf_polars/cudf_polars/streaming/actor_graph/over.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,7 @@
shutdown_on_error,
)
from cudf_polars.streaming.over import Over, _build_over_groupby_irs
from cudf_polars.utils.cuda_stream import stream_ordered_after

if TYPE_CHECKING:
from rapidsmpf.communicator.communicator import Communicator
Expand Down Expand Up @@ -161,16 +162,17 @@ def _evaluate_ir_broadcast_sync(
ir: Over,
global_agg_df: DataFrame,
key_names: tuple[str, ...],
ir_context: IRExecutionContext,
br: BufferResource,
) -> TableChunk:
"""Map the per-group aggregate onto a chunk's rows to produce its Over output."""
chunk_df = chunk_to_frame(chunk, ir.children[0])
# global_agg_df and chunk_df may live on different streams (the former from
# the upstream allgather/reduction on ir_context's stream, the latter from
# the input message). Join them so the broadcast kernels read global_agg_df
# safely.
with ir_context.stream_ordered_after(chunk_df, global_agg_df) as stream:
# 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
Expand Down Expand Up @@ -211,7 +213,6 @@ async def _evaluate_broadcast_chunk(
ir,
global_agg_df,
key_names,
ir_context,
context.br(),
)

Expand Down
Loading