Support streaming "hint_sorted" metadata propagation

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

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

@@ -906,6 +906,7 @@ async def groupby_actor(
                 ch_out,
                 ch_in,
                 metadata_out,
+                input_metadata=metadata_in,
                 tracer=tracer,
             )
             return

python/cudf_polars/cudf_polars/streaming/actor_graph/hint_sorted.py

@@ -0,0 +1,203 @@
+# 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, TypeAlias, 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
+
+
+HintSortedOptions: TypeAlias = tuple[
+    tuple[str, ...], tuple[bool, ...], tuple[bool, ...]
+]
+
+
+def _hint_sorted_options(
+    ir: MapFunction,
+) -> HintSortedOptions:
+    """Return normalized ``hint_sorted`` options."""
+    assert ir.name == "hint_sorted"
+    return cast("HintSortedOptions", 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.
+    # 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.
+    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

python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py

@@ -18,14 +18,15 @@
 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,
 )
 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,
     gather_in_task_group,
@@ -42,6 +43,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 +97,7 @@ async def default_node_single(
             ch_out,
             ch_in,
             metadata_out,
+            input_metadata=metadata_in,
             handle_empty_input=True,
             tracer=tracer,
         )
@@ -136,9 +139,13 @@ 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)
+        )
+        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)
             # Set "duplicated" to False as soon as we
@@ -209,13 +216,17 @@ 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,
+                    ordering_metadata=child_ordering_metadata,
+                )
+                for chunk, child, child_ordering_metadata in zip(
+                    ready_chunks,
+                    ir.children,
+                    child_ordering_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(

python/cudf_polars/cudf_polars/streaming/actor_graph/over.py

@@ -767,6 +767,7 @@ async def over_actor(
                 ch_out,
                 ch_in,
                 metadata_out,
+                input_metadata=metadata_in,
                 tracer=tracer,
             )
             return