Fix order preservation in grouped-window evaluation

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()

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: