From d6e20f692dfd57a956c2c43b9810b1a0dbb4c01d Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 15 May 2026 07:22:21 -0700 Subject: [PATCH 1/2] add LocalRepartitioner.repartition_by_orderscheme --- .../rapidsmpf/collectives/shuffle.py | 58 +++++++++++- .../tests/experimental/test_shuffler.py | 93 +++++++++++++++++++ 2 files changed, 149 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/shuffle.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/shuffle.py index d3952657135f..6f308cb05129 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/shuffle.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/shuffle.py @@ -25,9 +25,13 @@ ) from rapidsmpf.streaming.cudf.table_chunk import TableChunk +import polars as pl + import pylibcudf as plc import pylibcudf.partitioning +import pylibcudf.search +from cudf_polars.containers import DataFrame, DataType from cudf_polars.dsl.expr import Col from cudf_polars.experimental.rapidsmpf.dispatch import ( generate_ir_sub_network, @@ -40,14 +44,15 @@ send_metadata, ) from cudf_polars.experimental.shuffle import Shuffle -from cudf_polars.utils.cuda_stream import stream_ordered_after +from cudf_polars.utils.cuda_stream import join_cuda_streams, stream_ordered_after if TYPE_CHECKING: - from collections.abc import Generator + from collections.abc import Generator, Sequence from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.memory.packed_data import PackedData from rapidsmpf.streaming.core.channel import Channel + from rapidsmpf.streaming.cudf.channel_metadata import OrderScheme from rmm.pylibrmm.stream import Stream @@ -322,6 +327,55 @@ async def repartition_by_index( ), ) + async def repartition_by_orderscheme( + self, + *, + scheme: OrderScheme, + stream: Stream, + key_column_indices: Sequence[int] | None = None, + ) -> None: + """ + Re-partition items by range using the boundaries in an OrderScheme. + + Parameters + ---------- + scheme + The ``OrderScheme`` to use for re-partitioning. + stream + The stream to use for the unpack operation. + key_column_indices + The indices of the columns to use for re-partitioning. + If ``None``, column indices in ``scheme`` are used. + """ + column_order = [k.order for k in scheme.keys] + null_order = [k.null_order for k in scheme.keys] + if key_column_indices is None: + key_column_indices = [k.column_index for k in scheme.keys] + boundaries, boundary_stream = scheme.get_boundaries() + join_cuda_streams(downstreams=(stream,), upstreams=(boundary_stream,)) + async with self._local_shuffle.inserting() as inserter: + for table in self._iter_chunks(stream): + key_table = plc.Table([table.columns()[i] for i in key_column_indices]) + split_col = plc.search.lower_bound( + key_table, boundaries, column_order, null_order, stream=stream + ) + splits = ( + DataFrame.from_table( + plc.Table([split_col]), + ["split"], + [DataType(pl.Int32())], + stream=stream, + ) + .to_polars()["split"] + .to_list() + ) + inserter.insert_split( + TableChunk.from_pylibcudf_table( + table, stream, exclusive_view=True, br=self._br + ), + splits, + ) + def local_partitions(self) -> list[int]: """Return the local partition IDs.""" return self._local_shuffle.local_partitions() diff --git a/python/cudf_polars/tests/experimental/test_shuffler.py b/python/cudf_polars/tests/experimental/test_shuffler.py index 84818741815c..f146f5bf2ce9 100644 --- a/python/cudf_polars/tests/experimental/test_shuffler.py +++ b/python/cudf_polars/tests/experimental/test_shuffler.py @@ -9,12 +9,16 @@ from rapidsmpf.streaming.cudf.channel_metadata import ( ChannelMetadata, HashScheme, + OrderKey, + OrderScheme, Partitioning, ) from rapidsmpf.streaming.cudf.table_chunk import TableChunk import polars as pl +import pylibcudf as plc + from cudf_polars.containers import DataFrame, DataType from cudf_polars.experimental.rapidsmpf.collectives.common import reserve_op_id from cudf_polars.experimental.rapidsmpf.collectives.shuffle import ( @@ -270,3 +274,92 @@ async def _run() -> None: engine=spmd_engine, local_df=local_df, op_id=op_id ) assert global_df.height == 12 * comm.nranks + + +@pytest.mark.spmd +@pytest.mark.parametrize("local_count", [1, 2, 3]) +def test_local_repartitioner_boundaries(spmd_engine, local_count) -> None: + context = spmd_engine.context + comm = spmd_engine.comm + + # Each rank holds a sorted, contiguous key slice of [0, n_total) + n_per_rank = 6 + key_start = comm.rank * n_per_rank + pl_df = pl.DataFrame( + { + "key": pl.Series( + list(range(key_start, key_start + n_per_rank)), dtype=pl.Int32() + ), + "val": pl.Series(list(range(n_per_rank)), dtype=pl.Int32()), + } + ) + col_names = pl_df.columns + dtypes = [DataType(dt) for dt in pl_df.dtypes] + + results: list[tuple[int, pl.DataFrame]] = [] + + async def _run() -> None: + stream = context.get_stream_from_pool() + cudf_df = DataFrame.from_polars(pl_df, stream) + + # Build an OrderScheme whose boundaries evenly split this rank's key range + boundary_keys = [ + key_start + n_per_rank * i // local_count for i in range(1, local_count) + ] + boundary_df = DataFrame.from_polars( + pl.DataFrame({"key": pl.Series(boundary_keys, dtype=pl.Int32())}), stream + ) + scheme = OrderScheme( + [OrderKey(0, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE)], + TableChunk.from_pylibcudf_table( + boundary_df.table, stream, exclusive_view=False, br=context.br() + ), + strict_boundaries=True, + ) + + with reserve_op_id() as op_id: + shuffle = ShuffleManager( + context, comm, num_partitions=comm.nranks, collective_id=op_id + ) + # Route each rank's sorted slice only to its own partition + splits = [0] * comm.rank + [n_per_rank] * (comm.nranks - comm.rank - 1) + async with shuffle.inserting() as inserter: + inserter.insert_split( + TableChunk.from_pylibcudf_table( + cudf_df.table, stream, exclusive_view=True, br=context.br() + ), + splits, + ) + + local = LocalRepartitioner(shuffle, local_count=local_count) + await local.repartition_by_orderscheme( + scheme=scheme, + stream=stream, + key_column_indices=(0,), + ) + + for pid in local.local_partitions(): + tbl = local.extract_chunk(pid, stream) + results.append( + ( + pid, + DataFrame.from_table( + tbl, col_names, dtypes, stream + ).to_polars(), + ) + ) + + asyncio.run(_run()) + + assert len(results) == local_count + + # All rows present + all_keys = sorted(pl.concat([df for _, df in results])["key"].to_list()) + assert all_keys == list(range(key_start, key_start + n_per_rank)) + + # Each partition holds keys in the correct range + for pid, df in results: + lo = key_start + n_per_rank * pid // local_count + hi = key_start + n_per_rank * (pid + 1) // local_count + for key_val in df["key"].to_list(): + assert lo <= key_val < hi From 951283e2645688931ed589f4717e6692e610c677 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 15 May 2026 11:52:02 -0700 Subject: [PATCH 2/2] align with rapidsmpf change --- .../experimental/rapidsmpf/collectives/shuffle.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/shuffle.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/shuffle.py index 6f308cb05129..32c555fe9ed3 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/shuffle.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/shuffle.py @@ -351,13 +351,17 @@ async def repartition_by_orderscheme( null_order = [k.null_order for k in scheme.keys] if key_column_indices is None: key_column_indices = [k.column_index for k in scheme.keys] - boundaries, boundary_stream = scheme.get_boundaries() - join_cuda_streams(downstreams=(stream,), upstreams=(boundary_stream,)) + boundary_chunk = scheme.get_boundaries(self._br) + join_cuda_streams(downstreams=(stream,), upstreams=(boundary_chunk.stream,)) async with self._local_shuffle.inserting() as inserter: for table in self._iter_chunks(stream): key_table = plc.Table([table.columns()[i] for i in key_column_indices]) split_col = plc.search.lower_bound( - key_table, boundaries, column_order, null_order, stream=stream + key_table, + boundary_chunk.table_view(), + column_order, + null_order, + stream=stream, ) splits = ( DataFrame.from_table(