diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py index 805c00c29e92..7710cf9eb01a 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py @@ -10,6 +10,7 @@ import polars as pl import pylibcudf as plc +from cudf_streaming.channel_metadata import Ordering from cudf_streaming.partition_utils import ( packed_data_from_cudf_packed_columns, unpack_and_concat, @@ -29,7 +30,6 @@ from cudf_polars.utils.cuda_stream import stream_ordered_after 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.memory.packed_data import PackedData @@ -44,6 +44,15 @@ _PartitionRange = tuple[int, int] +def get_strict_ordering(ordering: Ordering, br: BufferResource) -> Ordering: + """Return an equivalent Ordering with strict boundaries.""" + return Ordering( + ordering.keys, + ordering.get_boundaries(br), + strict_boundaries=True, + ) + + @dataclass(frozen=True) class _RoutingPlan: """Static ownership and exchange plan for one ordering adjustment.""" diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py index e8de685d45cc..c24cf2a11a1c 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py @@ -608,7 +608,10 @@ async def sort_actor( partitioning = NormalizedPartitioning.from_keys( metadata_in.partitioning, comm.nranks, keys=order_keys ) - if partitioning.is_strictly_sorted(order_keys): + if partitioning.is_ordered( + order_keys, + level="local" if metadata_in.duplicated else "flat", + ): if tracer is not None: tracer.decision = "already_sorted" await chunkwise_evaluate( 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 e1c93620d1f0..f7678a538a2f 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py @@ -11,6 +11,10 @@ from cudf_streaming.channel_metadata import ( ChannelMetadata, HashScheme, + OrderKey, + OrderScheme, + Ordering, + Partitioning, ) from cudf_streaming.table_chunk import TableChunk from rapidsmpf.communicator.single import new_communicator as single_comm @@ -22,6 +26,11 @@ from cudf_polars.dsl.expr import Col, NamedExpr from cudf_polars.dsl.ir import IR, Distinct, GroupBy, Select from cudf_polars.dsl.utils.naming import names_to_indices, unique_names +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 ( generate_ir_sub_network, @@ -38,10 +47,12 @@ empty_table_chunk, evaluate_batch, evaluate_chunk, + gather_in_task_group, maybe_remap_partitioning, process_children, recv_metadata, send_metadata, + shutdown_channels_on_error, shutdown_on_error, ) from cudf_polars.streaming.groupby import _has_stable_sorted_agg, combine, decompose @@ -55,6 +66,7 @@ from cudf_polars.dsl.ir import IRExecutionContext from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator from cudf_polars.streaming.actor_graph.tracing import ActorTracer + from cudf_polars.streaming.actor_graph.utils import PartitioningLevel from cudf_polars.typing import Schema @@ -445,6 +457,178 @@ async def _shuffle_reduce( await ch_out.drain(context) +def _remap_ordering_keys( + ordering: Ordering, + column_indices: tuple[int, ...], + br: BufferResource, +) -> Ordering: + """Return ``ordering`` with keys remapped to another schema.""" + return Ordering( + [ + 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, + ) + + +def _groupby_output_metadata( + ir: GroupBy | Distinct, + decomposed: DecomposedGroupBy, + local_count: int, + partitioning: Partitioning, + duplicated: bool, # noqa: FBT001 + *, + context: Context, +) -> ChannelMetadata: + """Return groupby output metadata after final reduction/select.""" + partitioning = maybe_remap_partitioning( + decomposed.reduction_ir, + partitioning, + child_ir=decomposed.reduction_ir, + context=context, + ) + if decomposed.select_ir is not None: + partitioning = maybe_remap_partitioning( + decomposed.select_ir, + partitioning, + child_ir=decomposed.reduction_ir, + context=context, + ) + else: + partitioning = maybe_remap_partitioning( + ir, + partitioning, + child_ir=ir.children[0], + context=context, + ) + return ChannelMetadata( + local_count=local_count, + partitioning=partitioning, + duplicated=duplicated, + ) + + +async def _send_locally_aggregated_chunks( + context: Context, + decomposed: DecomposedGroupBy, + ir_context: IRExecutionContext, + ch_out: Channel[TableChunk], + ch_in: Channel[TableChunk], + target_partition_size: int, + *, + aggregated: TableChunk, + input_drained: bool, +) -> None: + """Send locally aggregated chunks, then drain the remaining input.""" + seq_num = 0 + while True: + await send_chunk( + context, + ch_out, + _enforce_schema(aggregated, decomposed.reduction_ir.schema, context.br()), + seq_num, + tracer=None, + ) + seq_num += 1 + del aggregated + if input_drained: + break + aggregated, input_drained, _ = await _local_aggregation( + context, + decomposed, + ir_context, + ch_in, + target_partition_size, + ) + await ch_out.drain(context) + + +async def _ordered_adjust_reduce( + context: Context, + comm: Communicator, + decomposed: DecomposedGroupBy, + ir_context: IRExecutionContext, + ch_out: Channel[TableChunk], + ch_in: Channel[TableChunk], + metadata_in: ChannelMetadata, + collective_id: int, + target_partition_size: int, + *, + aggregated: TableChunk, + input_drained: bool, + input_ordering: Ordering, + tracer: ActorTracer | None = None, +) -> None: + """Adjust locally aggregated data to strict ordering boundaries.""" + 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()) + ch_local = context.create_channel() + ch_adjusted = context.create_channel() + adjusted_metadata = _adjusted_ordering_metadata( + comm, metadata_in, partial_output_ordering + ) + metadata_out = _groupby_output_metadata( + decomposed.ir, + decomposed, + adjusted_metadata.local_count, + adjusted_metadata.partitioning, + adjusted_metadata.duplicated, + context=context, + ) + if tracer is not None: + tracer.decision = "adjust_ordering" + + await send_metadata(ch_out, context, metadata_out) + if tracer is not None and metadata_out.duplicated: + tracer.set_duplicated() + + async def reduce_adjusted_chunks() -> None: + extract_irs = [decomposed.reduction_ir] + ( + [decomposed.select_ir] if decomposed.select_ir else [] + ) + while (msg := await ch_adjusted.recv(context)) is not None: + chunk = await evaluate_chunk( + context, + TableChunk.from_message(msg, br=context.br()), + *extract_irs, + ir_context=ir_context, + ) + await send_chunk(context, ch_out, chunk, msg.sequence_number, tracer=tracer) + await ch_out.drain(context) + + async with shutdown_channels_on_error(context, ch_local, ch_adjusted): + await gather_in_task_group( + _send_locally_aggregated_chunks( + context, + decomposed, + ir_context, + ch_local, + ch_in, + target_partition_size, + aggregated=aggregated, + input_drained=input_drained, + ), + adjust_ordering( + context, + comm, + decomposed.reduction_ir, + ir_context, + ch_adjusted, + ch_local, + partial_input_ordering, + partial_output_ordering, + collective_id=collective_id, + ), + reduce_adjusted_chunks(), + ) + + def _enforce_schema( chunk: TableChunk, canonical_schema: dict[str, Any], @@ -506,6 +690,27 @@ def _maintain_order(ir: GroupBy | Distinct) -> bool: ) +def _partition_count_for_rank(rank: int, nranks: int, npartitions: int) -> int: + """Return the contiguous output-partition count owned by one rank.""" + start, stop = _partition_range(rank, nranks, npartitions) + return stop - start + + +def _adjusted_ordering_metadata( + comm: Communicator, + metadata_in: ChannelMetadata, + output_ordering: Ordering, +) -> ChannelMetadata: + """Return metadata for data adjusted to strict ordering boundaries.""" + return ChannelMetadata( + local_count=_partition_count_for_rank( + comm.rank, comm.nranks, output_ordering.num_boundaries + 1 + ), + partitioning=Partitioning(OrderScheme([output_ordering]), "inherit"), + duplicated=metadata_in.duplicated, + ) + + async def _choose_strategy( context: Context, comm: Communicator, @@ -660,13 +865,19 @@ async def groupby_actor( metadata_in = await recv_metadata(ch_in, context) nranks = comm.nranks + group_keys = _key_indices(ir, ir.children[0].schema, concrete_prefix=True) partitioning = NormalizedPartitioning.from_keys( metadata_in.partitioning, nranks, - keys=_key_indices(ir, ir.children[0].schema, concrete_prefix=True), + keys=group_keys, + ) + partitioning_level: PartitioningLevel = ( + "local" if metadata_in.duplicated else "flat" ) maintain_order = _maintain_order(ir) - fully_partitioned = partitioning.is_strictly_partitioned() + fully_partitioned = partitioning.is_strictly_partitioned( + level=partitioning_level, + ) fallback_case = ( # NOTE: This criteria means that we fell back # to one partition at lowering time. @@ -741,6 +952,32 @@ async def groupby_actor( aggregated=aggregated, tracer=tracer, ) + elif ( + # adjust_ordering requires row-ordered chunks. maintain_order=True + # preserves key order through local aggregation for ordered input. + maintain_order + and not metadata_in.duplicated + and partitioning.is_ordered( + group_keys, + level="flat", + ) + ): + assert isinstance(partitioning.inter_rank_scheme, OrderScheme) + await _ordered_adjust_reduce( + context, + comm, + decomposed, + ir_context, + ch_out, + ch_in, + metadata_in, + collective_ids.pop(), + target_partition_size, + aggregated=aggregated, + input_drained=input_drained, + input_ordering=partitioning.inter_rank_scheme.orderings[0], + tracer=tracer, + ) else: await _shuffle_reduce( context, diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py index 8be6f71bdeda..469ff880206b 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py @@ -750,7 +750,9 @@ async def over_actor( keys=ir.key_indices, allow_subset=True, ) - if partitioning.is_strictly_partitioned(): + if partitioning.is_strictly_partitioned( + level="local" if metadata_in.duplicated else "flat", + ): metadata_out = ChannelMetadata( local_count=metadata_in.local_count, partitioning=maybe_remap_partitioning( diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py index 1b7153954f2f..ba04e92e47a8 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -75,6 +75,12 @@ InterRankScheme: TypeAlias = HashScheme | OrderScheme | None PartitioningScheme: TypeAlias = InterRankScheme | Literal["inherit"] +# Partitioning-level predicates: +# - "flat": inter-rank scheme with local layout inherited from it. +# - "inter_rank": inter-rank scheme only. +# - "local": explicit local scheme only. +PartitioningLevel: TypeAlias = Literal["flat", "inter_rank", "local"] + # cuDF column/concatenate row limit (int32) CUDF_ROW_LIMIT = 2**31 - 1 # Stay well below the cuDF row limit when forming a single table/partition. @@ -1271,29 +1277,70 @@ def __eq__(self, other: object) -> bool: and self.local_scheme == other.local_scheme ) - def is_strictly_partitioned(self) -> bool: - """True if data is strictly partitioned with no boundary straddling.""" - if not self: + def _scheme_for_level(self, level: PartitioningLevel) -> PartitioningScheme: + """Return the scheme relevant to the requested partitioning level.""" + match level: + case "flat": + if self.local_scheme != "inherit": + return None + return self.inter_rank_scheme + case "inter_rank": + return self.inter_rank_scheme + case "local": + return self.local_scheme + + @staticmethod + def _scheme_is_strict(scheme: PartitioningScheme) -> bool: + """True when one scheme proves strict partitioning.""" + if scheme is None or scheme == "inherit": return False - for scheme in [self.inter_rank_scheme, self.local_scheme]: - if isinstance(scheme, OrderScheme): - ordering = scheme.orderings[0] - if ordering.strict_boundaries: - continue - return False + if isinstance(scheme, OrderScheme): + return scheme.orderings[0].strict_boundaries return True - def is_strictly_sorted(self, order_keys: Sequence[OrderKey]) -> bool: - """True if the selected ordering proves sortedness for order_keys.""" - if not self or not isinstance(self.inter_rank_scheme, OrderScheme): - return False - ordering = self.inter_rank_scheme.orderings[0] + @staticmethod + def _ordering_covers_keys( + ordering: Ordering, + order_keys: Sequence[int | OrderKey], + ) -> bool: + """True when an ordering covers the requested sort keys.""" + ordering_keys = ordering.keys if len(ordering.keys) < len(order_keys): # If we are only sorted on a subset of the keys, we need strict # boundaries to know later keys cannot interleave across chunks. - return ordering.strict_boundaries + if not ordering.strict_boundaries: + return False + order_keys = order_keys[: len(ordering.keys)] + else: + ordering_keys = ordering.keys[: len(order_keys)] + for current, target in zip(ordering_keys, order_keys, strict=True): + if isinstance(target, OrderKey): + if current != target: + return False + elif current.column_index != target: + return False return True + def is_strictly_partitioned( + self, + *, + level: PartitioningLevel = "flat", + ) -> bool: + """True if data is strictly partitioned at the requested level.""" + return self._scheme_is_strict(self._scheme_for_level(level)) + + def is_ordered( + self, + order_keys: Sequence[int | OrderKey], + *, + level: PartitioningLevel = "flat", + ) -> bool: + """True if the selected ordering covers order_keys.""" + scheme = self._scheme_for_level(level) + if not isinstance(scheme, OrderScheme): + return False + return self._ordering_covers_keys(scheme.orderings[0], order_keys) + def is_aligned_with( self, other: NormalizedPartitioning, br: BufferResource ) -> bool: @@ -1317,9 +1364,17 @@ def _schemes_aligned( ) return lhs == "inherit" and rhs == "inherit" + def _local_schemes_strict( + lhs: PartitioningScheme, rhs: PartitioningScheme + ) -> bool: + return (lhs == "inherit" and rhs == "inherit") or ( + self._scheme_is_strict(lhs) and self._scheme_is_strict(rhs) + ) + return ( - self.is_strictly_partitioned() - and other.is_strictly_partitioned() + self.is_strictly_partitioned(level="inter_rank") + and other.is_strictly_partitioned(level="inter_rank") + and _local_schemes_strict(self.local_scheme, other.local_scheme) and _schemes_aligned(self.inter_rank_scheme, other.inter_rank_scheme) and _schemes_aligned(self.local_scheme, other.local_scheme) ) diff --git a/python/cudf_polars/tests/streaming/test_groupby.py b/python/cudf_polars/tests/streaming/test_groupby.py index da6db146ca68..65dd21ffbce8 100644 --- a/python/cudf_polars/tests/streaming/test_groupby.py +++ b/python/cudf_polars/tests/streaming/test_groupby.py @@ -83,6 +83,25 @@ async def fake_allgather_reduce(_context, _comm, _op_id, *local_values): assert tracer.decision == "shuffle" +@pytest.mark.parametrize( + "nranks,npartitions,expected", + [ + (2, 5, [3, 2]), + (3, 5, [2, 2, 1]), + (4, 10, [3, 2, 3, 2]), + ], +) +def test_partition_count_for_rank_uses_contiguous_ownership( + nranks, npartitions, expected +): + """GroupBy metadata uses the same uneven partition ownership as adjust_ordering.""" + counts = [ + groupby_actor_graph._partition_count_for_rank(rank, nranks, npartitions) + for rank in range(nranks) + ] + assert counts == expected + + @pytest.mark.parametrize("keys", [("key",), ("key", "key2")]) @pytest.mark.parametrize("agg", ["sum", "mean", "len", "min", "max"]) def test_dynamic_groupby_basic(df, streaming_engine, keys, agg): @@ -111,6 +130,27 @@ def test_dynamic_groupby_shuffle_strategy(streaming_engine_factory): assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) +@pytest.mark.parametrize("group_keys", [("key", "subkey"), ("key",)]) +def test_dynamic_groupby_after_sort_on_group_keys(spmd_engine_factory, group_keys): + """Group sorted data by the full sort key set or a sorted-key prefix.""" + streaming_engine = spmd_engine_factory( + StreamingOptions(target_partition_size=128), + ) + df = pl.LazyFrame( + { + "key": [0] * 16 + [1] * 16 + [2] * 16 + [3] * 16, + "subkey": ([0] * 8 + [1] * 8) * 4, + "value": range(64), + } + ) + q = ( + df.sort("key", "subkey") + .group_by(*group_keys, maintain_order=True) + .agg(pl.col("value").sum()) + ) + assert_gpu_result_equal(q, engine=streaming_engine) + + def test_dynamic_groupby_single_group(streaming_engine): """Test dynamic groupby where all rows have the same key.""" df = pl.LazyFrame({"key": [1] * 100, "value": range(100)}) diff --git a/python/cudf_polars/tests/streaming/test_metadata.py b/python/cudf_polars/tests/streaming/test_metadata.py index c9a17fce6ee8..b312ff8f0d5c 100644 --- a/python/cudf_polars/tests/streaming/test_metadata.py +++ b/python/cudf_polars/tests/streaming/test_metadata.py @@ -547,8 +547,21 @@ def test_is_strictly_partitioned_order_scheme(spmd_engine): strict = _make_order_scheme(spmd_engine.context, strict=True) non_strict = _make_order_scheme(spmd_engine.context, strict=False) assert NormalizedPartitioning(strict, "inherit").is_strictly_partitioned() + assert not NormalizedPartitioning(strict, "inherit").is_strictly_partitioned( + level="local" + ) assert not NormalizedPartitioning(non_strict, "inherit").is_strictly_partitioned() assert not NormalizedPartitioning(strict, non_strict).is_strictly_partitioned() + assert NormalizedPartitioning(strict, non_strict).is_strictly_partitioned( + level="inter_rank" + ) + assert ( + NormalizedPartitioning(strict, non_strict).is_strictly_partitioned( + level="local" + ) + is False + ) + assert NormalizedPartitioning(strict, strict).is_strictly_partitioned(level="local") def test_is_aligned_with_order_scheme(spmd_engine): @@ -898,15 +911,15 @@ def test_sort_output_metadata(spmd_engine_factory, by, descending, nulls_last) - @pytest.mark.parametrize( - "scheme_key_count,strict,expected", + "scheme_key_count,strict_boundaries,expected", [ - (1, True, True), # prefix match + strict → skip - (1, False, False), # prefix match + non-strict → no skip - (2, True, True), # exact match + strict → skip - (2, False, True), # exact match + non-strict → skip (strict irrelevant) + (1, True, True), # prefix match + strict → sorted + (1, False, False), # prefix match + non-strict → not sorted + (2, True, True), # exact match + strict → sorted + (2, False, True), # exact match + non-strict → sorted ], ) -def test_is_strictly_sorted(spmd_engine, scheme_key_count, strict, expected) -> None: +def test_is_ordered(spmd_engine, scheme_key_count, strict_boundaries, expected) -> None: df_lf = pl.LazyFrame({"x": list(range(5)), "y": list(range(5))}) base_ir = Translator(df_lf._ldf.visit(), spmd_engine).translate_ir() asc, before = plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE @@ -936,7 +949,9 @@ def test_is_strictly_sorted(spmd_engine, scheme_key_count, strict, expected) -> exclusive_view=False, br=ctx.br(), ) - scheme = OrderScheme([Ordering(keys, boundary_chunk, strict_boundaries=strict)]) + scheme = OrderScheme( + [Ordering(keys, boundary_chunk, strict_boundaries=strict_boundaries)] + ) meta = ChannelMetadata( 3, partitioning=Partitioning(inter_rank=scheme, local="inherit") ) @@ -945,4 +960,23 @@ def test_is_strictly_sorted(spmd_engine, scheme_key_count, strict, expected) -> partitioning = NormalizedPartitioning.from_keys( meta.partitioning, nranks=1, keys=order_keys ) - assert partitioning.is_strictly_sorted(order_keys) is expected + assert partitioning.is_ordered(order_keys) is expected + assert partitioning.is_ordered(order_keys, level="flat") is expected + assert not partitioning.is_ordered(order_keys, level="local") + + nested = NormalizedPartitioning(scheme, scheme) + assert not nested.is_ordered(order_keys) + assert nested.is_ordered(order_keys, level="inter_rank") is expected + assert nested.is_ordered(order_keys, level="local") is expected + + if scheme_key_count == 2: + desc, after = plc.types.Order.DESCENDING, plc.types.NullOrder.AFTER + mismatched_order_keys = [ + (OrderKey(1, asc, before), OrderKey(0, asc, before)), + (OrderKey(0, desc, before), OrderKey(1, asc, before)), + (OrderKey(0, asc, after), OrderKey(1, asc, before)), + ] + for mismatched_keys in mismatched_order_keys: + assert not partitioning.is_ordered(mismatched_keys) + assert not nested.is_ordered(mismatched_keys, level="inter_rank") + assert not nested.is_ordered(mismatched_keys, level="local")