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 eb1ef2bbff6a..8be6f71bdeda 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py @@ -300,10 +300,15 @@ def _evaluate_window_with_stamps( 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, + plc.Table( + [ + columns[n_child + 2], # origin rank + columns[n_child], # origin chunk index (local) + columns[n_child + 1], # origin row index (in chunk) + ] + ), + [plc.types.Order.ASCENDING] * 3, + [plc.types.NullOrder.AFTER] * 3, stream=stream, ) columns = table.columns() diff --git a/python/cudf_polars/tests/streaming/test_spmd.py b/python/cudf_polars/tests/streaming/test_spmd.py index eb9ef7703b32..326d39421265 100644 --- a/python/cudf_polars/tests/streaming/test_spmd.py +++ b/python/cudf_polars/tests/streaming/test_spmd.py @@ -29,9 +29,12 @@ ) from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id from cudf_polars.testing.asserts import assert_gpu_result_equal +from cudf_polars.testing.io import make_partitioned_source from cudf_polars.utils.config import MemoryResourceConfig if TYPE_CHECKING: + from pathlib import Path + from rapidsmpf.communicator.communicator import Communicator pytestmark = pytest.mark.spmd @@ -703,6 +706,44 @@ def test_over_shared_group_ordering_multirank( assert global_result["result"].to_list() == expected_values +def test_over_preserves_input_order_within_source_chunk( + comm: Communicator, tmp_path: Path +) -> None: + n_rows = 64 + make_partitioned_source( + pl.DataFrame( + { + "g": [0] * n_rows, + "x": list(range(n_rows)), + } + ), + tmp_path, + "parquet", + row_group_size=2, + ) + + with SPMDEngine( + comm=comm, + executor_options={ + "max_rows_per_partition": 2, + "dynamic_planning": {}, + "fallback_mode": "raise", + }, + ) as engine: + if engine.nranks != 1: + pytest.skip("expected values are defined for exactly 1 rank") + + q = ( + pl.scan_parquet(tmp_path) + .sort("x") + .select( + "x", + pl.col("x").cum_sum().over("g").alias("result"), + ) + ) + assert_gpu_result_equal(q, engine=engine) + + def test_over_nonscalar_duplicated_input( comm: Communicator, ) -> None: