diff --git a/python/cudf_polars/cudf_polars/callback.py b/python/cudf_polars/cudf_polars/callback.py index dcdac8fdfe36..2c254720925a 100644 --- a/python/cudf_polars/cudf_polars/callback.py +++ b/python/cudf_polars/cudf_polars/callback.py @@ -24,7 +24,9 @@ from rmm._cuda import gpu import cudf_polars.dsl.tracing -from cudf_polars.dsl.ir import IRExecutionContext +from cudf_polars.dsl.ir import ( + IRExecutionContext, +) from cudf_polars.dsl.tracing import CUDF_POLARS_NVTX_DOMAIN from cudf_polars.dsl.translate import Translator from cudf_polars.utils.config import ( diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index b645a6b8b330..68b6da4a1f16 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -20,6 +20,7 @@ import itertools import json import random +import reprlib import time import uuid from collections.abc import Sized @@ -66,6 +67,7 @@ ) if TYPE_CHECKING: + import concurrent.futures from collections.abc import ( Callable, Generator, @@ -74,7 +76,6 @@ Iterator, Sequence, ) - from concurrent.futures import ThreadPoolExecutor from typing import Literal, Self from polars import polars # type: ignore[attr-defined] @@ -82,6 +83,7 @@ from rmm.pylibrmm.stream import Stream from cudf_polars.containers.dataframe import NamedColumn + from cudf_polars.dsl.utils.io import CachedParquetInfo from cudf_polars.streaming.rank_aware_source import RankAwareSource from cudf_polars.typing import CSECache, ClosedInterval, Schema, Slice as Zlice from cudf_polars.utils.config import ParquetOptions @@ -138,7 +140,7 @@ class IRExecutionContext: Identifier for the query being executed. """ - py_executor: ThreadPoolExecutor | None = field(default=None) + py_executor: concurrent.futures.ThreadPoolExecutor | None = field(default=None) get_cuda_stream: Callable[[], Stream] = field(default=get_cuda_stream) query_id: uuid.UUID = field(default_factory=uuid.uuid4) @@ -570,6 +572,8 @@ def do_evaluate( def _parquet_physical_types( paths: list[str], columns: list[str] | None ) -> dict[str, plc.DataType]: + # TODO: Use prefetched metadata + # https://github.com/rapidsai/cudf/issues/22940 metadata = plc.io.parquet_metadata.read_parquet_metadata(plc.io.SourceInfo(paths)) column_types = metadata.schema().column_types() @@ -637,6 +641,7 @@ class Scan(IR): """Input from files.""" __slots__ = ( + "cached_parquet_info", "cloud_options", "include_file_paths", "n_rows", @@ -663,7 +668,7 @@ class Scan(IR): "predicate", "parquet_options", ) - _n_non_child_args = 11 + _n_non_child_args = 12 typ: str """What type of file are we reading? Parquet, CSV, etc...""" reader_options: dict[str, Any] @@ -686,6 +691,8 @@ class Scan(IR): """Mask to apply to the read dataframe.""" parquet_options: ParquetOptions """Parquet-specific options.""" + cached_parquet_info: list[CachedParquetInfo] | None + """Cached parquet file metadata.""" PARQUET_DEFAULT_CHUNK_SIZE: int = 0 # unlimited PARQUET_DEFAULT_PASS_LIMIT: int = 16 * 1024**3 # 16GiB @@ -704,6 +711,7 @@ def __init__( include_file_paths: str | None, predicate: expr.NamedExpr | None, parquet_options: ParquetOptions, + cached_parquet_info: list[CachedParquetInfo] | None = None, ): self.schema = schema self.typ = typ @@ -728,9 +736,14 @@ def __init__( include_file_paths, predicate, parquet_options, + cached_parquet_info, ) self.children = () self.parquet_options = parquet_options + self.cached_parquet_info = cached_parquet_info + + Scan._validate_cached_parquet_info(self.paths, self.cached_parquet_info) + if self.typ not in ("csv", "parquet", "ndjson"): # pragma: no cover # This line is unhittable ATM since IPC/Anonymous scan raise # on the polars side @@ -810,6 +823,21 @@ def __init__( "Reading only parquet metadata to produce row index." ) + @staticmethod + def _validate_cached_parquet_info( + paths: list[str], + cached_parquet_info: list[CachedParquetInfo] | None, + ) -> None: + if cached_parquet_info is not None and paths != [ + info.path for info in cached_parquet_info + ]: + missing = reprlib.repr( + set(paths) - {info.path for info in cached_parquet_info} + ) + raise AssertionError( + f"Paths do not match cached parquet info. Missing paths: {missing}" + ) + def get_hashable(self) -> Hashable: """ Hashable representation of the node. @@ -866,12 +894,34 @@ def add_file_paths( @staticmethod @nvtx_annotate_cudf_polars(message="Scan._get_parquet_row_count_from_metadata") def _get_parquet_row_count_from_metadata( - paths: list[str], skip_rows: int, n_rows: int + paths: list[str], + skip_rows: int, + n_rows: int, + parquet_options: ParquetOptions, + cached_parquet_info: list[CachedParquetInfo] | None, ) -> int: # Zero-width parquet files lose their row count when read through # pylibcudf. See https://github.com/rapidsai/cudf/issues/21428 - meta = plc.io.parquet_metadata.read_parquet_metadata(plc.io.SourceInfo(paths)) - num_rows = meta.num_rows() - skip_rows + if parquet_options.prefetch_file_metadata: + if cached_parquet_info is None: + raise AssertionError( + "Cached parquet info is required when prefetching file metadata is enabled" + ) + + Scan._validate_cached_parquet_info(paths, cached_parquet_info) + parquet_metadatas = [ + info.file_metadata for info in cached_parquet_info + ] # pragma: no cover + num_rows = sum( + metadata.num_rows for metadata in parquet_metadatas + ) # pragma: no cover + else: + meta = plc.io.parquet_metadata.read_parquet_metadata( + plc.io.SourceInfo(paths) + ) + num_rows = meta.num_rows() + + num_rows -= skip_rows if n_rows != -1: num_rows = min(num_rows, n_rows) return max(num_rows, 0) @@ -892,6 +942,7 @@ def do_evaluate( include_file_paths: str | None, predicate: expr.NamedExpr | None, parquet_options: ParquetOptions, + cached_parquet_info: list[CachedParquetInfo] | None, *, context: IRExecutionContext, ) -> DataFrame: @@ -1007,6 +1058,24 @@ def read_csv_header( df, ) elif typ == "parquet": + if parquet_options.prefetch_file_metadata: + if cached_parquet_info is None: + raise AssertionError( + "Cached parquet info is required when prefetching file metadata is enabled" + ) + Scan._validate_cached_parquet_info(paths, cached_parquet_info) + filepath_sources = [] + parquet_metadatas = [] + for info in cached_parquet_info: + filepath_sources.append( + plc.io.types.FilepathSource(info.path, info.size) + ) + parquet_metadatas.append(info.file_metadata) + source_info = plc.io.SourceInfo(filepath_sources) + else: + parquet_metadatas = None + source_info = plc.io.SourceInfo(paths) + filters = None if predicate is not None and row_index is None: # Can't apply filters during read if we have a row index. @@ -1016,9 +1085,7 @@ def read_csv_header( ), stream=stream, ) - builder = plc.io.parquet.ParquetReaderOptions.builder( - plc.io.SourceInfo(paths) - ) + builder = plc.io.parquet.ParquetReaderOptions.builder(source_info) if filters is not None and parquet_options.use_jit_filter: builder.use_jit_filter(use_jit_filter=True) parquet_reader_options = builder.decimal_width( @@ -1038,6 +1105,7 @@ def read_csv_header( parquet_reader_options, chunk_read_limit=parquet_options.chunk_read_limit, pass_read_limit=parquet_options.pass_read_limit, + parquet_metadatas=parquet_metadatas, stream=stream, ) chunk = reader.read_chunk() @@ -1053,7 +1121,9 @@ def read_csv_header( [concatenated_columns[i], columns.pop()], stream=stream ) num_rows = ( - cls._get_parquet_row_count_from_metadata(paths, skip_rows, n_rows) + cls._get_parquet_row_count_from_metadata( + paths, skip_rows, n_rows, parquet_options, cached_parquet_info + ) if not names else None ) @@ -1070,12 +1140,16 @@ def read_csv_header( ) else: tbl_w_meta = plc.io.parquet.read_parquet( - parquet_reader_options, stream=stream + parquet_reader_options, + parquet_metadatas=parquet_metadatas, + stream=stream, ) # TODO: consider nested column names? col_names = tbl_w_meta.column_names(include_children=False) num_rows = ( - cls._get_parquet_row_count_from_metadata(paths, skip_rows, n_rows) + cls._get_parquet_row_count_from_metadata( + paths, skip_rows, n_rows, parquet_options, cached_parquet_info + ) if not col_names else None ) @@ -1747,7 +1821,11 @@ def evaluate( stream = context.get_cuda_stream() scan = self.children[0] effective_rows = Scan._get_parquet_row_count_from_metadata( - scan.paths, scan.skip_rows, scan.n_rows + scan.paths, + scan.skip_rows, + scan.n_rows, + scan.parquet_options, + None, ) dtype = DataType(pl.UInt32()) col = Column( diff --git a/python/cudf_polars/cudf_polars/dsl/translate.py b/python/cudf_polars/cudf_polars/dsl/translate.py index 341438442bb9..b46514ff8f82 100644 --- a/python/cudf_polars/cudf_polars/dsl/translate.py +++ b/python/cudf_polars/cudf_polars/dsl/translate.py @@ -466,6 +466,7 @@ def _(node: plrs._ir_nodes.Scan, translator: Translator, schema: Schema) -> ir.I ) ), parquet_options, + cached_parquet_info=None, ) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py new file mode 100644 index 000000000000..cc49522d498a --- /dev/null +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Utilities for IR nodes.""" + +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 + +if TYPE_CHECKING: + from cudf_polars.dsl.ir import IR + from cudf_polars.streaming.base import StatsCollector + + +@dataclass(frozen=True) +class CachedParquetInfo: + """ + Metadata for a parquet file. + + File metadata is only cached when the setting + ``ParquetOptions.prefetch_file_metadata`` is ``True``. Metadata is cached + for the duration of the query. + + Parameters + ---------- + path + The path of an individual parquet file. This is one element of a + ``paths`` tuple in a ``Scan`` node. + size + The size of the parquet file, in bytes. This is typically only set + for remote URLs, since it allows skipping subsequent HTTP HEAD requests + made by kvikio on operations involving that file. + file_metadata + The ``FileMetaData`` object for the parquet file returned from + ``read_parquet_footers``. + """ + + path: str + size: int + file_metadata: plc.io.parquet_metadata.FileMetaData + + +@nvtx_annotate_cudf_polars(message="fetch_parquet_footers_for_paths") +def _prefetch_parquet_footers_for_paths(paths: list[str]) -> list[CachedParquetInfo]: + """ + Prefetch parquet footers for a list of paths. + + This is typically executed concurrently with prefetch operations for other + path groups for other parquet scan nodes. + + Parameters + ---------- + paths + The paths to prefetch. + + Returns + ------- + paths + The original input ``paths``. + metadata + The list of ``FileMetaData`` objects for the ``paths``. + """ + # TODO: https://github.com/rapidsai/cudf/issues/22734, use object metadata from polars + # For now, we'll just use kvikio to explicitly get the size. + sizes = [] + + try: # pragma: no cover; kvikio is optional + import kvikio + except ImportError: + kvikio = None + + for path in paths: + if ( + paths and kvikio is not None and plc.io.SourceInfo._is_remote_uri(path) + ): # pragma: no cover; kvikio is optional + # We're OK to use `kvikio.RemoteFile.open` here. It does make an HTTP HEAD + # request for S3/HTTP endpoints, but that's the entire reason we're running + # this code. So long as it makes just *one* HTTP request, there's no advantage + # to inferring the endpoint type. + with kvikio.RemoteFile.open(path) as remote_file: + sizes.append(remote_file.nbytes()) + else: + sizes.append(None) + + metadata = plc.io.parquet_metadata.read_parquet_footers( + plc.io.types.SourceInfo( + [ + plc.io.types.FilepathSource(path, size) + for path, size in zip(paths, sizes, strict=True) + ] + ) + ) + + return [ + CachedParquetInfo(path, size, file_metadata) + for path, size, file_metadata in zip(paths, sizes, metadata, strict=True) + ] + + +@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, +) -> 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.") + + 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(): + 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) diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 97b6dba20f3f..bc82c3bc2ec5 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -26,6 +26,10 @@ 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 @@ -408,14 +412,12 @@ def _find_memory_error(exc: BaseException) -> MemoryError | None: def execute_ir_on_rank( ctx: Context, comm: Communicator, - py_executor: ThreadPoolExecutor, ir: IR, + ir_context: IRExecutionContext, partition_info: MutableMapping[IR, PartitionInfo], config_options: ConfigOptions[StreamingExecutor], stats: StatsCollector, collective_id_map: dict[IR, list[int]], - *, - query_id: uuid.UUID, ) -> tuple[pl.DataFrame, list[ChannelMetadata]]: """ Execute a Polars IR query on a single rank's GPU. @@ -430,10 +432,10 @@ def execute_ir_on_rank( The active RapidsMPF streaming context for this rank. comm The active RapidsMPF communicator for this rank. - py_executor - Thread-pool executor used to drive the actor network. ir Root IR node describing the query. + ir_context + Execution context reused across scan-task execution. partition_info Per-node partition metadata. config_options @@ -442,8 +444,6 @@ def execute_ir_on_rank( Statistics collector. collective_id_map Mapping from IR nodes to their pre-allocated collective operation IDs. - query_id - Unique identifier for the query, propagated into actor traces. Returns ------- @@ -452,9 +452,6 @@ def execute_ir_on_rank( metadata Collected channel metadata. """ - ir_context = IRExecutionContext( - py_executor, get_cuda_stream=ctx.br().stream_pool.get_stream, query_id=query_id - ) metadata_collector: list[ChannelMetadata] = [] nodes, output = generate_network( @@ -695,15 +692,26 @@ def evaluate_on_rank( # so we only log it once. log_query_plan(ir, config_options) + ir_context = IRExecutionContext( + 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, comm, - py_executor, ir, + ir_context, partition_info, config_options, stats, collective_id_map, - query_id=query_id, ) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 53bcb27ee6e5..b3b4438812c5 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -43,7 +43,7 @@ from cudf_polars.containers import DataFrame, DataType from cudf_polars.dsl.expr import NamedExpr - from cudf_polars.dsl.ir import IRExecutionContext + from cudf_polars.dsl.ir import CachedParquetInfo, IRExecutionContext from cudf_polars.streaming.base import ( DataSourceInfo, SerializedDataSourceInfo, @@ -194,6 +194,7 @@ class SplitScan(IR): __slots__ = ( "base_scan", + "cached_parquet_info", "parquet_options", "paths", "schema", @@ -219,6 +220,7 @@ class SplitScan(IR): """Total number of splits.""" parquet_options: ParquetOptions """Parquet-specific options.""" + cached_parquet_info: list[CachedParquetInfo] | None def __init__( self, @@ -228,6 +230,7 @@ def __init__( split_index: int, total_splits: int, parquet_options: ParquetOptions, + cached_parquet_info: list[CachedParquetInfo] | None = None, ): self.schema = schema self.base_scan = base_scan @@ -248,8 +251,10 @@ def __init__( base_scan.include_file_paths, base_scan.predicate, parquet_options, + cached_parquet_info, ) self.parquet_options = parquet_options + self.cached_parquet_info = cached_parquet_info self.children = () if base_scan.typ not in ("parquet",): # pragma: no cover raise NotImplementedError( @@ -284,6 +289,7 @@ def do_evaluate( include_file_paths: str | None, predicate: NamedExpr | None, parquet_options: ParquetOptions, + cached_parquet_info: list[CachedParquetInfo] | None, *, context: IRExecutionContext, ) -> DataFrame: @@ -303,10 +309,24 @@ def do_evaluate( # - We can use all this information to calculate the # "skip_rows" and "n_rows" options to use locally. - rowgroup_metadata = plc.io.parquet_metadata.read_parquet_metadata( - plc.io.SourceInfo(paths) - ).rowgroup_metadata() - total_row_groups = len(rowgroup_metadata) + if cached_parquet_info is not None: + parquet_metadatas = [info.file_metadata for info in cached_parquet_info] + + row_group_num_rows = [ + num_rows + for metadata in parquet_metadatas + for num_rows in metadata.row_group_num_rows + ] + + else: + row_group_num_rows = [ + rg["num_rows"] + for rg in plc.io.parquet_metadata.read_parquet_metadata( + plc.io.SourceInfo(paths) + ).rowgroup_metadata() + ] + + total_row_groups = len(row_group_num_rows) if total_splits <= total_row_groups: # We have enough row-groups in the file to align # all "total_splits" of our reads with row-group @@ -315,17 +335,14 @@ def do_evaluate( # the row-group indices to "skip_rows" and "n_rows". rg_stride = total_row_groups // total_splits skip_rgs = rg_stride * split_index - skip_rows = sum(rg["num_rows"] for rg in rowgroup_metadata[:skip_rgs]) - n_rows = sum( - rg["num_rows"] - for rg in rowgroup_metadata[skip_rgs : skip_rgs + rg_stride] - ) + skip_rows = sum(row_group_num_rows[:skip_rgs]) + n_rows = sum(row_group_num_rows[skip_rgs : skip_rgs + rg_stride]) else: # There are not enough row-groups to align # all "total_splits" of our reads with row-group # boundaries. Use metadata to directly calculate # "skip_rows" and "n_rows" for the current read. - total_rows = sum(rg["num_rows"] for rg in rowgroup_metadata) + total_rows = sum(row_group_num_rows) n_rows = total_rows // total_splits skip_rows = n_rows * split_index @@ -349,6 +366,7 @@ def do_evaluate( include_file_paths, predicate, parquet_options, + cached_parquet_info, context=context, ) @@ -363,6 +381,7 @@ class FusedScan(IR): __slots__ = ( "base_scan", + "cached_parquet_info", "parquet_options", "paths", "schema", @@ -380,6 +399,8 @@ class FusedScan(IR): """File paths assigned to this task.""" parquet_options: ParquetOptions """Parquet-specific options.""" + cached_parquet_info: list[CachedParquetInfo] | None + """Cached parquet metadata.""" def __init__( self, @@ -387,11 +408,13 @@ def __init__( base_scan: Scan, paths: list[str], parquet_options: ParquetOptions, + cached_parquet_info: list[CachedParquetInfo] | None = None, ): self.schema = schema self.base_scan = base_scan self.paths = paths self.parquet_options = parquet_options + self.cached_parquet_info = cached_parquet_info self._non_child_args = ( base_scan.schema, base_scan.typ, @@ -404,6 +427,7 @@ def __init__( base_scan.include_file_paths, base_scan.predicate, parquet_options, + cached_parquet_info, ) self.children = () @@ -431,6 +455,7 @@ def do_evaluate( include_file_paths: str | None, predicate: NamedExpr | None, parquet_options: ParquetOptions, + cached_parquet_info: list[CachedParquetInfo] | None, *, context: IRExecutionContext, ) -> DataFrame: @@ -448,6 +473,7 @@ def do_evaluate( include_file_paths, predicate, parquet_options, + cached_parquet_info, context=context, ) @@ -538,7 +564,7 @@ def can_use_native_parquet_node( @lower_ir_node.register(Scan) def _( ir: Scan, rec: LowerIRTransformer -) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: +) -> tuple[StreamingScan, MutableMapping[IR, PartitionInfo]]: config_options = rec.state["config_options"] parquet_options = config_options.parquet_options if ( @@ -590,24 +616,30 @@ class StreamingScan(IR): __slots__ = ( "base_scan", + "scan_type", "scans", "schema", ) _non_child = ( "scans", "base_scan", + "scan_type", ) - _n_non_child_args = 2 + _n_non_child_args = 3 scans: Sequence[SplitScan] | Sequence[FusedScan] base_scan: Scan def __init__( - self, scans: Sequence[SplitScan] | Sequence[FusedScan], base_scan: Scan + self, + scans: Sequence[SplitScan] | Sequence[FusedScan], + base_scan: Scan, + scan_type: Literal["split", "fused"], ): self.scans = scans self.base_scan = base_scan self.schema = base_scan.schema - self._non_child_args = (scans, base_scan) + self.scan_type = scan_type + self._non_child_args = (scans, base_scan, scan_type) self.children = () @classmethod @@ -639,12 +671,13 @@ def for_split_files( sindex, plan.factor, parquet_options, + None, ) ) sindex += 1 splits_created += 1 sindex = 0 - return cls(scans, base_scan) + return cls(scans, base_scan, "split") @classmethod def for_fused_files( @@ -667,11 +700,12 @@ def for_fused_files( base_scan, base_scan.paths[offset : offset + plan.factor], parquet_options, + None, ) for offset in range(paths_start, paths_end, plan.factor) if base_scan.paths[offset : offset + plan.factor] ] - return cls(scans, base_scan) + return cls(scans, base_scan, "fused") def get_hashable(self) -> Hashable: """Hashable representation of the node.""" @@ -856,6 +890,20 @@ def _sink_to_file( return True +def _columnchunk_metadata_from_footers( + footers: list[plc.io.parquet_metadata.FileMetaData], +) -> dict[str, list[int]]: + columnchunk_metadata: dict[str, list[int]] = {} + for fmd in footers: + for rg in fmd.row_groups: + for col in rg.columns: + name = ".".join(col.meta_data.path_in_schema) + columnchunk_metadata.setdefault(name, []).append( + col.meta_data.total_uncompressed_size + ) + return columnchunk_metadata + + class ParquetMetadata: """ Parquet metadata container. @@ -869,6 +917,7 @@ class ParquetMetadata: """ __slots__ = ( + "cached_parquet_info", "column_names", "max_footer_samples", "mean_size_per_file", @@ -876,6 +925,8 @@ class ParquetMetadata: "paths", "row_count", "sample_paths", + "sampled_file_count", + "total_file_count", ) paths: tuple[str, ...] @@ -892,18 +943,27 @@ class ParquetMetadata: """All column names found it the dataset.""" sample_paths: tuple[str, ...] """Sampled file paths.""" + cached_parquet_info: list[CachedParquetInfo] | None + """Cached parquet info for the sampled paths. Only set if all files were sampled.""" @nvtx_annotate_cudf_polars(message="ParquetMetadata") def __init__(self, paths: tuple[str, ...], max_footer_samples: int): + from cudf_polars.dsl.utils.io import _prefetch_parquet_footers_for_paths + self.paths = paths self.max_footer_samples = max_footer_samples self.row_count = None self.num_row_groups_per_file = () self.mean_size_per_file = {} self.column_names = () - stride = ( - max(1, int(len(paths) / max_footer_samples)) if max_footer_samples else 1 - ) + self.cached_parquet_info = None + self.total_file_count = len(self.paths) + self.sampled_file_count = 0 + if max_footer_samples <= 0: + self.sample_paths = () + return + + stride = max(1, int(len(paths) / max_footer_samples)) self.sample_paths = paths[: stride * max_footer_samples : stride] if not self.sample_paths: @@ -911,21 +971,24 @@ def __init__(self, paths: tuple[str, ...], max_footer_samples: int): # TODO: This requires row_count to be nullable. Why do we allow empty paths? return - total_file_count = len(self.paths) sampled_file_count = len(self.sample_paths) - sample_metadata = plc.io.parquet_metadata.read_parquet_metadata( - plc.io.SourceInfo(list(self.sample_paths)) + + sample_parquet_info = _prefetch_parquet_footers_for_paths( + list(self.sample_paths) ) + sample_footers = [info.file_metadata for info in sample_parquet_info] - if total_file_count == sampled_file_count: - row_count = sample_metadata.num_rows() + sampled_row_count = sum(fmd.num_rows for fmd in sample_footers) + if self.total_file_count == sampled_file_count: + row_count = sampled_row_count + self.cached_parquet_info = sample_parquet_info else: - num_rows_per_sampled_file = int( - sample_metadata.num_rows() / sampled_file_count - ) - row_count = num_rows_per_sampled_file * total_file_count + num_rows_per_sampled_file = int(sampled_row_count / sampled_file_count) + row_count = num_rows_per_sampled_file * self.total_file_count - num_row_groups_per_sampled_file = sample_metadata.num_rowgroups_per_file() + num_row_groups_per_sampled_file = [ + len(fmd.row_groups) for fmd in sample_footers + ] rowgroup_offsets_per_file = list( itertools.accumulate(num_row_groups_per_sampled_file, initial=0) ) @@ -935,7 +998,9 @@ def __init__(self, paths: tuple[str, ...], max_footer_samples: int): sum(uncompressed_sizes[start:end]) for (start, end) in itertools.pairwise(rowgroup_offsets_per_file) ] - for name, uncompressed_sizes in sample_metadata.columnchunk_metadata().items() + for name, uncompressed_sizes in _columnchunk_metadata_from_footers( + sample_footers + ).items() } self.column_names = tuple(column_sizes_per_file) @@ -945,6 +1010,7 @@ def __init__(self, paths: tuple[str, ...], max_footer_samples: int): } self.num_row_groups_per_file = tuple(num_row_groups_per_sampled_file) self.row_count = row_count + self.sampled_file_count = sampled_file_count @nvtx_annotate_cudf_polars(message="_sample_rg_sizes") @@ -1017,13 +1083,19 @@ class ParquetSourceInfo: type: Literal["parquet"] = "parquet" def __init__( - self, row_count: int | None, per_file_means: dict[str, int] | None = None + self, + row_count: int | None, + per_file_means: dict[str, int] | None = None, + *, + # TODO: change this to cached_parquet_info + cached_parquet_info: list[CachedParquetInfo] | None = None, ): if per_file_means is None: per_file_means = {} self.row_count = row_count self.per_file_means = per_file_means + self.cached_parquet_info = cached_parquet_info @classmethod def from_paths( @@ -1082,7 +1154,15 @@ def from_paths( else max(footer_mean, decoded_floor) ) - return cls(row_count, per_file_means) + cached_parquet_info: list[CachedParquetInfo] | None + if ( + metadata.sampled_file_count == metadata.total_file_count + and metadata.cached_parquet_info is not None + ): + cached_parquet_info = list(metadata.cached_parquet_info) + else: + cached_parquet_info = None + return cls(row_count, per_file_means, cached_parquet_info=cached_parquet_info) def column_storage_size(self, column: str) -> int | None: """Return the average storage size for a single column in one file.""" diff --git a/python/cudf_polars/cudf_polars/streaming/select.py b/python/cudf_polars/cudf_polars/streaming/select.py index e1466f0b6610..a4bb5f2ad0bf 100644 --- a/python/cudf_polars/cudf_polars/streaming/select.py +++ b/python/cudf_polars/cudf_polars/streaming/select.py @@ -1,9 +1,10 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Parallel Select Logic.""" from __future__ import annotations +import dataclasses from collections import defaultdict from typing import TYPE_CHECKING @@ -431,8 +432,17 @@ def _( if scan_child and scan_child.predicate is None and scan_child.typ == "parquet": # Special Case: Fast count. + # We can't use prefetched file metadata here, because we're in lowering, + # not execution, so we don't have an IRExecutionContext with the prefetched + # file metadata yet. count = Scan._get_parquet_row_count_from_metadata( - scan_child.paths, scan_child.skip_rows, scan_child.n_rows + scan_child.paths, + scan_child.skip_rows, + scan_child.n_rows, + dataclasses.replace( + scan_child.parquet_options, prefetch_file_metadata=False + ), + None, ) dtype = ir.exprs[0].value.dtype diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 8e3bcdcc3d9f..d536fa4aa9c4 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -215,6 +215,9 @@ class ParquetOptions: Whether to use the native rapidsmpf node for parquet reading. This option is only used by the streaming executor. Default is False. + prefetch_file_metadata + Whether to prefetch parquet file metadata and pass it through + `parquet_metadatas` to avoid rereading file footers. use_jit_filter Whether to use JIT compilation for post-read filtering in Parquet scans. When enabled, filter predicates are JIT-compiled to CUDA kernels for @@ -261,6 +264,13 @@ class ParquetOptions: default=False, ) ) + prefetch_file_metadata: bool = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__PREFETCH_FILE_METADATA", + _bool_converter, + default=False, + ) + ) use_jit_filter: bool = dataclasses.field( default_factory=_make_default_factory( f"{_env_prefix}__USE_JIT_FILTER", @@ -284,6 +294,13 @@ def __post_init__(self) -> None: # noqa: D105 raise TypeError("max_row_group_samples must be an int") if not isinstance(self.use_rapidsmpf_native, bool): raise TypeError("use_rapidsmpf_native must be a bool") + if not isinstance(self.prefetch_file_metadata, bool): + raise TypeError("prefetch_file_metadata must be a bool") + + if self.use_rapidsmpf_native and self.prefetch_file_metadata: + raise NotImplementedError( + "'use_rapidsmpf_native=True' does not currently support 'prefetch_file_metadata=True'" + ) if not isinstance(self.use_jit_filter, bool): raise TypeError("use_jit_filter must be a bool") @@ -880,6 +897,11 @@ def from_polars_engine( user_parquet_options = engine.config.get("parquet_options", {}) if user_parquet_options is None: user_parquet_options = {} + + if isinstance(user_parquet_options, dict): + parquet_options = ParquetOptions(**user_parquet_options) + else: + parquet_options = user_parquet_options # This is set in polars, and so can't be overridden by the environment user_raise_on_fail = engine.config.get("raise_on_fail", False) user_memory_resource_config = engine.config.get("memory_resource_config", None) @@ -907,6 +929,10 @@ def from_polars_engine( match user_executor: case "in-memory": executor = InMemoryExecutor(**user_executor_options) + if parquet_options.prefetch_file_metadata: + raise NotImplementedError( + "Prefetching is not supported for the in-memory executor." + ) case "streaming": user_executor_options = user_executor_options.copy() if "min_device_size" not in user_executor_options: @@ -929,7 +955,7 @@ def from_polars_engine( kwargs = { "raise_on_fail": user_raise_on_fail, - "parquet_options": ParquetOptions(**user_parquet_options), + "parquet_options": parquet_options, "executor": executor, "device": engine.device, "memory_resource_config": user_memory_resource_config, diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index fe535e25c474..13e88ead7731 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -4,7 +4,7 @@ from __future__ import annotations import math -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast import pytest @@ -12,7 +12,15 @@ from cudf_polars import Translator from cudf_polars.containers import DataType -from cudf_polars.dsl.ir import IRExecutionContext, Scan +from cudf_polars.dsl.ir import ( + Empty, + IRExecutionContext, + Scan, +) +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.base import ( DataSourceInfo, @@ -36,10 +44,14 @@ if TYPE_CHECKING: import concurrent.futures + from collections.abc import Callable from pathlib import Path from typing import Any, Literal + import pylibcudf as plc + import cudf_polars.engine.core + from cudf_polars.engine.core import StreamingEngine @pytest.fixture(scope="module") @@ -92,6 +104,54 @@ def test_scan_parquet_use_rapidsmpf_native(tmp_path, df, streaming_engine_factor assert_gpu_result_equal(pl.scan_parquet(tmp_path), engine=streaming_engine) +@pytest.mark.parametrize( + "target_partition_size_and_n_files", [(1_000, 1), (1_000, 2), (1_000_000, 5)] +) +def test_scan_parquet_prefetch_file_metadata( + tmp_path: Path, + target_partition_size_and_n_files: tuple[int, int], + df: pl.DataFrame, + streaming_engine_factory: Callable[..., StreamingEngine], +): + target_partition_size, n_files = target_partition_size_and_n_files + streaming_engine = streaming_engine_factory( + StreamingOptions( + target_partition_size=target_partition_size, + parquet_options={"prefetch_file_metadata": True}, + ), + ) + make_partitioned_source(df, tmp_path, "parquet", n_files=n_files) + 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) + + +def test_prefetch_parquet_file_metadata_no_parquet_scans() -> None: + result = prefetch_parquet_file_metadata_for_ir( + Empty({}), py_executor=None, stats=None + ) + assert result == {} + + +def test_prefetch_file_metadata_select_fast_count( + df: pl.DataFrame, + streaming_engine_factory: Callable[..., 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) + + # --------------------------------------------------------------------------- # Tests migrated from tests/streaming/test_scan.py # --------------------------------------------------------------------------- @@ -187,7 +247,10 @@ def test_scan_union(engine: pl.GPUEngine, tmp_path: Path) -> None: assert_gpu_result_equal(q, engine=engine) -def _make_parquet_scan(paths: list[str]) -> Scan: +def _make_parquet_scan( + paths: list[str], parquet_options: ParquetOptions | None = None +) -> Scan: + parquet_options = parquet_options or ParquetOptions() return Scan( {"x": DataType(pl.Int64())}, "parquet", @@ -200,7 +263,8 @@ def _make_parquet_scan(paths: list[str]) -> Scan: None, None, None, - ParquetOptions(), + parquet_options, + None, ) @@ -289,12 +353,262 @@ def test_expand_scan_for_rank_split_files( def test_streaming_scan_raises() -> None: # This isn't reachable by normal cudf-polars usage. scan = _make_parquet_scan(["file.parquet"]) - fused = FusedScan(scan.schema, scan, scan.paths, scan.parquet_options) + fused = FusedScan(scan.schema, scan, scan.paths, scan.parquet_options, []) ctx = IRExecutionContext() with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): StreamingScan.do_evaluate([fused], scan, context=ctx) +def test_scan_missing_prefetch_metadata_raises() -> None: + # This isn't reachable by polars' public API, so we test it directly. + scan = _make_parquet_scan( + ["file.parquet"], parquet_options=ParquetOptions(prefetch_file_metadata=True) + ) + ctx = IRExecutionContext() + + with pytest.raises( + AssertionError, + match=r"Cached parquet info is required", + ): + Scan.do_evaluate( + scan.schema, + scan.typ, + scan.reader_options, + scan.paths, + scan.with_columns, + scan.skip_rows, + scan.n_rows, + scan.row_index, + scan.include_file_paths, + scan.predicate, + scan.parquet_options, + None, + context=ctx, + ) + + with pytest.raises( + AssertionError, + match=r"Paths do not match cached parquet info", + ): + Scan.do_evaluate( + scan.schema, + scan.typ, + scan.reader_options, + scan.paths, + scan.with_columns, + scan.skip_rows, + scan.n_rows, + scan.row_index, + scan.include_file_paths, + scan.predicate, + scan.parquet_options, + [], + context=ctx, + ) + + +def test_streaming_scan_missing_prefetch_metadata_raises() -> None: + # This isn't reachable by polars' public API, so we test it directly. + scan = _make_parquet_scan( + ["file.parquet"], parquet_options=ParquetOptions(prefetch_file_metadata=True) + ) + fused = FusedScan(scan.schema, scan, scan.paths, scan.parquet_options, []) + + ctx = IRExecutionContext() + with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): + StreamingScan.do_evaluate([fused], scan, context=ctx) + + +def test_split_scan_do_evaluate_missing_prefetch_metadata() -> None: + paths = ["/some/missing/file.parquet"] + parquet_options = ParquetOptions(prefetch_file_metadata=True) + context = IRExecutionContext() + schema = {"x": DataType(pl.Int64())} + + with pytest.raises( + AssertionError, + match=(r"Paths do not match cached parquet info."), + ): + SplitScan.do_evaluate( + 0, + 4, + schema, + "parquet", + {}, + paths, + None, + 0, + -1, + None, + None, + None, + parquet_options, + [], + context=context, + ) + + +def test_prefetch_file_metadata_join( + tmp_path: Path, streaming_engine_factory: Callable[..., 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) + + +def _make_cached_parquet_info( + paths: list[str], size: int = 10 +) -> list[CachedParquetInfo]: + return [ + # `file_metadata` is not used by identity/hash tests. + # It only needs to be a stable value for equality checks. + CachedParquetInfo( + path=path, + size=size, + file_metadata=cast("plc.io.parquet_metadata.FileMetaData", path), + ) + for path in paths + ] + + +def test_prefetch_file_metadata_with_cached_scan_parent_nodes( + tmp_path: Path, streaming_engine_factory: Callable[..., StreamingEngine] +) -> None: + # Regression test for replace not replacing StreamingScan nodes with their prefetched variants. + source = tmp_path / "data.parquet" + pl.DataFrame( + { + "k": [1, 1, 2, 2, 3, 3], + "v": [10, 11, 20, 21, 30, 31], + } + ).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) + + +def test_fused_scan_identity_equality() -> None: + base = _make_parquet_scan(["a.parquet", "b.parquet"]) + paths = ["a.parquet"] + info = _make_cached_parquet_info(paths) + + a = FusedScan(base.schema, base, paths, base.parquet_options, info) + b = FusedScan(base.schema, base, paths, base.parquet_options, info.copy()) + c = FusedScan(base.schema, base, ["b.parquet"], base.parquet_options, info) + + assert a == b + assert hash(a) == hash(b) + assert a != c + + +def test_split_scan_identity_equality() -> None: + base = _make_parquet_scan(["a.parquet"]) + info = _make_cached_parquet_info(base.paths) + + a = SplitScan(base.schema, base, base.paths, 0, 4, base.parquet_options, info) + b = SplitScan( + base.schema, base, base.paths, 0, 4, base.parquet_options, info.copy() + ) + c = SplitScan(base.schema, base, base.paths, 1, 4, base.parquet_options, info) + + assert a == b + assert hash(a) == hash(b) + assert a != c + + +def test_streaming_scan_identity_equality() -> None: + base = _make_parquet_scan(["a.parquet"]) + split = SplitScan( + base.schema, + base, + base.paths, + 0, + 2, + base.parquet_options, + _make_cached_parquet_info(base.paths, size=10), + ) + split_same = SplitScan( + base.schema, + base, + base.paths, + 0, + 2, + base.parquet_options, + _make_cached_parquet_info(base.paths, size=11), + ) + split_diff = SplitScan( + base.schema, + base, + base.paths, + 1, + 2, + base.parquet_options, + _make_cached_parquet_info(base.paths, size=10), + ) + + a = StreamingScan([split], base, "split") + b = StreamingScan([split_same], base, "split") + c = StreamingScan([split_diff], base, "split") + + assert a == b + assert hash(a) == hash(b) + assert a != c + + +def test_cached_parquet_info_excluded_from_identity() -> None: + base = _make_parquet_scan(["a.parquet"]) + info = _make_cached_parquet_info(base.paths) + + scan_without = _make_parquet_scan(base.paths) + scan_with = Scan( + base.schema, + "parquet", + {}, + None, + base.paths, + None, + 0, + -1, + None, + None, + None, + base.parquet_options, + info, + ) + assert scan_without == scan_with + assert hash(scan_without) == hash(scan_with) + + split_without = SplitScan( + base.schema, base, base.paths, 0, 4, base.parquet_options, None + ) + split_with = SplitScan( + base.schema, base, base.paths, 0, 4, base.parquet_options, info + ) + assert split_without == split_with + assert hash(split_without) == hash(split_with) + + fused_without = FusedScan(base.schema, base, base.paths, base.parquet_options, None) + fused_with = FusedScan(base.schema, base, base.paths, base.parquet_options, info) + assert fused_without == fused_with + assert hash(fused_without) == hash(fused_with) + + class FooSource(DataSourceInfo): def __init__(self, size: int): self._size = size diff --git a/python/cudf_polars/tests/streaming/test_stats.py b/python/cudf_polars/tests/streaming/test_stats.py index 79cc1d5a57de..22c7a1795ee1 100644 --- a/python/cudf_polars/tests/streaming/test_stats.py +++ b/python/cudf_polars/tests/streaming/test_stats.py @@ -12,16 +12,23 @@ import polars as pl import pylibcudf as plc +import rmm.pylibrmm.stream +import cudf_polars.containers import cudf_polars.streaming.io as streaming_io from cudf_polars import Translator from cudf_polars.containers import DataType -from cudf_polars.dsl.ir import Empty, Projection +from cudf_polars.dsl.ir import ( + Empty, + Projection, +) from cudf_polars.engine.options import StreamingOptions from cudf_polars.streaming.base import SerializedDataSourceInfo, StatsCollector from cudf_polars.streaming.io import ( DataFrameSourceInfo, + ParquetMetadata, ParquetSourceInfo, + _build_parquet_source, _clear_source_info_cache, ) from cudf_polars.streaming.statistics import collect_statistics @@ -33,16 +40,22 @@ import concurrent.futures import pathlib + from cudf_polars.typing import Schema + @pytest.fixture(scope="module") -def df(): - return pl.DataFrame( +def df_and_schema() -> tuple[pl.DataFrame, Schema]: + stream = rmm.pylibrmm.stream.Stream() + df = pl.DataFrame( { "x": range(3_000), "y": ["cat", "dog", "fish"] * 1_000, "z": [1.0, 2.0, 3.0, 4.0, 5.0] * 600, } ) + df_ = cudf_polars.containers.DataFrame.from_polars(df, stream=stream) + schema = {column.name: column.dtype for column in df_.columns} + return df, schema # Simple engine for IR translation / stats collection only (no actual GPU execution) @@ -59,8 +72,11 @@ def stats_engine(): def test_base_stats_dataframescan( - df, stats_engine, parquet_stats_executor: concurrent.futures.ThreadPoolExecutor + df_and_schema: tuple[pl.DataFrame, Schema], + stats_engine, + parquet_stats_executor: concurrent.futures.ThreadPoolExecutor, ): + df, _schema = df_and_schema row_count = df.height q = pl.LazyFrame(df) ir = Translator(q._ldf.visit(), stats_engine).translate_ir() @@ -81,7 +97,7 @@ def test_base_stats_dataframescan( @pytest.mark.parametrize("max_row_group_samples", [1, 0]) def test_base_stats_parquet( tmp_path, - df, + df_and_schema: tuple[pl.DataFrame, Schema], n_files, row_group_size, max_footer_samples, @@ -89,6 +105,7 @@ def test_base_stats_parquet( parquet_stats_executor: concurrent.futures.ThreadPoolExecutor, ): _clear_source_info_cache() + df, _schema = df_and_schema make_partitioned_source( df, tmp_path, @@ -147,6 +164,8 @@ class FakeParquetMetadata: def __init__(self, paths: tuple[str, ...], max_footer_samples: int) -> None: self.paths = paths self.max_footer_samples = max_footer_samples + self.sampled_file_count = 1 + self.total_file_count = len(paths) sampled_cols: list[str] = [] @@ -228,6 +247,63 @@ def test_parquet_round_trip_empty() -> None: assert restored.per_file_means == {} +def test_parquet_source_info_stores_footers_when_all_files_sampled( + tmp_path: pathlib.Path, + df_and_schema: tuple[pl.DataFrame, Schema], +) -> None: + _clear_source_info_cache() + df, schema = df_and_schema + make_partitioned_source(df, tmp_path, "parquet", n_files=2) + paths = tuple(str(p) for p in sorted(tmp_path.iterdir())) + info = _build_parquet_source( + paths, + frozenset(df.columns), + tuple(schema.items()), + max_footer_samples=10, + max_row_group_samples=0, + ) + + assert info.cached_parquet_info is not None + assert len(info.cached_parquet_info) == len(paths) + assert ( + sum(cached.file_metadata.num_rows for cached in info.cached_parquet_info) + == df.height + ) + + +def test_parquet_source_info_omits_footers_when_paths_are_sampled( + tmp_path: pathlib.Path, + df_and_schema: tuple[pl.DataFrame, Schema], +) -> None: + _clear_source_info_cache() + df, schema = df_and_schema + make_partitioned_source(df, tmp_path, "parquet", n_files=5) + paths = tuple(str(p) for p in sorted(tmp_path.iterdir())) + info = _build_parquet_source( + paths, + frozenset(df.columns), + tuple(schema.items()), + max_footer_samples=2, + max_row_group_samples=0, + ) + + assert info.cached_parquet_info is None + + +def test_parquet_metadata_reads_footers( + tmp_path: pathlib.Path, + df_and_schema: tuple[pl.DataFrame, Schema], +) -> None: + df, _schema = df_and_schema + make_partitioned_source(df, tmp_path, "parquet", n_files=1) + path = next(tmp_path.iterdir()) + metadata = ParquetMetadata((str(path),), max_footer_samples=1) + + assert metadata.cached_parquet_info is not None + assert len(metadata.cached_parquet_info) == 1 + assert metadata.row_count == df.height + + def test_dataframe_round_trip() -> None: info = DataFrameSourceInfo(2500) data = info.serialize() @@ -328,10 +404,11 @@ def test_serialize_stats_roundtrip_dataframescan( def test_serialize_stats_roundtrip_parquet( tmp_path: pathlib.Path, - df: pl.DataFrame, + df_and_schema: tuple[pl.DataFrame, Schema], parquet_stats_executor: concurrent.futures.ThreadPoolExecutor, ) -> None: _clear_source_info_cache() + df, _schema = df_and_schema make_partitioned_source(df, tmp_path, "parquet", n_files=3) engine = pl.GPUEngine( raise_on_fail=True, diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 7a9811c7d146..cc3d34f1d7ae 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -331,6 +331,7 @@ def test_parquet_options_from_env(monkeypatch: pytest.MonkeyPatch) -> None: m.setenv("CUDF_POLARS__PARQUET_OPTIONS__MAX_FOOTER_SAMPLES", "0") m.setenv("CUDF_POLARS__PARQUET_OPTIONS__MAX_ROW_GROUP_SAMPLES", "0") m.setenv("CUDF_POLARS__PARQUET_OPTIONS__USE_RAPIDSMPF_NATIVE", "0") + m.setenv("CUDF_POLARS__PARQUET_OPTIONS__PREFETCH_FILE_METADATA", "1") m.setenv("CUDF_POLARS__PARQUET_OPTIONS__USE_JIT_FILTER", "1") # Test default @@ -343,6 +344,7 @@ def test_parquet_options_from_env(monkeypatch: pytest.MonkeyPatch) -> None: assert config.parquet_options.max_footer_samples == 0 assert config.parquet_options.max_row_group_samples == 0 assert config.parquet_options.use_rapidsmpf_native is False + assert config.parquet_options.prefetch_file_metadata is True assert config.parquet_options.use_jit_filter is True with monkeypatch.context() as m: @@ -420,6 +422,7 @@ def test_fallback_mode_default(monkeypatch: pytest.MonkeyPatch) -> None: "max_footer_samples", "max_row_group_samples", "use_rapidsmpf_native", + "prefetch_file_metadata", "use_jit_filter", ], ) @@ -433,6 +436,22 @@ def test_validate_parquet_options(option: str) -> None: ) +def test_prefetch_and_use_rapidsmpf_native_raises() -> None: + with pytest.raises( + NotImplementedError, + match="'use_rapidsmpf_native=True' does not currently support 'prefetch_file_metadata=True'", + ): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + parquet_options={ + "use_rapidsmpf_native": True, + "prefetch_file_metadata": True, + }, + ) + ) + + def test_validate_raise_on_fail() -> None: with pytest.raises(TypeError, match="'raise_on_fail' must be"): ConfigOptions.from_polars_engine( diff --git a/python/cudf_polars/tests/test_scan.py b/python/cudf_polars/tests/test_scan.py index ccbd68050588..6b0185f41d1f 100644 --- a/python/cudf_polars/tests/test_scan.py +++ b/python/cudf_polars/tests/test_scan.py @@ -16,12 +16,15 @@ import polars as pl +from cudf_polars.containers import DataType +from cudf_polars.dsl.ir import IRExecutionContext, Scan from cudf_polars.testing.asserts import ( assert_gpu_result_equal, assert_ir_translation_raises, ) from cudf_polars.testing.engine_utils import is_streaming_engine from cudf_polars.testing.io import make_partitioned_source +from cudf_polars.utils.config import ConfigOptions, ParquetOptions from cudf_polars.utils.versions import ( POLARS_VERSION_LT_138, POLARS_VERSION_LT_139, @@ -169,6 +172,46 @@ def test_negative_slice_pushdown_raises(engine: pl.GPUEngine, tmp_path): assert_ir_translation_raises(q, engine, NotImplementedError) +def test_scan_parquet_prefetch_file_metadata_in_memory_raises(): + with pytest.raises( + NotImplementedError, + match=r"Prefetching is not supported for the in-memory executor.", + ): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="in-memory", + parquet_options=ParquetOptions(prefetch_file_metadata=True), + ) + ) + + +def test_scan_do_evaluate_missing_prefetch_metadata() -> None: + paths = ["/some/missing/file.parquet"] + parquet_options = ParquetOptions(prefetch_file_metadata=True) + context = IRExecutionContext() + schema = {"a": DataType(pl.Int64())} + + with pytest.raises( + AssertionError, + match=(r"Paths do not match cached parquet info."), + ): + Scan.do_evaluate( + schema, + "parquet", + {}, + paths, + None, + 0, + -1, + None, + None, + None, + parquet_options, + [], + context=context, + ) + + def test_scan_unsupported_raises(engine: pl.GPUEngine, tmp_path): df = pl.DataFrame({"a": [1, 2, 3]}) diff --git a/python/cudf_polars/tests/test_select.py b/python/cudf_polars/tests/test_select.py index f37c2d195d92..435bec5d2031 100644 --- a/python/cudf_polars/tests/test_select.py +++ b/python/cudf_polars/tests/test_select.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -8,10 +8,12 @@ import polars as pl +from cudf_polars.dsl.ir import Scan from cudf_polars.testing.asserts import ( assert_gpu_result_equal, assert_ir_translation_raises, ) +from cudf_polars.utils.config import ParquetOptions def test_select(engine: pl.GPUEngine): @@ -147,3 +149,54 @@ def test_select_fast_count_parquet_skip_rows( q = pl.scan_parquet(file).slice(1, 5).select(pl.len()) assert_gpu_result_equal(q, engine=engine) + + +PARQUET_FAST_COUNT_ROWS = 10 + + +@pytest.fixture(scope="module") +def parquet_fast_count_df() -> pl.DataFrame: + return pl.DataFrame({"a": range(PARQUET_FAST_COUNT_ROWS)}) + + +@pytest.fixture( + params=[ + pytest.param({"skip_rows": 0, "n_rows": None}, id="all_rows"), + pytest.param({"skip_rows": 3, "n_rows": None}, id="skip_rows"), + pytest.param({"skip_rows": 2, "n_rows": 4}, id="skip_rows_and_limit"), + pytest.param({"skip_rows": 0, "n_rows": 5}, id="n_rows"), + pytest.param({"skip_rows": 8, "n_rows": 10}, id="skip_near_end"), + pytest.param( + {"skip_rows": PARQUET_FAST_COUNT_ROWS, "n_rows": None}, + id="skip_all", + ), + ], +) +def parquet_scan_row_bounds(request) -> dict[str, int | None]: + return request.param + + +def test_get_parquet_row_count_from_metadata_raises() -> None: + paths = ["/some/missing/file.parquet"] + parquet_options = ParquetOptions(prefetch_file_metadata=True) + + with pytest.raises(AssertionError, match=r"Cached parquet info is required"): + Scan._get_parquet_row_count_from_metadata( + paths, + skip_rows=0, + n_rows=-1, + parquet_options=parquet_options, + cached_parquet_info=None, + ) + + with pytest.raises( + AssertionError, + match=(r"Paths do not match cached parquet info."), + ): + Scan._get_parquet_row_count_from_metadata( + paths, + skip_rows=0, + n_rows=-1, + parquet_options=parquet_options, + cached_parquet_info=[], + )