diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 30b111e8c8c1..2334b4069533 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -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( @@ -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: @@ -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 @@ -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, + ) + with context.stream_ordered_after(left, right) as stream: null_equality = ( plc.types.NullEquality.EQUAL if nulls_equal @@ -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 ) @@ -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): diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py index ba1799f8daba..0322645d02d1 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py @@ -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 @@ -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 @@ -211,7 +213,6 @@ async def _evaluate_broadcast_chunk( ir, global_agg_df, key_names, - ir_context, context.br(), )