Improve Ordering reuse for groupby and join

python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py

@@ -701,6 +701,18 @@ def _sort_to_order_keys(ir: Sort) -> list[OrderKey]:
     ]
 
 
+def _can_sort_chunkwise(
+    ordering: Ordering | None, order_keys: Sequence[OrderKey]
+) -> bool:
+    """Return true when ordering avoids a global sort."""
+    if ordering is None:
+        return False
+    keys = tuple(ordering.keys)
+    return keys == tuple(order_keys[: len(keys)]) and (
+        len(keys) == len(order_keys) or ordering.strict_boundaries
+    )
+
+
 def _build_order_scheme(
     context: Context,
     order_keys: list[OrderKey],
@@ -832,10 +844,10 @@ async def sort_actor(
         partitioning = NormalizedPartitioning.from_keys(
             metadata_in.partitioning, comm.nranks, keys=order_keys
         )
-        if partitioning.is_ordered(
-            order_keys,
-            level="local" if metadata_in.duplicated else "flat",
-        ):
+        ordering = partitioning.get_ordering(
+            level="local" if metadata_in.duplicated else "flat"
+        )
+        if _can_sort_chunkwise(ordering, order_keys):
             if tracer is not None:
                 tracer.decision = "already_sorted"
             await chunkwise_evaluate(

python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py

@@ -891,6 +891,13 @@ async def groupby_actor(
         fully_partitioned = partitioning.is_strictly_partitioned(
             level=partitioning_level,
         )
+        input_ordering = None if metadata_in.duplicated else partitioning.get_ordering()
+        can_adjust_preserving_order = (
+            isinstance(ir, GroupBy)
+            and preserves_output_order
+            and input_ordering is not None
+            and not _has_stable_sorted_agg(ir.agg_requests)
+        )
         fallback_case = (
             # NOTE: This criteria means that we fell back
             # to one partition at lowering time.
@@ -936,7 +943,7 @@ async def groupby_actor(
             ir_context,
             ch_in,
             target_partition_size,
-            allow_early_exit=not maintain_order,
+            allow_early_exit=not maintain_order or can_adjust_preserving_order,
         )
 
         skip_global_comm = metadata_in.duplicated or isinstance(
@@ -952,7 +959,7 @@ async def groupby_actor(
             collective_ids,
             target_partition_size,
             skip_global_comm,
-            maintain_order,
+            maintain_order and not can_adjust_preserving_order,
             tracer,
         )
 
@@ -969,11 +976,7 @@ async def groupby_actor(
                 aggregated=aggregated,
                 tracer=tracer,
             )
-        elif not metadata_in.duplicated and partitioning.is_ordered(
-            group_keys,
-            level="flat",
-        ):
-            assert isinstance(partitioning.inter_rank_scheme, OrderScheme)
+        elif input_ordering is not None:
             await _ordered_adjust_reduce(
                 context,
                 comm,
@@ -986,7 +989,7 @@ async def groupby_actor(
                 target_partition_size,
                 aggregated=aggregated,
                 input_drained=input_drained,
-                input_ordering=partitioning.inter_rank_scheme.orderings[0],
+                input_ordering=input_ordering,
                 preserves_output_order=preserves_output_order,
                 tracer=tracer,
             )

python/cudf_polars/cudf_polars/streaming/actor_graph/join.py

@@ -139,6 +139,12 @@ class OrderedJoinStrategy:
 JoinStrategy: TypeAlias = (
     BroadcastJoinStrategy | ShuffleJoinStrategy | OrderedJoinStrategy
 )
+OrderedJoinDecision: TypeAlias = Literal[
+    "ordered_aligned",
+    "ordered_adjust_left",
+    "ordered_adjust_right",
+    "ordered_adjust_both",
+]
 
 
 @define_actor()
@@ -538,6 +544,17 @@ def _make_ordered_strategy(
     if not _ordering_prefix_matches(right_ordering, reference, right_key_indices):
         return None
 
+    left_output_ordering = _update_ordering_indices(reference, left_key_indices)
+    output_ordering = _update_ordering_indices(reference, output_key_indices)
+    if not (
+        left_ordering.locally_ordered
+        and join_preserves_side_order(ir.options[5], "left")
+    ):
+        left_output_ordering = left_output_ordering.with_locally_ordered(
+            locally_ordered=False
+        )
+        output_ordering = output_ordering.with_locally_ordered(locally_ordered=False)
+
     return OrderedJoinStrategy(
         output_indices=output_key_indices,
         left_indices=left_key_indices,
@@ -546,9 +563,11 @@ def _make_ordered_strategy(
         right_keys=right_keys[:reference_key_count],
         left_input_ordering=left_ordering,
         right_input_ordering=right_ordering,
-        left_output_ordering=_update_ordering_indices(reference, left_key_indices),
-        right_output_ordering=_update_ordering_indices(reference, right_key_indices),
-        output_ordering=_update_ordering_indices(reference, output_key_indices),
+        left_output_ordering=left_output_ordering,
+        right_output_ordering=_update_ordering_indices(
+            reference, right_key_indices
+        ).with_locally_ordered(locally_ordered=False),
+        output_ordering=output_ordering,
     )
 
 
@@ -744,6 +763,19 @@ def _local_count_for_ordering(comm: Communicator, ordering: Ordering) -> int:
     return stop - start
 
 
+def _ordered_join_decision(
+    *, left_aligned: bool, right_aligned: bool
+) -> OrderedJoinDecision:
+    """Return the trace decision for ordered join-side alignment."""
+    if left_aligned and right_aligned:
+        return "ordered_aligned"
+    if left_aligned:
+        return "ordered_adjust_right"
+    if right_aligned:
+        return "ordered_adjust_left"
+    return "ordered_adjust_both"
+
+
 async def _adjust_ordered_join_side(
     context: Context,
     comm: Communicator,
@@ -754,21 +786,30 @@ async def _adjust_ordered_join_side(
     input_ordering: Ordering,
     output_ordering: Ordering,
     *,
+    already_aligned: bool,
     collective_id: int,
 ) -> None:
     """Send metadata, then align one join side to output_ordering."""
-    await send_metadata(
-        ch_out,
-        context,
-        ChannelMetadata(
-            local_count=_local_count_for_ordering(comm, output_ordering),
-            partitioning=Partitioning(
-                OrderScheme([output_ordering]),
-                local="inherit",
-            ),
-            duplicated=False,
+    output_metadata = ChannelMetadata(
+        local_count=_local_count_for_ordering(comm, output_ordering),
+        partitioning=Partitioning(
+            OrderScheme([output_ordering]),
+            local="inherit",
         ),
+        duplicated=False,
     )
+    if already_aligned:
+        await replay_buffered_channel(
+            context,
+            ch_out,
+            ch_in,
+            (),
+            output_metadata,
+            trace_ir=schema_ir,
+        )
+        return
+
+    await send_metadata(ch_out, context, output_metadata)
     await adjust_ordering(
         context,
         comm,
@@ -796,6 +837,12 @@ async def _ordered_join(
     tracer: ActorTracer | None,
 ) -> None:
     """Align ordered inputs to common boundaries, then join partition-wise."""
+    left_boundaries_aligned = strategy.left_input_ordering.boundaries_aligned_with(
+        strategy.left_output_ordering, context.br()
+    )
+    right_boundaries_aligned = strategy.right_input_ordering.boundaries_aligned_with(
+        strategy.right_output_ordering, context.br()
+    )
     metadata_out = ChannelMetadata(
         local_count=_local_count_for_ordering(comm, strategy.output_ordering),
         partitioning=Partitioning(
@@ -806,10 +853,26 @@ async def _ordered_join(
     )
     await send_metadata(ch_out, context, metadata_out)
 
-    await gather_in_task_group(
+    left_metadata, right_metadata = await gather_in_task_group(
         recv_metadata(ch_left, context),
         recv_metadata(ch_right, context),
     )
+    left_target_count = _local_count_for_ordering(comm, strategy.left_output_ordering)
+    right_target_count = _local_count_for_ordering(comm, strategy.right_output_ordering)
+    left_aligned = (
+        left_boundaries_aligned
+        and not left_metadata.duplicated
+        and left_metadata.local_count == left_target_count
+    )
+    right_aligned = (
+        right_boundaries_aligned
+        and not right_metadata.duplicated
+        and right_metadata.local_count == right_target_count
+    )
+    if tracer is not None:
+        tracer.decision = _ordered_join_decision(
+            left_aligned=left_aligned, right_aligned=right_aligned
+        )
     ch_left_adjusted = context.create_channel()
     ch_right_adjusted = context.create_channel()
     async with shutdown_on_error(
@@ -829,6 +892,7 @@ async def _ordered_join(
                 ch_left,
                 strategy.left_input_ordering,
                 strategy.left_output_ordering,
+                already_aligned=left_aligned,
                 collective_id=collective_ids.pop(0),
             ),
             _adjust_ordered_join_side(
@@ -840,6 +904,7 @@ async def _ordered_join(
                 ch_right,
                 strategy.right_input_ordering,
                 strategy.right_output_ordering,
+                already_aligned=right_aligned,
                 collective_id=collective_ids.pop(0),
             ),
             _join_chunks(

python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py

@@ -54,6 +54,7 @@
         Callable,
         Coroutine,
         Generator,
+        Iterable,
         Iterator,
         Sequence,
     )
@@ -1278,7 +1279,7 @@ async def replay_buffered_channel(
     context: Context,
     ch_out: Channel[TableChunk],
     ch_in: Channel[TableChunk],
-    buffered_chunks: ChunkStore,
+    buffered_chunks: Iterable[Message],
     metadata: ChannelMetadata,
     *,
     trace_ir: IR,
@@ -1295,7 +1296,7 @@ async def replay_buffered_channel(
     ch_in
         The buffered input channel.
     buffered_chunks
-        The buffered chunks to yield first.
+        Buffered messages to yield first. May be empty.
     metadata
         The metadata to send to the output channel.
     trace_ir
@@ -1358,29 +1359,6 @@ def _scheme_is_strict(scheme: PartitioningScheme) -> bool:
             return scheme.orderings[0].strict_boundaries
         return True
 
-    @staticmethod
-    def _ordering_covers_keys(
-        ordering: Ordering,
-        order_keys: Sequence[int | OrderKey],
-    ) -> bool:
-        """True when an ordering covers the requested sort keys."""
-        ordering_keys = ordering.keys
-        if len(ordering.keys) < len(order_keys):
-            # If we are only sorted on a subset of the keys, we need strict
-            # boundaries to know later keys cannot interleave across chunks.
-            if not ordering.strict_boundaries:
-                return False
-            order_keys = order_keys[: len(ordering.keys)]
-        else:
-            ordering_keys = ordering.keys[: len(order_keys)]
-        for current, target in zip(ordering_keys, order_keys, strict=True):
-            if isinstance(target, OrderKey):
-                if current != target:
-                    return False
-            elif current.column_index != target:
-                return False
-        return True
-
     def is_strictly_partitioned(
         self,
         *,
@@ -1389,17 +1367,10 @@ def is_strictly_partitioned(
         """True if data is strictly partitioned at the requested level."""
         return self._scheme_is_strict(self._scheme_for_level(level))
 
-    def is_ordered(
-        self,
-        order_keys: Sequence[int | OrderKey],
-        *,
-        level: PartitioningLevel = "flat",
-    ) -> bool:
-        """True if the selected ordering covers order_keys."""
+    def get_ordering(self, *, level: PartitioningLevel = "flat") -> Ordering | None:
+        """Return the normalized ordering for the requested partitioning level."""
         scheme = self._scheme_for_level(level)
-        if not isinstance(scheme, OrderScheme):
-            return False
-        return self._ordering_covers_keys(scheme.orderings[0], order_keys)
+        return scheme.orderings[0] if isinstance(scheme, OrderScheme) else None
 
     def is_aligned_with(
         self, other: NormalizedPartitioning, br: BufferResource