From e904d2ff5e4504de74e7befb7955b12a1af178ac Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 14 May 2026 17:57:33 -0700 Subject: [PATCH 1/3] basic get_boundaries method in Python --- .../rapidsmpf/streaming/cudf/channel_metadata.pyi | 2 ++ .../rapidsmpf/streaming/cudf/channel_metadata.pyx | 12 +++++++++++- .../tests/streaming/test_channel_metadata.py | 15 +++++++++++++++ 3 files changed, 28 insertions(+), 1 deletion(-) diff --git a/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyi b/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyi index 12f47c73a..30669b3eb 100644 --- a/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyi +++ b/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyi @@ -9,6 +9,7 @@ from dataclasses import dataclass from typing import Literal, Self import pylibcudf as plc +from rmm.pylibrmm.stream import Stream from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.streaming.core.message import Message @@ -45,6 +46,7 @@ class OrderScheme: def strict_boundaries(self) -> bool: ... @property def num_boundaries(self) -> int: ... + def get_boundaries(self) -> tuple[plc.Table, Stream]: ... def with_keys(self, new_keys: Sequence[OrderKey]) -> OrderScheme: ... def boundaries_aligned_with( self, other: OrderScheme, br: BufferResource diff --git a/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyx b/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyx index c2a927b47..92e09ee5d 100644 --- a/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyx +++ b/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyx @@ -9,10 +9,12 @@ from libcpp.utility cimport move from libcpp.vector cimport vector from pylibcudf.libcudf.types cimport null_order as cpp_null_order from pylibcudf.libcudf.types cimport order as cpp_order +from pylibcudf.table cimport Table +from rmm.pylibrmm.stream cimport Stream from rapidsmpf.memory.buffer_resource cimport BufferResource from rapidsmpf.streaming.core.message cimport Message -from rapidsmpf.streaming.cudf.table_chunk cimport TableChunk +from rapidsmpf.streaming.cudf.table_chunk cimport TableChunk, cpp_TableChunk cdef extern from * nogil: @@ -182,6 +184,14 @@ cdef class OrderScheme: """Number of boundary rows (N-1 for N partitions).""" return self._handle.boundaries.get().shape().first + def get_boundaries(self) -> tuple[Table, Stream]: + """Return the boundary rows and their CUDA stream as a zero-copy view.""" + cdef const cpp_TableChunk* chunk = self._handle.boundaries.get() + cdef Stream stream = Stream._from_cudaStream_t(chunk.stream().value()) + return Table.from_table_view_of_arbitrary( + chunk.table_view(), owner=self, stream=stream + ), stream + def with_keys(self, object new_keys) -> OrderScheme: """Return a new ``OrderScheme`` with updated key column indices.""" cdef vector[cpp_OrderKey] cpp_keys diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_channel_metadata.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_channel_metadata.py index abb843b8b..c15ca73e2 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_channel_metadata.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/test_channel_metadata.py @@ -109,6 +109,21 @@ def test_order_scheme(context: Context) -> None: ) +def test_order_scheme_get_boundaries(context: Context) -> None: + scheme = _two_key_order_scheme(context) + tbl, stream = scheme.get_boundaries() + assert tbl.num_columns() == 2 + assert tbl.num_rows() == 1 + scheme2 = OrderScheme( + scheme.keys, + TableChunk.from_pylibcudf_table( + tbl, stream, exclusive_view=False, br=context.br() + ), + strict_boundaries=scheme.strict_boundaries, + ) + assert scheme2.boundaries_aligned_with(scheme, context.br()) + + def test_order_scheme_with_keys(context: Context) -> None: """with_keys shares boundaries and updates column indices.""" o1 = _two_key_order_scheme(context) From ba6619dad0d1fe493eb47b6e573387bcb92b90eb Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 15 May 2026 11:45:55 -0700 Subject: [PATCH 2/3] return single TableChunk instead of tuple[Table, Stream] --- .../rapidsmpf/streaming/cudf/channel_metadata.pyi | 3 +-- .../rapidsmpf/streaming/cudf/channel_metadata.pyx | 11 +++++++---- .../tests/streaming/test_channel_metadata.py | 10 ++++------ 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyi b/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyi index 30669b3eb..09b075ad0 100644 --- a/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyi +++ b/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyi @@ -9,7 +9,6 @@ from dataclasses import dataclass from typing import Literal, Self import pylibcudf as plc -from rmm.pylibrmm.stream import Stream from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.streaming.core.message import Message @@ -46,7 +45,7 @@ class OrderScheme: def strict_boundaries(self) -> bool: ... @property def num_boundaries(self) -> int: ... - def get_boundaries(self) -> tuple[plc.Table, Stream]: ... + def get_boundaries(self, br: BufferResource) -> TableChunk: ... def with_keys(self, new_keys: Sequence[OrderKey]) -> OrderScheme: ... def boundaries_aligned_with( self, other: OrderScheme, br: BufferResource diff --git a/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyx b/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyx index 92e09ee5d..1363f7844 100644 --- a/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyx +++ b/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyx @@ -184,13 +184,16 @@ cdef class OrderScheme: """Number of boundary rows (N-1 for N partitions).""" return self._handle.boundaries.get().shape().first - def get_boundaries(self) -> tuple[Table, Stream]: - """Return the boundary rows and their CUDA stream as a zero-copy view.""" + def get_boundaries(self, BufferResource br not None) -> TableChunk: + """Return the boundary rows as a zero-copy TableChunk view.""" cdef const cpp_TableChunk* chunk = self._handle.boundaries.get() cdef Stream stream = Stream._from_cudaStream_t(chunk.stream().value()) - return Table.from_table_view_of_arbitrary( + tbl = Table.from_table_view_of_arbitrary( chunk.table_view(), owner=self, stream=stream - ), stream + ) + return TableChunk.from_pylibcudf_table( + tbl, stream, exclusive_view=False, br=br + ) def with_keys(self, object new_keys) -> OrderScheme: """Return a new ``OrderScheme`` with updated key column indices.""" diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_channel_metadata.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_channel_metadata.py index c15ca73e2..5ea474781 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_channel_metadata.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/test_channel_metadata.py @@ -111,14 +111,12 @@ def test_order_scheme(context: Context) -> None: def test_order_scheme_get_boundaries(context: Context) -> None: scheme = _two_key_order_scheme(context) - tbl, stream = scheme.get_boundaries() - assert tbl.num_columns() == 2 - assert tbl.num_rows() == 1 + chunk = scheme.get_boundaries(context.br()) + assert chunk.table_view().num_columns() == 2 + assert chunk.table_view().num_rows() == 1 scheme2 = OrderScheme( scheme.keys, - TableChunk.from_pylibcudf_table( - tbl, stream, exclusive_view=False, br=context.br() - ), + chunk, strict_boundaries=scheme.strict_boundaries, ) assert scheme2.boundaries_aligned_with(scheme, context.br()) From 3576b55c51a8d401de9e35b43eaadb57eadc85fb Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 22 May 2026 09:56:13 -0700 Subject: [PATCH 3/3] update docstring --- .../rapidsmpf/streaming/cudf/channel_metadata.pyx | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyx b/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyx index 1363f7844..28753f1f6 100644 --- a/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyx +++ b/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyx @@ -185,7 +185,19 @@ cdef class OrderScheme: return self._handle.boundaries.get().shape().first def get_boundaries(self, BufferResource br not None) -> TableChunk: - """Return the boundary rows as a zero-copy TableChunk view.""" + """ + Return the boundary rows. + + Parameters + ---------- + br + Buffer resource to associate with the returned table chunk. + + Returns + ------- + TableChunk + A non-exclusive view of the boundary rows owned by this scheme. + """ cdef const cpp_TableChunk* chunk = self._handle.boundaries.get() cdef Stream stream = Stream._from_cudaStream_t(chunk.stream().value()) tbl = Table.from_table_view_of_arbitrary(