align with main

Author: rjzamora

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

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 """OrderScheme adjustment utilities for the RapidsMPF streaming runtime."""
 
@@ -10,12 +10,12 @@
 import polars as pl
 
 import pylibcudf as plc
+from cudf_streaming.partition_utils import unpack_and_concat
+from cudf_streaming.table_chunk import TableChunk
 from pylibcudf.contiguous_split import pack
-from rapidsmpf.integrations.cudf.partition import unpack_and_concat
 from rapidsmpf.memory.packed_data import PackedData
 from rapidsmpf.streaming.coll.sparse_alltoall import SparseAlltoall
 from rapidsmpf.streaming.core.message import Message
-from rapidsmpf.streaming.cudf.table_chunk import TableChunk
 
 from cudf_polars.containers import DataFrame, DataType
 from cudf_polars.streaming.actor_graph.utils import (
@@ -27,11 +27,11 @@
 from cudf_polars.utils.cuda_stream import stream_ordered_after
 
 if TYPE_CHECKING:
+    from cudf_streaming.channel_metadata import OrderScheme, Ordering
     from rapidsmpf.communicator.communicator import Communicator
     from rapidsmpf.memory.buffer_resource import BufferResource
     from rapidsmpf.streaming.core.channel import Channel
     from rapidsmpf.streaming.core.context import Context
-    from rapidsmpf.streaming.cudf.channel_metadata import OrderScheme
     from rmm.pylibrmm.stream import Stream
 
     from cudf_polars.dsl.ir import IR, IRExecutionContext
@@ -41,6 +41,12 @@
 _PID_PLC_DTYPE = plc.DataType(plc.TypeId.INT32)
 
 
+def _primary_ordering(scheme: OrderScheme) -> Ordering:
+    """Return the single ordering supported by adjust_orderscheme for now."""
+    (ordering,) = scheme.orderings
+    return ordering
+
+
 def _contiguous_owner(pid: int, nranks: int, npartitions: int) -> int:
     """Return the rank owning *pid* under contiguous partition assignment."""
     return pid * nranks // npartitions
@@ -81,10 +87,12 @@ def _contiguous_owners(
 
 def _validate_schemes(input_scheme: OrderScheme, output_scheme: OrderScheme) -> None:
     """Validate the first-pass flat OrderScheme adjustment contract."""
-    if not output_scheme.strict_boundaries:
+    input_ordering = _primary_ordering(input_scheme)
+    output_ordering = _primary_ordering(output_scheme)
+    if not output_ordering.strict_boundaries:
         raise ValueError("adjust_orderscheme requires a strict output OrderScheme.")
-    prefix_len = len(output_scheme.keys)
-    if input_scheme.keys[:prefix_len] != output_scheme.keys:
+    prefix_len = len(output_ordering.keys)
+    if input_ordering.keys[:prefix_len] != output_ordering.keys:
         raise NotImplementedError(
             "adjust_orderscheme currently requires the output OrderScheme keys "
             "to be a prefix of the input OrderScheme keys."
@@ -94,18 +102,18 @@ def _validate_schemes(input_scheme: OrderScheme, output_scheme: OrderScheme) ->
 def _split_points(
     table: plc.Table,
     boundary_table: plc.Table,
-    scheme: OrderScheme,
+    ordering: Ordering,
     stream: Stream,
 ) -> list[int]:
     """Return row split points that partition *table* by *scheme* boundaries."""
     if boundary_table.num_rows() == 0:
         return []
-    key_table = plc.Table([table.columns()[key.column_index] for key in scheme.keys])
+    key_table = plc.Table([table.columns()[key.column_index] for key in ordering.keys])
     split_col = plc.search.lower_bound(
         key_table,
         boundary_table,
-        [key.order for key in scheme.keys],
-        [key.null_order for key in scheme.keys],
+        [key.order for key in ordering.keys],
+        [key.null_order for key in ordering.keys],
         stream=stream,
     )
     return (
@@ -133,16 +141,16 @@ def _append_partition_id(table: plc.Table, pid: int, stream: Stream) -> plc.Tabl
 def _boundary_search_positions(
     input_boundary_table: plc.Table,
     output_boundary_table: plc.Table,
-    output_scheme: OrderScheme,
+    output_ordering: Ordering,
     stream: Stream,
 ) -> tuple[list[int], list[int]]:
     """Search output boundary positions for projected input boundary rows."""
     if input_boundary_table.num_rows() == 0:
         return [], []
-    prefix_len = len(output_scheme.keys)
+    prefix_len = len(output_ordering.keys)
     input_prefix_boundaries = plc.Table(input_boundary_table.columns()[:prefix_len])
-    orders = [key.order for key in output_scheme.keys]
-    null_orders = [key.null_order for key in output_scheme.keys]
+    orders = [key.order for key in output_ordering.keys]
+    null_orders = [key.null_order for key in output_ordering.keys]
     lower_col = plc.search.lower_bound(
         output_boundary_table,
         input_prefix_boundaries,
@@ -169,16 +177,16 @@ def _boundary_search_positions(
 def _peer_ranks(
     rank: int,
     nranks: int,
-    input_scheme: OrderScheme,
-    output_scheme: OrderScheme,
+    input_ordering: Ordering,
+    output_ordering: Ordering,
     lower_positions: list[int],
     upper_positions: list[int],
 ) -> tuple[list[int], list[int]]:
     """Return source and destination ranks needed for OrderScheme adjustment."""
-    input_npartitions = input_scheme.num_boundaries + 1
-    output_npartitions = output_scheme.num_boundaries + 1
-    output_prefix_only = len(output_scheme.keys) < len(input_scheme.keys)
-    include_upper_boundary = output_prefix_only or not input_scheme.strict_boundaries
+    input_npartitions = input_ordering.num_boundaries + 1
+    output_npartitions = output_ordering.num_boundaries + 1
+    output_prefix_only = len(output_ordering.keys) < len(input_ordering.keys)
+    include_upper_boundary = output_prefix_only or not input_ordering.strict_boundaries
 
     def dsts_for_source(source_rank: int) -> list[int]:
         input_start, input_stop = _partition_range(
@@ -265,10 +273,10 @@ async def _adjust_orderscheme_local(
     ir_context: IRExecutionContext,
     ch_out: Channel[TableChunk],
     ch_in: Channel[TableChunk],
-    output_scheme: OrderScheme,
+    output_ordering: Ordering,
 ) -> None:
-    npartitions = output_scheme.num_boundaries + 1
-    boundary_chunk = output_scheme.get_boundaries(context.br())
+    npartitions = output_ordering.num_boundaries + 1
+    boundary_chunk = output_ordering.get_boundaries(context.br())
     boundary_table = boundary_chunk.table_view()
     pending_pid: int | None = None
     pending_chunks: ChunkStore | None = None
@@ -294,11 +302,11 @@ async def emit_pending(pid: int) -> None:
         if chunk.table_view().num_rows() == 0:
             continue
         with stream_ordered_after(
-            context.get_stream_from_pool,
+            context.br().stream_pool.get_stream,
             upstreams=(chunk.stream, boundary_chunk.stream),
         ) as stream:
             table = chunk.table_view()
-            splits = _split_points(table, boundary_table, output_scheme, stream)
+            splits = _split_points(table, boundary_table, output_ordering, stream)
             for pid, piece in enumerate(
                 plc.copying.split(table, splits, stream=stream)
             ):
@@ -372,7 +380,9 @@ async def adjust_orderscheme(
     sortedness is not checked here.
     """
     _validate_schemes(input_scheme, output_scheme)
-    npartitions = output_scheme.num_boundaries + 1
+    input_ordering = _primary_ordering(input_scheme)
+    output_ordering = _primary_ordering(output_scheme)
+    npartitions = output_ordering.num_boundaries + 1
     local_pids = _local_partitions(comm.rank, comm.nranks, npartitions)
 
     if comm.nranks > 1 and collective_id is None:
@@ -386,31 +396,31 @@ async def adjust_orderscheme(
                 ir_context,
                 ch_out,
                 ch_in,
-                output_scheme,
+                output_ordering,
             )
             return
 
-        input_boundary_chunk = input_scheme.get_boundaries(context.br())
-        boundary_chunk = output_scheme.get_boundaries(context.br())
+        input_boundary_chunk = input_ordering.get_boundaries(context.br())
+        boundary_chunk = output_ordering.get_boundaries(context.br())
         boundary_table = boundary_chunk.table_view()
         srcs: list[int] = []
         dsts: list[int] = []
         if comm.nranks > 1:
             with stream_ordered_after(
-                context.get_stream_from_pool,
+                context.br().stream_pool.get_stream,
                 upstreams=(input_boundary_chunk.stream, boundary_chunk.stream),
             ) as stream:
                 lower_positions, upper_positions = _boundary_search_positions(
                     input_boundary_chunk.table_view(),
                     boundary_table,
-                    output_scheme,
+                    output_ordering,
                     stream,
                 )
             srcs, dsts = _peer_ranks(
                 comm.rank,
                 comm.nranks,
-                input_scheme,
-                output_scheme,
+                input_ordering,
+                output_ordering,
                 lower_positions,
                 upper_positions,
             )
@@ -429,11 +439,13 @@ async def adjust_orderscheme(
                 if chunk.table_view().num_rows() == 0:
                     continue
                 with stream_ordered_after(
-                    context.get_stream_from_pool,
+                    context.br().stream_pool.get_stream,
                     upstreams=(chunk.stream, boundary_chunk.stream),
                 ) as stream:
                     table = chunk.table_view()
-                    splits = _split_points(table, boundary_table, output_scheme, stream)
+                    splits = _split_points(
+                        table, boundary_table, output_ordering, stream
+                    )
                     for pid, piece in enumerate(
                         plc.copying.split(table, splits, stream=stream)
                     ):
@@ -473,7 +485,7 @@ async def adjust_orderscheme(
                     output_chunks[pid].extend(chunks)
             else:
                 assert exchange is not None
-                stream = context.get_stream_from_pool()
+                stream = context.br().stream_pool.get_stream()
                 for packed in exchange.extract(source_rank):
                     remote_piece = _unpack_remote_piece(packed, stream, context.br())
                     if remote_piece is None:

python/cudf_polars/tests/streaming/test_adjust_orderscheme.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from __future__ import annotations
@@ -12,12 +12,13 @@
 import polars as pl
 
 import pylibcudf as plc
-from rapidsmpf.streaming.core.message import Message
-from rapidsmpf.streaming.cudf.channel_metadata import (
+from cudf_streaming.channel_metadata import (
     OrderKey,
     OrderScheme,
+    Ordering,
 )
-from rapidsmpf.streaming.cudf.table_chunk import TableChunk
+from cudf_streaming.table_chunk import TableChunk
+from rapidsmpf.streaming.core.message import Message
 
 from cudf_polars.containers import DataFrame, DataType
 from cudf_polars.dsl.ir import Empty, IRExecutionContext
@@ -71,16 +72,24 @@ def _make_scheme(
     )
     return OrderScheme(
         [
-            OrderKey(index, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE)
-            for index in key_indices
-        ],
-        TableChunk.from_pylibcudf_table(
-            boundary_df.table,
-            stream,
-            exclusive_view=True,
-            br=context.br(),
-        ),
-        strict_boundaries=strict,
+            Ordering(
+                [
+                    OrderKey(
+                        index,
+                        plc.types.Order.ASCENDING,
+                        plc.types.NullOrder.BEFORE,
+                    )
+                    for index in key_indices
+                ],
+                TableChunk.from_pylibcudf_table(
+                    boundary_df.table,
+                    stream,
+                    exclusive_view=True,
+                    br=context.br(),
+                ),
+                strict_boundaries=strict,
+            )
+        ]
     )
 
 
@@ -118,7 +127,7 @@ async def _adjust_and_collect(
     """Run adjustment and collect output chunks by partition ID."""
     ch_in = context.create_channel()
     ch_out = context.create_channel()
-    stream = context.get_stream_from_pool()
+    stream = context.br().stream_pool.get_stream()
     output: dict[int, pl.DataFrame] = {}
 
     async def _produce() -> None:
@@ -147,7 +156,7 @@ async def _consume() -> None:
 
     with ThreadPoolExecutor(max_workers=1) as executor:
         ir_context = IRExecutionContext(
-            executor, get_cuda_stream=context.get_stream_from_pool
+            executor, get_cuda_stream=context.br().stream_pool.get_stream
         )
         await gather_in_task_group(
             _produce(),
@@ -190,7 +199,7 @@ async def _adjust_direct(
     ch_out = context.create_channel()
     with ThreadPoolExecutor(max_workers=1) as executor:
         ir_context = IRExecutionContext(
-            executor, get_cuda_stream=context.get_stream_from_pool
+            executor, get_cuda_stream=context.br().stream_pool.get_stream
         )
         await adjust_orderscheme(
             context,
@@ -222,7 +231,7 @@ def test_adjust_orderscheme_rejects_invalid_schemes(
     match: str,
 ) -> None:
     context = spmd_engine.context
-    stream = context.get_stream_from_pool()
+    stream = context.br().stream_pool.get_stream()
     input_scheme = _make_scheme(context, 4, key_indices=input_keys, stream=stream)
     output_scheme = _make_scheme(
         context,
@@ -247,7 +256,7 @@ def test_adjust_orderscheme_requires_collective_id(
     if comm.nranks == 1:
         pytest.skip("collective_id is only required for multi-rank runs.")
 
-    stream = context.get_stream_from_pool()
+    stream = context.br().stream_pool.get_stream()
     input_scheme = _make_scheme(context, 4, stream=stream)
     output_scheme = _make_scheme(context, 4, stream=stream)
 
@@ -274,7 +283,7 @@ def test_adjust_orderscheme_sparse_boundary_shift(
         pytest.skip("This test expects exactly two ranks.")
 
     keys = list(range(4)) if comm.rank == 0 else list(range(4, 8))
-    stream = context.get_stream_from_pool()
+    stream = context.br().stream_pool.get_stream()
     # Input sorted on (key, val) is also sorted on the target key prefix.
     input_scheme = _make_scheme(context, (4, 4), key_indices=(0, 1), stream=stream)
     output_scheme = _make_scheme(context, target_boundary, stream=stream)
@@ -304,7 +313,7 @@ def test_adjust_orderscheme_emits_empty_owned_partitions(
         pytest.skip("This test expects exactly two ranks.")
 
     keys = [0, 1, 2] if comm.rank == 0 else [5, 8]
-    stream = context.get_stream_from_pool()
+    stream = context.br().stream_pool.get_stream()
     input_scheme = _make_scheme(context, 5, stream=stream)
     output_scheme = _make_scheme(context, [3, 5, 7], stream=stream)
 
@@ -331,13 +340,14 @@ def test_adjust_orderscheme_emits_empty_owned_partitions(
 def test_adjust_orderscheme_all_empty_input(spmd_engine: SPMDEngine) -> None:
     context = spmd_engine.context
     comm = spmd_engine.comm
-    stream = context.get_stream_from_pool()
+    stream = context.br().stream_pool.get_stream()
     input_scheme = _make_scheme(context, 5, stream=stream)
     output_scheme = _make_scheme(context, [3, 5, 7], stream=stream)
+    output_npartitions = output_scheme.orderings[0].num_boundaries + 1
     expected: _ExpectedPartitions = {
         pid: []
-        for pid in range(output_scheme.num_boundaries + 1)
-        if pid * comm.nranks // (output_scheme.num_boundaries + 1) == comm.rank
+        for pid in range(output_npartitions)
+        if pid * comm.nranks // output_npartitions == comm.rank
     }
 
     if comm.nranks == 1:
@@ -384,7 +394,7 @@ def test_adjust_orderscheme_single_rank_no_collective(
     if comm.nranks != 1:
         pytest.skip("This test covers the single-rank path.")
 
-    stream = context.get_stream_from_pool()
+    stream = context.br().stream_pool.get_stream()
     input_scheme = _make_scheme(context, 4, stream=stream)
     output_scheme = _make_scheme(context, target_boundary, stream=stream)
     output_by_pid = asyncio.run(
@@ -407,7 +417,7 @@ def test_adjust_orderscheme_multi_chunk_input(spmd_engine: SPMDEngine) -> None:
     if comm.nranks != 1:
         pytest.skip("This test covers local chunk accumulation.")
 
-    stream = context.get_stream_from_pool()
+    stream = context.br().stream_pool.get_stream()
     input_scheme = _make_scheme(context, 4, stream=stream)
     output_scheme = _make_scheme(context, 4, stream=stream)
     output = asyncio.run(