Preserve input order for grouped window evaluation

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

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

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(