From bf53cd8ed4f3b89fb506480fc87fbee87a71bd3b Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 21 Aug 2026 07:00:58 -0700 Subject: [PATCH 1/4] Update cython to simplify Ordering usage --- .../streaming/actor_graph/groupby.py | 19 ++++++++----------- .../cudf_streaming/channel_metadata.pxd | 3 +++ .../cudf_streaming/channel_metadata.pyi | 1 + .../cudf_streaming/channel_metadata.pyx | 6 ++++++ .../tests/test_channel_metadata.py | 11 +++++++++++ 5 files changed, 29 insertions(+), 11 deletions(-) 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 f7678a538a2f..af14755cb25d 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..e13718a0cc6b 100644 --- a/python/cudf_streaming/cudf_streaming/tests/test_channel_metadata.py +++ b/python/cudf_streaming/cudf_streaming/tests/test_channel_metadata.py @@ -231,6 +231,17 @@ def test_ordering_with_keys(context: Context) -> None: assert ordering.boundaries_aligned_with(ordering2, context.br()) +def test_ordering_as_strict(context: Context) -> None: + """as_strict shares boundaries and marks them strict.""" + o1 = _two_key_order_scheme(context) + ordering = o1.orderings[0] + strict_ordering = ordering.as_strict() + assert strict_ordering.keys == ordering.keys + assert strict_ordering.num_boundaries == ordering.num_boundaries + assert strict_ordering.strict_boundaries + assert ordering.boundaries_aligned_with(strict_ordering, context.br()) + + def test_ordering_boundaries_aligned_with(context: Context) -> None: """Boundary comparison ignores key indices but checks values and ordering.""" df = plc.Table( From 717b753ac9f88c9bf3dbc10abb577bf9acb70bc3 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 21 Aug 2026 08:07:33 -0700 Subject: [PATCH 2/4] address CI failures and coderabbit review --- .../tests/test_channel_metadata.py | 100 ++++++++++++------ 1 file changed, 69 insertions(+), 31 deletions(-) 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 e13718a0cc6b..27811909ed84 100644 --- a/python/cudf_streaming/cudf_streaming/tests/test_channel_metadata.py +++ b/python/cudf_streaming/cudf_streaming/tests/test_channel_metadata.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING +import pyarrow as pa import pylibcudf as plc import pytest @@ -34,39 +35,54 @@ def _make_boundaries(context: Context, table: plc.Table) -> TableChunk: ) +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( + [ + plc.Column.from_arrow( + pa.array(int_values, type=pa.int64()) + ), + plc.Column.from_arrow( + pa.array(string_values, type=pa.string()) + ), + ] + ), + ), + 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,15 +247,37 @@ def test_ordering_with_keys(context: Context) -> None: assert ordering.boundaries_aligned_with(ordering2, context.br()) -def test_ordering_as_strict(context: Context) -> None: +@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.""" - o1 = _two_key_order_scheme(context) - ordering = o1.orderings[0] + 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 ordering.boundaries_aligned_with(strict_ordering, context.br()) + 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: From e93c629a9527454cb6e994f51c2ef9375cf631bf Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 21 Aug 2026 09:10:59 -0700 Subject: [PATCH 3/4] remove arrow usage --- .../tests/test_channel_metadata.py | 22 ++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) 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 27811909ed84..736f8c8b96ff 100644 --- a/python/cudf_streaming/cudf_streaming/tests/test_channel_metadata.py +++ b/python/cudf_streaming/cudf_streaming/tests/test_channel_metadata.py @@ -6,7 +6,6 @@ from typing import TYPE_CHECKING -import pyarrow as pa import pylibcudf as plc import pytest @@ -35,6 +34,19 @@ 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], @@ -60,11 +72,11 @@ def _two_key_ordering_from_boundary_values( context, plc.Table( [ - plc.Column.from_arrow( - pa.array(int_values, type=pa.int64()) + _column_from_values( + int_values, plc.DataType(plc.TypeId.INT64), 0 ), - plc.Column.from_arrow( - pa.array(string_values, type=pa.string()) + _column_from_values( + string_values, plc.DataType(plc.TypeId.STRING), "" ), ] ), From e0ea6b9794c4d6b39a042e048c0d43fba88956c7 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 21 Aug 2026 10:26:54 -0700 Subject: [PATCH 4/4] fix test setup --- .../cudf_streaming/tests/test_channel_metadata.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 736f8c8b96ff..7f4fef0fc1b8 100644 --- a/python/cudf_streaming/cudf_streaming/tests/test_channel_metadata.py +++ b/python/cudf_streaming/cudf_streaming/tests/test_channel_metadata.py @@ -76,7 +76,7 @@ def _two_key_ordering_from_boundary_values( int_values, plc.DataType(plc.TypeId.INT64), 0 ), _column_from_values( - string_values, plc.DataType(plc.TypeId.STRING), "" + string_values, plc.DataType(plc.TypeId.STRING), "x" ), ] ),