diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py index 290865ae9087..ae2ce782dde3 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py +++ b/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( diff --git a/python/cudf_streaming/cudf_streaming/channel_metadata.pxd b/python/cudf_streaming/cudf_streaming/channel_metadata.pxd index f2037164a9ec..f59e8c8c8bef 100644 --- a/python/cudf_streaming/cudf_streaming/channel_metadata.pxd +++ b/python/cudf_streaming/cudf_streaming/channel_metadata.pxd @@ -41,6 +41,9 @@ cdef extern from "" \ 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 diff --git a/python/cudf_streaming/cudf_streaming/channel_metadata.pyi b/python/cudf_streaming/cudf_streaming/channel_metadata.pyi index 60599a2cf732..c08650585b4a 100644 --- a/python/cudf_streaming/cudf_streaming/channel_metadata.pyi +++ b/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 diff --git a/python/cudf_streaming/cudf_streaming/channel_metadata.pyx b/python/cudf_streaming/cudf_streaming/channel_metadata.pyx index 70d5ebfadc82..a40a4b1fee98 100644 --- a/python/cudf_streaming/cudf_streaming/channel_metadata.pyx +++ b/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.""" diff --git a/python/cudf_streaming/cudf_streaming/tests/test_channel_metadata.py b/python/cudf_streaming/cudf_streaming/tests/test_channel_metadata.py index 90dad3974400..7f4fef0fc1b8 100644 --- a/python/cudf_streaming/cudf_streaming/tests/test_channel_metadata.py +++ b/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(