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 d9811a73e620..eb1ef2bbff6a 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.streaming.utils import _contains_input_order_window_without_order_by from cudf_polars.utils.cuda_stream import stream_ordered_after if TYPE_CHECKING: @@ -287,12 +288,25 @@ def _evaluate_window_with_stamps( ir: Over, ir_context: IRExecutionContext, stamps: OriginStamps, + *, + sort_by_input_order: bool, ) -> DataFrame: """Evaluate *ir* on the un-stamped portion of *chunk*; reattach stamps after.""" child_schema = ir.children[0].schema stream = ir_context.get_cuda_stream() - columns = chunk.table_view().columns() n_child = len(child_schema) + table = chunk.table_view() + if sort_by_input_order: + columns = table.columns() + table = plc.sorting.stable_sort_by_key( + table, + # Sort by (rank, chunk_index) + plc.Table([columns[n_child + 2], columns[n_child]]), + [plc.types.Order.ASCENDING] * 2, + [plc.types.NullOrder.AFTER] * 2, + stream=stream, + ) + columns = table.columns() input_df = DataFrame.from_table( plc.Table(columns[:n_child]), @@ -511,6 +525,8 @@ async def _evaluate_and_route_to_origin( return_shuffle: ShuffleManager, num_ranks: int, stamps: OriginStamps, + *, + sort_by_input_order: bool, ) -> None: """Window-evaluate each local forward partition, then ship rows back to their origin.""" async with return_shuffle.inserting() as inserter: @@ -523,7 +539,12 @@ async def _evaluate_and_route_to_origin( extracted, stream, exclusive_view=True, br=context.br() ) evaluated = await ir_context.to_thread( - _evaluate_window_with_stamps, partition, ir, ir_context, stamps + _evaluate_window_with_stamps, + partition, + ir, + ir_context, + stamps, + sort_by_input_order=sort_by_input_order, ) routed, splits = await ir_context.to_thread( _partition_by_origin_rank, evaluated, num_ranks, context.br() @@ -633,6 +654,9 @@ async def _shuffle_and_reassemble( ) ch_replay = context.create_channel() + sort_by_input_order = _contains_input_order_window_without_order_by( + [ne.value for ne in ir.exprs] + ) sequence_numbers, _ = await gather_in_task_group( _distribute_by_group( context, @@ -656,6 +680,7 @@ async def _shuffle_and_reassemble( return_shuffle, comm.nranks, stamps, + sort_by_input_order=sort_by_input_order, ) await _reassemble_input_chunks( context, ch_out, ir_context, return_shuffle, sequence_numbers, ir, tracer diff --git a/python/cudf_polars/cudf_polars/streaming/select.py b/python/cudf_polars/cudf_polars/streaming/select.py index 53196c562e18..a4bb5f2ad0bf 100644 --- a/python/cudf_polars/cudf_polars/streaming/select.py +++ b/python/cudf_polars/cudf_polars/streaming/select.py @@ -25,7 +25,6 @@ from cudf_polars.streaming.over import _fuse_over_nodes from cudf_polars.streaming.repartition import Repartition from cudf_polars.streaming.utils import ( - _contains_input_order_window_without_order_by, _contains_unsupported_fill_strategy, _dynamic_planning_on, _lower_ir_fallback, @@ -414,19 +413,6 @@ def _( ), ) - if rec.state["nranks"] > 1 and _contains_input_order_window_without_order_by( - [e.value for e in ir.exprs] - ): - return _lower_ir_fallback( - ir.reconstruct([child]), - rec, - msg=( - "input-order-sensitive window expressions without order_by are " - "not supported across multiple ranks; falling back to a single " - "partition." - ), - ) - # Fast count optimization - reads parquet metadata only, works regardless of partitioning scan_child: Scan | None = None if Select._is_len_expr(ir.exprs): diff --git a/python/cudf_polars/tests/streaming/test_spmd.py b/python/cudf_polars/tests/streaming/test_spmd.py index 1d50107e663b..8858da773d47 100644 --- a/python/cudf_polars/tests/streaming/test_spmd.py +++ b/python/cudf_polars/tests/streaming/test_spmd.py @@ -563,7 +563,19 @@ def test_over_multirank( assert grp["result"].to_list() == [None, *expected_xs[:-1]] -def test_over_shift_without_order_by_multirank_raises(comm: Communicator) -> None: +@pytest.mark.parametrize( + "expr,expected", + [ + (pl.col("x").shift(1).over("g").alias("result"), "shift"), + (pl.col("x").cum_sum().over("g").alias("result"), "cum_sum"), + ], + ids=["shift", "cum_sum"], +) +def test_over_input_order_without_order_by_multirank( + comm: Communicator, + expr: pl.Expr, + expected: str, +) -> None: with SPMDEngine( comm=comm, executor_options={ @@ -576,18 +588,32 @@ def test_over_shift_without_order_by_multirank_raises(comm: Communicator) -> Non pytest.skip("requires multiple ranks") rank = engine.rank + local_xs = [rank * 3 + 1, rank * 3 + 2, rank * 3 + 3] lf = pl.LazyFrame( { "g": [0, 0, 0], - "x": [rank * 3 + 1, rank * 3 + 2, rank * 3 + 3], + "x": local_xs, } ) - q = lf.select(pl.col("x").shift(1).over("g")) - with pytest.raises( - NotImplementedError, - match=r"input-order-sensitive window expressions without order_by", - ): - q.collect(engine=engine) + local_result = lf.select(pl.col("x"), expr).collect(engine=engine) + assert local_result["x"].to_list() == local_xs + + with reserve_op_id() as op_id: + global_result = allgather_polars_dataframe( + engine=engine, local_df=local_result, op_id=op_id + ).sort("x") + + xs = list(range(1, 3 * engine.nranks + 1)) + assert global_result["x"].to_list() == xs + if expected == "shift": + expected_values = [None, *xs[:-1]] + else: + total = 0 + expected_values = [] + for x in xs: + total += x + expected_values.append(total) + assert global_result["result"].to_list() == expected_values def test_over_nonscalar_duplicated_input(