WIP: Fetch parquet metadata in streaming network

Author: TomAugspurger

python/cudf_polars/cudf_polars/dsl/utils/io.py

@@ -13,7 +13,7 @@
 
 from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars
 from cudf_polars.dsl.traversal import traversal
-from cudf_polars.streaming.io import Scan, StreamingScan
+from cudf_polars.streaming.io import Scan
 
 if TYPE_CHECKING:
     from cudf_polars.dsl.ir import IR
@@ -105,6 +105,70 @@ def _prefetch_parquet_footers_for_paths(paths: list[str]) -> list[CachedParquetI
     ]
 
 
+def _cached_parquet_info_from_stats(
+    stats: StatsCollector | None,
+) -> dict[str, CachedParquetInfo]:
+    """Return path -> cached parquet info seeded from statistics collection."""
+    from cudf_polars.streaming.io import ParquetSourceInfo, Scan
+
+    cached_parquet_info: dict[str, CachedParquetInfo] = {}
+    if stats is None:
+        return cached_parquet_info
+
+    for node, datasource_info in stats.scan_stats.items():
+        if (
+            isinstance(node, Scan)
+            and node.typ == "parquet"
+            and isinstance(datasource_info, ParquetSourceInfo)
+            and datasource_info.cached_parquet_info is not None
+        ):
+            for info in datasource_info.cached_parquet_info:
+                cached_parquet_info[info.path] = info
+    return cached_parquet_info
+
+
+@nvtx_annotate_cudf_polars(message="prefetch_cached_parquet_info_for_paths")
+def prefetch_cached_parquet_info_for_paths(
+    paths: list[str],
+    *,
+    stats: StatsCollector | None = None,
+    py_executor: concurrent.futures.Executor | None = None,
+) -> list[CachedParquetInfo]:
+    """
+    Prefetch parquet metadata for a path group.
+
+    Reuses footers already collected during statistics gathering when
+    available and fetches any remaining paths.
+
+    Parameters
+    ----------
+    paths
+        Ordered list of parquet file paths for one scan task group.
+    stats
+        Optional statistics collector with already-cached footers.
+    py_executor
+        Thread pool used when prefetch must run on a worker thread.
+
+    Returns
+    -------
+    Cached parquet metadata ordered to match ``paths``.
+    """
+    cached_by_path = _cached_parquet_info_from_stats(stats)
+    missing_paths = [path for path in paths if path not in cached_by_path]
+
+    if missing_paths:
+        if py_executor is None:
+            fetched = _prefetch_parquet_footers_for_paths(missing_paths)
+        else:
+            fetched = py_executor.submit(
+                _prefetch_parquet_footers_for_paths, missing_paths
+            ).result()
+        for info in fetched:
+            cached_by_path[info.path] = info
+
+    return [cached_by_path[path] for path in paths]
+
+
 @nvtx_annotate_cudf_polars(message="prefetch_parquet_file_metadata_for_ir")
 def prefetch_parquet_file_metadata_for_ir(
     root: IR,
@@ -130,7 +194,7 @@ def prefetch_parquet_file_metadata_for_ir(
     -------
     A dictionary mapping each individual path to its cached parquet metadata.
     """
-    from cudf_polars.streaming.io import ParquetSourceInfo, StreamingScan
+    from cudf_polars.streaming.io import StreamingScan
 
     all_paths: set[str] = set()
 
@@ -142,18 +206,7 @@ def prefetch_parquet_file_metadata_for_ir(
         elif isinstance(node, Scan) and node.typ == "parquet":  # pragma: no cover
             raise RuntimeError("Unexpected parquet 'Scan' node in lowered IR graph.")
 
-    cached_parquet_info: dict[str, CachedParquetInfo] = {}
-    if stats is not None:
-        for node, datasource_info in stats.scan_stats.items():
-            if (
-                isinstance(node, Scan)
-                and node.typ == "parquet"
-                and isinstance(datasource_info, ParquetSourceInfo)
-                and datasource_info.cached_parquet_info is not None
-            ):
-                for info in datasource_info.cached_parquet_info:
-                    cached_parquet_info[info.path] = info
-
+    cached_parquet_info = _cached_parquet_info_from_stats(stats)
     missing_paths = all_paths - set(cached_parquet_info.keys())
     cm: contextlib.AbstractContextManager[concurrent.futures.Executor | None]
 
@@ -173,28 +226,3 @@ def prefetch_parquet_file_metadata_for_ir(
             for info in future.result():
                 cached_parquet_info[info.path] = info
     return cached_parquet_info
-
-
-def attach_cached_parquet_metadata(
-    root: IR,
-    cached_parquet_info_map: dict[str, CachedParquetInfo],
-) -> None:
-    """
-    Attach prefetched metadata to scan nodes.
-
-    This is an optimization only and does not affect IR identity.
-
-    Parameters
-    ----------
-    root
-        Root of the IR graph to update.
-    cached_parquet_info_map
-        Mapping from file paths to cached parquet metadata.
-    """
-    for node in traversal([root]):
-        if isinstance(node, StreamingScan):
-            for scan in node.scans:
-                cached = [cached_parquet_info_map[path] for path in scan.paths]
-                Scan._validate_cached_parquet_info(scan.paths, cached)
-                scan.cached_parquet_info = cached
-                scan._non_child_args = (*scan._non_child_args[:-1], cached)

python/cudf_polars/cudf_polars/engine/core.py

@@ -26,10 +26,6 @@
 
 from cudf_polars.containers import DataFrame
 from cudf_polars.dsl.ir import IRExecutionContext
-from cudf_polars.dsl.utils.io import (
-    attach_cached_parquet_metadata,
-    prefetch_parquet_file_metadata_for_ir,
-)
 from cudf_polars.streaming.actor_graph.collectives import ReserveOpIDs
 from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id
 from cudf_polars.streaming.actor_graph.core import generate_network
@@ -696,14 +692,6 @@ def evaluate_on_rank(
         py_executor, get_cuda_stream=ctx.br().stream_pool.get_stream, query_id=query_id
     )
 
-    if config_options.parquet_options.prefetch_file_metadata:
-        cached_parquet_info_map = prefetch_parquet_file_metadata_for_ir(
-            ir,
-            ir_context.py_executor,
-            stats=stats,
-        )
-        attach_cached_parquet_metadata(ir, cached_parquet_info_map)
-
     with ReserveOpIDs(ir, config_options) as collective_id_map:
         return execute_ir_on_rank(
             ctx,

python/cudf_polars/cudf_polars/streaming/actor_graph/core.py

@@ -19,11 +19,12 @@
 )
 from cudf_polars.dsl.traversal import CachingVisitor, traversal
 from cudf_polars.streaming.actor_graph.dispatch import FanoutInfo
+from cudf_polars.streaming.actor_graph.io import parquet_metadata_prefetch_node
 from cudf_polars.streaming.actor_graph.nodes import (
     generate_ir_sub_network_wrapper,
     metadata_drain_node,
 )
-from cudf_polars.streaming.io import StreamingScan
+from cudf_polars.streaming.io import StreamingScan, can_use_native_parquet_node
 from cudf_polars.streaming.over import Over
 from cudf_polars.utils.config import SPMDContext
 
@@ -35,6 +36,7 @@
     from cudf_streaming.channel_metadata import ChannelMetadata
     from cudf_streaming.table_chunk import TableChunk
     from rapidsmpf.communicator.communicator import Communicator
+    from rapidsmpf.streaming.chunks.arbitrary import ArbitraryChunk
     from rapidsmpf.streaming.core.channel import Channel
     from rapidsmpf.streaming.core.context import Context
     from rapidsmpf.streaming.core.leaf_actor import DeferredMessages
@@ -44,6 +46,7 @@
         GenState,
         SubNetGenerator,
     )
+    from cudf_polars.streaming.actor_graph.io import MetadataMessagePayload
     from cudf_polars.streaming.base import PartitionInfo, StatsCollector
     from cudf_polars.streaming.parallel import ConfigOptions
     from cudf_polars.utils.config import StreamingExecutor
@@ -200,6 +203,66 @@ def _mark_children_unbounded(node: IR) -> None:
     return fanout_nodes
 
 
+def _collect_scan_metadata_groups(
+    ir: IR,
+    *,
+    partition_info: MutableMapping[IR, PartitionInfo],
+    config_options: ConfigOptions,
+    nranks: int,
+) -> dict[tuple[str, ...], set[StreamingScan]]:
+    groups: defaultdict[tuple[str, ...], set[StreamingScan]] = defaultdict(set)
+    if not config_options.parquet_options.prefetch_file_metadata:
+        return {}
+
+    for node in traversal([ir]):
+        if not isinstance(node, StreamingScan):
+            continue
+        if node.base_scan.typ != "parquet":
+            continue
+        node_partition_info = partition_info[node]
+        assert node_partition_info.io_plan is not None, (
+            "Scan node must have a partition plan"
+        )
+        use_native = can_use_native_parquet_node(
+            node.base_scan,
+            plan=node_partition_info.io_plan,
+            count=node_partition_info.count,
+            nranks=nranks,
+            parquet_options=config_options.parquet_options,
+            config_options=config_options,
+        )
+        if use_native:
+            continue
+        for scan in node.scans:
+            groups[tuple(scan.paths)].add(node)
+    return dict(groups)
+
+
+def _build_scan_metadata_channels(
+    context: Context,
+    metadata_scan_groups: dict[tuple[str, ...], set[StreamingScan]],
+) -> tuple[
+    dict[tuple[str, ...], tuple[Channel[ArbitraryChunk[MetadataMessagePayload]], ...]],
+    dict[
+        StreamingScan,
+        dict[tuple[str, ...], Channel[ArbitraryChunk[MetadataMessagePayload]]],
+    ],
+]:
+    metadata_group_channels: dict[
+        tuple[str, ...], tuple[Channel[ArbitraryChunk[MetadataMessagePayload]], ...]
+    ] = {}
+    metadata_channels_by_scan: dict[
+        StreamingScan,
+        dict[tuple[str, ...], Channel[ArbitraryChunk[MetadataMessagePayload]]],
+    ] = {}
+    for key, scans in metadata_scan_groups.items():
+        channels = tuple(context.create_channel() for _ in scans)
+        metadata_group_channels[key] = channels
+        for scan, channel in zip(sorted(scans, key=id), channels, strict=True):
+            metadata_channels_by_scan.setdefault(scan, {})[key] = channel
+    return metadata_group_channels, metadata_channels_by_scan
+
+
 def generate_network(
     context: Context,
     comm: Communicator,
@@ -257,6 +320,15 @@ def generate_network(
     # Get max_io_threads from config (default: 2)
     max_io_threads_global = config_options.executor.max_io_threads
     max_io_threads_local = max(1, max_io_threads_global // max(1, num_io_nodes))
+    metadata_scan_groups = _collect_scan_metadata_groups(
+        ir,
+        partition_info=partition_info,
+        config_options=config_options,
+        nranks=comm.nranks,
+    )
+    metadata_group_channels, metadata_channels_by_scan = _build_scan_metadata_channels(
+        context, metadata_scan_groups
+    )
 
     # Generate the network
     state: GenState = {
@@ -269,12 +341,26 @@ def generate_network(
         "max_io_threads": max_io_threads_local,
         "stats": stats,
         "collective_id_map": collective_id_map,
+        "metadata_scan_groups": metadata_scan_groups,
+        "metadata_group_channels": metadata_group_channels,
+        "metadata_channels_by_scan": metadata_channels_by_scan,
     }
     mapper: SubNetGenerator = CachingVisitor(
         generate_ir_sub_network_wrapper, state=state
     )
     nodes_dict, channels = mapper(ir)
     ch_out = channels[ir].reserve_output_slot()
+    metadata_nodes = [
+        parquet_metadata_prefetch_node(
+            context,
+            ir_context,
+            key,
+            state["metadata_group_channels"][key],
+            stats,
+            sorted(scans, key=id)[0],
+        )
+        for key, scans in state["metadata_scan_groups"].items()
+    ]
 
     # Add node to drain metadata before pull_from_channel
     # (since pull_from_channel doesn't handle metadata messages)
@@ -294,6 +380,7 @@ def generate_network(
 
     # Flatten the nodes dictionary into a list for run_actor_network
     nodes: list[Any] = [node for node_list in nodes_dict.values() for node in node_list]
+    nodes.extend(metadata_nodes)
     nodes.extend([drain_node, output_node])
 
     # Return network and output hook

python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.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
 """Dispatching for the RapidsMPF streaming runtime."""
 
@@ -13,14 +13,15 @@
     from collections.abc import MutableMapping
 
     from rapidsmpf.communicator.communicator import Communicator
+    from rapidsmpf.streaming.chunks.arbitrary import ArbitraryChunk
+    from rapidsmpf.streaming.core.channel import Channel
     from rapidsmpf.streaming.core.context import Context
 
     from cudf_polars.dsl.ir import IR, IRExecutionContext
+    from cudf_polars.streaming.actor_graph.io import MetadataMessagePayload
     from cudf_polars.streaming.actor_graph.utils import ChannelManager
-    from cudf_polars.streaming.base import (
-        PartitionInfo,
-        StatsCollector,
-    )
+    from cudf_polars.streaming.base import PartitionInfo, StatsCollector
+    from cudf_polars.streaming.io import StreamingScan
     from cudf_polars.utils.config import ConfigOptions, StreamingExecutor
 
 
@@ -58,6 +59,14 @@ class GenState(TypedDict):
         Statistics collector.
     collective_id_map
         The mapping of IR nodes to lists of collective IDs.
+    metadata_scan_groups
+        Mapping from parquet metadata group key to dependent StreamingScan nodes.
+    metadata_group_channels
+        Mapping from parquet metadata group key to output channels for each
+        dependent StreamingScan node.
+    metadata_channels_by_scan
+        Mapping from each StreamingScan node to its parquet metadata input
+        channels, keyed by parquet metadata group key.
     """
 
     context: Context
@@ -69,6 +78,14 @@ class GenState(TypedDict):
     max_io_threads: int
     stats: StatsCollector
     collective_id_map: dict[IR, list[int]]
+    metadata_scan_groups: dict[tuple[str, ...], set[StreamingScan]]
+    metadata_group_channels: dict[
+        tuple[str, ...], tuple[Channel[ArbitraryChunk[MetadataMessagePayload]], ...]
+    ]
+    metadata_channels_by_scan: dict[
+        StreamingScan,
+        dict[tuple[str, ...], Channel[ArbitraryChunk[MetadataMessagePayload]]],
+    ]
 
 
 SubNetGenerator: TypeAlias = GenericTransformer[