Add Ordering.as_strict helper

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

@@ -13,7 +13,6 @@
     HashScheme,
     OrderKey,
     OrderScheme,
-    Ordering,
     Partitioning,
 )
 from cudf_streaming.table_chunk import TableChunk
@@ -29,7 +28,6 @@
 from cudf_polars.streaming.actor_graph.collectives.ordering import (
     _partition_range,
     adjust_ordering,
-    get_strict_ordering,
 )
 from cudf_polars.streaming.actor_graph.collectives.shuffle import ShuffleManager
 from cudf_polars.streaming.actor_graph.dispatch import (
@@ -59,6 +57,7 @@
 from cudf_polars.streaming.repartition import Repartition
 
 if TYPE_CHECKING:
+    from cudf_streaming.channel_metadata import Ordering
     from rapidsmpf.communicator.communicator import Communicator
     from rapidsmpf.memory.buffer_resource import BufferResource
     from rapidsmpf.streaming.core.channel import Channel
@@ -460,16 +459,15 @@ async def _shuffle_reduce(
 def _remap_ordering_keys(
     ordering: Ordering,
     column_indices: tuple[int, ...],
-    br: BufferResource,
 ) -> Ordering:
     """Return ``ordering`` with keys remapped to another schema."""
-    return Ordering(
-        [
+    return ordering.with_keys(
+        tuple(
             OrderKey(index, key.order, key.null_order)
-            for key, index in zip(ordering.keys, column_indices, strict=True)
-        ],
-        ordering.get_boundaries(br),
-        strict_boundaries=ordering.strict_boundaries,
+            for key, index in zip(
+                ordering.keys[: len(column_indices)], column_indices, strict=True
+            )
+        )
     )
 
 
@@ -565,9 +563,8 @@ async def _ordered_adjust_reduce(
     partial_input_ordering = _remap_ordering_keys(
         input_ordering,
         decomposed.shuffle_indices[: len(input_ordering.keys)],
-        context.br(),
     )
-    partial_output_ordering = get_strict_ordering(partial_input_ordering, context.br())
+    partial_output_ordering = partial_input_ordering.as_strict()
     ch_local = context.create_channel()
     ch_adjusted = context.create_channel()
     adjusted_metadata = _adjusted_ordering_metadata(

python/cudf_streaming/cudf_streaming/channel_metadata.pxd

@@ -41,6 +41,9 @@ cdef extern from "<cudf_streaming/channel_metadata.hpp>" \
         cpp_Ordering(
             vector[cpp_OrderKey], unique_ptr[cpp_TableChunk], bool_t
         ) except +ex_handler
+        cpp_Ordering(
+            vector[cpp_OrderKey], shared_ptr[cpp_TableChunk], bool_t
+        ) except +ex_handler
         vector[cpp_OrderKey] keys
         shared_ptr[cpp_TableChunk] boundaries
         bool_t strict_boundaries

python/cudf_streaming/cudf_streaming/channel_metadata.pyi

@@ -50,6 +50,7 @@ class Ordering:
     @property
     def num_boundaries(self) -> int: ...
     def get_boundaries(self, br: BufferResource) -> TableChunk: ...
+    def as_strict(self) -> Ordering: ...
     def with_keys(self, new_keys: Sequence[OrderKey]) -> Ordering: ...
     def boundaries_aligned_with(
         self, other: Ordering, br: BufferResource

python/cudf_streaming/cudf_streaming/channel_metadata.pyx

@@ -150,6 +150,12 @@ cdef class Ordering:
         ret._handle = move(ordering)
         return ret
 
+    def as_strict(self) -> Ordering:
+        """Return an equivalent ``Ordering`` with strict boundaries."""
+        return Ordering.from_cpp(
+            cpp_Ordering(self._handle.keys, self._handle.boundaries, True)
+        )
+
     @property
     def keys(self) -> tuple:
         """Sort keys, one per sort column."""

python/cudf_streaming/cudf_streaming/tests/test_channel_metadata.py

@@ -34,39 +34,67 @@ def _make_boundaries(context: Context, table: plc.Table) -> TableChunk:
     )
 
 
+def _column_from_values(
+    values: list[int | str | None],
+    dtype: plc.DataType,
+    non_null_value: int | str,
+) -> plc.Column:
+    if values and all(value is None for value in values):
+        return plc.Column.all_null_like(
+            plc.Column.from_iterable_of_py([non_null_value], dtype),
+            len(values),
+        )
+    return plc.Column.from_iterable_of_py(values, dtype)
+
+
+def _two_key_ordering_from_boundary_values(
+    context: Context,
+    int_values: list[int | None],
+    string_values: list[str | None],
+    *,
+    strict_boundaries: bool = False,
+) -> Ordering:
+    """Two-key Ordering with mixed INT64/STRING boundary columns."""
+    return Ordering(
+        [
+            OrderKey(
+                0,
+                plc.types.Order.ASCENDING,
+                plc.types.NullOrder.BEFORE,
+            ),
+            OrderKey(
+                1,
+                plc.types.Order.DESCENDING,
+                plc.types.NullOrder.AFTER,
+            ),
+        ],
+        _make_boundaries(
+            context,
+            plc.Table(
+                [
+                    _column_from_values(
+                        int_values, plc.DataType(plc.TypeId.INT64), 0
+                    ),
+                    _column_from_values(
+                        string_values, plc.DataType(plc.TypeId.STRING), "x"
+                    ),
+                ]
+            ),
+        ),
+        strict_boundaries=strict_boundaries,
+    )
+
+
 def _two_key_order_scheme(
     context: Context, *, strict_boundaries: bool = False
 ) -> OrderScheme:
     """Two-key OrderScheme with a 1-row boundary table (2 partitions)."""
-    boundaries = _make_boundaries(
-        context,
-        plc.Table(
-            [
-                plc.Column.from_iterable_of_py(
-                    [100], plc.DataType(plc.TypeId.INT64)
-                ),
-                plc.Column.from_iterable_of_py(
-                    ["abc"], plc.DataType(plc.TypeId.STRING)
-                ),
-            ]
-        ),
-    )
     return OrderScheme(
         [
-            Ordering(
-                [
-                    OrderKey(
-                        0,
-                        plc.types.Order.ASCENDING,
-                        plc.types.NullOrder.BEFORE,
-                    ),
-                    OrderKey(
-                        1,
-                        plc.types.Order.DESCENDING,
-                        plc.types.NullOrder.AFTER,
-                    ),
-                ],
-                boundaries,
+            _two_key_ordering_from_boundary_values(
+                context,
+                [100],
+                ["abc"],
                 strict_boundaries=strict_boundaries,
             )
         ]
@@ -231,6 +259,39 @@ def test_ordering_with_keys(context: Context) -> None:
     assert ordering.boundaries_aligned_with(ordering2, context.br())
 
 
+@pytest.mark.parametrize(
+    "int_values,string_values",
+    [
+        ([100], ["abc"]),
+        ([], []),
+        ([None], [None]),
+    ],
+)
+def test_ordering_as_strict(
+    context: Context,
+    int_values: list[int | None],
+    string_values: list[str | None],
+) -> None:
+    """as_strict shares boundaries and marks them strict."""
+    ordering = _two_key_ordering_from_boundary_values(
+        context, int_values, string_values
+    )
+    strict_ordering = ordering.as_strict()
+    strict_ordering2 = _two_key_ordering_from_boundary_values(
+        context,
+        int_values,
+        string_values,
+        strict_boundaries=True,
+    )
+    assert strict_ordering.keys == ordering.keys
+    assert strict_ordering.num_boundaries == ordering.num_boundaries
+    assert strict_ordering.strict_boundaries
+    assert not ordering.boundaries_aligned_with(strict_ordering, context.br())
+    assert strict_ordering.boundaries_aligned_with(
+        strict_ordering2, context.br()
+    )
+
+
 def test_ordering_boundaries_aligned_with(context: Context) -> None:
     """Boundary comparison ignores key indices but checks values and ordering."""
     df = plc.Table(