From fd0b5af89c538152d5107d74821366e982592781 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 20 Aug 2026 08:43:49 -0700 Subject: [PATCH 1/4] add initial hint_sorted_actor --- .../streaming/actor_graph/__init__.py | 3 +- .../streaming/actor_graph/groupby.py | 1 + .../streaming/actor_graph/hint_sorted.py | 197 ++++++++++++++++++ .../streaming/actor_graph/nodes.py | 24 ++- .../cudf_polars/streaming/actor_graph/over.py | 1 + .../streaming/actor_graph/utils.py | 81 ++++++- .../tests/streaming/test_metadata.py | 122 ++++++++++- 7 files changed, 407 insertions(+), 22 deletions(-) create mode 100644 python/cudf_polars/cudf_polars/streaming/actor_graph/hint_sorted.py diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py index a4dd049a0eae..c78c8084ae54 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """RapidsMPF streaming-engine support.""" @@ -12,6 +12,7 @@ # ``@generate_ir_sub_network.register(...)`` handlers at import time so the # dispatch table is populated before any query is evaluated. import cudf_polars.streaming.actor_graph.groupby +import cudf_polars.streaming.actor_graph.hint_sorted import cudf_polars.streaming.actor_graph.io import cudf_polars.streaming.actor_graph.join import cudf_polars.streaming.actor_graph.over 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..9eff98df753d 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py @@ -695,6 +695,7 @@ async def groupby_actor( ch_out, ch_in, metadata_out, + input_metadata=metadata_in, tracer=tracer, ) return diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/hint_sorted.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/hint_sorted.py new file mode 100644 index 000000000000..42b78852f2a7 --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/hint_sorted.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Streaming actor for ``MapFunction("hint_sorted")``.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, cast + +import pylibcudf as plc +from cudf_streaming.channel_metadata import ( + ChannelMetadata, + OrderKey, + OrderScheme, + Ordering, + Partitioning, +) +from cudf_streaming.table_chunk import TableChunk +from rapidsmpf.streaming.core.actor import define_actor + +from cudf_polars.dsl.ir import IR, MapFunction +from cudf_polars.dsl.utils.naming import names_to_indices +from cudf_polars.streaming.actor_graph.dispatch import generate_ir_sub_network +from cudf_polars.streaming.actor_graph.utils import ( + ChannelManager, + process_children, + recv_metadata, + send_metadata, + shutdown_on_error, +) +from cudf_polars.utils import sorting +from cudf_polars.utils.dtypes import make_empty_column + +if TYPE_CHECKING: + from rapidsmpf.communicator.communicator import Communicator + from rapidsmpf.streaming.core.channel import Channel + from rapidsmpf.streaming.core.context import Context + + from cudf_polars.dsl.ir import IRExecutionContext + from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator + + +def _hint_sorted_options( + ir: MapFunction, +) -> tuple[tuple[str, ...], tuple[bool, ...], tuple[bool, ...]]: + """Return normalized ``hint_sorted`` options.""" + return cast( + "tuple[tuple[str, ...], tuple[bool, ...], tuple[bool, ...]]", ir.options + ) + + +def _hint_sorted_order_keys(ir: MapFunction) -> list[OrderKey]: + """Convert ``MapFunction("hint_sorted")`` options to ordering keys.""" + column_names, descending, nulls_last = _hint_sorted_options(ir) + orders, null_orders = sorting.sort_order( + descending, nulls_last=nulls_last, num_keys=len(column_names) + ) + return [ + OrderKey(index, order, null_order) + for index, order, null_order in zip( + names_to_indices(column_names, ir.schema), + orders, + null_orders, + strict=True, + ) + ] + + +def _order_scheme_has_keys(scheme: OrderScheme, keys: list[OrderKey]) -> bool: + """Check for an exact ordering match.""" + return any(list(ordering.keys) == keys for ordering in scheme.orderings) + + +def _metadata_satisfies_hint(metadata: ChannelMetadata, keys: list[OrderKey]) -> bool: + """Check whether existing metadata already advertises the requested ordering.""" + if metadata.partitioning is None: + return False + scheme = metadata.partitioning.inter_rank + return isinstance(scheme, OrderScheme) and _order_scheme_has_keys(scheme, keys) + + +def _trivial_ordering_metadata( + context: Context, + comm: Communicator, + ir: MapFunction, + metadata: ChannelMetadata, + keys: list[OrderKey], +) -> ChannelMetadata | None: + """Temporary policy: attach ordering only when boundaries are trivial.""" + if comm.nranks != 1 or metadata.local_count > 1: + return None + + partitioning = metadata.partitioning + existing_orderings: list[Ordering] = [] + local = "inherit" + if partitioning is not None: + local = partitioning.local + if isinstance(partitioning.inter_rank, OrderScheme): + existing_orderings = list(partitioning.inter_rank.orderings) + + column_names = _hint_sorted_options(ir)[0] + stream = context.br().stream_pool.get_stream() + boundaries = TableChunk.from_pylibcudf_table( + plc.Table( + [make_empty_column(ir.schema[name], stream) for name in column_names] + ), + stream, + exclusive_view=False, + br=context.br(), + ) + ordering = Ordering(keys, boundaries, strict_boundaries=True) + return ChannelMetadata( + local_count=metadata.local_count, + partitioning=Partitioning( + OrderScheme([*existing_orderings, ordering]), + local, + ), + duplicated=metadata.duplicated, + ) + + +async def extract_hint_sorted_metadata( + context: Context, + comm: Communicator, + ir: MapFunction, + ir_context: IRExecutionContext, + metadata: ChannelMetadata, + ch_in: Channel[TableChunk], + ch_replay: Channel[TableChunk], +) -> tuple[ChannelMetadata, Channel[TableChunk]]: + """Resolve output metadata and the channel to forward for ``hint_sorted``.""" + keys = _hint_sorted_order_keys(ir) + if not keys or _metadata_satisfies_hint(metadata, keys): + return metadata, ch_in + + # Future policy hook: use downstream partitioning hints to decide whether + # to extract real boundaries. The extraction path will consume ``ch_in``, + # replay consumed data through ``ch_replay``, and return ``ch_replay`` as + # the forwarding channel. + + # For now, only the trivial single-partition case can synthesize correct + # strict boundaries without a collective. + trivial_metadata = _trivial_ordering_metadata(context, comm, ir, metadata, keys) + return (metadata if trivial_metadata is None else trivial_metadata), ch_in + + +@define_actor() +async def hint_sorted_actor( + context: Context, + comm: Communicator, + ir: MapFunction, + ir_context: IRExecutionContext, + ch_out: Channel[TableChunk], + ch_in: Channel[TableChunk], + ch_replay: Channel[TableChunk], +) -> None: + """Forward data and attach safe ordering metadata for ``hint_sorted``.""" + async with shutdown_on_error( + context, ch_in, ch_replay, ch_out, trace_ir=ir, ir_context=ir_context + ): + metadata = await recv_metadata(ch_in, context) + metadata, ch_forward = await extract_hint_sorted_metadata( + context, + comm, + ir, + ir_context, + metadata, + ch_in, + ch_replay, + ) + await send_metadata(ch_out, context, metadata) + while (msg := await ch_forward.recv(context)) is not None: + await ch_out.send(context, msg) + await ch_out.drain(context) + + +@generate_ir_sub_network.register(MapFunction) +def _( + ir: MapFunction, rec: SubNetGenerator +) -> tuple[dict[IR, list[Any]], dict[IR, ChannelManager]]: + if ir.name != "hint_sorted": + return generate_ir_sub_network.dispatch(IR)(ir, rec) + + nodes, channels = process_children(ir, rec) + channels[ir] = ChannelManager(rec.state["context"]) + ch_replay = rec.state["context"].create_channel() + nodes[ir] = [ + hint_sorted_actor( + rec.state["context"], + rec.state["comm"], + ir, + rec.state["ir_context"], + channels[ir].reserve_input_slot(), + channels[ir.children[0]].reserve_output_slot(), + ch_replay, + ) + ] + return nodes, channels diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py index b0ca603bec47..e575ad1c7605 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py @@ -18,7 +18,6 @@ from rapidsmpf.streaming.core.message import Message from rapidsmpf.streaming.core.spillable_messages import SpillableMessages -from cudf_polars.containers import DataFrame from cudf_polars.dsl.ir import IR, Empty from cudf_polars.streaming.actor_graph.dispatch import ( generate_ir_sub_network, @@ -26,6 +25,7 @@ from cudf_polars.streaming.actor_graph.tracing import send_chunk from cudf_polars.streaming.actor_graph.utils import ( ChannelManager, + chunk_to_frame, chunkwise_evaluate, empty_table_chunk, gather_in_task_group, @@ -42,6 +42,7 @@ from rapidsmpf.streaming.core.channel import Channel from rapidsmpf.streaming.core.context import Context + from cudf_polars.containers import DataFrame from cudf_polars.dsl.ir import IRExecutionContext from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator @@ -95,6 +96,7 @@ async def default_node_single( ch_out, ch_in, metadata_out, + input_metadata=metadata_in, handle_empty_input=True, tracer=tracer, ) @@ -136,9 +138,10 @@ async def default_node_multi( local_count = 1 duplicated = True partitioning = None - for idx, md_child in enumerate( - await gather_in_task_group(*(recv_metadata(ch, context) for ch in chs_in)) - ): + child_metadatas = await gather_in_task_group( + *(recv_metadata(ch, context) for ch in chs_in) + ) + for idx, md_child in enumerate(child_metadatas): # Use simple "max" rule to determine counts. local_count = max(md_child.local_count, local_count) # Set "duplicated" to False as soon as we @@ -209,13 +212,14 @@ async def default_node_multi( net_memory_delta=0, ) dfs = [ - DataFrame.from_table( - chunk.table_view(), # type: ignore[union-attr] - list(child.schema.keys()), - list(child.schema.values()), - chunk.stream, # type: ignore[union-attr] + chunk_to_frame( + cast("TableChunk", chunk), + child, + metadata=child_metadata, + ) + for chunk, child, child_metadata in zip( + ready_chunks, ir.children, child_metadatas, strict=True ) - for chunk, child in zip(ready_chunks, ir.children, strict=True) ] with opaque_memory_usage(extra): df = await ir_context.to_thread( 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..dd3fba6cf41f 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py @@ -765,6 +765,7 @@ async def over_actor( ch_out, ch_in, metadata_out, + input_metadata=metadata_in, tracer=tracer, ) return 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..40079e704e5b 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -725,6 +725,8 @@ def _evaluate_chunk_sync( ir: IR, ir_context: IRExecutionContext, br: BufferResource, + *, + metadata: ChannelMetadata | None = None, ) -> TableChunk: """ Apply an IR node's do_evaluate to a table chunk (synchronous). @@ -742,17 +744,17 @@ def _evaluate_chunk_sync( The IR execution context. br The buffer resource for lifetime tracking. + metadata + Optional channel metadata to synthesize local DataFrame metadata from. Returns ------- The resulting table chunk after evaluation. """ - input_schema = ir.children[0].schema - names = list(input_schema.keys()) - dtypes = list(input_schema.values()) + df_in = chunk_to_frame(chunk, ir.children[0], metadata=metadata) df = ir.do_evaluate( *ir._non_child_args, - DataFrame.from_table(chunk.table_view(), names, dtypes, chunk.stream), + df_in, context=ir_context, ) return TableChunk.from_pylibcudf_table( @@ -765,6 +767,7 @@ async def evaluate_chunk( chunk: TableChunk, *irs: IR, ir_context: IRExecutionContext, + metadata: ChannelMetadata | None = None, ) -> TableChunk: """ Make chunk available, reserve memory, and evaluate. @@ -780,6 +783,9 @@ async def evaluate_chunk( in order within a single memory reservation. ir_context The IR execution context. + metadata + Optional channel metadata to synthesize local DataFrame metadata from + during the first evaluation. Returns ------- @@ -795,8 +801,14 @@ async def evaluate_chunk( with opaque_memory_usage(extra): for single_ir in irs: chunk = await ir_context.to_thread( - _evaluate_chunk_sync, chunk, single_ir, ir_context, context.br() + _evaluate_chunk_sync, + chunk, + single_ir, + ir_context, + context.br(), + metadata=metadata, ) + metadata = None return chunk @@ -935,11 +947,12 @@ async def chunkwise_evaluate( ch_in: Channel[TableChunk], metadata: ChannelMetadata, *, + input_metadata: ChannelMetadata | None = None, handle_empty_input: bool = False, tracer: ActorTracer | None = None, ) -> None: """ - Apply IR evaluation chunk-by-chunk, preserving partitioning. + Apply IR evaluation chunk-by-chunk. Use when data is already partitioned on the relevant keys and each chunk can be processed independently. @@ -957,7 +970,10 @@ async def chunkwise_evaluate( ch_in The input channel. metadata - The channel metadata to forward (partitioning preserved). + The channel metadata to forward. + input_metadata + The input metadata to synthesize local DataFrame metadata from. + Defaults to ``metadata`` for partition-preserving callers. handle_empty_input If True and no chunks are received, create an empty chunk and evaluate it. Use for operations like aggregations that always produce output. @@ -983,13 +999,20 @@ async def chunkwise_evaluate( TableChunk.from_message(msg, br=context.br()), ir, ir_context=ir_context, + metadata=metadata if input_metadata is None else input_metadata, ) del msg, cd await send_chunk(context, ch_out, result, seq_num, tracer=tracer) if handle_empty_input and not received_any: chunk = empty_table_chunk(ir.children[0], context, ir_context.get_cuda_stream()) - result = await evaluate_chunk(context, chunk, ir, ir_context=ir_context) + result = await evaluate_chunk( + context, + chunk, + ir, + ir_context=ir_context, + metadata=metadata if input_metadata is None else input_metadata, + ) del chunk await send_chunk(context, ch_out, result, 0, tracer=tracer) @@ -1550,7 +1573,42 @@ def empty_table_chunk(ir: IR, context: Context, stream: Stream) -> TableChunk: ) -def chunk_to_frame(chunk: TableChunk, ir: IR) -> DataFrame: +def _leading_order_keys(metadata: ChannelMetadata | None) -> dict[int, OrderKey]: + """Return unambiguous leading order keys implied by channel metadata.""" + if metadata is None or metadata.partitioning is None: + return {} + + candidates: dict[int, OrderKey | None] = {} + for scheme in (metadata.partitioning.inter_rank, metadata.partitioning.local): + if not isinstance(scheme, OrderScheme): + continue + for ordering in scheme.orderings: + if not ordering.keys: + continue + key = ordering.keys[0] + current = candidates.get(key.column_index, key) + candidates[key.column_index] = key if current == key else None + return {index: key for index, key in candidates.items() if key is not None} + + +def _apply_ordering_metadata( + df: DataFrame, metadata: ChannelMetadata | None +) -> DataFrame: + """Apply safe column-level sortedness metadata implied by ``metadata``.""" + for index, key in _leading_order_keys(metadata).items(): + if index >= df.num_columns: + continue + df.columns[index].set_sorted( + is_sorted=plc.types.Sorted.YES, + order=key.order, + null_order=key.null_order, + ) + return df + + +def chunk_to_frame( + chunk: TableChunk, ir: IR, *, metadata: ChannelMetadata | None = None +) -> DataFrame: """ Convert a TableChunk to a DataFrame. @@ -1560,17 +1618,20 @@ def chunk_to_frame(chunk: TableChunk, ir: IR) -> DataFrame: The TableChunk to convert. ir The IR node to use for the schema. + metadata + Optional channel metadata to synthesize local DataFrame metadata from. Returns ------- A DataFrame. """ - return DataFrame.from_table( + df = DataFrame.from_table( chunk.table_view(), list(ir.schema.keys()), list(ir.schema.values()), chunk.stream, ) + return _apply_ordering_metadata(df, metadata) def _is_already_partitioned( diff --git a/python/cudf_polars/tests/streaming/test_metadata.py b/python/cudf_polars/tests/streaming/test_metadata.py index c9a17fce6ee8..5914b358efc8 100644 --- a/python/cudf_polars/tests/streaming/test_metadata.py +++ b/python/cudf_polars/tests/streaming/test_metadata.py @@ -5,6 +5,8 @@ from __future__ import annotations +import asyncio + import pytest import polars as pl @@ -23,17 +25,29 @@ from cudf_polars import Translator from cudf_polars.containers import DataFrame, DataType from cudf_polars.dsl import expr -from cudf_polars.dsl.ir import GroupBy, HStack, Projection, Select, Sort +from cudf_polars.dsl.ir import ( + DataFrameScan, + GroupBy, + HStack, + IRExecutionContext, + MapFunction, + Projection, + Select, + Sort, +) from cudf_polars.engine.options import StreamingOptions from cudf_polars.streaming.actor_graph.collectives.sort import ( _sort_to_order_keys, ) from cudf_polars.streaming.actor_graph.core import evaluate_logical_plan +from cudf_polars.streaming.actor_graph.hint_sorted import extract_hint_sorted_metadata from cudf_polars.streaming.actor_graph.utils import ( NormalizedPartitioning, + _apply_ordering_metadata, maybe_remap_partitioning, ) from cudf_polars.utils.config import ConfigOptions +from cudf_polars.utils.dtypes import make_empty_column @pytest.fixture(scope="module") @@ -508,6 +522,112 @@ def _make_order_scheme(context, *, key_indices=(0,), values=(100, 200), strict=F ) +def _hint_sorted_ir() -> MapFunction: + schema = {"a": DataType(pl.Int64()), "b": DataType(pl.Int64())} + child = DataFrameScan(schema, pl.DataFrame({"a": [1], "b": [2]})._df, None) + return MapFunction(schema, "hint_sorted", [[("a", False, False)]], child) + + +async def _extract_hint_sorted_metadata(spmd_engine, metadata: ChannelMetadata): + context = spmd_engine.context + ch_in = context.create_channel() + ch_replay = context.create_channel() + result = await extract_hint_sorted_metadata( + context, + spmd_engine.comm, + _hint_sorted_ir(), + IRExecutionContext(), + metadata, + ch_in, + ch_replay, + ) + return (*result, ch_in, ch_replay) + + +def test_hint_sorted_metadata_attaches_single_partition_ordering(spmd_engine) -> None: + metadata = ChannelMetadata(local_count=1) + + result, ch_forward, ch_in, ch_replay = asyncio.run( + _extract_hint_sorted_metadata(spmd_engine, metadata) + ) + + assert result.partitioning is not None + assert isinstance(result.partitioning.inter_rank, OrderScheme) + (ordering,) = result.partitioning.inter_rank.orderings + assert list(ordering.keys) == [ + OrderKey(0, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE) + ] + assert ordering.strict_boundaries is True + assert ch_forward is ch_in + assert ch_forward is not ch_replay + + +def test_hint_sorted_metadata_ignores_multi_partition_without_boundaries( + spmd_engine, +) -> None: + metadata = ChannelMetadata(local_count=2) + + result, ch_forward, ch_in, ch_replay = asyncio.run( + _extract_hint_sorted_metadata(spmd_engine, metadata) + ) + + assert result is metadata + assert ch_forward is ch_in + assert ch_forward is not ch_replay + + +def test_apply_ordering_metadata_marks_leading_key_only(spmd_engine) -> None: + stream = spmd_engine.context.br().stream_pool.get_stream() + df = DataFrame.from_polars(pl.DataFrame({"a": [1, 1, 2], "b": [2, 1, 3]}), stream) + scheme = _make_order_scheme(spmd_engine.context, key_indices=(0, 1)) + metadata = ChannelMetadata( + local_count=1, + partitioning=Partitioning(scheme, local="inherit"), + ) + + result = _apply_ordering_metadata(df, metadata) + + assert result.column_map["a"].is_sorted == plc.types.Sorted.YES + assert result.column_map["b"].is_sorted == plc.types.Sorted.NO + + +def test_apply_ordering_metadata_skips_conflicting_keys(spmd_engine) -> None: + stream = spmd_engine.context.br().stream_pool.get_stream() + df = DataFrame.from_polars(pl.DataFrame({"a": [1, 2, 3]}), stream) + before = plc.types.NullOrder.BEFORE + + def empty_boundaries() -> TableChunk: + return TableChunk.from_pylibcudf_table( + plc.Table([make_empty_column(DataType(pl.Int64()), stream)]), + stream, + exclusive_view=False, + br=spmd_engine.context.br(), + ) + + scheme = OrderScheme( + [ + Ordering( + [OrderKey(0, plc.types.Order.ASCENDING, before)], + empty_boundaries(), + strict_boundaries=True, + ), + Ordering( + [OrderKey(0, plc.types.Order.DESCENDING, before)], + empty_boundaries(), + strict_boundaries=True, + ), + ] + ) + metadata = ChannelMetadata( + local_count=1, + partitioning=Partitioning(scheme, local="inherit"), + ) + + result = _apply_ordering_metadata(df, metadata) + + assert result.column_map["a"].is_sorted == plc.types.Sorted.NO + + @pytest.mark.parametrize( "keys,strict,should_match", [ From 2709d137d229ef1c2c7c012fb207455907b38a9b Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 20 Aug 2026 08:52:18 -0700 Subject: [PATCH 2/4] update comment --- .../cudf_polars/streaming/actor_graph/hint_sorted.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/hint_sorted.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/hint_sorted.py index 42b78852f2a7..a669e863b6d6 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/hint_sorted.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/hint_sorted.py @@ -43,6 +43,7 @@ def _hint_sorted_options( ir: MapFunction, ) -> tuple[tuple[str, ...], tuple[bool, ...], tuple[bool, ...]]: """Return normalized ``hint_sorted`` options.""" + assert ir.name == "hint_sorted" return cast( "tuple[tuple[str, ...], tuple[bool, ...], tuple[bool, ...]]", ir.options ) @@ -136,6 +137,8 @@ async def extract_hint_sorted_metadata( # to extract real boundaries. The extraction path will consume ``ch_in``, # replay consumed data through ``ch_replay``, and return ``ch_replay`` as # the forwarding channel. + # TODO: Integrate replay-capable ``extract_orderscheme_partitioning``. + # See https://github.com/NVIDIA/cudf/pull/22526. # For now, only the trivial single-partition case can synthesize correct # strict boundaries without a collective. From 8e493aa8006833f67654108b2dfd195fd641ea29 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 20 Aug 2026 11:38:31 -0700 Subject: [PATCH 3/4] address coderabbit review --- .../streaming/actor_graph/utils.py | 19 +++++++++++-------- .../tests/streaming/test_metadata.py | 19 +++++++++++++++++++ 2 files changed, 30 insertions(+), 8 deletions(-) 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 40079e704e5b..f7353d049910 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -1578,16 +1578,19 @@ def _leading_order_keys(metadata: ChannelMetadata | None) -> dict[int, OrderKey] if metadata is None or metadata.partitioning is None: return {} + scheme = metadata.partitioning.local + if scheme == "inherit": + scheme = metadata.partitioning.inter_rank + if not isinstance(scheme, OrderScheme): + return {} + candidates: dict[int, OrderKey | None] = {} - for scheme in (metadata.partitioning.inter_rank, metadata.partitioning.local): - if not isinstance(scheme, OrderScheme): + for ordering in scheme.orderings: + if not ordering.keys: continue - for ordering in scheme.orderings: - if not ordering.keys: - continue - key = ordering.keys[0] - current = candidates.get(key.column_index, key) - candidates[key.column_index] = key if current == key else None + key = ordering.keys[0] + current = candidates.get(key.column_index, key) + candidates[key.column_index] = key if current == key else None return {index: key for index, key in candidates.items() if key is not None} diff --git a/python/cudf_polars/tests/streaming/test_metadata.py b/python/cudf_polars/tests/streaming/test_metadata.py index 5914b358efc8..5e72be237cb6 100644 --- a/python/cudf_polars/tests/streaming/test_metadata.py +++ b/python/cudf_polars/tests/streaming/test_metadata.py @@ -591,6 +591,25 @@ def test_apply_ordering_metadata_marks_leading_key_only(spmd_engine) -> None: assert result.column_map["b"].is_sorted == plc.types.Sorted.NO +def test_apply_ordering_metadata_ignores_inter_rank_ordering_if_local_hash( + spmd_engine, +) -> None: + stream = spmd_engine.context.br().stream_pool.get_stream() + df = DataFrame.from_polars(pl.DataFrame({"a": [2, 1, 3]}), stream) + scheme = _make_order_scheme(spmd_engine.context, key_indices=(0,)) + metadata = ChannelMetadata( + local_count=1, + partitioning=Partitioning( + inter_rank=scheme, + local=HashScheme((0,), 1), + ), + ) + + result = _apply_ordering_metadata(df, metadata) + + assert result.column_map["a"].is_sorted == plc.types.Sorted.NO + + def test_apply_ordering_metadata_skips_conflicting_keys(spmd_engine) -> None: stream = spmd_engine.context.br().stream_pool.get_stream() df = DataFrame.from_polars(pl.DataFrame({"a": [1, 2, 3]}), stream) From 1fc293c7ec814901ee8849336b1c8ff044890950 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 21 Aug 2026 07:53:04 -0700 Subject: [PATCH 4/4] address review --- .../streaming/actor_graph/hint_sorted.py | 13 +++-- .../streaming/actor_graph/nodes.py | 13 +++-- .../streaming/actor_graph/utils.py | 54 ++++++++++++------- .../tests/streaming/test_metadata.py | 7 +-- 4 files changed, 56 insertions(+), 31 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/hint_sorted.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/hint_sorted.py index a669e863b6d6..3da749ab6677 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/hint_sorted.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/hint_sorted.py @@ -4,7 +4,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, TypeAlias, cast import pylibcudf as plc from cudf_streaming.channel_metadata import ( @@ -39,14 +39,17 @@ from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator +HintSortedOptions: TypeAlias = tuple[ + tuple[str, ...], tuple[bool, ...], tuple[bool, ...] +] + + def _hint_sorted_options( ir: MapFunction, -) -> tuple[tuple[str, ...], tuple[bool, ...], tuple[bool, ...]]: +) -> HintSortedOptions: """Return normalized ``hint_sorted`` options.""" assert ir.name == "hint_sorted" - return cast( - "tuple[tuple[str, ...], tuple[bool, ...], tuple[bool, ...]]", ir.options - ) + return cast("HintSortedOptions", ir.options) def _hint_sorted_order_keys(ir: MapFunction) -> list[OrderKey]: diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py index e575ad1c7605..732ca77bf82a 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py @@ -25,6 +25,7 @@ from cudf_polars.streaming.actor_graph.tracing import send_chunk from cudf_polars.streaming.actor_graph.utils import ( ChannelManager, + _leading_order_keys, chunk_to_frame, chunkwise_evaluate, empty_table_chunk, @@ -141,6 +142,9 @@ async def default_node_multi( child_metadatas = await gather_in_task_group( *(recv_metadata(ch, context) for ch in chs_in) ) + child_ordering_metadatas = [ + _leading_order_keys(md_child) for md_child in child_metadatas + ] for idx, md_child in enumerate(child_metadatas): # Use simple "max" rule to determine counts. local_count = max(md_child.local_count, local_count) @@ -215,10 +219,13 @@ async def default_node_multi( chunk_to_frame( cast("TableChunk", chunk), child, - metadata=child_metadata, + ordering_metadata=child_ordering_metadata, ) - for chunk, child, child_metadata in zip( - ready_chunks, ir.children, child_metadatas, strict=True + for chunk, child, child_ordering_metadata in zip( + ready_chunks, + ir.children, + child_ordering_metadatas, + strict=True, ) ] with opaque_memory_usage(extra): 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 f7353d049910..a63028d8c063 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -74,6 +74,7 @@ InterRankScheme: TypeAlias = HashScheme | OrderScheme | None PartitioningScheme: TypeAlias = InterRankScheme | Literal["inherit"] +OrderingMetadata: TypeAlias = dict[int, OrderKey] # cuDF column/concatenate row limit (int32) CUDF_ROW_LIMIT = 2**31 - 1 @@ -726,7 +727,7 @@ def _evaluate_chunk_sync( ir_context: IRExecutionContext, br: BufferResource, *, - metadata: ChannelMetadata | None = None, + ordering_metadata: OrderingMetadata | None = None, ) -> TableChunk: """ Apply an IR node's do_evaluate to a table chunk (synchronous). @@ -744,14 +745,15 @@ def _evaluate_chunk_sync( The IR execution context. br The buffer resource for lifetime tracking. - metadata - Optional channel metadata to synthesize local DataFrame metadata from. + ordering_metadata + Optional precomputed ordering metadata to synthesize local DataFrame + metadata from. Returns ------- The resulting table chunk after evaluation. """ - df_in = chunk_to_frame(chunk, ir.children[0], metadata=metadata) + df_in = chunk_to_frame(chunk, ir.children[0], ordering_metadata=ordering_metadata) df = ir.do_evaluate( *ir._non_child_args, df_in, @@ -767,7 +769,7 @@ async def evaluate_chunk( chunk: TableChunk, *irs: IR, ir_context: IRExecutionContext, - metadata: ChannelMetadata | None = None, + ordering_metadata: OrderingMetadata | None = None, ) -> TableChunk: """ Make chunk available, reserve memory, and evaluate. @@ -783,9 +785,9 @@ async def evaluate_chunk( in order within a single memory reservation. ir_context The IR execution context. - metadata - Optional channel metadata to synthesize local DataFrame metadata from - during the first evaluation. + ordering_metadata + Optional precomputed ordering metadata to synthesize local DataFrame + metadata from during the first evaluation. Returns ------- @@ -806,9 +808,9 @@ async def evaluate_chunk( single_ir, ir_context, context.br(), - metadata=metadata, + ordering_metadata=ordering_metadata, ) - metadata = None + ordering_metadata = None return chunk @@ -984,6 +986,10 @@ async def chunkwise_evaluate( if tracer is not None and metadata.duplicated: tracer.set_duplicated() + input_ordering_metadata = _leading_order_keys( + metadata if input_metadata is None else input_metadata + ) + received_any = False while (msg := await ch_in.recv(context)) is not None: received_any = True @@ -999,7 +1005,7 @@ async def chunkwise_evaluate( TableChunk.from_message(msg, br=context.br()), ir, ir_context=ir_context, - metadata=metadata if input_metadata is None else input_metadata, + ordering_metadata=input_ordering_metadata, ) del msg, cd await send_chunk(context, ch_out, result, seq_num, tracer=tracer) @@ -1011,7 +1017,7 @@ async def chunkwise_evaluate( chunk, ir, ir_context=ir_context, - metadata=metadata if input_metadata is None else input_metadata, + ordering_metadata=input_ordering_metadata, ) del chunk await send_chunk(context, ch_out, result, 0, tracer=tracer) @@ -1573,7 +1579,7 @@ def empty_table_chunk(ir: IR, context: Context, stream: Stream) -> TableChunk: ) -def _leading_order_keys(metadata: ChannelMetadata | None) -> dict[int, OrderKey]: +def _leading_order_keys(metadata: ChannelMetadata | None) -> OrderingMetadata: """Return unambiguous leading order keys implied by channel metadata.""" if metadata is None or metadata.partitioning is None: return {} @@ -1595,10 +1601,10 @@ def _leading_order_keys(metadata: ChannelMetadata | None) -> dict[int, OrderKey] def _apply_ordering_metadata( - df: DataFrame, metadata: ChannelMetadata | None + df: DataFrame, ordering_metadata: OrderingMetadata ) -> DataFrame: - """Apply safe column-level sortedness metadata implied by ``metadata``.""" - for index, key in _leading_order_keys(metadata).items(): + """Apply precomputed safe column-level sortedness metadata.""" + for index, key in ordering_metadata.items(): if index >= df.num_columns: continue df.columns[index].set_sorted( @@ -1610,7 +1616,10 @@ def _apply_ordering_metadata( def chunk_to_frame( - chunk: TableChunk, ir: IR, *, metadata: ChannelMetadata | None = None + chunk: TableChunk, + ir: IR, + *, + ordering_metadata: OrderingMetadata | None = None, ) -> DataFrame: """ Convert a TableChunk to a DataFrame. @@ -1621,8 +1630,9 @@ def chunk_to_frame( The TableChunk to convert. ir The IR node to use for the schema. - metadata - Optional channel metadata to synthesize local DataFrame metadata from. + ordering_metadata + Optional precomputed ordering metadata to synthesize local DataFrame + metadata from. Returns ------- @@ -1634,7 +1644,11 @@ def chunk_to_frame( list(ir.schema.values()), chunk.stream, ) - return _apply_ordering_metadata(df, metadata) + return ( + df + if ordering_metadata is None + else _apply_ordering_metadata(df, ordering_metadata) + ) def _is_already_partitioned( diff --git a/python/cudf_polars/tests/streaming/test_metadata.py b/python/cudf_polars/tests/streaming/test_metadata.py index 5e72be237cb6..db45c72f8bc8 100644 --- a/python/cudf_polars/tests/streaming/test_metadata.py +++ b/python/cudf_polars/tests/streaming/test_metadata.py @@ -44,6 +44,7 @@ from cudf_polars.streaming.actor_graph.utils import ( NormalizedPartitioning, _apply_ordering_metadata, + _leading_order_keys, maybe_remap_partitioning, ) from cudf_polars.utils.config import ConfigOptions @@ -585,7 +586,7 @@ def test_apply_ordering_metadata_marks_leading_key_only(spmd_engine) -> None: partitioning=Partitioning(scheme, local="inherit"), ) - result = _apply_ordering_metadata(df, metadata) + result = _apply_ordering_metadata(df, _leading_order_keys(metadata)) assert result.column_map["a"].is_sorted == plc.types.Sorted.YES assert result.column_map["b"].is_sorted == plc.types.Sorted.NO @@ -605,7 +606,7 @@ def test_apply_ordering_metadata_ignores_inter_rank_ordering_if_local_hash( ), ) - result = _apply_ordering_metadata(df, metadata) + result = _apply_ordering_metadata(df, _leading_order_keys(metadata)) assert result.column_map["a"].is_sorted == plc.types.Sorted.NO @@ -642,7 +643,7 @@ def empty_boundaries() -> TableChunk: partitioning=Partitioning(scheme, local="inherit"), ) - result = _apply_ordering_metadata(df, metadata) + result = _apply_ordering_metadata(df, _leading_order_keys(metadata)) assert result.column_map["a"].is_sorted == plc.types.Sorted.NO