diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index f61a738d0bbd..d7e1d81f750e 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -1711,6 +1711,17 @@ def is_equal(self, other: Self) -> bool: ) ) + def with_prefetched_metadata( + self, + cached_parquet_info: list[CachedParquetInfo] | None, + ) -> tuple[Any, ...]: + """ + Return ``self.args`` with cached_parquet_info inserted. + + This is a noop for DataFrameScan, which doesn't use parquet metadata. + """ + return self._non_child_args + @classmethod @log_do_evaluate @nvtx_annotate_cudf_polars(message="DataFrameScan") diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index 288dcd79f460..403dd8358770 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -4,19 +4,15 @@ from __future__ import annotations -import concurrent.futures -import contextlib from dataclasses import dataclass from typing import TYPE_CHECKING import pylibcudf as plc 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 from cudf_polars.streaming.base import StatsCollector @@ -105,96 +101,52 @@ def _prefetch_parquet_footers_for_paths(paths: list[str]) -> list[CachedParquetI ] -@nvtx_annotate_cudf_polars(message="prefetch_parquet_file_metadata_for_ir") -def prefetch_parquet_file_metadata_for_ir( - root: IR, - py_executor: concurrent.futures.Executor | None, - stats: StatsCollector | None = None, +def cached_parquet_info_from_stats( + stats: StatsCollector, ) -> dict[str, CachedParquetInfo]: - """ - Prefetch parquet metadata for all parquet scans in an IR graph. - - Parameters - ---------- - root - The root of the IR graph, which will be traversed. - py_executor - The thread pool executor to use for fetching parquet metadata concurrently. - stats - The stats collector. The file metadata might have already been - prefetched during statistics collection, when the number of files - sampled equals the total number of files. Providing ``stats`` here will - skip rereading metadata for those files. - - Returns - ------- - A dictionary mapping each individual path to its cached parquet metadata. - """ - from cudf_polars.streaming.io import ParquetSourceInfo, StreamingScan - - all_paths: set[str] = set() - - for node in traversal([root]): - if isinstance(node, StreamingScan): - for scan in node.scans: - for path in scan.paths: - all_paths.add(path) - elif isinstance(node, Scan) and node.typ == "parquet": # pragma: no cover - raise RuntimeError("Unexpected parquet 'Scan' node in lowered IR graph.") + """Return path -> cached parquet info seeded from statistics collection.""" + from cudf_polars.streaming.io import ParquetSourceInfo 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 - - missing_paths = all_paths - set(cached_parquet_info.keys()) - cm: contextlib.AbstractContextManager[concurrent.futures.Executor | None] - - if py_executor is None: - cm = py_executor = concurrent.futures.ThreadPoolExecutor() - else: - # We didn't create the executor, so we don't close it. - cm = contextlib.nullcontext() - - with cm: - futures = [ - py_executor.submit(_prefetch_parquet_footers_for_paths, [path]) - for path in missing_paths - ] - - for future in concurrent.futures.as_completed(futures): - for info in future.result(): + 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 -def attach_cached_parquet_metadata( - root: IR, - cached_parquet_info_map: dict[str, CachedParquetInfo], -) -> None: +@nvtx_annotate_cudf_polars(message="prefetch_cached_parquet_info_for_paths") +def prefetch_cached_parquet_info_for_paths( + paths: list[str], stats: StatsCollector +) -> list[CachedParquetInfo]: """ - Attach prefetched metadata to scan nodes. + Prefetch parquet metadata for a path group. - This is an optimization only and does not affect IR identity. + Reuses footers already collected during statistics gathering when + available and fetches any remaining paths. Parameters ---------- - root - Root of the IR graph to update. - cached_parquet_info_map - Mapping from file paths to cached parquet metadata. + paths + Ordered list of parquet file paths for one scan task group. + stats + Optional statistics collector with already-cached footers. + + Returns + ------- + Cached parquet metadata ordered to match ``paths``. """ - 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) + 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: + fetched = _prefetch_parquet_footers_for_paths(missing_paths) + for info in fetched: + cached_by_path[info.path] = info + + return [cached_by_path[path] for path in paths] diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index e87728b4d01f..5ff6372cf35e 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -27,10 +27,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.quent._plan import build_plan from cudf_polars.streaming.actor_graph.collectives import ReserveOpIDs from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id @@ -774,14 +770,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, diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py index c923d342db99..2c1d7977b636 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py @@ -19,6 +19,11 @@ ) 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 ( + ParquetMetadataCache, + collect_metadata_scans, + parquet_metadata_prefetch_node, +) from cudf_polars.streaming.actor_graph.nodes import ( generate_ir_sub_network_wrapper, metadata_drain_node, @@ -264,6 +269,16 @@ 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_scans = collect_metadata_scans( + ir, + partition_info=partition_info, + config_options=config_options, + nranks=comm.nranks, + ) + metadata_channel_by_scan = { + scan: context.create_channel() for scan in metadata_scans + } + metadata_cache = ParquetMetadataCache(stats) # Generate the network state: GenState = { @@ -276,12 +291,24 @@ def generate_network( "max_io_threads": max_io_threads_local, "stats": stats, "collective_id_map": collective_id_map, + "metadata_scans": metadata_scans, + "metadata_channel_by_scan": metadata_channel_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, + scan, + metadata_channel_by_scan[scan], + metadata_cache, + ) + for scan in metadata_scans + ] # Add node to drain metadata before pull_from_channel # (since pull_from_channel doesn't handle metadata messages) @@ -301,6 +328,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 diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py index 2554d95fe750..19c3597b8552 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py +++ b/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,11 @@ class GenState(TypedDict): Statistics collector. collective_id_map The mapping of IR nodes to lists of collective IDs. + metadata_scans + Non-native parquet StreamingScan nodes that need metadata prefetch. + metadata_channel_by_scan + Mapping from each eligible StreamingScan node to its single metadata + input channel. """ context: Context @@ -69,6 +75,10 @@ class GenState(TypedDict): max_io_threads: int stats: StatsCollector collective_id_map: dict[IR, list[int]] + metadata_scans: list[StreamingScan] + metadata_channel_by_scan: dict[ + StreamingScan, Channel[ArbitraryChunk[MetadataMessagePayload]] + ] SubNetGenerator: TypeAlias = GenericTransformer[ diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py index 4c5773874aa5..82be404cecb6 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -5,10 +5,13 @@ from __future__ import annotations import asyncio +import contextlib import functools import io import math -from typing import TYPE_CHECKING, Any, cast +import reprlib +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, TypeAlias, cast import polars as pl @@ -16,6 +19,8 @@ from cudf_streaming.channel_metadata import ChannelMetadata from cudf_streaming.table_chunk import TableChunk from rapidsmpf.memory.memory_reservation import opaque_memory_usage +from rapidsmpf.streaming.chunks.arbitrary import ArbitraryChunk +from rapidsmpf.streaming.core.channel import Channel from rapidsmpf.streaming.core.memory_reserve_or_wait import ( reserve_memory, ) @@ -30,6 +35,8 @@ _prepare_parquet_predicate, ) from cudf_polars.dsl.to_ast import to_parquet_filter +from cudf_polars.dsl.traversal import traversal +from cudf_polars.dsl.utils.io import prefetch_cached_parquet_info_for_paths from cudf_polars.streaming.actor_graph.dispatch import ( generate_ir_sub_network, ) @@ -58,13 +65,13 @@ from cudf_polars.streaming.rank_aware_source import RankAwareSource if TYPE_CHECKING: - from collections.abc import Callable, Sequence + from collections.abc import Callable, MutableMapping, Sequence 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 IR, IRExecutionContext, Scan + from cudf_polars.dsl.utils.io import CachedParquetInfo from cudf_polars.streaming.actor_graph.core import SubNetGenerator from cudf_polars.streaming.actor_graph.tracing import ActorTracer from cudf_polars.streaming.base import ( @@ -73,7 +80,91 @@ StatsCollector, ) from cudf_polars.streaming.io import FusedScan, SplitScan - from cudf_polars.utils.config import ParquetOptions + from cudf_polars.utils.config import ConfigOptions, ParquetOptions + + +MetadataChannel: TypeAlias = Channel[ArbitraryChunk["MetadataMessagePayload"]] +MetadataChannelByScan: TypeAlias = dict[StreamingScan, MetadataChannel] + + +@dataclass(frozen=True) +class MetadataMessagePayload: + """Parquet metadata payload sent to scan actors.""" + + group_key: tuple[str, ...] + cached_parquet_info: list[CachedParquetInfo] + + +class ParquetMetadataCache: + """ + Query-scoped cache for prefetched parquet metadata. + + Coordinates footer reads across concurrent metadata prefetch actors so each + distinct ``scan.paths`` tuple is fetched at most once per query/rank. + """ + + def __init__( + self, + stats: StatsCollector, + fetch: Callable[ + [list[str], StatsCollector], list[CachedParquetInfo] + ] = prefetch_cached_parquet_info_for_paths, + ) -> None: + self._stats = stats + self._cached_by_key: dict[tuple[str, ...], list[CachedParquetInfo]] = {} + self._pending_by_key: dict[ + tuple[str, ...], asyncio.Future[list[CachedParquetInfo]] + ] = {} + self._lock = asyncio.Lock() + self._fetch = fetch + + async def get( + self, + paths: list[str], + ir_context: IRExecutionContext, + ) -> list[CachedParquetInfo]: + """ + Return cached parquet metadata for ``paths``, fetching on first use. + + Concurrent callers with identical ``paths`` share a single in-flight fetch. + """ + key = tuple(paths) + async with self._lock: + if key in self._cached_by_key: + return self._cached_by_key[key] + if key in self._pending_by_key: + future = self._pending_by_key[key] + should_fetch = False + else: + loop = asyncio.get_running_loop() + future = loop.create_future() + self._pending_by_key[key] = future + should_fetch = True + + if should_fetch: + try: + result = await ir_context.to_thread( + self._fetch, + list(key), + self._stats, + ) + except BaseException as exc: + async with self._lock: + self._pending_by_key.pop(key, None) + if not future.done(): + if isinstance(exc, asyncio.CancelledError): + future.cancel() + else: + future.set_exception(exc) + raise + async with self._lock: + self._cached_by_key[key] = result + self._pending_by_key.pop(key) + if not future.done(): + future.set_result(result) + return result + + return await future class Lineariser: @@ -200,7 +291,7 @@ async def dataframescan_node( ) # Build list of IR slices to read - ir_slices = [] + ir_slices: list[DataFrameScan] = [] # Partial workaround for # https://github.com/pola-rs/polars/issues/23214 If a struct column # has nulls and is sliced then polars exports invalid validity @@ -502,11 +593,12 @@ def _( async def read_chunk( context: Context, - scan: IR, + scan: DataFrameScan | SplitScan | FusedScan, seq_num: int, ch_out: Channel[TableChunk], ir_context: IRExecutionContext, estimated_chunk_bytes: int, + cached_parquet_info: list[CachedParquetInfo] | None = None, tracer: ActorTracer | None = None, ) -> None: """ @@ -527,17 +619,25 @@ async def read_chunk( estimated_chunk_bytes Estimated size of the chunk in bytes. Used for memory reservation with block spilling to avoid thrashing. + cached_parquet_info + Optional prefetched parquet metadata for parquet scans. tracer The actor tracer for collecting runtime statistics. """ + args = scan.with_prefetched_metadata(cached_parquet_info) + + # Help mypy with the type inference of the scan.do_evaluate method. + # DataFrameScan, SplitScan, and FusedScan have different signatures for the + # do_evaluate method, but we promise that calling with `args` is fine. + do_evaluate: Callable[..., DataFrame] = scan.do_evaluate with opaque_memory_usage( await reserve_memory( context, size=estimated_chunk_bytes, net_memory_delta=estimated_chunk_bytes ) ): df = await ir_context.to_thread( - scan.do_evaluate, - *scan._non_child_args, + do_evaluate, + *args, context=ir_context, ) chunk = TableChunk.from_pylibcudf_table( @@ -549,12 +649,118 @@ async def read_chunk( await send_chunk(context, ch_out, chunk, seq_num, tracer=tracer) +@define_actor() +async def parquet_metadata_prefetch_node( + context: Context, + ir_context: IRExecutionContext, + ir: StreamingScan, + ch_out: Channel[ArbitraryChunk[MetadataMessagePayload]], + metadata_cache: ParquetMetadataCache, +) -> None: + """ + Fetch parquet metadata for each scan task and send it to the paired scan actor. + + Parameters + ---------- + context + The rapidsmpf context. + ir_context + The execution context for the IR node. Prefetching is offloaded to a thread from + its thread pool. + ir + The StreamingScan node. This actor will send one a message per scan task in this + streaming scan node. + ch_out + The output channel. The Scan actor generated for this StreamingScan node will + read messages from this channel. + metadata_cache + Shared query-scoped cache for prefetched parquet metadata. + + Notes + ----- + This actor emits one message per SplitScan / FusedScan in the streaming scan. + The messages are sent in the order of the scans. + """ + async with shutdown_on_error(context, ch_out, trace_ir=ir, ir_context=ir_context): + for scan in ir.scans: + scan = cast("SplitScan | FusedScan", scan) + key = tuple[str, ...](scan.paths) + cached_parquet_info = await metadata_cache.get(list(key), ir_context) + payload = MetadataMessagePayload( + group_key=key, + cached_parquet_info=cached_parquet_info, + ) + await ch_out.send_metadata( + context, + Message(0, ArbitraryChunk(payload)), + ) + await ch_out.drain(context) + + +def _start_metadata_receiver( + context: Context, + ch_metadata: Channel[ArbitraryChunk[MetadataMessagePayload]], + scans: Sequence[SplitScan | FusedScan], +) -> tuple[list[asyncio.Future[list[CachedParquetInfo]]], asyncio.Task[None]]: + """ + Receive metadata messages sequentially and expose one Future per scan task. + + A single receiver task preserves channel order while allowing concurrent + producers to await only the metadata for their assigned task index. + """ + loop = asyncio.get_running_loop() + futures = [loop.create_future() for _ in range(len(scans))] + + async def _receive() -> None: + try: + for task_idx, scan in enumerate(scans): + msg = await ch_metadata.recv_metadata(context) + cached = recv_prefetched_parquet_metadata_handler( + msg, tuple(scan.paths) + ) + futures[task_idx].set_result(cached) + except BaseException as exc: + # Propagate the failure (including cancellation / premature stop) to + # every unfinished future so downstream awaiters fail fast instead of + # hanging on metadata that will never arrive. + for future in futures: + if not future.done(): + if isinstance(exc, asyncio.CancelledError): + future.cancel() + else: + future.set_exception(exc) + raise + + receiver_task = asyncio.create_task(_receive()) + return futures, receiver_task + + +def recv_prefetched_parquet_metadata_handler( + msg: Message[ArbitraryChunk[MetadataMessagePayload]] | None, + group_key: tuple[str, ...], +) -> list[CachedParquetInfo]: + """Synchronous handler for prefetched parquet metadata messages.""" + if msg is None: + raise AssertionError( + f"Missing parquet metadata message for paths: {reprlib.repr(group_key)}" + ) + payload = ArbitraryChunk[MetadataMessagePayload].from_message(msg).release() + if payload.group_key != group_key: + difference = set(group_key) ^ set(payload.group_key) + raise AssertionError( + "Unexpected parquet metadata key on scan input channel. " + f"{reprlib.repr(difference)}" + ) + return payload.cached_parquet_info + + @define_actor() async def scan_node( context: Context, ir: StreamingScan, ir_context: IRExecutionContext, ch_out: Channel[TableChunk], + ch_metadata: Channel[ArbitraryChunk[MetadataMessagePayload]] | None, *, num_producers: int, estimated_chunk_bytes: int, @@ -572,81 +778,118 @@ async def scan_node( The execution context for the IR node. ch_out The output Channel[TableChunk]. + ch_metadata + Optional channel carrying prefetched parquet metadata messages, one + per `SplitScan`/`FusedScan` in `ir.scans` order. num_producers The number of producers to use for the scan node. estimated_chunk_bytes Estimated size of each chunk in bytes. Used for memory reservation with block spilling to avoid thrashing. """ - scans: Sequence[SplitScan] | Sequence[FusedScan] = ir.scans - - async with shutdown_on_error( - context, ch_out, trace_ir=ir, ir_context=ir_context - ) as tracer: - # Send basic metadata - await send_metadata( - ch_out, - context, - ChannelMetadata(local_count=len(scans)), + scans = cast("Sequence[SplitScan | FusedScan]", ir.scans) + metadata_futures: list[asyncio.Future[list[CachedParquetInfo]]] | None = None + _metadata_receiver_task: asyncio.Task[None] | None = None + if ch_metadata is not None: + metadata_futures, _metadata_receiver_task = _start_metadata_receiver( + context, ch_metadata, scans ) - # If there is nothing to scan, drain the channel and return - if len(scans) == 0: - await ch_out.drain(context) - return + async def cached_parquet_info_for_task( + task_idx: int, + ) -> list[CachedParquetInfo] | None: + if metadata_futures is None: + return None + return await metadata_futures[task_idx] - # If there is only one scan or one producer, we can - # skip the lineariser and read the chunks directly - if len(scans) == 1 or num_producers == 1: - for seq_num, scan in enumerate(scans): - await read_chunk( - context, - scan, - seq_num, - ch_out, - ir_context, - estimated_chunk_bytes, - tracer=tracer, - ) - await ch_out.drain(context) - return + shutdown_channels: list[Channel[Any]] = [ch_out] + if ch_metadata is not None: + shutdown_channels.append(ch_metadata) - # Use Lineariser to ensure ordered delivery - num_producers = min(num_producers, len(scans)) - lineariser = Lineariser(context, ch_out, num_producers) - - # Assign tasks to producers using round-robin - producer_tasks: list[list[tuple[int, SplitScan | FusedScan]]] = [ - [] for _ in range(num_producers) - ] - for task_idx, scan in enumerate(scans): - producer_id = task_idx % num_producers - # mypy resolves __iter__ on union-of-sequences to the common base (IR) - producer_tasks[producer_id].append((task_idx, scan)) # type: ignore[arg-type] + try: + async with shutdown_on_error( + context, + *shutdown_channels, + trace_ir=ir, + ir_context=ir_context, + ) as tracer: + # Send basic metadata + await send_metadata( + ch_out, + context, + ChannelMetadata(local_count=len(scans)), + ) - async def _producer(producer_id: int, ch_out: Channel) -> None: - for task_idx, scan in producer_tasks[producer_id]: - await read_chunk( - context, - scan, - task_idx, - ch_out, - ir_context, - estimated_chunk_bytes, - tracer=tracer, + # If there is nothing to scan, drain the channel and return + if len(scans) == 0: + await ch_out.drain(context) + return + + # If there is only one scan or one producer, we can + # skip the lineariser and read the chunks directly + if len(scans) == 1 or num_producers == 1: + for seq_num, scan in enumerate(scans): + cached_parquet_info = await cached_parquet_info_for_task(seq_num) + await read_chunk( + context, + scan, + seq_num, + ch_out, + ir_context, + estimated_chunk_bytes, + cached_parquet_info, + tracer=tracer, + ) + await ch_out.drain(context) + return + + # Use Lineariser to ensure ordered delivery + num_producers = min(num_producers, len(scans)) + lineariser = Lineariser(context, ch_out, num_producers) + + # Assign tasks to producers using round-robin + producer_tasks: list[list[tuple[int, SplitScan | FusedScan]]] = [ + [] for _ in range(num_producers) + ] + for task_idx, scan in enumerate(scans): + producer_id = task_idx % num_producers + producer_tasks[producer_id].append((task_idx, scan)) + + async def _producer(producer_id: int, ch_out: Channel) -> None: + for task_idx, scan in producer_tasks[producer_id]: + cached_parquet_info = await cached_parquet_info_for_task(task_idx) + await read_chunk( + context, + scan, + task_idx, + ch_out, + ir_context, + estimated_chunk_bytes, + cached_parquet_info, + tracer=tracer, + ) + await ch_out.drain(context) + + async with ( + shutdown_on_error(context, *lineariser.input_channels, trace_ir=ir), + ): + await gather_in_task_group( + lineariser.drain(), + *( + _producer(i, ch_in) + for i, ch_in in enumerate(lineariser.input_channels) + ), ) - await ch_out.drain(context) - - async with ( - shutdown_on_error(context, *lineariser.input_channels, trace_ir=ir), - ): - await gather_in_task_group( - lineariser.drain(), - *( - _producer(i, ch_in) - for i, ch_in in enumerate(lineariser.input_channels) - ), - ) + finally: + # Always finalize the background metadata receiver, even on early + # return or failure, so it is never left orphaned. + if _metadata_receiver_task is not None: + _metadata_receiver_task.cancel() + # Awaiting also retrieves any exception the receiver raised (already + # surfaced to producers via the per-task futures), avoiding a stray + # "Task exception was never retrieved" warning. + with contextlib.suppress(BaseException): + await _metadata_receiver_task def make_rapidsmpf_read_parquet_node( @@ -821,6 +1064,7 @@ def _( ir, rec.state["ir_context"], ch_out, + rec.state["metadata_channel_by_scan"].get(ir), num_producers=num_producers, estimated_chunk_bytes=( plan.estimated_chunk_bytes or executor.target_partition_size @@ -960,3 +1204,40 @@ def _( ] return nodes, channels + + +def collect_metadata_scans( + ir: IR, + *, + partition_info: MutableMapping[IR, PartitionInfo], + config_options: ConfigOptions, + nranks: int, +) -> list[StreamingScan]: + """Return non-native parquet StreamingScan nodes that need metadata prefetch.""" + if not config_options.parquet_options.prefetch_file_metadata: + return [] + + metadata_scans: list[StreamingScan] = [] + for node in traversal([ir]): + if not isinstance(node, StreamingScan): + continue + if node.base_scan.typ != "parquet": + continue + if not node.scans: + 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 + metadata_scans.append(node) + return metadata_scans diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index b3b4438812c5..de61e16c75e1 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -273,6 +273,15 @@ def get_hashable(self) -> Hashable: self.parquet_options, ) + def with_prefetched_metadata( + self, + cached_parquet_info: list[CachedParquetInfo] | None, + ) -> tuple[Any, ...]: + """Return ``do_evaluate`` args, substituting prefetched parquet metadata when provided.""" + if cached_parquet_info is None: + return self._non_child_args + return (*self._non_child_args[:-1], cached_parquet_info) + @classmethod def do_evaluate( cls, @@ -441,6 +450,15 @@ def get_hashable(self) -> Hashable: self.parquet_options, ) + def with_prefetched_metadata( + self, + cached_parquet_info: list[CachedParquetInfo] | None, + ) -> tuple[Any, ...]: + """Return ``do_evaluate`` args, substituting prefetched parquet metadata when provided.""" + if cached_parquet_info is None: + return self._non_child_args + return (*self._non_child_args[:-1], cached_parquet_info) + @classmethod def do_evaluate( cls, diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 13e88ead7731..213cbe45aabd 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -3,29 +3,41 @@ from __future__ import annotations +import asyncio import math +from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, cast import pytest import polars as pl +from rapidsmpf.streaming.chunks.arbitrary import ArbitraryChunk +from rapidsmpf.streaming.core.message import Message + from cudf_polars import Translator from cudf_polars.containers import DataType from cudf_polars.dsl.ir import ( - Empty, + DataFrameScan, IRExecutionContext, Scan, + Union, ) from cudf_polars.dsl.utils.io import ( CachedParquetInfo, - prefetch_parquet_file_metadata_for_ir, ) from cudf_polars.engine.options import StreamingOptions +from cudf_polars.streaming.actor_graph.io import ( + MetadataMessagePayload, + ParquetMetadataCache, + collect_metadata_scans, + recv_prefetched_parquet_metadata_handler, +) from cudf_polars.streaming.base import ( DataSourceInfo, IOPartitionFlavor, IOPartitionPlan, + PartitionInfo, StatsCollector, ) from cudf_polars.streaming.io import ( @@ -51,6 +63,7 @@ import pylibcudf as plc import cudf_polars.engine.core + from cudf_polars.dsl.ir import IR from cudf_polars.engine.core import StreamingEngine @@ -65,6 +78,16 @@ def df(): ) +@pytest.fixture +def prefetch_file_metadata_engine( + streaming_engine_factory: Callable[..., StreamingEngine], +): + """Streaming Engine fixture with parquet metadata prefetching enabled.""" + return streaming_engine_factory( + StreamingOptions(parquet_options={"prefetch_file_metadata": True}), + ) + + @pytest.mark.parametrize( "fmt, scan_fn", [ @@ -124,32 +147,61 @@ def test_scan_parquet_prefetch_file_metadata( assert_gpu_result_equal(pl.scan_parquet(tmp_path), engine=streaming_engine) -def test_prefetch_file_metadata_non_parquet_scan(df, streaming_engine_factory) -> None: - streaming_engine = streaming_engine_factory( - StreamingOptions(parquet_options={"prefetch_file_metadata": True}), - ) - assert_gpu_result_equal(df.lazy().select("x"), engine=streaming_engine) +@pytest.mark.timeout(90) +def test_scan_parquet_prefetch_metadata_shared_scan_paths( + tmp_path: Path, + df: pl.DataFrame, + prefetch_file_metadata_engine: StreamingEngine, +): + # The spmd-small case creates *many* partitions with the length-3000 df. + # A smaller dataframe gives us sufficient test coverage, and runs much faster. + if ( + prefetch_file_metadata_engine.config["executor_options"][ + "max_rows_per_partition" + ] + == SMALL_MAX_ROWS_PER_PARTITION + ): + df = df.head(40) + make_partitioned_source(df, tmp_path, "parquet", n_files=2) + scan = pl.scan_parquet(tmp_path) + query = pl.concat([scan.select("x"), scan.select("x")]) + assert_gpu_result_equal(query, engine=prefetch_file_metadata_engine) + + +def test_scan_parquet_prefetch_metadata_disjoint_scan_paths( + tmp_path: Path, + prefetch_file_metadata_engine: StreamingEngine, +): + left = pl.DataFrame({"x": [1, 2, 3]}) + right = pl.DataFrame({"x": [4, 5, 6]}) + left.write_parquet(tmp_path / "left.parquet") + right.write_parquet(tmp_path / "right.parquet") -def test_prefetch_parquet_file_metadata_no_parquet_scans() -> None: - result = prefetch_parquet_file_metadata_for_ir( - Empty({}), py_executor=None, stats=None + query = pl.concat( + [ + pl.scan_parquet(tmp_path / "left.parquet"), + pl.scan_parquet(tmp_path / "right.parquet"), + ] ) - assert result == {} + assert_gpu_result_equal(query, engine=prefetch_file_metadata_engine) + + +def test_prefetch_file_metadata_non_parquet_scan( + df: pl.DataFrame, prefetch_file_metadata_engine: StreamingEngine +) -> None: + assert_gpu_result_equal(df.lazy().select("x"), engine=prefetch_file_metadata_engine) def test_prefetch_file_metadata_select_fast_count( df: pl.DataFrame, - streaming_engine_factory: Callable[..., StreamingEngine], + prefetch_file_metadata_engine: StreamingEngine, tmp_path: Path, ) -> None: - streaming_engine = streaming_engine_factory( - StreamingOptions(parquet_options={"prefetch_file_metadata": True}), - ) source = tmp_path / "data.parquet" df.write_parquet(source) q = pl.scan_parquet(source).select(pl.len()) - assert_gpu_result_equal(q, engine=streaming_engine) + assert_gpu_result_equal(q, engine=prefetch_file_metadata_engine) # --------------------------------------------------------------------------- @@ -449,19 +501,15 @@ def test_split_scan_do_evaluate_missing_prefetch_metadata() -> None: def test_prefetch_file_metadata_join( - tmp_path: Path, streaming_engine_factory: Callable[..., StreamingEngine] + tmp_path: Path, prefetch_file_metadata_engine: StreamingEngine ) -> None: p1 = tmp_path / "f1.parquet" p2 = tmp_path / "f2.parquet" pl.DataFrame({"k": [1, 2, 3], "a": [4, 5, 6]}).write_parquet(p1) pl.DataFrame({"k": [1, 2, 3], "b": [7, 8, 9]}).write_parquet(p2) - engine = streaming_engine_factory( - StreamingOptions(parquet_options={"prefetch_file_metadata": True}), - ) - q = pl.scan_parquet(p1).join(pl.scan_parquet(p2), on="k") - q.collect(engine=engine) + q.collect(engine=prefetch_file_metadata_engine) def _make_cached_parquet_info( @@ -480,7 +528,7 @@ def _make_cached_parquet_info( def test_prefetch_file_metadata_with_cached_scan_parent_nodes( - tmp_path: Path, streaming_engine_factory: Callable[..., StreamingEngine] + tmp_path: Path, prefetch_file_metadata_engine: StreamingEngine ) -> None: # Regression test for replace not replacing StreamingScan nodes with their prefetched variants. source = tmp_path / "data.parquet" @@ -491,16 +539,35 @@ def test_prefetch_file_metadata_with_cached_scan_parent_nodes( } ).write_parquet(source) - engine = streaming_engine_factory( - StreamingOptions(parquet_options={"prefetch_file_metadata": True}), - ) - cached_scan = pl.scan_parquet(source).cache() left = cached_scan.group_by("k").agg(pl.col("v").sum().alias("sum_v")) right = cached_scan.group_by("k").agg(pl.len().alias("n")) q = left.join(right, on="k").sort("k") - assert_gpu_result_equal(q, engine=engine) + assert_gpu_result_equal(q, engine=prefetch_file_metadata_engine) + + +def test_with_prefetched_metadata() -> None: + base = _make_parquet_scan(["a.parquet"]) + info = _make_cached_parquet_info(base.paths) + + dfs = DataFrameScan(base.schema, pl.DataFrame({"x": [1]})._df, None) + assert dfs.with_prefetched_metadata(info) == dfs._non_child_args + assert dfs.with_prefetched_metadata(None) == dfs._non_child_args + + split = SplitScan(base.schema, base, base.paths, 0, 4, base.parquet_options, None) + assert split.with_prefetched_metadata(None) == split._non_child_args + assert split.with_prefetched_metadata(info) == ( + *split._non_child_args[:-1], + info, + ) + + fused = FusedScan(base.schema, base, base.paths, base.parquet_options, None) + assert fused.with_prefetched_metadata(None) == fused._non_child_args + assert fused.with_prefetched_metadata(info) == ( + *fused._non_child_args[:-1], + info, + ) def test_fused_scan_identity_equality() -> None: @@ -647,6 +714,16 @@ def _make_config(target: int) -> ConfigOptions: return ConfigOptions.from_polars_engine(engine) +def _make_prefetch_config(target: int) -> ConfigOptions: + engine = pl.GPUEngine( + raise_on_fail=True, + executor="streaming", + executor_options={"target_partition_size": target}, + parquet_options={"prefetch_file_metadata": True}, + ) + return ConfigOptions.from_polars_engine(engine) + + @pytest.mark.parametrize( "file_size,n_paths,expected_factor,expected_flavor", [ @@ -673,3 +750,144 @@ def test_scan_partition_plan_nearest( plan = scan_partition_plan(scan, FooStats(scan, file_size), _make_config(10)) assert plan.factor == expected_factor assert plan.flavor == expected_flavor + + +def test_collect_metadata_scans_one_actor_per_streaming_scan() -> None: + parquet_options = ParquetOptions(prefetch_file_metadata=True) + paths = [f"part.{i}.parquet" for i in range(6)] + base_scan = _make_parquet_scan(paths, parquet_options) + plan = IOPartitionPlan(9, IOPartitionFlavor.SPLIT_FILES) + partition_count = plan.factor * len(paths) + streaming_scan = expand_scan_for_rank( + base_scan, + plan, + partition_count, + rank=0, + nranks=1, + parquet_options=parquet_options, + ) + assert len(streaming_scan.scans) == partition_count + assert len({tuple(scan.paths) for scan in streaming_scan.scans}) == len(paths) + + config_options = _make_prefetch_config(873_630_000) + partition_info: dict[IR, PartitionInfo] = { + streaming_scan: PartitionInfo(count=partition_count, io_plan=plan), + } + metadata_scans = collect_metadata_scans( + streaming_scan, + partition_info=partition_info, + config_options=config_options, + nranks=1, + ) + assert metadata_scans == [streaming_scan] + + +def test_collect_metadata_scans_union_disjoint_paths() -> None: + parquet_options = ParquetOptions(prefetch_file_metadata=True) + plan = IOPartitionPlan(1, IOPartitionFlavor.FUSED_FILES) + left = expand_scan_for_rank( + _make_parquet_scan(["left.parquet"], parquet_options), + plan, + 1, + rank=0, + nranks=1, + parquet_options=parquet_options, + ) + right = expand_scan_for_rank( + _make_parquet_scan(["right.parquet"], parquet_options), + plan, + 1, + rank=0, + nranks=1, + parquet_options=parquet_options, + ) + union = Union(left.schema, None, False, left, right) # noqa: FBT003 + config_options = _make_prefetch_config(10_000) + partition_info: dict[IR, PartitionInfo] = { + left: PartitionInfo(count=1, io_plan=plan), + right: PartitionInfo(count=1, io_plan=plan), + union: PartitionInfo(count=2), + } + metadata_scans = collect_metadata_scans( + union, + partition_info=partition_info, + config_options=config_options, + nranks=1, + ) + assert metadata_scans == [left, right] + + +def test_collect_metadata_scans_skips_empty_rank() -> None: + parquet_options = ParquetOptions(prefetch_file_metadata=True) + plan = IOPartitionPlan(3, IOPartitionFlavor.SINGLE_READ) + paths = ["a.parquet", "b.parquet", "c.parquet"] + streaming_scan = expand_scan_for_rank( + _make_parquet_scan(paths, parquet_options), + plan, + 1, + rank=1, + nranks=2, + parquet_options=parquet_options, + ) + assert len(streaming_scan.scans) == 0 + config_options = _make_prefetch_config(10_000) + partition_info: dict[IR, PartitionInfo] = { + streaming_scan: PartitionInfo(count=0, io_plan=plan), + } + metadata_scans = collect_metadata_scans( + streaming_scan, + partition_info=partition_info, + config_options=config_options, + nranks=2, + ) + assert metadata_scans == [] + + +def test_recv_prefetched_parquet_metadata_handler_errors() -> None: + with pytest.raises( + AssertionError, match=r"Missing parquet metadata message for paths: .*" + ): + recv_prefetched_parquet_metadata_handler(None, ("file.parquet",)) + + msg = Message( + 0, + ArbitraryChunk( + MetadataMessagePayload( + group_key=("file.parquet",), + cached_parquet_info=[ + # We don't use file_metadata, so just lie about it. + CachedParquetInfo(path="file.parquet", size=10, file_metadata=None) # type: ignore[arg-type] + ], + ) + ), + ) + with pytest.raises( + AssertionError, + match=r"Unexpected parquet metadata key on scan input channel. .*", + ): + recv_prefetched_parquet_metadata_handler(msg, ("file2.parquet",)) + + +def test_parquet_metadata_cache_dedupes_identical_paths() -> None: + fetch_count = 0 + + def mock_fetch(paths: list[str], stats: StatsCollector) -> list[CachedParquetInfo]: + nonlocal fetch_count + fetch_count += 1 + return _make_cached_parquet_info(paths) + + cache = ParquetMetadataCache(StatsCollector(), fetch=mock_fetch) + paths = ["a.parquet", "b.parquet"] + + async def run() -> list[list[CachedParquetInfo]]: + with ThreadPoolExecutor(max_workers=2) as executor: + ir_context = IRExecutionContext(executor) + async with asyncio.TaskGroup() as tg: + first = tg.create_task(cache.get(paths, ir_context)) + second = tg.create_task(cache.get(paths, ir_context)) + return [first.result(), second.result()] + + results = asyncio.run(run()) + assert fetch_count == 1 + assert results[0] == results[1] + assert [info.path for info in results[0]] == paths