Add OrderScheme.get_boundaries API

python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyi

@@ -45,6 +45,7 @@ class OrderScheme:
     def strict_boundaries(self) -> bool: ...
     @property
     def num_boundaries(self) -> int: ...
+    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

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,29 @@ cdef class OrderScheme:
         """Number of boundary rows (N-1 for N partitions)."""
         return self._handle.boundaries.get().shape().first
 
+    def get_boundaries(self, BufferResource br not None) -> TableChunk:
+        """
+        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(
+            chunk.table_view(), owner=self, 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."""
         cdef vector[cpp_OrderKey] cpp_keys

python/rapidsmpf/rapidsmpf/tests/streaming/test_channel_metadata.py

@@ -109,6 +109,19 @@ def test_order_scheme(context: Context) -> None:
         )
 
 
+def test_order_scheme_get_boundaries(context: Context) -> None:
+    scheme = _two_key_order_scheme(context)
+    chunk = scheme.get_boundaries(context.br())
+    assert chunk.table_view().num_columns() == 2
+    assert chunk.table_view().num_rows() == 1
+    scheme2 = OrderScheme(
+        scheme.keys,
+        chunk,
+        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)