From 80d92b79a0430d352d7cdf894b0b9f24e0545328 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 22 May 2026 12:15:11 -0700 Subject: [PATCH 01/51] Prefetch parquet metadata for scan tasks Add a parquet metadata prefetch cache in IR execution context and plumb it through in-memory and streaming scan paths to reuse footer metadata across related reads. --- python/cudf_polars/cudf_polars/callback.py | 9 +- python/cudf_polars/cudf_polars/dsl/ir.py | 102 +++++++++++++++++- python/cudf_polars/cudf_polars/engine/core.py | 26 ++--- .../cudf_polars/cudf_polars/streaming/io.py | 37 +++++-- .../cudf_polars/cudf_polars/utils/config.py | 13 +++ .../cudf_polars/tests/streaming/test_scan.py | 11 ++ python/cudf_polars/tests/test_config.py | 3 + python/cudf_polars/tests/test_scan.py | 17 +++ 8 files changed, 190 insertions(+), 28 deletions(-) diff --git a/python/cudf_polars/cudf_polars/callback.py b/python/cudf_polars/cudf_polars/callback.py index 93dd0790bd59..a91e63b1423e 100644 --- a/python/cudf_polars/cudf_polars/callback.py +++ b/python/cudf_polars/cudf_polars/callback.py @@ -24,7 +24,10 @@ from rmm._cuda import gpu import cudf_polars.dsl.tracing -from cudf_polars.dsl.ir import IRExecutionContext +from cudf_polars.dsl.ir import ( + IRExecutionContext, + prefetch_parquet_file_metadata_for_ir, +) from cudf_polars.dsl.tracing import CUDF_POLARS_NVTX_DOMAIN from cudf_polars.dsl.translate import Translator from cudf_polars.utils.config import ( @@ -301,6 +304,10 @@ def _callback( ): if config_options.executor.name == "in-memory": context = IRExecutionContext() + prefetch_parquet_file_metadata_for_ir( + ir, + context, + ) df = ir.evaluate(cache={}, timer=timer, context=context).to_polars() if timer is None: return df diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 8f2e8fe686b1..e6e9532e8555 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -14,6 +14,7 @@ from __future__ import annotations import asyncio +import concurrent.futures import contextlib import contextvars import functools @@ -47,7 +48,10 @@ from cudf_polars.dsl.expressions.base import ExecutionContext from cudf_polars.dsl.nodebase import Node from cudf_polars.dsl.to_ast import _DECIMAL_IDS, to_ast, to_parquet_filter -from cudf_polars.dsl.tracing import log_do_evaluate, nvtx_annotate_cudf_polars +from cudf_polars.dsl.tracing import ( + log_do_evaluate, + nvtx_annotate_cudf_polars, +) from cudf_polars.dsl.utils.reshape import broadcast from cudf_polars.dsl.utils.windows import ( offsets_to_windows, @@ -66,7 +70,6 @@ if TYPE_CHECKING: from collections.abc import Callable, Generator, Hashable, Iterable, Sequence - from concurrent.futures import ThreadPoolExecutor from typing import Literal, Self from polars import polars # type: ignore[attr-defined] @@ -129,9 +132,12 @@ 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) + parquet_file_metadata: dict[ + tuple[str, ...], list[plc.io.parquet_metadata.FileMetaData] + ] = field(default_factory=dict) async def to_thread( self, func: Callable[P, T], /, *args: P.args, **kwargs: P.kwargs @@ -185,6 +191,81 @@ def stream_ordered_after(self, *dfs: DataFrame) -> Generator[Stream, None, None] yield result_stream +@nvtx_annotate_cudf_polars(message="PrefetchParquetFootersForPaths") +def _fetch_parquet_footers_for_paths( + paths: tuple[str, ...], +) -> tuple[tuple[str, ...], list[plc.io.parquet_metadata.FileMetaData]]: + metadata = plc.io.parquet_metadata.read_parquet_footers( + plc.io.SourceInfo(list(paths)) + ) + return paths, metadata + + +@nvtx_annotate_cudf_polars(message="PrefetchParquetFileMetadataForIR") +def prefetch_parquet_file_metadata_for_ir( + root: IR, + context: IRExecutionContext, +) -> None: + """ + Prefetch parquet metadata for all parquet scans in an IR graph. + + Parameters + ---------- + root + The root of the IR graph, which will be traversed. + context + The IR execution context. Its ``py_executor`` is used to fetch + metadata concurrently, its ``parquet_file_metadata`` is mutated + to cache the newly read parquet metadata. + """ + from cudf_polars.dsl.traversal import traversal + + groups = { + tuple(node.paths) + for node in traversal([root]) + if isinstance(node, Scan) and node.typ == "parquet" and node.paths + } + if not groups: + return + + missing_files = { + (path,) + for group in groups + for path in group + if (path,) not in context.parquet_file_metadata + } + + cm: contextlib.AbstractContextManager[concurrent.futures.Executor | None] + + if context.py_executor is None: + cm = executor = concurrent.futures.ThreadPoolExecutor() + # We didn't create the executor, so we don't close it. + else: + cm = contextlib.nullcontext() + executor = context.py_executor + + if len(missing_files) > 0: + with cm: + futures = [ + executor.submit(_fetch_parquet_footers_for_paths, missing_file) + for missing_file in missing_files + ] + + for future in concurrent.futures.as_completed(futures): + paths, metadata = future.result() + context.parquet_file_metadata.setdefault(paths, metadata) + + for group in groups: + context.parquet_file_metadata.setdefault( + group, + list( + itertools.chain.from_iterable( + context.parquet_file_metadata[(path,)] for path in group + ) + ), + ) + + _BINOPS = { plc.binaryop.BinaryOperator.EQUAL, plc.binaryop.BinaryOperator.NOT_EQUAL, @@ -802,6 +883,16 @@ def read_csv_header( df, ) elif typ == "parquet": + if parquet_options.prefetch_file_metadata: + try: + parquet_metadatas = context.parquet_file_metadata[tuple(paths)] + except KeyError as e: + raise AssertionError( + f"Parquet file metadata was not prefetched for paths: {list(paths)}." + ) from e + else: + parquet_metadatas = None + filters = None if predicate is not None and row_index is None: # Can't apply filters during read if we have a row index. @@ -830,6 +921,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() @@ -862,7 +954,9 @@ 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) diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 5fddabcf78df..505944894748 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -24,7 +24,7 @@ import polars as pl from cudf_polars.containers import DataFrame -from cudf_polars.dsl.ir import IRExecutionContext +from cudf_polars.dsl.ir import IRExecutionContext, 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 @@ -407,14 +407,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. @@ -429,10 +427,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 @@ -441,8 +439,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 ------- @@ -451,9 +447,6 @@ def execute_ir_on_rank( metadata Collected channel metadata. """ - ir_context = IRExecutionContext( - py_executor, get_cuda_stream=ctx.get_stream_from_pool, query_id=query_id - ) metadata_collector: list[ChannelMetadata] = [] nodes, output = generate_network( @@ -686,6 +679,14 @@ def evaluate_on_rank( """ stats = allgather_stats(comm, ctx.br(), ir, config_options, py_executor) ir, partition_info = lower_ir_graph(ir, config_options, stats) + ir_context = IRExecutionContext( + py_executor, get_cuda_stream=ctx.get_stream_from_pool, query_id=query_id + ) + if config_options.parquet_options.prefetch_file_metadata: + prefetch_parquet_file_metadata_for_ir( + ir, + ir_context, + ) if comm.rank == 0: # At least for now, the query plan is identical on all ranks, @@ -696,11 +697,10 @@ def evaluate_on_rank( 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 4472735ac356..e93782435dd5 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -206,10 +206,30 @@ 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) + # row_group_num_rows: list[int] + if parquet_options.prefetch_file_metadata: + try: + parquet_metadatas = context.parquet_file_metadata[tuple(paths)] + except KeyError as e: + raise AssertionError( + f"Parquet file metadata was not prefetched for paths: {list(paths)}." + ) from e + + 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 @@ -218,17 +238,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 diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index eefe8d618c88..7122a2a88102 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -206,6 +206,10 @@ 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. + Default is False. """ _env_prefix = "CUDF_POLARS__PARQUET_OPTIONS" @@ -247,6 +251,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, + ) + ) def __post_init__(self) -> None: # noqa: D105 if not isinstance(self.chunked, bool): @@ -263,6 +274,8 @@ 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") def default_target_partition_size(min_device_size: int | None) -> int: diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 11afaac838ac..2a46442a63e7 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -58,6 +58,17 @@ 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) +def test_scan_parquet_prefetch_file_metadata(tmp_path, df, streaming_engine_factory): + streaming_engine = streaming_engine_factory( + StreamingOptions( + target_partition_size=1_000, + parquet_options={"prefetch_file_metadata": True}, + ), + ) + make_partitioned_source(df, tmp_path, "parquet", n_files=2) + assert_gpu_result_equal(pl.scan_parquet(tmp_path), engine=streaming_engine) + + # --------------------------------------------------------------------------- # Tests migrated from tests/streaming/test_scan.py # --------------------------------------------------------------------------- diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 19bede15d290..6dc45ccfa8ed 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -327,6 +327,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") # Test default engine = pl.GPUEngine() @@ -338,6 +339,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 with monkeypatch.context() as m: m.setenv("CUDF_POLARS__PARQUET_OPTIONS__CHUNKED", "foo") @@ -416,6 +418,7 @@ def test_fallback_mode_default(monkeypatch: pytest.MonkeyPatch) -> None: "max_footer_samples", "max_row_group_samples", "use_rapidsmpf_native", + "prefetch_file_metadata", ], ) def test_validate_parquet_options(option: str) -> None: diff --git a/python/cudf_polars/tests/test_scan.py b/python/cudf_polars/tests/test_scan.py index 1dca93dd38ae..bc71e1541f86 100644 --- a/python/cudf_polars/tests/test_scan.py +++ b/python/cudf_polars/tests/test_scan.py @@ -167,6 +167,23 @@ def test_negative_slice_pushdown_raises(engine: pl.GPUEngine, tmp_path): assert_ir_translation_raises(q, engine, NotImplementedError) +@pytest.mark.parametrize("chunked", [False, True], ids=["single_read", "chunked"]) +def test_scan_parquet_prefetch_file_metadata( + tmp_path: Path, df: pl.DataFrame, *, chunked: bool +): + make_partitioned_source(df, tmp_path / "file", "parquet") + q = pl.scan_parquet(tmp_path / "file") + engine = pl.GPUEngine( + executor="in-memory", + raise_on_fail=True, + parquet_options={ + "chunked": chunked, + "prefetch_file_metadata": True, + }, + ) + assert_gpu_result_equal(q, engine=engine) + + def test_scan_unsupported_raises(engine: pl.GPUEngine, tmp_path): df = pl.DataFrame({"a": [1, 2, 3]}) From 109ad43f1fdc62510800bdc561178db007e08aa4 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 2 Jun 2026 07:02:16 -0700 Subject: [PATCH 02/51] Use FilepathSource --- python/cudf_polars/cudf_polars/dsl/ir.py | 84 ++++++++++++------- .../cudf_polars/cudf_polars/streaming/io.py | 2 +- 2 files changed, 57 insertions(+), 29 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index e6e9532e8555..6d795f4d3c6e 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -113,6 +113,13 @@ ] +@dataclass(frozen=True) +class ParquetFileMetadata: + # TODO: rename this; it's too close to plc.io.parquet_metadata.FileMetadata + source_info: plc.io.types.SourceInfo + metadata: list[plc.io.parquet_metadata.FileMetaData] + + @dataclass(frozen=True) class IRExecutionContext: """ @@ -130,14 +137,24 @@ class IRExecutionContext: A zero-argument callable that returns a CUDA stream. query_id Identifier for the query being executed. + parquet_file_metadata + Cache of parquet file metadata. The keys here are the ``paths`` of Scan nodes + with a ``parquet`` type. The values are ``ParquetFileMetadata`` objects. The + ``source_info`` is the ``SourceInfo`` object that was used to read the metadata, + which will contain the known file size for remote files. ``metadata`` is the list + of ``FileMetaData`` objects that were read. Both ``source_info`` and ``metadata`` + should be used in ``read_parquet`` calls later on. """ 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) - parquet_file_metadata: dict[ - tuple[str, ...], list[plc.io.parquet_metadata.FileMetaData] - ] = field(default_factory=dict) + + # This should cache a tuple(size, FileMetadata) + # maybe a dataclass + parquet_file_metadata: dict[tuple[str, ...], ParquetFileMetadata] = field( + default_factory=dict + ) async def to_thread( self, func: Callable[P, T], /, *args: P.args, **kwargs: P.kwargs @@ -194,11 +211,24 @@ def stream_ordered_after(self, *dfs: DataFrame) -> Generator[Stream, None, None] @nvtx_annotate_cudf_polars(message="PrefetchParquetFootersForPaths") def _fetch_parquet_footers_for_paths( paths: tuple[str, ...], -) -> tuple[tuple[str, ...], list[plc.io.parquet_metadata.FileMetaData]]: - metadata = plc.io.parquet_metadata.read_parquet_footers( - plc.io.SourceInfo(list(paths)) - ) - return paths, metadata +) -> tuple[tuple[str, ...], ParquetFileMetadata]: + # https://github.com/rapidsai/cudf/issues/22734 will provide this metadata, allowing + # us to skip the HEAD request below. + import kvikio # TODO: is this a required dependency? + + filepath_sources = [] + for path in paths: + if path.startswith(("s3://", "https://", "http://")): + # this should be the one and only HEAD request we perform. + with kvikio.RemoteFile.open_s3_url(path) as remote_file: + size = remote_file.nbytes() + else: + size = None + filepath_sources.append(plc.io.types.FilepathSource(path, size)) # type: ignore[attr-defined] + + source_info = plc.io.types.SourceInfo(filepath_sources) + metadata = plc.io.parquet_metadata.read_parquet_footers(source_info) + return paths, ParquetFileMetadata(source_info, metadata) @nvtx_annotate_cudf_polars(message="PrefetchParquetFileMetadataForIR") @@ -228,11 +258,14 @@ def prefetch_parquet_file_metadata_for_ir( if not groups: return - missing_files = { - (path,) - for group in groups - for path in group - if (path,) not in context.parquet_file_metadata + # missing_files = { + # (path,) + # for group in groups + # for path in group + # if (path,) not in context.parquet_file_metadata + # } + missing_paths = { + paths for paths in groups if paths not in context.parquet_file_metadata } cm: contextlib.AbstractContextManager[concurrent.futures.Executor | None] @@ -244,27 +277,17 @@ def prefetch_parquet_file_metadata_for_ir( cm = contextlib.nullcontext() executor = context.py_executor - if len(missing_files) > 0: + if len(missing_paths) > 0: with cm: futures = [ - executor.submit(_fetch_parquet_footers_for_paths, missing_file) - for missing_file in missing_files + executor.submit(_fetch_parquet_footers_for_paths, paths) + for paths in missing_paths ] for future in concurrent.futures.as_completed(futures): paths, metadata = future.result() context.parquet_file_metadata.setdefault(paths, metadata) - for group in groups: - context.parquet_file_metadata.setdefault( - group, - list( - itertools.chain.from_iterable( - context.parquet_file_metadata[(path,)] for path in group - ) - ), - ) - _BINOPS = { plc.binaryop.BinaryOperator.EQUAL, @@ -885,17 +908,22 @@ def read_csv_header( elif typ == "parquet": if parquet_options.prefetch_file_metadata: try: - parquet_metadatas = context.parquet_file_metadata[tuple(paths)] + cached_metadata = context.parquet_file_metadata[tuple(paths)] + + source_info = cached_metadata.source_info + parquet_metadatas = cached_metadata.metadata except KeyError as e: raise AssertionError( f"Parquet file metadata was not prefetched for paths: {list(paths)}." ) from e else: + source_info = plc.io.types.SourceInfo(paths) parquet_metadatas = None filters = None if predicate is not None and row_index is None: # Can't apply filters during read if we have a row index. + # TODO: check if this does I/O. filters = to_parquet_filter( _prepare_parquet_predicate( predicate.value, paths, schema, with_columns @@ -903,7 +931,7 @@ def read_csv_header( stream=stream, ) parquet_reader_options = ( - plc.io.parquet.ParquetReaderOptions.builder(plc.io.SourceInfo(paths)) + plc.io.parquet.ParquetReaderOptions.builder(source_info) .decimal_width(plc.TypeId.DECIMAL128) .build() ) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index e93782435dd5..2353a1d0f08f 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -217,7 +217,7 @@ def do_evaluate( row_group_num_rows = [ num_rows - for metadata in parquet_metadatas + for metadata in parquet_metadatas.metadata for num_rows in metadata.row_group_num_rows ] From 1b28cbdbaf42ff7ac3009ac157615241e1597ac2 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 2 Jun 2026 10:58:16 -0700 Subject: [PATCH 03/51] Revert "Use FilepathSource" This reverts commit 109ad43f1fdc62510800bdc561178db007e08aa4. --- python/cudf_polars/cudf_polars/dsl/ir.py | 84 +++++++------------ .../cudf_polars/cudf_polars/streaming/io.py | 2 +- 2 files changed, 29 insertions(+), 57 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index da0bf74831f0..bd174e58310e 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -113,13 +113,6 @@ ] -@dataclass(frozen=True) -class ParquetFileMetadata: - # TODO: rename this; it's too close to plc.io.parquet_metadata.FileMetadata - source_info: plc.io.types.SourceInfo - metadata: list[plc.io.parquet_metadata.FileMetaData] - - @dataclass(frozen=True) class IRExecutionContext: """ @@ -137,24 +130,14 @@ class IRExecutionContext: A zero-argument callable that returns a CUDA stream. query_id Identifier for the query being executed. - parquet_file_metadata - Cache of parquet file metadata. The keys here are the ``paths`` of Scan nodes - with a ``parquet`` type. The values are ``ParquetFileMetadata`` objects. The - ``source_info`` is the ``SourceInfo`` object that was used to read the metadata, - which will contain the known file size for remote files. ``metadata`` is the list - of ``FileMetaData`` objects that were read. Both ``source_info`` and ``metadata`` - should be used in ``read_parquet`` calls later on. """ 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) - - # This should cache a tuple(size, FileMetadata) - # maybe a dataclass - parquet_file_metadata: dict[tuple[str, ...], ParquetFileMetadata] = field( - default_factory=dict - ) + parquet_file_metadata: dict[ + tuple[str, ...], list[plc.io.parquet_metadata.FileMetaData] + ] = field(default_factory=dict) async def to_thread( self, func: Callable[P, T], /, *args: P.args, **kwargs: P.kwargs @@ -211,24 +194,11 @@ def stream_ordered_after(self, *dfs: DataFrame) -> Generator[Stream, None, None] @nvtx_annotate_cudf_polars(message="PrefetchParquetFootersForPaths") def _fetch_parquet_footers_for_paths( paths: tuple[str, ...], -) -> tuple[tuple[str, ...], ParquetFileMetadata]: - # https://github.com/rapidsai/cudf/issues/22734 will provide this metadata, allowing - # us to skip the HEAD request below. - import kvikio # TODO: is this a required dependency? - - filepath_sources = [] - for path in paths: - if path.startswith(("s3://", "https://", "http://")): - # this should be the one and only HEAD request we perform. - with kvikio.RemoteFile.open_s3_url(path) as remote_file: - size = remote_file.nbytes() - else: - size = None - filepath_sources.append(plc.io.types.FilepathSource(path, size)) # type: ignore[attr-defined] - - source_info = plc.io.types.SourceInfo(filepath_sources) - metadata = plc.io.parquet_metadata.read_parquet_footers(source_info) - return paths, ParquetFileMetadata(source_info, metadata) +) -> tuple[tuple[str, ...], list[plc.io.parquet_metadata.FileMetaData]]: + metadata = plc.io.parquet_metadata.read_parquet_footers( + plc.io.SourceInfo(list(paths)) + ) + return paths, metadata @nvtx_annotate_cudf_polars(message="PrefetchParquetFileMetadataForIR") @@ -258,14 +228,11 @@ def prefetch_parquet_file_metadata_for_ir( if not groups: return - # missing_files = { - # (path,) - # for group in groups - # for path in group - # if (path,) not in context.parquet_file_metadata - # } - missing_paths = { - paths for paths in groups if paths not in context.parquet_file_metadata + missing_files = { + (path,) + for group in groups + for path in group + if (path,) not in context.parquet_file_metadata } cm: contextlib.AbstractContextManager[concurrent.futures.Executor | None] @@ -277,17 +244,27 @@ def prefetch_parquet_file_metadata_for_ir( cm = contextlib.nullcontext() executor = context.py_executor - if len(missing_paths) > 0: + if len(missing_files) > 0: with cm: futures = [ - executor.submit(_fetch_parquet_footers_for_paths, paths) - for paths in missing_paths + executor.submit(_fetch_parquet_footers_for_paths, missing_file) + for missing_file in missing_files ] for future in concurrent.futures.as_completed(futures): paths, metadata = future.result() context.parquet_file_metadata.setdefault(paths, metadata) + for group in groups: + context.parquet_file_metadata.setdefault( + group, + list( + itertools.chain.from_iterable( + context.parquet_file_metadata[(path,)] for path in group + ) + ), + ) + _BINOPS = { plc.binaryop.BinaryOperator.EQUAL, @@ -897,22 +874,17 @@ def read_csv_header( elif typ == "parquet": if parquet_options.prefetch_file_metadata: try: - cached_metadata = context.parquet_file_metadata[tuple(paths)] - - source_info = cached_metadata.source_info - parquet_metadatas = cached_metadata.metadata + parquet_metadatas = context.parquet_file_metadata[tuple(paths)] except KeyError as e: raise AssertionError( f"Parquet file metadata was not prefetched for paths: {list(paths)}." ) from e else: - source_info = plc.io.types.SourceInfo(paths) parquet_metadatas = None filters = None if predicate is not None and row_index is None: # Can't apply filters during read if we have a row index. - # TODO: check if this does I/O. filters = to_parquet_filter( _prepare_parquet_predicate( predicate.value, paths, schema, with_columns @@ -920,7 +892,7 @@ def read_csv_header( stream=stream, ) parquet_reader_options = ( - plc.io.parquet.ParquetReaderOptions.builder(source_info) + plc.io.parquet.ParquetReaderOptions.builder(plc.io.SourceInfo(paths)) .decimal_width(plc.TypeId.DECIMAL128) .build() ) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 2353a1d0f08f..e93782435dd5 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -217,7 +217,7 @@ def do_evaluate( row_group_num_rows = [ num_rows - for metadata in parquet_metadatas.metadata + for metadata in parquet_metadatas for num_rows in metadata.row_group_num_rows ] From 3bff8b1e9283eea43af589d5f655c13e4079318f Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 2 Jun 2026 11:18:28 -0700 Subject: [PATCH 04/51] fixes --- python/cudf_polars/cudf_polars/dsl/ir.py | 56 ++++++++++++++---------- 1 file changed, 32 insertions(+), 24 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index bd174e58310e..b3978f42d01e 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -52,6 +52,7 @@ log_do_evaluate, nvtx_annotate_cudf_polars, ) +from cudf_polars.dsl.traversal import traversal from cudf_polars.dsl.utils.reshape import broadcast from cudf_polars.dsl.utils.windows import ( offsets_to_windows, @@ -130,6 +131,10 @@ class IRExecutionContext: A zero-argument callable that returns a CUDA stream. query_id Identifier for the query being executed. + parquet_file_metadata + A cache of parquet file metadata. The keys are the ``paths`` of Scan nodes + with a ``parquet`` type. The values are a list of ``FileMetaData`` objects + associated with those ``paths``. """ py_executor: concurrent.futures.ThreadPoolExecutor | None = field(default=None) @@ -191,17 +196,36 @@ def stream_ordered_after(self, *dfs: DataFrame) -> Generator[Stream, None, None] yield result_stream -@nvtx_annotate_cudf_polars(message="PrefetchParquetFootersForPaths") -def _fetch_parquet_footers_for_paths( +@nvtx_annotate_cudf_polars(message="fetch_parquet_footers_for_paths") +def _prefetch_parquet_footers_for_paths( paths: tuple[str, ...], ) -> tuple[tuple[str, ...], list[plc.io.parquet_metadata.FileMetaData]]: + """ + Prefetch parquet footers for a list of paths. + + This is typically executed concurrently with prefetch operations for other + groups of ``paths`` for other ``Scan`` nodes. + + Parameters + ---------- + paths + The tuple of paths to prefetch. These correspond to ``paths`` in a ``Scan`` node. + + Returns + ------- + paths + The original input ``paths``. Useful for associating the result with the metadata + when executing out of order concurrently. + metadata + The list of ``FileMetaData`` objects for the ``paths``. + """ metadata = plc.io.parquet_metadata.read_parquet_footers( plc.io.SourceInfo(list(paths)) ) return paths, metadata -@nvtx_annotate_cudf_polars(message="PrefetchParquetFileMetadataForIR") +@nvtx_annotate_cudf_polars(message="prefetch_parquet_file_metadata_for_ir") def prefetch_parquet_file_metadata_for_ir( root: IR, context: IRExecutionContext, @@ -218,8 +242,6 @@ def prefetch_parquet_file_metadata_for_ir( metadata concurrently, its ``parquet_file_metadata`` is mutated to cache the newly read parquet metadata. """ - from cudf_polars.dsl.traversal import traversal - groups = { tuple(node.paths) for node in traversal([root]) @@ -228,13 +250,9 @@ def prefetch_parquet_file_metadata_for_ir( if not groups: return - missing_files = { - (path,) - for group in groups - for path in group - if (path,) not in context.parquet_file_metadata + missing_paths = { + paths for paths in groups if paths not in context.parquet_file_metadata } - cm: contextlib.AbstractContextManager[concurrent.futures.Executor | None] if context.py_executor is None: @@ -244,27 +262,17 @@ def prefetch_parquet_file_metadata_for_ir( cm = contextlib.nullcontext() executor = context.py_executor - if len(missing_files) > 0: + if missing_paths: with cm: futures = [ - executor.submit(_fetch_parquet_footers_for_paths, missing_file) - for missing_file in missing_files + executor.submit(_prefetch_parquet_footers_for_paths, paths) + for paths in missing_paths ] for future in concurrent.futures.as_completed(futures): paths, metadata = future.result() context.parquet_file_metadata.setdefault(paths, metadata) - for group in groups: - context.parquet_file_metadata.setdefault( - group, - list( - itertools.chain.from_iterable( - context.parquet_file_metadata[(path,)] for path in group - ) - ), - ) - _BINOPS = { plc.binaryop.BinaryOperator.EQUAL, From f3fa17f499276d31ce441daa1d9091758c836d64 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Thu, 4 Jun 2026 10:31:05 -0700 Subject: [PATCH 05/51] use streamingscan --- python/cudf_polars/cudf_polars/dsl/ir.py | 18 +++++++++++++----- python/cudf_polars/cudf_polars/engine/core.py | 5 +++-- python/cudf_polars/cudf_polars/streaming/io.py | 2 +- 3 files changed, 17 insertions(+), 8 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 5bd196c7ae84..0fa03f23970d 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -242,11 +242,19 @@ def prefetch_parquet_file_metadata_for_ir( metadata concurrently, its ``parquet_file_metadata`` is mutated to cache the newly read parquet metadata. """ - groups = { - tuple(node.paths) - for node in traversal([root]) - if isinstance(node, Scan) and node.typ == "parquet" and node.paths - } + from cudf_polars.streaming.io import SplitScan, StreamingScan + + groups = set() + for node in traversal([root]): + if isinstance(node, StreamingScan): + for child in node.children: + if isinstance(child, Scan) and child.typ == "parquet": + groups.add(tuple(child.paths)) + elif isinstance(node, Scan) and node.typ == "parquet": + groups.add(tuple(node.paths)) + elif isinstance(node, SplitScan) and node.base_scan.typ == "parquet": + groups.add(tuple(node.base_scan.paths)) + if not groups: return diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 8793f4466bd3..7239e472041c 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -678,7 +678,9 @@ def evaluate_on_rank( Collected channel metadata. """ stats = allgather_stats(comm, ctx.br(), ir, config_options, py_executor) - ir, partition_info = lower_ir_graph(ir, config_options, stats) + ir, partition_info = lower_ir_graph( + ir, config_options, stats, rank=comm.rank, nranks=comm.nranks + ) if comm.rank == 0: # At least for now, the query plan is identical on all ranks, @@ -688,7 +690,6 @@ def evaluate_on_rank( ir_context = IRExecutionContext( py_executor, get_cuda_stream=ctx.get_stream_from_pool, query_id=query_id ) - if config_options.parquet_options.prefetch_file_metadata: prefetch_parquet_file_metadata_for_ir( ir, diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 4c4c9b132226..d9be4fc38aa6 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -441,7 +441,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 ( From 7f19690cf0915da595f1ad042fa734ed68478515 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Thu, 4 Jun 2026 10:40:35 -0700 Subject: [PATCH 06/51] use streamingscan --- python/cudf_polars/cudf_polars/dsl/ir.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 0fa03f23970d..beba45df1e48 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -250,6 +250,8 @@ def prefetch_parquet_file_metadata_for_ir( for child in node.children: if isinstance(child, Scan) and child.typ == "parquet": groups.add(tuple(child.paths)) + elif isinstance(child, SplitScan) and child.base_scan.typ == "parquet": + groups.add(tuple(child.base_scan.paths)) elif isinstance(node, Scan) and node.typ == "parquet": groups.add(tuple(node.paths)) elif isinstance(node, SplitScan) and node.base_scan.typ == "parquet": From cf23d6155774d69de2ca7082f96a9d4627e25548 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Thu, 4 Jun 2026 17:41:26 -0700 Subject: [PATCH 07/51] Fix the test --- python/cudf_polars/cudf_polars/dsl/ir.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index beba45df1e48..a419a240f322 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -247,11 +247,11 @@ def prefetch_parquet_file_metadata_for_ir( groups = set() for node in traversal([root]): if isinstance(node, StreamingScan): - for child in node.children: - if isinstance(child, Scan) and child.typ == "parquet": - groups.add(tuple(child.paths)) - elif isinstance(child, SplitScan) and child.base_scan.typ == "parquet": - groups.add(tuple(child.base_scan.paths)) + for scan in node.scans: + if isinstance(scan, Scan) and scan.typ == "parquet": + groups.add(tuple(scan.paths)) + elif isinstance(scan, SplitScan) and scan.base_scan.typ == "parquet": + groups.add(tuple(scan.base_scan.paths)) elif isinstance(node, Scan) and node.typ == "parquet": groups.add(tuple(node.paths)) elif isinstance(node, SplitScan) and node.base_scan.typ == "parquet": From 0533d4fe79eb80b88724476ab4072832433ed356 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Thu, 4 Jun 2026 17:59:10 -0700 Subject: [PATCH 08/51] Review --- python/cudf_polars/cudf_polars/dsl/ir.py | 11 +++++++---- python/cudf_polars/cudf_polars/streaming/io.py | 8 +++++--- 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index a419a240f322..4c5d4bf59dae 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -267,10 +267,10 @@ def prefetch_parquet_file_metadata_for_ir( if context.py_executor is None: cm = executor = concurrent.futures.ThreadPoolExecutor() - # We didn't create the executor, so we don't close it. else: - cm = contextlib.nullcontext() executor = context.py_executor + # We didn't create the executor, so we don't close it. + cm = contextlib.nullcontext() if missing_paths: with cm: @@ -892,9 +892,12 @@ def read_csv_header( try: parquet_metadatas = context.parquet_file_metadata[tuple(paths)] except KeyError as e: - raise AssertionError( + msg = ( f"Parquet file metadata was not prefetched for paths: {list(paths)}." - ) from e + "Please report this as a bug to cudf-polars. You can work around it " + "by setting 'CUDF_POLARS__PARQUET_OPTIONS__PREFETCH_FILE_METADATA=0'." + ) + raise AssertionError(msg) from e else: parquet_metadatas = None diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index d9be4fc38aa6..34bcf34bac01 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -306,14 +306,16 @@ def do_evaluate( # - We can use all this information to calculate the # "skip_rows" and "n_rows" options to use locally. - # row_group_num_rows: list[int] if parquet_options.prefetch_file_metadata: try: parquet_metadatas = context.parquet_file_metadata[tuple(paths)] except KeyError as e: - raise AssertionError( + msg = ( f"Parquet file metadata was not prefetched for paths: {list(paths)}." - ) from e + "Please report this as a bug to cudf-polars. You can work around it " + "by setting 'CUDF_POLARS__PARQUET_OPTIONS__PREFETCH_FILE_METADATA=0'." + ) + raise AssertionError(msg) from e row_group_num_rows = [ num_rows From ac0d7995da85a8b1f012b2ed6f95ba52839bfc40 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 5 Jun 2026 04:34:34 -0700 Subject: [PATCH 09/51] Tests --- python/cudf_polars/cudf_polars/dsl/ir.py | 37 ++++++++++++--- python/cudf_polars/tests/test_select.py | 60 ++++++++++++++++++++++++ 2 files changed, 91 insertions(+), 6 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 4c5d4bf59dae..b789887aaa21 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -451,6 +451,7 @@ def __init__(self, schema: Schema, options: Any, predicate: expr.NamedExpr | Non def _parquet_physical_types( paths: list[str], columns: list[str] | None ) -> dict[str, plc.DataType]: + # This may not be able use prefetched metadata, since we don't (currently) have a Schema. metadata = plc.io.parquet_metadata.read_parquet_metadata(plc.io.SourceInfo(paths)) column_types = metadata.schema().column_types() @@ -747,12 +748,32 @@ 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, + context: IRExecutionContext, ) -> 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: + try: + parquet_metadatas = context.parquet_file_metadata[tuple(paths)] + except KeyError as e: + msg = ( + f"Parquet file metadata was not prefetched for paths: {list(paths)}." + "Please report this as a bug to cudf-polars. You can work around it " + "by setting 'CUDF_POLARS__PARQUET_OPTIONS__PREFETCH_FILE_METADATA=0'." + ) + raise AssertionError(msg) from e + num_rows = sum(metadata.num_rows for metadata in parquet_metadatas) + 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) @@ -945,7 +966,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, context + ) if not names else None ) @@ -969,7 +992,9 @@ def read_csv_header( # 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, context + ) if not col_names else None ) @@ -1647,7 +1672,7 @@ 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, context ) dtype = DataType(pl.UInt32()) col = Column( diff --git a/python/cudf_polars/tests/test_select.py b/python/cudf_polars/tests/test_select.py index f37c2d195d92..7cc2223ed268 100644 --- a/python/cudf_polars/tests/test_select.py +++ b/python/cudf_polars/tests/test_select.py @@ -147,3 +147,63 @@ 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 +def prefetch_engine() -> pl.GPUEngine: + return pl.GPUEngine( + executor="in-memory", + raise_on_fail=True, + parquet_options={"prefetch_file_metadata": True}, + ) + + +@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_select_fast_count_parquet_prefetch_metadata( + tmp_path, + parquet_fast_count_df: pl.DataFrame, + prefetch_engine: pl.GPUEngine, + parquet_scan_row_bounds: dict[str, int | None], +) -> None: + skip_rows = parquet_scan_row_bounds["skip_rows"] + assert skip_rows is not None + n_rows = parquet_scan_row_bounds["n_rows"] + + file = tmp_path / "data.parquet" + parquet_fast_count_df.write_parquet(file) + + if skip_rows == 0 and n_rows is None: + q = pl.scan_parquet(file) + elif skip_rows == 0: + q = pl.scan_parquet(file, n_rows=n_rows) + elif n_rows is None: + q = pl.scan_parquet(file).slice(skip_rows) + else: + q = pl.scan_parquet(file).slice(skip_rows, n_rows) + + q = q.select(pl.len()) + assert_gpu_result_equal(q, engine=prefetch_engine) From 5d10839cd620919542ce3cdc12552ee428defbbd Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 5 Jun 2026 05:53:45 -0700 Subject: [PATCH 10/51] tests --- python/cudf_polars/cudf_polars/dsl/ir.py | 4 ++-- .../cudf_polars/streaming/select.py | 9 +++++++- python/cudf_polars/tests/test_select.py | 23 +++++++++++++++++++ 3 files changed, 33 insertions(+), 3 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index b789887aaa21..72a1874b6946 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -752,11 +752,11 @@ def _get_parquet_row_count_from_metadata( skip_rows: int, n_rows: int, parquet_options: ParquetOptions, - context: IRExecutionContext, + context: IRExecutionContext | None, ) -> int: # Zero-width parquet files lose their row count when read through # pylibcudf. See https://github.com/rapidsai/cudf/issues/21428 - if parquet_options.prefetch_file_metadata: + if parquet_options.prefetch_file_metadata and context is not None: try: parquet_metadatas = context.parquet_file_metadata[tuple(paths)] except KeyError as e: diff --git a/python/cudf_polars/cudf_polars/streaming/select.py b/python/cudf_polars/cudf_polars/streaming/select.py index e1466f0b6610..81bfec030bc9 100644 --- a/python/cudf_polars/cudf_polars/streaming/select.py +++ b/python/cudf_polars/cudf_polars/streaming/select.py @@ -431,8 +431,15 @@ 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, + scan_child.parquet_options, + context=None, ) dtype = ir.exprs[0].value.dtype diff --git a/python/cudf_polars/tests/test_select.py b/python/cudf_polars/tests/test_select.py index 7cc2223ed268..b5ab87d26042 100644 --- a/python/cudf_polars/tests/test_select.py +++ b/python/cudf_polars/tests/test_select.py @@ -8,10 +8,12 @@ import polars as pl +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.utils.config import ParquetOptions def test_select(engine: pl.GPUEngine): @@ -207,3 +209,24 @@ def test_select_fast_count_parquet_prefetch_metadata( q = q.select(pl.len()) assert_gpu_result_equal(q, engine=prefetch_engine) + + +def test_get_parquet_row_count_from_metadata_missing_prefetch() -> None: + paths = ["/some/missing/file.parquet"] + parquet_options = ParquetOptions(prefetch_file_metadata=True) + context = IRExecutionContext() + + with pytest.raises( + AssertionError, + match=( + r"Parquet file metadata was not prefetched for paths: " + r"\['/some/missing/file\.parquet'\]\." + ), + ): + Scan._get_parquet_row_count_from_metadata( + paths, + skip_rows=0, + n_rows=-1, + parquet_options=parquet_options, + context=context, + ) From 2decc397d7a157189e07ab7265ed34f60d920c96 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 5 Jun 2026 07:26:31 -0700 Subject: [PATCH 11/51] test coverage --- .../cudf_polars/tests/streaming/test_scan.py | 161 +++++++++++++++++- python/cudf_polars/tests/test_scan.py | 32 ++++ 2 files changed, 192 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 1c0e9f7e31d3..1bc51e3d9209 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -11,7 +11,13 @@ 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, + prefetch_parquet_file_metadata_for_ir, +) +from cudf_polars.dsl.traversal import traversal from cudf_polars.engine.options import StreamingOptions from cudf_polars.streaming.base import IOPartitionFlavor, IOPartitionPlan from cudf_polars.streaming.io import SplitScan, StreamingScan, expand_scan_for_rank @@ -74,6 +80,128 @@ def test_scan_parquet_prefetch_file_metadata(tmp_path, df, streaming_engine_fact assert_gpu_result_equal(pl.scan_parquet(tmp_path), engine=streaming_engine) +def test_prefetch_parquet_file_metadata_streaming_scan_children( + tmp_path, + df, + parquet_stats_executor: concurrent.futures.ThreadPoolExecutor, +) -> None: + # Small files fused into Scan children (not SplitScan) inside StreamingScan. + make_partitioned_source(df, tmp_path, "parquet", n_files=3) + engine = pl.GPUEngine( + raise_on_fail=True, + executor="streaming", + executor_options={"target_partition_size": 1_000_000}, + parquet_options={"prefetch_file_metadata": True}, + ) + q = pl.scan_parquet(tmp_path) + qir = Translator(q._ldf.visit(), engine).translate_ir() + config_options = ConfigOptions.from_polars_engine(engine) + stats = collect_statistics(qir, config_options, parquet_stats_executor) + ir, partition_info = lower_ir_graph( + qir, + config_options, + stats, + rank=0, + nranks=1, + ) + + streaming_scan = next( + node for node in traversal([ir]) if isinstance(node, StreamingScan) + ) + parquet_scans = [ + scan + for scan in streaming_scan.scans + if isinstance(scan, Scan) and scan.typ == "parquet" + ] + assert parquet_scans + io_plan = partition_info[streaming_scan].io_plan + assert io_plan is not None + assert io_plan.flavor != IOPartitionFlavor.SPLIT_FILES + + context = IRExecutionContext() + prefetch_parquet_file_metadata_for_ir(ir, context) + + for scan in parquet_scans: + assert tuple(scan.paths) in context.parquet_file_metadata + assert context.parquet_file_metadata[tuple(scan.paths)] + + +def test_prefetch_parquet_file_metadata_streaming_split_scan_children( + tmp_path, + df, + parquet_stats_executor: concurrent.futures.ThreadPoolExecutor, +) -> None: + # Large file split into SplitScan children inside StreamingScan. + make_partitioned_source(df, tmp_path, "parquet", n_files=1) + engine = pl.GPUEngine( + raise_on_fail=True, + executor="streaming", + executor_options={"target_partition_size": 1_000}, + parquet_options={"prefetch_file_metadata": True}, + ) + q = pl.scan_parquet(tmp_path) + qir = Translator(q._ldf.visit(), engine).translate_ir() + config_options = ConfigOptions.from_polars_engine(engine) + stats = collect_statistics(qir, config_options, parquet_stats_executor) + ir, partition_info = lower_ir_graph( + qir, + config_options, + stats, + rank=0, + nranks=1, + ) + + streaming_scan = next( + node for node in traversal([ir]) if isinstance(node, StreamingScan) + ) + split_scans = [ + scan + for scan in streaming_scan.scans + if isinstance(scan, SplitScan) and scan.base_scan.typ == "parquet" + ] + assert split_scans + io_plan = partition_info[streaming_scan].io_plan + assert io_plan is not None + assert io_plan.flavor == IOPartitionFlavor.SPLIT_FILES + + context = IRExecutionContext() + prefetch_parquet_file_metadata_for_ir(ir, context) + + for scan in split_scans: + paths = tuple(scan.base_scan.paths) + assert paths in context.parquet_file_metadata + assert context.parquet_file_metadata[paths] + + +def test_prefetch_parquet_file_metadata_split_scan_root(tmp_path, df) -> None: + # SplitScan is stored in StreamingScan.scans, not IR children, so prefetch + # must also handle a SplitScan node passed as the traversal root. + make_partitioned_source(df, tmp_path, "parquet", n_files=1) + path = str(next(tmp_path.glob("*.parquet"))) + split_scans = expand_scan_for_rank( + _make_parquet_scan([path]), + IOPartitionPlan(4, IOPartitionFlavor.SPLIT_FILES), + rank=0, + nranks=1, + parquet_options=ParquetOptions(), + ) + split_scan = split_scans[0] + assert isinstance(split_scan, SplitScan) + + context = IRExecutionContext() + prefetch_parquet_file_metadata_for_ir(split_scan, context) + + paths = tuple(split_scan.base_scan.paths) + assert paths in context.parquet_file_metadata + assert context.parquet_file_metadata[paths] + + +def test_prefetch_parquet_file_metadata_no_parquet_scans() -> None: + context = IRExecutionContext() + prefetch_parquet_file_metadata_for_ir(Empty({}), context) + assert context.parquet_file_metadata == {} + + # --------------------------------------------------------------------------- # Tests migrated from tests/streaming/test_scan.py # --------------------------------------------------------------------------- @@ -265,3 +393,34 @@ def test_streaming_scan_raises() -> None: ctx = IRExecutionContext() with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): StreamingScan.do_evaluate([scan], 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"Parquet file metadata was not prefetched for paths: " + r"\['/some/missing/file\.parquet'\]\." + ), + ): + SplitScan.do_evaluate( + 0, + 4, + schema, + "parquet", + {}, + paths, + None, + 0, + -1, + None, + None, + None, + parquet_options, + context=context, + ) diff --git a/python/cudf_polars/tests/test_scan.py b/python/cudf_polars/tests/test_scan.py index eceb440f34eb..6fc63aed9de9 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 ParquetOptions from cudf_polars.utils.versions import ( POLARS_VERSION_LT_138, POLARS_VERSION_LT_139, @@ -185,6 +188,35 @@ def test_scan_parquet_prefetch_file_metadata( assert_gpu_result_equal(q, engine=engine) +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"Parquet file metadata was not prefetched for paths: " + r"\['/some/missing/file\.parquet'\]\." + ), + ): + 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]}) From 46bf9e2a2aa627cb9382a4d31a5c594c647d48bf Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 5 Jun 2026 07:50:30 -0700 Subject: [PATCH 12/51] Test revamp --- .../cudf_polars/tests/streaming/test_scan.py | 180 ++++++------------ 1 file changed, 58 insertions(+), 122 deletions(-) diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 1bc51e3d9209..0a58543d06a6 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -11,13 +11,7 @@ from cudf_polars import Translator from cudf_polars.containers import DataType -from cudf_polars.dsl.ir import ( - Empty, - IRExecutionContext, - Scan, - prefetch_parquet_file_metadata_for_ir, -) -from cudf_polars.dsl.traversal import traversal +from cudf_polars.dsl.ir import IRExecutionContext, Scan from cudf_polars.engine.options import StreamingOptions from cudf_polars.streaming.base import IOPartitionFlavor, IOPartitionPlan from cudf_polars.streaming.io import SplitScan, StreamingScan, expand_scan_for_rank @@ -80,126 +74,37 @@ def test_scan_parquet_prefetch_file_metadata(tmp_path, df, streaming_engine_fact assert_gpu_result_equal(pl.scan_parquet(tmp_path), engine=streaming_engine) -def test_prefetch_parquet_file_metadata_streaming_scan_children( - tmp_path, - df, - parquet_stats_executor: concurrent.futures.ThreadPoolExecutor, +def test_scan_parquet_prefetch_file_metadata_fused_files( + tmp_path, df, streaming_engine_factory ) -> None: - # Small files fused into Scan children (not SplitScan) inside StreamingScan. - make_partitioned_source(df, tmp_path, "parquet", n_files=3) - engine = pl.GPUEngine( - raise_on_fail=True, - executor="streaming", - executor_options={"target_partition_size": 1_000_000}, - parquet_options={"prefetch_file_metadata": True}, - ) - q = pl.scan_parquet(tmp_path) - qir = Translator(q._ldf.visit(), engine).translate_ir() - config_options = ConfigOptions.from_polars_engine(engine) - stats = collect_statistics(qir, config_options, parquet_stats_executor) - ir, partition_info = lower_ir_graph( - qir, - config_options, - stats, - rank=0, - nranks=1, - ) - - streaming_scan = next( - node for node in traversal([ir]) if isinstance(node, StreamingScan) + streaming_engine = streaming_engine_factory( + StreamingOptions( + target_partition_size=1_000_000, + parquet_options={"prefetch_file_metadata": True}, + ), ) - parquet_scans = [ - scan - for scan in streaming_scan.scans - if isinstance(scan, Scan) and scan.typ == "parquet" - ] - assert parquet_scans - io_plan = partition_info[streaming_scan].io_plan - assert io_plan is not None - assert io_plan.flavor != IOPartitionFlavor.SPLIT_FILES - - context = IRExecutionContext() - prefetch_parquet_file_metadata_for_ir(ir, context) - - for scan in parquet_scans: - assert tuple(scan.paths) in context.parquet_file_metadata - assert context.parquet_file_metadata[tuple(scan.paths)] + make_partitioned_source(df, tmp_path, "parquet", n_files=3) + assert_gpu_result_equal(pl.scan_parquet(tmp_path), engine=streaming_engine) -def test_prefetch_parquet_file_metadata_streaming_split_scan_children( - tmp_path, - df, - parquet_stats_executor: concurrent.futures.ThreadPoolExecutor, +def test_scan_parquet_prefetch_file_metadata_split_files( + tmp_path, df, streaming_engine_factory ) -> None: - # Large file split into SplitScan children inside StreamingScan. - make_partitioned_source(df, tmp_path, "parquet", n_files=1) - engine = pl.GPUEngine( - raise_on_fail=True, - executor="streaming", - executor_options={"target_partition_size": 1_000}, - parquet_options={"prefetch_file_metadata": True}, - ) - q = pl.scan_parquet(tmp_path) - qir = Translator(q._ldf.visit(), engine).translate_ir() - config_options = ConfigOptions.from_polars_engine(engine) - stats = collect_statistics(qir, config_options, parquet_stats_executor) - ir, partition_info = lower_ir_graph( - qir, - config_options, - stats, - rank=0, - nranks=1, - ) - - streaming_scan = next( - node for node in traversal([ir]) if isinstance(node, StreamingScan) + streaming_engine = streaming_engine_factory( + StreamingOptions( + target_partition_size=1_000, + parquet_options={"prefetch_file_metadata": True}, + ), ) - split_scans = [ - scan - for scan in streaming_scan.scans - if isinstance(scan, SplitScan) and scan.base_scan.typ == "parquet" - ] - assert split_scans - io_plan = partition_info[streaming_scan].io_plan - assert io_plan is not None - assert io_plan.flavor == IOPartitionFlavor.SPLIT_FILES - - context = IRExecutionContext() - prefetch_parquet_file_metadata_for_ir(ir, context) - - for scan in split_scans: - paths = tuple(scan.base_scan.paths) - assert paths in context.parquet_file_metadata - assert context.parquet_file_metadata[paths] - - -def test_prefetch_parquet_file_metadata_split_scan_root(tmp_path, df) -> None: - # SplitScan is stored in StreamingScan.scans, not IR children, so prefetch - # must also handle a SplitScan node passed as the traversal root. make_partitioned_source(df, tmp_path, "parquet", n_files=1) - path = str(next(tmp_path.glob("*.parquet"))) - split_scans = expand_scan_for_rank( - _make_parquet_scan([path]), - IOPartitionPlan(4, IOPartitionFlavor.SPLIT_FILES), - rank=0, - nranks=1, - parquet_options=ParquetOptions(), - ) - split_scan = split_scans[0] - assert isinstance(split_scan, SplitScan) - - context = IRExecutionContext() - prefetch_parquet_file_metadata_for_ir(split_scan, context) - - paths = tuple(split_scan.base_scan.paths) - assert paths in context.parquet_file_metadata - assert context.parquet_file_metadata[paths] + assert_gpu_result_equal(pl.scan_parquet(tmp_path), engine=streaming_engine) -def test_prefetch_parquet_file_metadata_no_parquet_scans() -> None: - context = IRExecutionContext() - prefetch_parquet_file_metadata_for_ir(Empty({}), context) - assert context.parquet_file_metadata == {} +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) # --------------------------------------------------------------------------- @@ -297,7 +202,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", @@ -310,7 +218,7 @@ def _make_parquet_scan(paths: list[str]) -> Scan: None, None, None, - ParquetOptions(), + parquet_options, ) @@ -387,9 +295,37 @@ def test_expand_scan_for_rank_split_files( assert scan.base_scan.paths == ["file.parquet"] -def test_streaming_scan_raises() -> None: +def test_scan_missing_prefetch_metadata_raises() -> None: + # This isn't reachable by normal cudf-polars usage. + scan = _make_parquet_scan( + ["file.parquet"], parquet_options=ParquetOptions(prefetch_file_metadata=True) + ) + ctx = IRExecutionContext() + with pytest.raises( + AssertionError, + match=r"Parquet file metadata was not prefetched for paths: \['file\.parquet'\]\.", + ): + 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 normal cudf-polars usage. - scan = _make_parquet_scan(["file.parquet"]) + scan = _make_parquet_scan( + ["file.parquet"], parquet_options=ParquetOptions(prefetch_file_metadata=True) + ) ctx = IRExecutionContext() with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): StreamingScan.do_evaluate([scan], scan, context=ctx) From 55a84d2ab3be0eddce28f72f16cddc92abd15922 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 5 Jun 2026 08:17:54 -0700 Subject: [PATCH 13/51] remove unused code --- python/cudf_polars/cudf_polars/dsl/ir.py | 2 -- python/cudf_polars/tests/streaming/test_scan.py | 13 ++++++++++++- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 72a1874b6946..06058867eaa7 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -254,8 +254,6 @@ def prefetch_parquet_file_metadata_for_ir( groups.add(tuple(scan.base_scan.paths)) elif isinstance(node, Scan) and node.typ == "parquet": groups.add(tuple(node.paths)) - elif isinstance(node, SplitScan) and node.base_scan.typ == "parquet": - groups.add(tuple(node.base_scan.paths)) if not groups: return diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 0a58543d06a6..675003e634cf 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -11,7 +11,12 @@ 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, + prefetch_parquet_file_metadata_for_ir, +) from cudf_polars.engine.options import StreamingOptions from cudf_polars.streaming.base import IOPartitionFlavor, IOPartitionPlan from cudf_polars.streaming.io import SplitScan, StreamingScan, expand_scan_for_rank @@ -107,6 +112,12 @@ def test_prefetch_file_metadata_non_parquet_scan(df, streaming_engine_factory) - assert_gpu_result_equal(df.lazy().select("x"), engine=streaming_engine) +def test_prefetch_parquet_file_metadata_no_parquet_scans() -> None: + context = IRExecutionContext() + prefetch_parquet_file_metadata_for_ir(Empty({}), context) + assert context.parquet_file_metadata == {} + + # --------------------------------------------------------------------------- # Tests migrated from tests/streaming/test_scan.py # --------------------------------------------------------------------------- From d000c0143a4cdd425598298cfaa15ccaff05cacc Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 5 Jun 2026 08:35:39 -0700 Subject: [PATCH 14/51] conditional prefetch --- python/cudf_polars/cudf_polars/callback.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/python/cudf_polars/cudf_polars/callback.py b/python/cudf_polars/cudf_polars/callback.py index a91e63b1423e..533eb8245623 100644 --- a/python/cudf_polars/cudf_polars/callback.py +++ b/python/cudf_polars/cudf_polars/callback.py @@ -304,10 +304,11 @@ def _callback( ): if config_options.executor.name == "in-memory": context = IRExecutionContext() - prefetch_parquet_file_metadata_for_ir( - ir, - context, - ) + if config_options.parquet_options.prefetch_file_metadata: + prefetch_parquet_file_metadata_for_ir( + ir, + context, + ) df = ir.evaluate(cache={}, timer=timer, context=context).to_polars() if timer is None: return df From c3b4f76f6a23c964e4003db8f46ac1debcbbe04d Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 9 Jun 2026 10:37:29 -0700 Subject: [PATCH 15/51] review --- python/cudf_polars/cudf_polars/dsl/ir.py | 3 +++ python/cudf_polars/tests/streaming/test_scan.py | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 06058867eaa7..b2d09e62a01d 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -135,6 +135,9 @@ class IRExecutionContext: A cache of parquet file metadata. The keys are the ``paths`` of Scan nodes with a ``parquet`` type. The values are a list of ``FileMetaData`` objects associated with those ``paths``. + + This cache lasts for the duration of the a single query's execution + (e.g. ``LazyFrame.collect()``). """ py_executor: concurrent.futures.ThreadPoolExecutor | None = field(default=None) diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 675003e634cf..17523ad3fbb3 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -307,7 +307,7 @@ def test_expand_scan_for_rank_split_files( def test_scan_missing_prefetch_metadata_raises() -> None: - # This isn't reachable by normal cudf-polars usage. + # 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) ) From bc0bde7ca2ac757dbd7bf667344e85727b10d5af Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 9 Jun 2026 11:29:21 -0700 Subject: [PATCH 16/51] combine tests --- .../cudf_polars/tests/streaming/test_scan.py | 44 +++++++------------ 1 file changed, 15 insertions(+), 29 deletions(-) diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 17523ad3fbb3..d7116db75e70 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -28,8 +28,11 @@ if TYPE_CHECKING: import concurrent.futures + from collections.abc import Callable from pathlib import Path + from cudf_polars.engine.core import StreamingEngine + @pytest.fixture(scope="module") def df(): @@ -68,40 +71,23 @@ 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) -def test_scan_parquet_prefetch_file_metadata(tmp_path, df, streaming_engine_factory): - streaming_engine = streaming_engine_factory( - StreamingOptions( - target_partition_size=1_000, - parquet_options={"prefetch_file_metadata": True}, - ), - ) - make_partitioned_source(df, tmp_path, "parquet", n_files=2) - assert_gpu_result_equal(pl.scan_parquet(tmp_path), engine=streaming_engine) - - -def test_scan_parquet_prefetch_file_metadata_fused_files( - tmp_path, df, streaming_engine_factory -) -> None: - streaming_engine = streaming_engine_factory( - StreamingOptions( - target_partition_size=1_000_000, - parquet_options={"prefetch_file_metadata": True}, - ), - ) - make_partitioned_source(df, tmp_path, "parquet", n_files=3) - assert_gpu_result_equal(pl.scan_parquet(tmp_path), engine=streaming_engine) - - -def test_scan_parquet_prefetch_file_metadata_split_files( - tmp_path, df, streaming_engine_factory -) -> None: +@pytest.mark.parametrize( + "target_partition_size_and_n_files", [(1_000, 1), (1_000, 2), (1_000_000, 3)] +) +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=1_000, + target_partition_size=target_partition_size, parquet_options={"prefetch_file_metadata": True}, ), ) - make_partitioned_source(df, tmp_path, "parquet", n_files=1) + make_partitioned_source(df, tmp_path, "parquet", n_files=n_files) assert_gpu_result_equal(pl.scan_parquet(tmp_path), engine=streaming_engine) From f82f7a9b130427eb7a9fccec8feadad6ddadc509 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 12 Jun 2026 13:16:01 -0700 Subject: [PATCH 17/51] fusedscan --- python/cudf_polars/cudf_polars/dsl/ir.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index b2d09e62a01d..95ef565ef4ad 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -245,16 +245,17 @@ def prefetch_parquet_file_metadata_for_ir( metadata concurrently, its ``parquet_file_metadata`` is mutated to cache the newly read parquet metadata. """ - from cudf_polars.streaming.io import SplitScan, StreamingScan + from cudf_polars.streaming.io import FusedScan, SplitScan, StreamingScan groups = set() for node in traversal([root]): if isinstance(node, StreamingScan): for scan in node.scans: - if isinstance(scan, Scan) and scan.typ == "parquet": + if ( + isinstance(scan, (SplitScan, FusedScan)) + and scan.base_scan.typ == "parquet" + ): groups.add(tuple(scan.paths)) - elif isinstance(scan, SplitScan) and scan.base_scan.typ == "parquet": - groups.add(tuple(scan.base_scan.paths)) elif isinstance(node, Scan) and node.typ == "parquet": groups.add(tuple(node.paths)) From 2b234ddcb6de210ead02895535e82d2db6ae9598 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 15 Jun 2026 09:29:00 -0700 Subject: [PATCH 18/51] Reuse parquet metadata --- python/cudf_polars/cudf_polars/dsl/ir.py | 34 ++++++++ python/cudf_polars/cudf_polars/engine/core.py | 7 +- .../cudf_polars/cudf_polars/streaming/io.py | 62 ++++++++++++--- .../cudf_polars/tests/streaming/test_stats.py | 79 ++++++++++++++++++- 4 files changed, 168 insertions(+), 14 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 95ef565ef4ad..c4832db0b679 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -78,6 +78,7 @@ from rmm.pylibrmm.stream import Stream from cudf_polars.containers.dataframe import NamedColumn + from cudf_polars.streaming.base import StatsCollector from cudf_polars.typing import CSECache, ClosedInterval, Schema, Slice as Zlice from cudf_polars.utils.config import ParquetOptions from cudf_polars.utils.timer import Timer @@ -286,6 +287,39 @@ def prefetch_parquet_file_metadata_for_ir( context.parquet_file_metadata.setdefault(paths, metadata) +def seed_parquet_file_metadata_from_stats( + stats: StatsCollector, + context: IRExecutionContext, +) -> None: + """ + Seed prefetched parquet footers from stats collection when available. + + Stats collection reads parquet footers for sampled paths. When the sample + covers all paths in a scan, those footers can be reused during metadata + prefetch to avoid rereading file footers. + + Parameters + ---------- + stats: StatsCollector + The stats collector. + context: IRExecutionContext + The execution context. Its ``parquet_file_metadata`` is mutated to + cache the parquet file metadata read during stats collection. + """ + from cudf_polars.streaming.io import ParquetSourceInfo + + for node, info in stats.scan_stats.items(): + if ( + isinstance(node, Scan) + and node.typ == "parquet" + and isinstance(info, ParquetSourceInfo) + and info.file_metadata is not None + ): + context.parquet_file_metadata.setdefault( + tuple(node.paths), info.file_metadata + ) + + _BINOPS = { plc.binaryop.BinaryOperator.EQUAL, plc.binaryop.BinaryOperator.NOT_EQUAL, diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 5d89b09cf097..67f85216b526 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -25,7 +25,11 @@ from rapidsmpf.streaming.core.actor import run_actor_network from cudf_polars.containers import DataFrame -from cudf_polars.dsl.ir import IRExecutionContext, prefetch_parquet_file_metadata_for_ir +from cudf_polars.dsl.ir import ( + IRExecutionContext, + prefetch_parquet_file_metadata_for_ir, + seed_parquet_file_metadata_from_stats, +) 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 @@ -692,6 +696,7 @@ def evaluate_on_rank( py_executor, get_cuda_stream=ctx.get_stream_from_pool, query_id=query_id ) if config_options.parquet_options.prefetch_file_metadata: + seed_parquet_file_metadata_from_stats(stats, ir_context) prefetch_parquet_file_metadata_for_ir( ir, ir_context, diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 8893b9006704..bc30313e3f23 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -792,6 +792,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. @@ -806,12 +820,15 @@ class ParquetMetadata: __slots__ = ( "column_names", + "file_metadata", "max_footer_samples", "mean_size_per_file", "num_row_groups_per_file", "paths", "row_count", "sample_paths", + "sampled_file_count", + "total_file_count", ) paths: tuple[str, ...] @@ -828,6 +845,8 @@ class ParquetMetadata: """All column names found it the dataset.""" sample_paths: tuple[str, ...] """Sampled file paths.""" + file_metadata: tuple[plc.io.parquet_metadata.FileMetaData, ...] | None + """Parquet footers read for ``sample_paths``. Populated only if all files were sampled.""" @nvtx_annotate_cudf_polars(message="ParquetMetadata") def __init__(self, paths: tuple[str, ...], max_footer_samples: int): @@ -837,9 +856,9 @@ def __init__(self, paths: tuple[str, ...], max_footer_samples: int): 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.file_metadata = None + max_footer_samples = max(1, max_footer_samples) + stride = max(1, int(len(paths) / max_footer_samples)) self.sample_paths = paths[: stride * max_footer_samples : stride] if not self.sample_paths: @@ -849,19 +868,21 @@ def __init__(self, paths: tuple[str, ...], max_footer_samples: int): total_file_count = len(self.paths) sampled_file_count = len(self.sample_paths) - sample_metadata = plc.io.parquet_metadata.read_parquet_metadata( + sample_footers = plc.io.parquet_metadata.read_parquet_footers( plc.io.SourceInfo(list(self.sample_paths)) ) + self.file_metadata = tuple(sample_footers) + sampled_row_count = sum(fmd.num_rows for fmd in sample_footers) if total_file_count == sampled_file_count: - row_count = sample_metadata.num_rows() + row_count = sampled_row_count else: - num_rows_per_sampled_file = int( - sample_metadata.num_rows() / sampled_file_count - ) + num_rows_per_sampled_file = int(sampled_row_count / sampled_file_count) row_count = num_rows_per_sampled_file * 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) ) @@ -871,7 +892,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) @@ -881,6 +904,8 @@ 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.total_file_count = total_file_count + self.sampled_file_count = sampled_file_count @nvtx_annotate_cudf_polars(message="_sample_rg_sizes") @@ -934,13 +959,18 @@ 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, + *, + file_metadata: list[plc.io.parquet_metadata.FileMetaData] | None = None, ): if per_file_means is None: per_file_means = {} self.row_count = row_count self.per_file_means = per_file_means + self.file_metadata = file_metadata @classmethod def from_paths( @@ -992,7 +1022,15 @@ def from_paths( for col in suspicious: per_file_means[col] = min_floor - return cls(row_count, per_file_means) + file_metadata: list[plc.io.parquet_metadata.FileMetaData] | None + if ( + metadata.sampled_file_count == metadata.total_file_count + and metadata.file_metadata is not None + ): + file_metadata = list(metadata.file_metadata) + else: + file_metadata = None + return cls(row_count, per_file_means, file_metadata=file_metadata) 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/tests/streaming/test_stats.py b/python/cudf_polars/tests/streaming/test_stats.py index 36bb7e3fc266..6a7f0aef66c3 100644 --- a/python/cudf_polars/tests/streaming/test_stats.py +++ b/python/cudf_polars/tests/streaming/test_stats.py @@ -13,12 +13,19 @@ 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, + IRExecutionContext, + Projection, + seed_parquet_file_metadata_from_stats, +) 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 @@ -159,6 +166,76 @@ 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: pl.DataFrame, +) -> None: + _clear_source_info_cache() + 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), max_footer_samples=10, max_row_group_samples=0 + ) + + assert info.file_metadata is not None + assert len(info.file_metadata) == len(paths) + assert sum(fmd.num_rows for fmd in info.file_metadata) == df.height + + +def test_parquet_source_info_omits_footers_when_paths_are_sampled( + tmp_path: pathlib.Path, + df: pl.DataFrame, +) -> None: + _clear_source_info_cache() + 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), max_footer_samples=2, max_row_group_samples=0 + ) + + assert info.file_metadata is None + + +def test_seed_parquet_file_metadata_from_stats( + tmp_path: pathlib.Path, + df: pl.DataFrame, + parquet_stats_executor: concurrent.futures.ThreadPoolExecutor, +) -> None: + _clear_source_info_cache() + make_partitioned_source(df, tmp_path, "parquet", n_files=2) + engine = pl.GPUEngine( + raise_on_fail=True, + executor="streaming", + parquet_options={"max_footer_samples": 10, "prefetch_file_metadata": True}, + ) + q = pl.scan_parquet(tmp_path) + ir = Translator(q._ldf.visit(), engine).translate_ir() + config = ConfigOptions.from_polars_engine(engine) + stats = collect_statistics(ir, config, parquet_stats_executor) + + context = IRExecutionContext() + seed_parquet_file_metadata_from_stats(stats, context) + + scan_node = next(node for node in stats.scan_stats if hasattr(node, "paths")) + assert ( + context.parquet_file_metadata[tuple(scan_node.paths)] + is stats.scan_stats[scan_node].file_metadata + ) + + +def test_parquet_metadata_reads_footers( + tmp_path: pathlib.Path, + df: pl.DataFrame, +) -> None: + 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.file_metadata is not None + assert len(metadata.file_metadata) == 1 + assert metadata.row_count == df.height + + def test_dataframe_round_trip() -> None: info = DataFrameSourceInfo(2500) data = info.serialize() From 759fb377e5aec3157aca7e4e55fd5a0e75daf0c7 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 15 Jun 2026 10:15:56 -0700 Subject: [PATCH 19/51] disable with 0 --- python/cudf_polars/cudf_polars/streaming/io.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index bc30313e3f23..2ef17ca2d16c 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -857,7 +857,12 @@ def __init__(self, paths: tuple[str, ...], max_footer_samples: int): self.mean_size_per_file = {} self.column_names = () self.file_metadata = None - max_footer_samples = max(1, max_footer_samples) + 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] @@ -866,7 +871,6 @@ 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_footers = plc.io.parquet_metadata.read_parquet_footers( plc.io.SourceInfo(list(self.sample_paths)) @@ -874,11 +878,11 @@ def __init__(self, paths: tuple[str, ...], max_footer_samples: int): self.file_metadata = tuple(sample_footers) sampled_row_count = sum(fmd.num_rows for fmd in sample_footers) - if total_file_count == sampled_file_count: + if self.total_file_count == sampled_file_count: row_count = sampled_row_count else: num_rows_per_sampled_file = int(sampled_row_count / sampled_file_count) - row_count = num_rows_per_sampled_file * total_file_count + row_count = num_rows_per_sampled_file * self.total_file_count num_row_groups_per_sampled_file = [ len(fmd.row_groups) for fmd in sample_footers @@ -904,7 +908,6 @@ 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.total_file_count = total_file_count self.sampled_file_count = sampled_file_count From 1d9e2147f8001a6e607da6f79c8125be5d1bb3da Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 22 Jun 2026 11:52:02 -0700 Subject: [PATCH 20/51] lint --- python/cudf_polars/cudf_polars/streaming/select.py | 2 +- python/cudf_polars/cudf_polars/utils/config.py | 2 +- python/cudf_polars/tests/test_scan.py | 2 +- python/cudf_polars/tests/test_select.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/select.py b/python/cudf_polars/cudf_polars/streaming/select.py index 81bfec030bc9..df946e653851 100644 --- a/python/cudf_polars/cudf_polars/streaming/select.py +++ b/python/cudf_polars/cudf_polars/streaming/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 """Parallel Select Logic.""" diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 8c7282991967..fbe127ef0098 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.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 """ diff --git a/python/cudf_polars/tests/test_scan.py b/python/cudf_polars/tests/test_scan.py index 6fc63aed9de9..60f3d487365e 100644 --- a/python/cudf_polars/tests/test_scan.py +++ b/python/cudf_polars/tests/test_scan.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 diff --git a/python/cudf_polars/tests/test_select.py b/python/cudf_polars/tests/test_select.py index b5ab87d26042..6365d34115af 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 From 8d0bc3a11369601391dfc190338025d9ba766e6c Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 22 Jun 2026 14:41:59 -0700 Subject: [PATCH 21/51] fixes --- .../cudf_polars/tests/streaming/test_scan.py | 1 - .../cudf_polars/tests/streaming/test_stats.py | 39 ++++++++++++++----- 2 files changed, 30 insertions(+), 10 deletions(-) diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index d6fced132a43..ee9f234ae16b 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -363,7 +363,6 @@ def test_streaming_scan_missing_prefetch_metadata_raises() -> None: scan = _make_parquet_scan( ["file.parquet"], parquet_options=ParquetOptions(prefetch_file_metadata=True) ) - scan = _make_parquet_scan(["file.parquet"]) fused = FusedScan(scan.schema, scan, scan.paths, scan.parquet_options) ctx = IRExecutionContext() diff --git a/python/cudf_polars/tests/streaming/test_stats.py b/python/cudf_polars/tests/streaming/test_stats.py index 1c83b81a5c8e..161caf904004 100644 --- a/python/cudf_polars/tests/streaming/test_stats.py +++ b/python/cudf_polars/tests/streaming/test_stats.py @@ -12,7 +12,9 @@ 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 @@ -40,16 +42,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) @@ -237,13 +245,18 @@ def test_parquet_round_trip_empty() -> None: def test_parquet_source_info_stores_footers_when_all_files_sampled( tmp_path: pathlib.Path, - df: pl.DataFrame, + 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), max_footer_samples=10, max_row_group_samples=0 + paths, + frozenset(df.columns), + tuple(schema.items()), + max_footer_samples=10, + max_row_group_samples=0, ) assert info.file_metadata is not None @@ -253,13 +266,18 @@ def test_parquet_source_info_stores_footers_when_all_files_sampled( def test_parquet_source_info_omits_footers_when_paths_are_sampled( tmp_path: pathlib.Path, - df: pl.DataFrame, + 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), max_footer_samples=2, max_row_group_samples=0 + paths, + frozenset(df.columns), + tuple(schema.items()), + max_footer_samples=2, + max_row_group_samples=0, ) assert info.file_metadata is None @@ -267,10 +285,11 @@ def test_parquet_source_info_omits_footers_when_paths_are_sampled( def test_seed_parquet_file_metadata_from_stats( 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=2) engine = pl.GPUEngine( raise_on_fail=True, @@ -294,8 +313,9 @@ def test_seed_parquet_file_metadata_from_stats( def test_parquet_metadata_reads_footers( tmp_path: pathlib.Path, - df: pl.DataFrame, + 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) @@ -405,10 +425,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, From 580c5a1222d523e8a4980d30d528f65c5b15d1db Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 23 Jun 2026 07:12:16 -0700 Subject: [PATCH 22/51] Error if both are provided --- python/cudf_polars/cudf_polars/utils/config.py | 5 +++++ python/cudf_polars/tests/streaming/test_scan.py | 2 +- python/cudf_polars/tests/test_config.py | 16 ++++++++++++++++ 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index fbe127ef0098..655248783a73 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -277,6 +277,11 @@ def __post_init__(self) -> None: # noqa: D105 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'" + ) + def default_target_partition_size(min_device_size: int | None) -> int: """Return the default target partition size.""" diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index ee9f234ae16b..2e19cfa86e78 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -359,7 +359,7 @@ def test_scan_missing_prefetch_metadata_raises() -> None: def test_streaming_scan_missing_prefetch_metadata_raises() -> None: - # This isn't reachable by normal cudf-polars usage. + # 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) ) diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index d02118daef03..05c5841a0d78 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -431,6 +431,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( From bb10333ec1df3b3cdbb44dfad89cf22d1db142b3 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 23 Jun 2026 07:13:08 -0700 Subject: [PATCH 23/51] remove redundant isinstance --- python/cudf_polars/cudf_polars/dsl/ir.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 7b6f127d0810..9efcf76568b6 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -246,17 +246,13 @@ def prefetch_parquet_file_metadata_for_ir( metadata concurrently, its ``parquet_file_metadata`` is mutated to cache the newly read parquet metadata. """ - from cudf_polars.streaming.io import FusedScan, SplitScan, StreamingScan + from cudf_polars.streaming.io import StreamingScan groups = set() for node in traversal([root]): if isinstance(node, StreamingScan): for scan in node.scans: - if ( - isinstance(scan, (SplitScan, FusedScan)) - and scan.base_scan.typ == "parquet" - ): - groups.add(tuple(scan.paths)) + groups.add(tuple(scan.paths)) elif isinstance(node, Scan) and node.typ == "parquet": groups.add(tuple(node.paths)) From 1f9be6bbb4d0ce27e0f7719d424d79ebd1daa5e1 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 23 Jun 2026 07:29:25 -0700 Subject: [PATCH 24/51] error message --- python/cudf_polars/cudf_polars/dsl/ir.py | 51 +++++++++++++++--------- 1 file changed, 33 insertions(+), 18 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 9efcf76568b6..1c9f179ab674 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -789,15 +789,7 @@ def _get_parquet_row_count_from_metadata( # Zero-width parquet files lose their row count when read through # pylibcudf. See https://github.com/rapidsai/cudf/issues/21428 if parquet_options.prefetch_file_metadata and context is not None: - try: - parquet_metadatas = context.parquet_file_metadata[tuple(paths)] - except KeyError as e: - msg = ( - f"Parquet file metadata was not prefetched for paths: {list(paths)}." - "Please report this as a bug to cudf-polars. You can work around it " - "by setting 'CUDF_POLARS__PARQUET_OPTIONS__PREFETCH_FILE_METADATA=0'." - ) - raise AssertionError(msg) from e + parquet_metadatas = Scan._lookup_parquet_metadatas(paths, context) num_rows = sum(metadata.num_rows for metadata in parquet_metadatas) else: meta = plc.io.parquet_metadata.read_parquet_metadata( @@ -942,15 +934,7 @@ def read_csv_header( ) elif typ == "parquet": if parquet_options.prefetch_file_metadata: - try: - parquet_metadatas = context.parquet_file_metadata[tuple(paths)] - except KeyError as e: - msg = ( - f"Parquet file metadata was not prefetched for paths: {list(paths)}." - "Please report this as a bug to cudf-polars. You can work around it " - "by setting 'CUDF_POLARS__PARQUET_OPTIONS__PREFETCH_FILE_METADATA=0'." - ) - raise AssertionError(msg) from e + parquet_metadatas = cls._lookup_parquet_metadatas(paths, context) else: parquet_metadatas = None @@ -1101,6 +1085,37 @@ def read_csv_header( ) return df.filter(mask) + @staticmethod + def _lookup_parquet_metadatas( + paths: list[str], context: IRExecutionContext + ) -> list[plc.io.parquet_metadata.FileMetaData]: + """ + Lookup parquet metadata from the prefetch metadata cache. + + Only call this when prefetching is enabled. + + Parameters + ---------- + paths + The list of paths from the Scan node's ``.paths``. + context + The IRExecutionContext with the prefetch parquet file metadata. + + Raises + ------ + AssertionError + If the parquet metadata for 'paths' is not found in the cache. + """ + try: + return context.parquet_file_metadata[tuple(paths)] + except KeyError as e: + msg = ( + f"Parquet file metadata was not prefetched for paths: {list(paths)}." + "Please report this as a bug to cudf-polars. You can work around it " + "by setting 'CUDF_POLARS__PARQUET_OPTIONS__PREFETCH_FILE_METADATA=0'." + ) + raise AssertionError(msg) from e + class Sink(IR): """Sink a dataframe to a file.""" From 611701434ac7f152afaca90cc259b8294b365a50 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 23 Jun 2026 11:07:08 -0700 Subject: [PATCH 25/51] Update the mock --- python/cudf_polars/tests/streaming/test_stats.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/tests/streaming/test_stats.py b/python/cudf_polars/tests/streaming/test_stats.py index 161caf904004..b616b16c6263 100644 --- a/python/cudf_polars/tests/streaming/test_stats.py +++ b/python/cudf_polars/tests/streaming/test_stats.py @@ -74,8 +74,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() @@ -96,7 +99,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, @@ -104,6 +107,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, @@ -162,6 +166,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] = [] From b66af9fa16c8fb9f73c9e4cd7296e77d4c664be1 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 23 Jun 2026 13:14:00 -0700 Subject: [PATCH 26/51] Switch to per-path caching --- python/cudf_polars/cudf_polars/dsl/ir.py | 69 +++++++++---------- .../cudf_polars/cudf_polars/streaming/io.py | 10 +-- .../cudf_polars/tests/streaming/test_stats.py | 13 ++-- 3 files changed, 44 insertions(+), 48 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 1c9f179ab674..fbef76459e8d 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -133,9 +133,8 @@ class IRExecutionContext: query_id Identifier for the query being executed. parquet_file_metadata - A cache of parquet file metadata. The keys are the ``paths`` of Scan nodes - with a ``parquet`` type. The values are a list of ``FileMetaData`` objects - associated with those ``paths``. + A cache of parquet file metadata keyed by file path for parquet scans. + The values are ``FileMetaData`` objects for each cached file. This cache lasts for the duration of the a single query's execution (e.g. ``LazyFrame.collect()``). @@ -144,9 +143,9 @@ class IRExecutionContext: 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) - parquet_file_metadata: dict[ - tuple[str, ...], list[plc.io.parquet_metadata.FileMetaData] - ] = field(default_factory=dict) + parquet_file_metadata: dict[str, plc.io.parquet_metadata.FileMetaData] = field( + default_factory=dict + ) async def to_thread( self, func: Callable[P, T], /, *args: P.args, **kwargs: P.kwargs @@ -202,30 +201,27 @@ def stream_ordered_after(self, *dfs: DataFrame) -> Generator[Stream, None, None] @nvtx_annotate_cudf_polars(message="fetch_parquet_footers_for_paths") def _prefetch_parquet_footers_for_paths( - paths: tuple[str, ...], -) -> tuple[tuple[str, ...], list[plc.io.parquet_metadata.FileMetaData]]: + paths: list[str], +) -> tuple[list[str], list[plc.io.parquet_metadata.FileMetaData]]: """ Prefetch parquet footers for a list of paths. This is typically executed concurrently with prefetch operations for other - groups of ``paths`` for other ``Scan`` nodes. + path groups for other parquet scan nodes. Parameters ---------- paths - The tuple of paths to prefetch. These correspond to ``paths`` in a ``Scan`` node. + The paths to prefetch. Returns ------- paths - The original input ``paths``. Useful for associating the result with the metadata - when executing out of order concurrently. + The original input ``paths``. metadata The list of ``FileMetaData`` objects for the ``paths``. """ - metadata = plc.io.parquet_metadata.read_parquet_footers( - plc.io.SourceInfo(list(paths)) - ) + metadata = plc.io.parquet_metadata.read_parquet_footers(plc.io.SourceInfo(paths)) return paths, metadata @@ -248,20 +244,27 @@ def prefetch_parquet_file_metadata_for_ir( """ from cudf_polars.streaming.io import StreamingScan - groups = set() + all_paths: list[str] = [] + seen_paths: set[str] = set() for node in traversal([root]): if isinstance(node, StreamingScan): for scan in node.scans: - groups.add(tuple(scan.paths)) + for path in scan.paths: + if path not in seen_paths: + seen_paths.add(path) + all_paths.append(path) elif isinstance(node, Scan) and node.typ == "parquet": - groups.add(tuple(node.paths)) + for path in node.paths: + if path not in seen_paths: + seen_paths.add(path) + all_paths.append(path) - if not groups: + if not all_paths: return - missing_paths = { - paths for paths in groups if paths not in context.parquet_file_metadata - } + missing_paths = [ + path for path in all_paths if path not in context.parquet_file_metadata + ] cm: contextlib.AbstractContextManager[concurrent.futures.Executor | None] if context.py_executor is None: @@ -273,14 +276,11 @@ def prefetch_parquet_file_metadata_for_ir( if missing_paths: with cm: - futures = [ - executor.submit(_prefetch_parquet_footers_for_paths, paths) - for paths in missing_paths - ] - - for future in concurrent.futures.as_completed(futures): - paths, metadata = future.result() - context.parquet_file_metadata.setdefault(paths, metadata) + paths, metadata = executor.submit( + _prefetch_parquet_footers_for_paths, missing_paths + ).result() + for path, file_metadata in zip(paths, metadata, strict=True): + context.parquet_file_metadata.setdefault(path, file_metadata) def seed_parquet_file_metadata_from_stats( @@ -311,9 +311,8 @@ def seed_parquet_file_metadata_from_stats( and isinstance(info, ParquetSourceInfo) and info.file_metadata is not None ): - context.parquet_file_metadata.setdefault( - tuple(node.paths), info.file_metadata - ) + for path, file_metadata in zip(node.paths, info.file_metadata, strict=True): + context.parquet_file_metadata.setdefault(path, file_metadata) _BINOPS = { @@ -1104,10 +1103,10 @@ def _lookup_parquet_metadatas( Raises ------ AssertionError - If the parquet metadata for 'paths' is not found in the cache. + If parquet metadata for any requested file path is not found in the cache. """ try: - return context.parquet_file_metadata[tuple(paths)] + return [context.parquet_file_metadata[path] for path in paths] except KeyError as e: msg = ( f"Parquet file metadata was not prefetched for paths: {list(paths)}." diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 0681de149f58..f18123feb1dc 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -289,15 +289,7 @@ def do_evaluate( # "skip_rows" and "n_rows" options to use locally. if parquet_options.prefetch_file_metadata: - try: - parquet_metadatas = context.parquet_file_metadata[tuple(paths)] - except KeyError as e: - msg = ( - f"Parquet file metadata was not prefetched for paths: {list(paths)}." - "Please report this as a bug to cudf-polars. You can work around it " - "by setting 'CUDF_POLARS__PARQUET_OPTIONS__PREFETCH_FILE_METADATA=0'." - ) - raise AssertionError(msg) from e + parquet_metadatas = Scan._lookup_parquet_metadatas(paths, context) row_group_num_rows = [ num_rows diff --git a/python/cudf_polars/tests/streaming/test_stats.py b/python/cudf_polars/tests/streaming/test_stats.py index b616b16c6263..e2d160e65440 100644 --- a/python/cudf_polars/tests/streaming/test_stats.py +++ b/python/cudf_polars/tests/streaming/test_stats.py @@ -22,6 +22,7 @@ Empty, IRExecutionContext, Projection, + Scan, seed_parquet_file_metadata_from_stats, ) from cudf_polars.engine.options import StreamingOptions @@ -310,11 +311,15 @@ def test_seed_parquet_file_metadata_from_stats( context = IRExecutionContext() seed_parquet_file_metadata_from_stats(stats, context) - scan_node = next(node for node in stats.scan_stats if hasattr(node, "paths")) - assert ( - context.parquet_file_metadata[tuple(scan_node.paths)] - is stats.scan_stats[scan_node].file_metadata + scan_node, parquet_info = next( + (node, info) + for node, info in stats.scan_stats.items() + if isinstance(node, Scan) and isinstance(info, ParquetSourceInfo) ) + file_metadata = parquet_info.file_metadata + assert file_metadata is not None + for path, metadata in zip(scan_node.paths, file_metadata, strict=True): + assert context.parquet_file_metadata[path] is metadata def test_parquet_metadata_reads_footers( From daa40782aeb741d3f5469a34d9405b3b8a24c0eb Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 23 Jun 2026 13:59:00 -0700 Subject: [PATCH 27/51] add a todo --- python/cudf_polars/cudf_polars/dsl/ir.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index fbef76459e8d..db1afb5b7588 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -276,6 +276,8 @@ def prefetch_parquet_file_metadata_for_ir( if missing_paths: with cm: + # TODO: "Consider batching footer reads for SplitScan" + # https://github.com/rapidsai/cudf/pull/22700#discussion_r3455262132 paths, metadata = executor.submit( _prefetch_parquet_footers_for_paths, missing_paths ).result() From 674209f0888a550b1c0a30784de80d126ad7acbb Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 9 Jun 2026 13:31:11 -0700 Subject: [PATCH 28/51] Followup to tom/cudf-sourceinfo-size (cherry picked from commit de8a7a9e36cfcdcecb33b77a810d615090e9cec4) (cherry picked from commit 83da3c4e31056ce9469c007f12265f4877819eb6) --- python/cudf_polars/cudf_polars/dsl/ir.py | 98 ++++++++++++++----- .../cudf_polars/cudf_polars/streaming/io.py | 37 ++++--- .../cudf_polars/tests/streaming/test_stats.py | 23 +++-- 3 files changed, 110 insertions(+), 48 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index db1afb5b7588..9d455f674172 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -115,6 +115,15 @@ ] +@dataclass(frozen=True) +class CachedParquetInfo: + """Metadata for a parquet file.""" + + path: str + size: int + file_metadata: plc.io.parquet_metadata.FileMetaData + + @dataclass(frozen=True) class IRExecutionContext: """ @@ -134,7 +143,9 @@ class IRExecutionContext: Identifier for the query being executed. parquet_file_metadata A cache of parquet file metadata keyed by file path for parquet scans. - The values are ``FileMetaData`` objects for each cached file. + The values are `CachedParquetInfo` objects storing + a ``SourceInfo`` with known size and a list of ``FileMetaData`` objects + associated with those ``paths``. This cache lasts for the duration of the a single query's execution (e.g. ``LazyFrame.collect()``). @@ -143,9 +154,7 @@ class IRExecutionContext: 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) - parquet_file_metadata: dict[str, plc.io.parquet_metadata.FileMetaData] = field( - default_factory=dict - ) + parquet_file_metadata: dict[str, CachedParquetInfo] = field(default_factory=dict) async def to_thread( self, func: Callable[P, T], /, *args: P.args, **kwargs: P.kwargs @@ -200,9 +209,7 @@ def stream_ordered_after(self, *dfs: DataFrame) -> Generator[Stream, None, None] @nvtx_annotate_cudf_polars(message="fetch_parquet_footers_for_paths") -def _prefetch_parquet_footers_for_paths( - paths: list[str], -) -> tuple[list[str], list[plc.io.parquet_metadata.FileMetaData]]: +def _prefetch_parquet_footers_for_paths(paths: list[str]) -> list[CachedParquetInfo]: """ Prefetch parquet footers for a list of paths. @@ -221,8 +228,39 @@ def _prefetch_parquet_footers_for_paths( metadata The list of ``FileMetaData`` objects for the ``paths``. """ - metadata = plc.io.parquet_metadata.read_parquet_footers(plc.io.SourceInfo(paths)) - return paths, metadata + # 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: + 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): + # 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") @@ -276,13 +314,16 @@ def prefetch_parquet_file_metadata_for_ir( if missing_paths: with cm: - # TODO: "Consider batching footer reads for SplitScan" - # https://github.com/rapidsai/cudf/pull/22700#discussion_r3455262132 - paths, metadata = executor.submit( - _prefetch_parquet_footers_for_paths, missing_paths - ).result() - for path, file_metadata in zip(paths, metadata, strict=True): - context.parquet_file_metadata.setdefault(path, file_metadata) + futures = [ + executor.submit(_prefetch_parquet_footers_for_paths, [path]) + for path in missing_paths + ] + + for future in concurrent.futures.as_completed(futures): + for cached_parquet_info in future.result(): + context.parquet_file_metadata.setdefault( + cached_parquet_info.path, cached_parquet_info + ) def seed_parquet_file_metadata_from_stats( @@ -311,10 +352,12 @@ def seed_parquet_file_metadata_from_stats( isinstance(node, Scan) and node.typ == "parquet" and isinstance(info, ParquetSourceInfo) - and info.file_metadata is not None + and info.cached_parquet_info is not None ): - for path, file_metadata in zip(node.paths, info.file_metadata, strict=True): - context.parquet_file_metadata.setdefault(path, file_metadata) + for cached_parquet_info in info.cached_parquet_info: + context.parquet_file_metadata.setdefault( + cached_parquet_info.path, cached_parquet_info + ) _BINOPS = { @@ -790,7 +833,8 @@ def _get_parquet_row_count_from_metadata( # Zero-width parquet files lose their row count when read through # pylibcudf. See https://github.com/rapidsai/cudf/issues/21428 if parquet_options.prefetch_file_metadata and context is not None: - parquet_metadatas = Scan._lookup_parquet_metadatas(paths, context) + cached_parquet_info = Scan._lookup_parquet_metadatas(paths, context) + parquet_metadatas = [info.file_metadata for info in cached_parquet_info] num_rows = sum(metadata.num_rows for metadata in parquet_metadatas) else: meta = plc.io.parquet_metadata.read_parquet_metadata( @@ -935,9 +979,17 @@ def read_csv_header( ) elif typ == "parquet": if parquet_options.prefetch_file_metadata: - parquet_metadatas = cls._lookup_parquet_metadatas(paths, context) + cached_parquet_info = cls._lookup_parquet_metadatas(paths, context) + source_info = plc.io.SourceInfo( + [ + plc.io.types.FilepathSource(info.path, info.size) + for info in cached_parquet_info + ] + ) + parquet_metadatas = [info.file_metadata for info in cached_parquet_info] else: parquet_metadatas = None + source_info = plc.io.SourceInfo(paths) filters = None if predicate is not None and row_index is None: @@ -949,7 +1001,7 @@ def read_csv_header( stream=stream, ) parquet_reader_options = ( - plc.io.parquet.ParquetReaderOptions.builder(plc.io.SourceInfo(paths)) + plc.io.parquet.ParquetReaderOptions.builder(source_info) .decimal_width(plc.TypeId.DECIMAL128) .build() ) @@ -1089,7 +1141,7 @@ def read_csv_header( @staticmethod def _lookup_parquet_metadatas( paths: list[str], context: IRExecutionContext - ) -> list[plc.io.parquet_metadata.FileMetaData]: + ) -> list[CachedParquetInfo]: """ Lookup parquet metadata from the prefetch metadata cache. diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 8e35b004d9a0..2ff8e7c7b9f4 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -23,6 +23,7 @@ Empty, Scan, Sink, + _prefetch_parquet_footers_for_paths, ) from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars from cudf_polars.streaming.base import ( @@ -42,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, @@ -289,7 +290,8 @@ def do_evaluate( # "skip_rows" and "n_rows" options to use locally. if parquet_options.prefetch_file_metadata: - parquet_metadatas = Scan._lookup_parquet_metadatas(paths, context) + cached_parquet_info = Scan._lookup_parquet_metadatas(paths, context) + parquet_metadatas = [info.file_metadata for info in cached_parquet_info] row_group_num_rows = [ num_rows @@ -864,8 +866,8 @@ class ParquetMetadata: """ __slots__ = ( + "cached_parquet_info", "column_names", - "file_metadata", "max_footer_samples", "mean_size_per_file", "num_row_groups_per_file", @@ -890,8 +892,8 @@ class ParquetMetadata: """All column names found it the dataset.""" sample_paths: tuple[str, ...] """Sampled file paths.""" - file_metadata: tuple[plc.io.parquet_metadata.FileMetaData, ...] | None - """Parquet footers read for ``sample_paths``. Populated only if all files were sampled.""" + 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): @@ -901,7 +903,7 @@ def __init__(self, paths: tuple[str, ...], max_footer_samples: int): self.num_row_groups_per_file = () self.mean_size_per_file = {} self.column_names = () - self.file_metadata = None + self.cached_parquet_info = None self.total_file_count = len(self.paths) self.sampled_file_count = 0 if max_footer_samples <= 0: @@ -917,14 +919,16 @@ def __init__(self, paths: tuple[str, ...], max_footer_samples: int): return sampled_file_count = len(self.sample_paths) - sample_footers = plc.io.parquet_metadata.read_parquet_footers( - plc.io.SourceInfo(list(self.sample_paths)) + + sample_parquet_info = _prefetch_parquet_footers_for_paths( + list(self.sample_paths) ) - self.file_metadata = tuple(sample_footers) + sample_footers = [info.file_metadata for info in sample_parquet_info] 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(sampled_row_count / sampled_file_count) row_count = num_rows_per_sampled_file * self.total_file_count @@ -1030,14 +1034,15 @@ def __init__( row_count: int | None, per_file_means: dict[str, int] | None = None, *, - file_metadata: list[plc.io.parquet_metadata.FileMetaData] | 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.file_metadata = file_metadata + self.cached_parquet_info = cached_parquet_info @classmethod def from_paths( @@ -1096,15 +1101,15 @@ def from_paths( else max(footer_mean, decoded_floor) ) - file_metadata: list[plc.io.parquet_metadata.FileMetaData] | None + cached_parquet_info: list[CachedParquetInfo] | None if ( metadata.sampled_file_count == metadata.total_file_count - and metadata.file_metadata is not None + and metadata.cached_parquet_info is not None ): - file_metadata = list(metadata.file_metadata) + cached_parquet_info = list(metadata.cached_parquet_info) else: - file_metadata = None - return cls(row_count, per_file_means, file_metadata=file_metadata) + 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/tests/streaming/test_stats.py b/python/cudf_polars/tests/streaming/test_stats.py index e2d160e65440..054582bacb5d 100644 --- a/python/cudf_polars/tests/streaming/test_stats.py +++ b/python/cudf_polars/tests/streaming/test_stats.py @@ -19,6 +19,7 @@ from cudf_polars import Translator from cudf_polars.containers import DataType from cudf_polars.dsl.ir import ( + CachedParquetInfo, Empty, IRExecutionContext, Projection, @@ -266,9 +267,12 @@ def test_parquet_source_info_stores_footers_when_all_files_sampled( max_row_group_samples=0, ) - assert info.file_metadata is not None - assert len(info.file_metadata) == len(paths) - assert sum(fmd.num_rows for fmd in info.file_metadata) == df.height + 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( @@ -287,7 +291,7 @@ def test_parquet_source_info_omits_footers_when_paths_are_sampled( max_row_group_samples=0, ) - assert info.file_metadata is None + assert info.cached_parquet_info is None def test_seed_parquet_file_metadata_from_stats( @@ -316,9 +320,10 @@ def test_seed_parquet_file_metadata_from_stats( for node, info in stats.scan_stats.items() if isinstance(node, Scan) and isinstance(info, ParquetSourceInfo) ) - file_metadata = parquet_info.file_metadata - assert file_metadata is not None - for path, metadata in zip(scan_node.paths, file_metadata, strict=True): + cached_parquet_info = parquet_info.cached_parquet_info + assert cached_parquet_info is not None + for path, metadata in zip(scan_node.paths, cached_parquet_info, strict=True): + assert isinstance(metadata, CachedParquetInfo) assert context.parquet_file_metadata[path] is metadata @@ -331,8 +336,8 @@ def test_parquet_metadata_reads_footers( path = next(tmp_path.iterdir()) metadata = ParquetMetadata((str(path),), max_footer_samples=1) - assert metadata.file_metadata is not None - assert len(metadata.file_metadata) == 1 + assert metadata.cached_parquet_info is not None + assert len(metadata.cached_parquet_info) == 1 assert metadata.row_count == df.height From 73a104087e98c7c0aa25b06620ab0739630b1349 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 30 Jun 2026 05:26:38 -0700 Subject: [PATCH 29/51] doc cachedparquetinfo --- python/cudf_polars/cudf_polars/dsl/ir.py | 26 +++++++++++++++++++----- 1 file changed, 21 insertions(+), 5 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index e82c5a3347e1..c4e2586bd7d8 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -48,10 +48,7 @@ from cudf_polars.dsl.expressions.base import ExecutionContext from cudf_polars.dsl.nodebase import Node from cudf_polars.dsl.to_ast import _DECIMAL_IDS, to_ast, to_parquet_filter -from cudf_polars.dsl.tracing import ( - log_do_evaluate, - nvtx_annotate_cudf_polars, -) +from cudf_polars.dsl.tracing import log_do_evaluate, nvtx_annotate_cudf_polars from cudf_polars.dsl.traversal import traversal from cudf_polars.dsl.utils.reshape import broadcast from cudf_polars.dsl.utils.windows import ( @@ -117,7 +114,26 @@ @dataclass(frozen=True) class CachedParquetInfo: - """Metadata for a parquet file.""" + """ + 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 From ab9d50083afc32e46e76546fa28bce729b7afec6 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 30 Jun 2026 10:56:10 -0700 Subject: [PATCH 30/51] WIP: POC for eliminating IRExecutionContext cache --- python/cudf_polars/cudf_polars/callback.py | 6 - python/cudf_polars/cudf_polars/dsl/ir.py | 219 +++++++++--------- .../cudf_polars/cudf_polars/dsl/translate.py | 1 + python/cudf_polars/cudf_polars/engine/core.py | 37 ++- .../cudf_polars/cudf_polars/streaming/io.py | 194 ++++++++++++++-- .../cudf_polars/cudf_polars/utils/config.py | 11 +- .../cudf_polars/tests/streaming/test_scan.py | 158 ++++++++++++- .../cudf_polars/tests/streaming/test_stats.py | 37 +-- python/cudf_polars/tests/test_scan.py | 80 +++++-- 9 files changed, 535 insertions(+), 208 deletions(-) diff --git a/python/cudf_polars/cudf_polars/callback.py b/python/cudf_polars/cudf_polars/callback.py index 7e0a3b99a36d..2c254720925a 100644 --- a/python/cudf_polars/cudf_polars/callback.py +++ b/python/cudf_polars/cudf_polars/callback.py @@ -26,7 +26,6 @@ import cudf_polars.dsl.tracing from cudf_polars.dsl.ir import ( IRExecutionContext, - prefetch_parquet_file_metadata_for_ir, ) from cudf_polars.dsl.tracing import CUDF_POLARS_NVTX_DOMAIN from cudf_polars.dsl.translate import Translator @@ -304,11 +303,6 @@ def _callback( ): if config_options.executor.name == "in-memory": context = IRExecutionContext() - if config_options.parquet_options.prefetch_file_metadata: - prefetch_parquet_file_metadata_for_ir( - ir, - context, - ) df = ir.evaluate(cache={}, timer=timer, context=context).to_polars() if timer is None: return df diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index c4e2586bd7d8..177cb5d79dd4 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -157,20 +157,11 @@ class IRExecutionContext: A zero-argument callable that returns a CUDA stream. query_id Identifier for the query being executed. - parquet_file_metadata - A cache of parquet file metadata keyed by file path for parquet scans. - The values are `CachedParquetInfo` objects storing - a ``SourceInfo`` with known size and a list of ``FileMetaData`` objects - associated with those ``paths``. - - This cache lasts for the duration of the a single query's execution - (e.g. ``LazyFrame.collect()``). """ 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) - parquet_file_metadata: dict[str, CachedParquetInfo] = field(default_factory=dict) async def to_thread( self, func: Callable[P, T], /, *args: P.args, **kwargs: P.kwargs @@ -282,8 +273,9 @@ 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, - context: IRExecutionContext, -) -> None: + py_executor: concurrent.futures.Executor | None, + stats: StatsCollector | None = None, +) -> dict[str, CachedParquetInfo]: """ Prefetch parquet metadata for all parquet scans in an IR graph. @@ -291,89 +283,62 @@ def prefetch_parquet_file_metadata_for_ir( ---------- root The root of the IR graph, which will be traversed. - context - The IR execution context. Its ``py_executor`` is used to fetch - metadata concurrently, its ``parquet_file_metadata`` is mutated - to cache the newly read parquet metadata. + 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 file paths to their cached parquet metadata. """ - from cudf_polars.streaming.io import StreamingScan + from cudf_polars.streaming.io import ParquetSourceInfo, StreamingScan + + all_paths: set[str] = set() - all_paths: list[str] = [] - seen_paths: set[str] = set() for node in traversal([root]): if isinstance(node, StreamingScan): for scan in node.scans: for path in scan.paths: - if path not in seen_paths: - seen_paths.add(path) - all_paths.append(path) + all_paths.add(path) elif isinstance(node, Scan) and node.typ == "parquet": for path in node.paths: - if path not in seen_paths: - seen_paths.add(path) - all_paths.append(path) + all_paths.add(path) - if not all_paths: - return + 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 = [ - path for path in all_paths if path not in context.parquet_file_metadata - ] + missing_paths = all_paths - set(cached_parquet_info.keys()) cm: contextlib.AbstractContextManager[concurrent.futures.Executor | None] - if context.py_executor is None: - cm = executor = concurrent.futures.ThreadPoolExecutor() + if py_executor is None: + cm = py_executor = concurrent.futures.ThreadPoolExecutor() else: - executor = context.py_executor # We didn't create the executor, so we don't close it. cm = contextlib.nullcontext() - if missing_paths: - with cm: - futures = [ - executor.submit(_prefetch_parquet_footers_for_paths, [path]) - for path in missing_paths - ] - - for future in concurrent.futures.as_completed(futures): - for cached_parquet_info in future.result(): - context.parquet_file_metadata.setdefault( - cached_parquet_info.path, cached_parquet_info - ) - - -def seed_parquet_file_metadata_from_stats( - stats: StatsCollector, - context: IRExecutionContext, -) -> None: - """ - Seed prefetched parquet footers from stats collection when available. - - Stats collection reads parquet footers for sampled paths. When the sample - covers all paths in a scan, those footers can be reused during metadata - prefetch to avoid rereading file footers. - - Parameters - ---------- - stats: StatsCollector - The stats collector. - context: IRExecutionContext - The execution context. Its ``parquet_file_metadata`` is mutated to - cache the parquet file metadata read during stats collection. - """ - from cudf_polars.streaming.io import ParquetSourceInfo + with cm: + futures = [ + py_executor.submit(_prefetch_parquet_footers_for_paths, [path]) + for path in missing_paths + ] - for node, info in stats.scan_stats.items(): - if ( - isinstance(node, Scan) - and node.typ == "parquet" - and isinstance(info, ParquetSourceInfo) - and info.cached_parquet_info is not None - ): - for cached_parquet_info in info.cached_parquet_info: - context.parquet_file_metadata.setdefault( - cached_parquet_info.path, cached_parquet_info - ) + for future in concurrent.futures.as_completed(futures): + for info in future.result(): + cached_parquet_info[info.path] = info + return cached_parquet_info _BINOPS = { @@ -611,6 +576,7 @@ class Scan(IR): """Input from files.""" __slots__ = ( + "cached_parquet_info", "cloud_options", "include_file_paths", "n_rows", @@ -636,8 +602,9 @@ class Scan(IR): "include_file_paths", "predicate", "parquet_options", + "cached_parquet_info", ) - _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] @@ -660,6 +627,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 @@ -678,6 +647,7 @@ def __init__( include_file_paths: str | None, predicate: expr.NamedExpr | None, parquet_options: ParquetOptions, + cached_parquet_info: list[CachedParquetInfo] | None, ): self.schema = schema self.typ = typ @@ -702,9 +672,11 @@ 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 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 @@ -784,6 +756,40 @@ def __init__( "Reading only parquet metadata to produce row index." ) + @classmethod + def with_prefetched_parquet_metadata( + cls, scan: Scan, cached_parquet_info: list[CachedParquetInfo] + ) -> Self: + """Create a new scan node, with prefetched parquet metadata set.""" + return cls( + scan.schema, + scan.typ, + scan.reader_options, + scan.cloud_options, + scan.paths, + scan.with_columns, + scan.skip_rows, + scan.n_rows, + scan.row_index, + scan.include_file_paths, + scan.predicate, + scan.parquet_options, + cached_parquet_info, + ) + + def is_equal(self, other: Self) -> bool: # noqa: D102 + # This needs to exclude 'cached_parquet_info' from the equality check. + if self is other: + return True + result = ( + self._ctor_arguments(self.children)[:-1] + == other._ctor_arguments(other.children)[:-1] + ) + # Eager CSE for nodes that match. + if result: + self.children = other.children + return result + def get_hashable(self) -> Hashable: """ Hashable representation of the node. @@ -791,6 +797,7 @@ def get_hashable(self) -> Hashable: The options dictionaries are serialised for hashing purposes as json strings. """ + # cached_parquet_info is deliberately not included in the hash data. schema_hash = tuple(self.schema.items()) return ( type(self), @@ -844,12 +851,19 @@ def _get_parquet_row_count_from_metadata( skip_rows: int, n_rows: int, parquet_options: ParquetOptions, - context: IRExecutionContext | None, + 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 - if parquet_options.prefetch_file_metadata and context is not None: - cached_parquet_info = Scan._lookup_parquet_metadatas(paths, context) + if parquet_options.prefetch_file_metadata: + if cached_parquet_info is None: + raise AssertionError( + "Prefetching is enabled, but no cached parquet info was provided." + ) + if paths != [info.path for info in cached_parquet_info]: + raise AssertionError( + f"Paths do not match cached parquet info. {paths} != {[info.path for info in cached_parquet_info]}" + ) parquet_metadatas = [info.file_metadata for info in cached_parquet_info] num_rows = sum(metadata.num_rows for metadata in parquet_metadatas) else: @@ -879,6 +893,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: @@ -995,7 +1010,14 @@ def read_csv_header( ) elif typ == "parquet": if parquet_options.prefetch_file_metadata: - cached_parquet_info = cls._lookup_parquet_metadatas(paths, context) + if cached_parquet_info is None: + raise AssertionError( + "Prefetching is enabled, but no cached parquet info was provided." + ) + if paths != [info.path for info in cached_parquet_info]: + raise AssertionError( + f"Paths do not match cached parquet info. {paths} != {[info.path for info in cached_parquet_info]}" + ) source_info = plc.io.SourceInfo( [ plc.io.types.FilepathSource(info.path, info.size) @@ -1052,7 +1074,7 @@ def read_csv_header( ) num_rows = ( cls._get_parquet_row_count_from_metadata( - paths, skip_rows, n_rows, parquet_options, context + paths, skip_rows, n_rows, parquet_options, cached_parquet_info ) if not names else None @@ -1154,37 +1176,6 @@ def read_csv_header( ) return df.filter(mask) - @staticmethod - def _lookup_parquet_metadatas( - paths: list[str], context: IRExecutionContext - ) -> list[CachedParquetInfo]: - """ - Lookup parquet metadata from the prefetch metadata cache. - - Only call this when prefetching is enabled. - - Parameters - ---------- - paths - The list of paths from the Scan node's ``.paths``. - context - The IRExecutionContext with the prefetch parquet file metadata. - - Raises - ------ - AssertionError - If parquet metadata for any requested file path is not found in the cache. - """ - try: - return [context.parquet_file_metadata[path] for path in paths] - except KeyError as e: - msg = ( - f"Parquet file metadata was not prefetched for paths: {list(paths)}." - "Please report this as a bug to cudf-polars. You can work around it " - "by setting 'CUDF_POLARS__PARQUET_OPTIONS__PREFETCH_FILE_METADATA=0'." - ) - raise AssertionError(msg) from e - class Sink(IR): """Sink a dataframe to a file.""" diff --git a/python/cudf_polars/cudf_polars/dsl/translate.py b/python/cudf_polars/cudf_polars/dsl/translate.py index d9f2b03c3f74..192c24931db0 100644 --- a/python/cudf_polars/cudf_polars/dsl/translate.py +++ b/python/cudf_polars/cudf_polars/dsl/translate.py @@ -429,6 +429,7 @@ def _(node: plrs._ir_nodes.Scan, translator: Translator, schema: Schema) -> ir.I else translate_predicate(translator, n=node.predicate, schema=schema) ), parquet_options, + cached_parquet_info=None, ) diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 25a2df20ef04..69241b602eba 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -27,9 +27,10 @@ from cudf_polars.containers import DataFrame from cudf_polars.dsl.ir import ( IRExecutionContext, + Scan, prefetch_parquet_file_metadata_for_ir, - seed_parquet_file_metadata_from_stats, ) +from cudf_polars.dsl.utils.replace import replace 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 @@ -682,6 +683,8 @@ def evaluate_on_rank( metadata Collected channel metadata. """ + from cudf_polars.dsl.traversal import traversal + stats = allgather_stats(comm, ctx.br(), ir, config_options, py_executor) ir, partition_info = lower_ir_graph( ir, config_options, stats, rank=comm.rank, nranks=comm.nranks @@ -697,11 +700,37 @@ def evaluate_on_rank( ) if config_options.parquet_options.prefetch_file_metadata: - seed_parquet_file_metadata_from_stats(stats, ir_context) - prefetch_parquet_file_metadata_for_ir( + cached_parquet_info_map = prefetch_parquet_file_metadata_for_ir( ir, - ir_context, + ir_context.py_executor, + stats=stats, ) + # Build a mapping {Scan: Scan} where the values have the prefetched metadata set. + # This is complicated by StreamingScan, SplitScan, and FusedScan, but whatever. + replacements: dict[IR, IR] = {} + from cudf_polars.streaming.io import FusedScan, SplitScan, StreamingScan + + for node in traversal([ir]): + if isinstance(node, Scan): + replacements[node] = Scan.with_prefetched_parquet_metadata( + node, [cached_parquet_info_map[path] for path in node.paths] + ) + elif isinstance(node, StreamingScan): + x = StreamingScan.with_prefetched_parquet_metadata( + node, cached_parquet_info_map + ) + assert node in partition_info + assert node.is_equal(x) + assert x in partition_info + replacements[node] = x + elif isinstance(node, (SplitScan, FusedScan)): + raise NotImplementedError( + f"Prefetching is not supported for StreamingScan, SplitScan, and FusedScan.: {type(node)}" + ) + + ir = replace([ir], replacements)[0] + # Now we want to rewrite the Scan nodes. We need to insert the prefetched + # parquet metadata for the paths in that node. with ReserveOpIDs(ir, config_options) as collective_id_map: return execute_ir_on_rank( diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 2ff8e7c7b9f4..a27455781735 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -180,6 +180,7 @@ class SplitScan(IR): __slots__ = ( "base_scan", + "cached_parquet_info", "parquet_options", "paths", "schema", @@ -193,6 +194,7 @@ class SplitScan(IR): "split_index", "total_splits", "parquet_options", + "cached_parquet_info", ) _n_non_child_args = 13 base_scan: Scan @@ -205,6 +207,7 @@ class SplitScan(IR): """Total number of splits.""" parquet_options: ParquetOptions """Parquet-specific options.""" + cached_parquet_info: list[CachedParquetInfo] | None def __init__( self, @@ -214,6 +217,7 @@ def __init__( split_index: int, total_splits: int, parquet_options: ParquetOptions, + cached_parquet_info: list[CachedParquetInfo] | None, ): self.schema = schema self.base_scan = base_scan @@ -234,14 +238,65 @@ 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( f"Unhandled Scan type for file splitting: {base_scan.typ}" ) + @classmethod + def with_prefetched_parquet_metadata( + cls, + node: SplitScan, + cached_parquet_info: list[CachedParquetInfo], + ) -> Self: + """ + Create a new SplitScan node, with prefetched parquet metadata set. + + Because SplitScan is a single-file scan, each composed Scan nodes will + use the same cached parquet metadata. + + Parameters + ---------- + node + The SplitScan node to create a new node from. + cached_parquet_info + The cached parquet metadata to set on the new node. This will be a + length-1 list, matching the length-1 ``path`` for the base scan node. + + Returns + ------- + The new SplitScan node. + """ + new_base = Scan.with_prefetched_parquet_metadata( + node.base_scan, cached_parquet_info + ) + # assert new_base.paths == [info.path for info in cached_parquet_info] + return cls( + node.schema, + new_base, + node.paths, + node.split_index, + node.total_splits, + node.parquet_options, + cached_parquet_info, + ) + + def is_equal(self, other: Self) -> bool: # noqa: D102 + # This needs to exclude 'cached_parquet_info' from the equality check. + if type(other) is not type(self): + return False + return self is other or ( + self.base_scan.is_equal(other.base_scan) + and self.paths == other.paths + and self.split_index == other.split_index + and self.total_splits == other.total_splits + ) + def get_hashable(self) -> Hashable: """Hashable representation of the node.""" return ( @@ -270,6 +325,7 @@ def do_evaluate( include_file_paths: str | None, predicate: NamedExpr | None, parquet_options: ParquetOptions, + cached_parquet_info: list[CachedParquetInfo] | None, *, context: IRExecutionContext, ) -> DataFrame: @@ -289,8 +345,7 @@ def do_evaluate( # - We can use all this information to calculate the # "skip_rows" and "n_rows" options to use locally. - if parquet_options.prefetch_file_metadata: - cached_parquet_info = Scan._lookup_parquet_metadatas(paths, context) + if cached_parquet_info is not None: parquet_metadatas = [info.file_metadata for info in cached_parquet_info] row_group_num_rows = [ @@ -347,6 +402,7 @@ def do_evaluate( include_file_paths, predicate, parquet_options, + cached_parquet_info, context=context, ) @@ -361,6 +417,7 @@ class FusedScan(IR): __slots__ = ( "base_scan", + "cached_parquet_info", "parquet_options", "paths", "schema", @@ -370,6 +427,7 @@ class FusedScan(IR): "base_scan", "paths", "parquet_options", + "cached_parquet_info", ) _n_non_child_args = 11 base_scan: Scan @@ -378,6 +436,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, @@ -385,11 +445,13 @@ def __init__( base_scan: Scan, paths: list[str], parquet_options: ParquetOptions, + cached_parquet_info: list[CachedParquetInfo] | 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, @@ -402,9 +464,40 @@ def __init__( base_scan.include_file_paths, base_scan.predicate, parquet_options, + cached_parquet_info, ) self.children = () + @classmethod + def with_prefetched_parquet_metadata( + cls, + node: FusedScan, + cached_parquet_info_map: dict[str, CachedParquetInfo], + ) -> Self: + """Create a new FusedScan node, with prefetched parquet metadata set.""" + cached_parquet_info = [cached_parquet_info_map[path] for path in node.paths] + if node.paths != [info.path for info in cached_parquet_info]: + raise AssertionError( + f"Paths do not match cached parquet info. {node.paths} != {[info.path for info in cached_parquet_info]}" + ) + return cls( + node.schema, + node.base_scan, + node.paths, + node.parquet_options, + cached_parquet_info, + ) + + def is_equal(self, other: Self) -> bool: # noqa: D102 + # This needs to exclude 'cached_parquet_info' from the equality check. + if type(other) is not type(self): + return False + return self is other or ( + self.base_scan.is_equal(other.base_scan) + and self.paths == other.paths + and self.parquet_options == other.parquet_options + ) + def get_hashable(self) -> Hashable: """Hashable representation of the node.""" return ( @@ -429,6 +522,7 @@ def do_evaluate( include_file_paths: str | None, predicate: NamedExpr | None, parquet_options: ParquetOptions, + cached_parquet_info: list[CachedParquetInfo] | None, *, context: IRExecutionContext, ) -> DataFrame: @@ -446,6 +540,7 @@ def do_evaluate( include_file_paths, predicate, parquet_options, + cached_parquet_info, context=context, ) @@ -573,26 +668,83 @@ 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 + def with_prefetched_parquet_metadata( + cls, + node: StreamingScan, + cached_parquet_info_map: dict[str, CachedParquetInfo], + ) -> Self: + """ + Create a new StreamingScan node, with prefetched parquet metadata set. + + Parameters + ---------- + node: StreamingScan + The StreamingScan node to create a new node from. + cached_parquet_info_map + The cached parquet metadata to set on the new node. + + Returns + ------- + Self: The new StreamingScan node. + """ + new_scans: list[SplitScan | FusedScan] = [] + if node.scan_type == "split": + new_scans = [] + for scan in node.scans: + new_parquet_info = [ + cached_parquet_info_map[path] for path in scan.paths + ] + # SplitScan should be generic / overload based on type. + new_scan = SplitScan.with_prefetched_parquet_metadata( + scan, # type: ignore[arg-type] + new_parquet_info, + ) + assert new_scan.cached_parquet_info is not None + assert new_scan.paths == [ + info.path for info in new_scan.cached_parquet_info + ] + new_scans.append(new_scan) + else: + new_scans = [ + FusedScan.with_prefetched_parquet_metadata( + scan, # type: ignore[arg-type] + cached_parquet_info_map, + ) + for scan in node.scans + ] + for scan in new_scans: + assert scan.cached_parquet_info is not None + assert scan.paths == [info.path for info in scan.cached_parquet_info] + + return cls(new_scans, node.base_scan, node.scan_type) # type: ignore[arg-type] + @classmethod def for_split_files( cls, @@ -614,6 +766,9 @@ def for_split_files( splits_created = 0 for path in local_paths: while sindex < plan.factor and splits_created < local_count: + # TODO: We should replace base_scan first, and then ensure cached_parquet_info + # is set properly from the start. Then we can remove the with_prefetched_parquet_metadata + # alternate constructor. scans.append( SplitScan( base_scan.schema, @@ -622,12 +777,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( @@ -644,17 +800,21 @@ def for_fused_files( local_offset, local_count = _rank_slice(partition_count, rank, nranks) paths_start = local_offset * plan.factor paths_end = paths_start + plan.factor * local_count - scans = [ - FusedScan( - base_scan.schema, - base_scan, - base_scan.paths[offset : offset + plan.factor], - parquet_options, - ) - for offset in range(paths_start, paths_end, plan.factor) - if base_scan.paths[offset : offset + plan.factor] - ] - return cls(scans, base_scan) + scans = [] + for offset in range(paths_start, paths_end, plan.factor): + paths = base_scan.paths[offset : offset + plan.factor] + if paths: + scans.append( + FusedScan( + base_scan.schema, + base_scan, + paths, + parquet_options, + None, + ) + ) + + return cls(scans, base_scan, "fused") def get_hashable(self) -> Hashable: """Hashable representation of the node.""" diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 1431acb90ea7..e37c03562a42 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -884,6 +884,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) @@ -911,6 +916,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: @@ -933,7 +942,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 2e19cfa86e78..d9d682bfc7e2 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -13,6 +13,7 @@ from cudf_polars import Translator from cudf_polars.containers import DataType from cudf_polars.dsl.ir import ( + CachedParquetInfo, Empty, IRExecutionContext, Scan, @@ -121,9 +122,10 @@ def test_prefetch_file_metadata_non_parquet_scan(df, streaming_engine_factory) - def test_prefetch_parquet_file_metadata_no_parquet_scans() -> None: - context = IRExecutionContext() - prefetch_parquet_file_metadata_for_ir(Empty({}), context) - assert context.parquet_file_metadata == {} + result = prefetch_parquet_file_metadata_for_ir( + Empty({}), py_executor=None, stats=None + ) + assert result == {} # --------------------------------------------------------------------------- @@ -238,6 +240,7 @@ def _make_parquet_scan( None, None, parquet_options, + None, ) @@ -326,7 +329,7 @@ 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) @@ -340,7 +343,7 @@ def test_scan_missing_prefetch_metadata_raises() -> None: ctx = IRExecutionContext() with pytest.raises( AssertionError, - match=r"Parquet file metadata was not prefetched for paths: \['file\.parquet'\]\.", + match=r"Paths do not match cached parquet info", ): Scan.do_evaluate( scan.schema, @@ -354,6 +357,7 @@ def test_scan_missing_prefetch_metadata_raises() -> None: scan.include_file_paths, scan.predicate, scan.parquet_options, + [], context=ctx, ) @@ -363,7 +367,7 @@ def test_streaming_scan_missing_prefetch_metadata_raises() -> None: scan = _make_parquet_scan( ["file.parquet"], parquet_options=ParquetOptions(prefetch_file_metadata=True) ) - 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"): @@ -378,10 +382,7 @@ def test_split_scan_do_evaluate_missing_prefetch_metadata() -> None: with pytest.raises( AssertionError, - match=( - r"Parquet file metadata was not prefetched for paths: " - r"\['/some/missing/file\.parquet'\]\." - ), + match=(r"Paths do not match cached parquet info."), ): SplitScan.do_evaluate( 0, @@ -397,5 +398,142 @@ def test_split_scan_do_evaluate_missing_prefetch_metadata() -> None: None, None, parquet_options, + [], context=context, ) + + +def test_split_scan_is_equal() -> None: + parquet_options = ParquetOptions(prefetch_file_metadata=True) + base_scan = _make_parquet_scan(["file.parquet"], parquet_options=parquet_options) + cached_info = [ + CachedParquetInfo(path="file.parquet", size=100, file_metadata=None) # type: ignore[arg-type] + ] + + scan_1 = SplitScan( + base_scan.schema, + base_scan, + ["file.parquet"], + 0, + 2, + parquet_options, + None, + ) + scan_2 = SplitScan( + base_scan.schema, + base_scan, + ["file.parquet"], + 0, + 2, + parquet_options, + cached_info, + ) + scan_3 = SplitScan( + base_scan.schema, + base_scan, + ["file.parquet"], + 1, + 2, + parquet_options, + None, + ) + + assert scan_1.is_equal(scan_2) + assert scan_1.get_hashable() == scan_2.get_hashable() + assert hash(scan_1) == hash(scan_2) + + assert not scan_1.is_equal(scan_3) + assert scan_1.get_hashable() != scan_3.get_hashable() + assert hash(scan_1) != hash(scan_3) + + +def test_fused_scan_is_equal() -> None: + parquet_options = ParquetOptions(prefetch_file_metadata=True) + base_scan = _make_parquet_scan( + ["file_1.parquet", "file_2.parquet"], parquet_options=parquet_options + ) + other_base_scan = _make_parquet_scan( + ["other.parquet"], parquet_options=parquet_options + ) + cached_info = [ + CachedParquetInfo(path="file_1.parquet", size=100, file_metadata=None), # type: ignore[arg-type] + CachedParquetInfo(path="file_2.parquet", size=200, file_metadata=None), # type: ignore[arg-type] + ] + + scan_1 = FusedScan( + base_scan.schema, + base_scan, + ["file_1.parquet", "file_2.parquet"], + parquet_options, + None, + ) + scan_2 = FusedScan( + base_scan.schema, + base_scan, + ["file_1.parquet", "file_2.parquet"], + parquet_options, + cached_info, + ) + scan_3 = FusedScan( + other_base_scan.schema, + other_base_scan, + ["other.parquet"], + parquet_options, + None, + ) + + assert scan_1.is_equal(scan_2) + assert scan_1.get_hashable() == scan_2.get_hashable() + assert hash(scan_1) == hash(scan_2) + + assert not scan_1.is_equal(scan_3) + assert scan_1.get_hashable() != scan_3.get_hashable() + assert hash(scan_1) != hash(scan_3) + + +def test_streaming_scan_is_equal() -> None: + parquet_options = ParquetOptions(prefetch_file_metadata=True) + base_scan = _make_parquet_scan(["file.parquet"], parquet_options=parquet_options) + cached_info = [ + CachedParquetInfo(path="file.parquet", size=100, file_metadata=None) # type: ignore[arg-type] + ] + + split_1 = SplitScan( + base_scan.schema, + base_scan, + ["file.parquet"], + 0, + 2, + parquet_options, + None, + ) + split_2 = SplitScan( + base_scan.schema, + base_scan, + ["file.parquet"], + 0, + 2, + parquet_options, + cached_info, + ) + split_3 = SplitScan( + base_scan.schema, + base_scan, + ["file.parquet"], + 1, + 2, + parquet_options, + None, + ) + + streaming_1 = StreamingScan([split_1], base_scan, "split") + streaming_2 = StreamingScan([split_2], base_scan, "split") + streaming_3 = StreamingScan([split_3], base_scan, "split") + + assert streaming_1.is_equal(streaming_2) + assert streaming_1.get_hashable() == streaming_2.get_hashable() + assert hash(streaming_1) == hash(streaming_2) + + assert not streaming_1.is_equal(streaming_3) + assert streaming_1.get_hashable() != streaming_3.get_hashable() + assert hash(streaming_1) != hash(streaming_3) diff --git a/python/cudf_polars/tests/streaming/test_stats.py b/python/cudf_polars/tests/streaming/test_stats.py index 054582bacb5d..507974fd2646 100644 --- a/python/cudf_polars/tests/streaming/test_stats.py +++ b/python/cudf_polars/tests/streaming/test_stats.py @@ -19,12 +19,8 @@ from cudf_polars import Translator from cudf_polars.containers import DataType from cudf_polars.dsl.ir import ( - CachedParquetInfo, Empty, - IRExecutionContext, Projection, - Scan, - seed_parquet_file_metadata_from_stats, ) from cudf_polars.engine.options import StreamingOptions from cudf_polars.streaming.base import SerializedDataSourceInfo, StatsCollector @@ -294,37 +290,8 @@ def test_parquet_source_info_omits_footers_when_paths_are_sampled( assert info.cached_parquet_info is None -def test_seed_parquet_file_metadata_from_stats( - tmp_path: pathlib.Path, - 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=2) - engine = pl.GPUEngine( - raise_on_fail=True, - executor="streaming", - parquet_options={"max_footer_samples": 10, "prefetch_file_metadata": True}, - ) - q = pl.scan_parquet(tmp_path) - ir = Translator(q._ldf.visit(), engine).translate_ir() - config = ConfigOptions.from_polars_engine(engine) - stats = collect_statistics(ir, config, parquet_stats_executor) - - context = IRExecutionContext() - seed_parquet_file_metadata_from_stats(stats, context) - - scan_node, parquet_info = next( - (node, info) - for node, info in stats.scan_stats.items() - if isinstance(node, Scan) and isinstance(info, ParquetSourceInfo) - ) - cached_parquet_info = parquet_info.cached_parquet_info - assert cached_parquet_info is not None - for path, metadata in zip(scan_node.paths, cached_parquet_info, strict=True): - assert isinstance(metadata, CachedParquetInfo) - assert context.parquet_file_metadata[path] is metadata +# TODO: Add a test that calls prefetch_parquet_file_metadata_for_ir +# and ensure that we don't refetch the metadata from stats def test_parquet_metadata_reads_footers( diff --git a/python/cudf_polars/tests/test_scan.py b/python/cudf_polars/tests/test_scan.py index 60f3d487365e..0ae0d6317816 100644 --- a/python/cudf_polars/tests/test_scan.py +++ b/python/cudf_polars/tests/test_scan.py @@ -17,14 +17,14 @@ import polars as pl from cudf_polars.containers import DataType -from cudf_polars.dsl.ir import IRExecutionContext, Scan +from cudf_polars.dsl.ir import CachedParquetInfo, 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 ParquetOptions +from cudf_polars.utils.config import ConfigOptions, ParquetOptions from cudf_polars.utils.versions import ( POLARS_VERSION_LT_138, POLARS_VERSION_LT_139, @@ -171,21 +171,17 @@ def test_negative_slice_pushdown_raises(engine: pl.GPUEngine, tmp_path): assert_ir_translation_raises(q, engine, NotImplementedError) -@pytest.mark.parametrize("chunked", [False, True], ids=["single_read", "chunked"]) -def test_scan_parquet_prefetch_file_metadata( - tmp_path: Path, df: pl.DataFrame, *, chunked: bool -): - make_partitioned_source(df, tmp_path / "file", "parquet") - q = pl.scan_parquet(tmp_path / "file") - engine = pl.GPUEngine( - executor="in-memory", - raise_on_fail=True, - parquet_options={ - "chunked": chunked, - "prefetch_file_metadata": True, - }, - ) - assert_gpu_result_equal(q, engine=engine) +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: @@ -196,10 +192,7 @@ def test_scan_do_evaluate_missing_prefetch_metadata() -> None: with pytest.raises( AssertionError, - match=( - r"Parquet file metadata was not prefetched for paths: " - r"\['/some/missing/file\.parquet'\]\." - ), + match=(r"Paths do not match cached parquet info."), ): Scan.do_evaluate( schema, @@ -213,6 +206,7 @@ def test_scan_do_evaluate_missing_prefetch_metadata() -> None: None, None, parquet_options, + [], context=context, ) @@ -861,3 +855,47 @@ def test_scan_parquet_is_between_literal_dtype_mismatch_22622( ) assert_gpu_result_equal(q, engine=engine) + + +def test_scan_is_equal() -> None: + parquet_options = ParquetOptions(prefetch_file_metadata=True) + kwargs = { + "schema": {"a": DataType(pl.Int64())}, + "typ": "parquet", + "reader_options": {}, + "cloud_options": {}, + "paths": ["file.parquet"], + "with_columns": None, + "skip_rows": 0, + "n_rows": -1, + "row_index": None, + "include_file_paths": None, + "predicate": None, + "parquet_options": parquet_options, + "cached_parquet_info": None, + } + + scan_1 = Scan(**kwargs) # type: ignore[arg-type] + scan_2 = Scan(**kwargs) # type: ignore[arg-type] + + assert scan_1.is_equal(scan_2) + assert hash(scan_1) == hash(scan_2) + + kwargs3 = { + **kwargs, + "cached_parquet_info": [ + CachedParquetInfo(path="file.parquet", size=100, file_metadata=None) # type: ignore[arg-type] + ], + } + + scan_3 = Scan(**kwargs3) # type: ignore[arg-type] + assert scan_1.is_equal(scan_3) + assert hash(scan_1) == hash(scan_3) + + kwargs4 = { + **kwargs, + "paths": ["file2.parquet"], + } + scan_4 = Scan(**kwargs4) # type: ignore[arg-type] + assert not scan_1.is_equal(scan_4) + assert hash(scan_1) != hash(scan_4) From d6c125d36ca29587f258adc2cff30c3180c439c9 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 30 Jun 2026 13:29:45 -0700 Subject: [PATCH 31/51] self review --- python/cudf_polars/cudf_polars/dsl/ir.py | 69 ++++----- python/cudf_polars/cudf_polars/engine/core.py | 27 ++-- .../cudf_polars/cudf_polars/streaming/io.py | 46 ++---- .../cudf_polars/tests/streaming/test_scan.py | 137 ------------------ python/cudf_polars/tests/test_scan.py | 46 +----- 5 files changed, 58 insertions(+), 267 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 9030b97b2ec5..1789945cb106 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -293,7 +293,7 @@ def prefetch_parquet_file_metadata_for_ir( Returns ------- - A dictionary mapping file paths to their cached parquet metadata. + A dictionary mapping each individual path to its cached parquet metadata. """ from cudf_polars.streaming.io import ParquetSourceInfo, StreamingScan @@ -677,6 +677,9 @@ def __init__( self.children = () self.parquet_options = parquet_options self.cached_parquet_info = cached_parquet_info + + # Scan._validate_cached_parquet_info(self.paths, self.parquet_options, 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 @@ -777,18 +780,22 @@ def with_prefetched_parquet_metadata( cached_parquet_info, ) - def is_equal(self, other: Self) -> bool: # noqa: D102 - # This needs to exclude 'cached_parquet_info' from the equality check. - if self is other: - return True - result = ( - self._ctor_arguments(self.children)[:-1] - == other._ctor_arguments(other.children)[:-1] - ) - # Eager CSE for nodes that match. - if result: - self.children = other.children - return result + @staticmethod + def _validate_cached_parquet_info( + paths: list[str], + parquet_options: ParquetOptions, + cached_parquet_info: list[CachedParquetInfo] | None, + ) -> None: + if parquet_options.prefetch_file_metadata and not cached_parquet_info: + raise AssertionError( + "Prefetching is enabled, but no cached parquet info was provided." + ) + elif cached_parquet_info is not None and paths != [ + info.path for info in cached_parquet_info + ]: + raise AssertionError( + f"Paths do not match cached parquet info. {paths} != {[info.path for info in cached_parquet_info]}" + ) def get_hashable(self) -> Hashable: """ @@ -856,15 +863,10 @@ def _get_parquet_row_count_from_metadata( # Zero-width parquet files lose their row count when read through # pylibcudf. See https://github.com/rapidsai/cudf/issues/21428 if parquet_options.prefetch_file_metadata: - if cached_parquet_info is None: - raise AssertionError( - "Prefetching is enabled, but no cached parquet info was provided." - ) - if paths != [info.path for info in cached_parquet_info]: - raise AssertionError( - f"Paths do not match cached parquet info. {paths} != {[info.path for info in cached_parquet_info]}" - ) - parquet_metadatas = [info.file_metadata for info in cached_parquet_info] + Scan._validate_cached_parquet_info( + paths, parquet_options, cached_parquet_info + ) + parquet_metadatas = [info.file_metadata for info in cached_parquet_info] # type: ignore[union-attr] num_rows = sum(metadata.num_rows for metadata in parquet_metadatas) else: meta = plc.io.parquet_metadata.read_parquet_metadata( @@ -1010,21 +1012,16 @@ def read_csv_header( ) elif typ == "parquet": if parquet_options.prefetch_file_metadata: - if cached_parquet_info is None: - raise AssertionError( - "Prefetching is enabled, but no cached parquet info was provided." - ) - if paths != [info.path for info in cached_parquet_info]: - raise AssertionError( - f"Paths do not match cached parquet info. {paths} != {[info.path for info in cached_parquet_info]}" - ) + Scan._validate_cached_parquet_info( + paths, parquet_options, cached_parquet_info + ) source_info = plc.io.SourceInfo( [ plc.io.types.FilepathSource(info.path, info.size) - for info in cached_parquet_info + for info in cached_parquet_info # type: ignore[union-attr] ] ) - parquet_metadatas = [info.file_metadata for info in cached_parquet_info] + parquet_metadatas = [info.file_metadata for info in cached_parquet_info] # type: ignore[union-attr] else: parquet_metadatas = None source_info = plc.io.SourceInfo(paths) @@ -1100,7 +1097,7 @@ def read_csv_header( col_names = tbl_w_meta.column_names(include_children=False) num_rows = ( cls._get_parquet_row_count_from_metadata( - paths, skip_rows, n_rows, parquet_options, context + paths, skip_rows, n_rows, parquet_options, cached_parquet_info ) if not col_names else None @@ -1779,7 +1776,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.parquet_options, context + 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/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 69241b602eba..ad661e7397f5 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -30,6 +30,7 @@ Scan, prefetch_parquet_file_metadata_for_ir, ) +from cudf_polars.dsl.traversal import traversal from cudf_polars.dsl.utils.replace import replace from cudf_polars.streaming.actor_graph.collectives import ReserveOpIDs from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id @@ -683,7 +684,7 @@ def evaluate_on_rank( metadata Collected channel metadata. """ - from cudf_polars.dsl.traversal import traversal + from cudf_polars.streaming.io import StreamingScan stats = allgather_stats(comm, ctx.br(), ir, config_options, py_executor) ir, partition_info = lower_ir_graph( @@ -705,32 +706,26 @@ def evaluate_on_rank( ir_context.py_executor, stats=stats, ) - # Build a mapping {Scan: Scan} where the values have the prefetched metadata set. - # This is complicated by StreamingScan, SplitScan, and FusedScan, but whatever. + # We'll replace scan nodes with variants that have the prefetched metadata set. + # We also update partition_info to point to the new nodes. replacements: dict[IR, IR] = {} - from cudf_polars.streaming.io import FusedScan, SplitScan, StreamingScan + new_node: Scan | StreamingScan for node in traversal([ir]): if isinstance(node, Scan): - replacements[node] = Scan.with_prefetched_parquet_metadata( + new_node = Scan.with_prefetched_parquet_metadata( node, [cached_parquet_info_map[path] for path in node.paths] ) + replacements[node] = new_node + partition_info[new_node] = partition_info[node] elif isinstance(node, StreamingScan): - x = StreamingScan.with_prefetched_parquet_metadata( + new_node = StreamingScan.with_prefetched_parquet_metadata( node, cached_parquet_info_map ) - assert node in partition_info - assert node.is_equal(x) - assert x in partition_info - replacements[node] = x - elif isinstance(node, (SplitScan, FusedScan)): - raise NotImplementedError( - f"Prefetching is not supported for StreamingScan, SplitScan, and FusedScan.: {type(node)}" - ) + replacements[node] = new_node + partition_info[new_node] = partition_info[node] ir = replace([ir], replacements)[0] - # Now we want to rewrite the Scan nodes. We need to insert the prefetched - # parquet metadata for the paths in that node. with ReserveOpIDs(ir, config_options) as collective_id_map: return execute_ir_on_rank( diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index a27455781735..a2233982f001 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -286,17 +286,6 @@ def with_prefetched_parquet_metadata( cached_parquet_info, ) - def is_equal(self, other: Self) -> bool: # noqa: D102 - # This needs to exclude 'cached_parquet_info' from the equality check. - if type(other) is not type(self): - return False - return self is other or ( - self.base_scan.is_equal(other.base_scan) - and self.paths == other.paths - and self.split_index == other.split_index - and self.total_splits == other.total_splits - ) - def get_hashable(self) -> Hashable: """Hashable representation of the node.""" return ( @@ -488,16 +477,6 @@ def with_prefetched_parquet_metadata( cached_parquet_info, ) - def is_equal(self, other: Self) -> bool: # noqa: D102 - # This needs to exclude 'cached_parquet_info' from the equality check. - if type(other) is not type(self): - return False - return self is other or ( - self.base_scan.is_equal(other.base_scan) - and self.paths == other.paths - and self.parquet_options == other.parquet_options - ) - def get_hashable(self) -> Hashable: """Hashable representation of the node.""" return ( @@ -800,20 +779,17 @@ def for_fused_files( local_offset, local_count = _rank_slice(partition_count, rank, nranks) paths_start = local_offset * plan.factor paths_end = paths_start + plan.factor * local_count - scans = [] - for offset in range(paths_start, paths_end, plan.factor): - paths = base_scan.paths[offset : offset + plan.factor] - if paths: - scans.append( - FusedScan( - base_scan.schema, - base_scan, - paths, - parquet_options, - None, - ) - ) - + scans = [ + FusedScan( + base_scan.schema, + 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, "fused") def get_hashable(self) -> Hashable: diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index d9d682bfc7e2..6c1cef391fcb 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -13,7 +13,6 @@ from cudf_polars import Translator from cudf_polars.containers import DataType from cudf_polars.dsl.ir import ( - CachedParquetInfo, Empty, IRExecutionContext, Scan, @@ -401,139 +400,3 @@ def test_split_scan_do_evaluate_missing_prefetch_metadata() -> None: [], context=context, ) - - -def test_split_scan_is_equal() -> None: - parquet_options = ParquetOptions(prefetch_file_metadata=True) - base_scan = _make_parquet_scan(["file.parquet"], parquet_options=parquet_options) - cached_info = [ - CachedParquetInfo(path="file.parquet", size=100, file_metadata=None) # type: ignore[arg-type] - ] - - scan_1 = SplitScan( - base_scan.schema, - base_scan, - ["file.parquet"], - 0, - 2, - parquet_options, - None, - ) - scan_2 = SplitScan( - base_scan.schema, - base_scan, - ["file.parquet"], - 0, - 2, - parquet_options, - cached_info, - ) - scan_3 = SplitScan( - base_scan.schema, - base_scan, - ["file.parquet"], - 1, - 2, - parquet_options, - None, - ) - - assert scan_1.is_equal(scan_2) - assert scan_1.get_hashable() == scan_2.get_hashable() - assert hash(scan_1) == hash(scan_2) - - assert not scan_1.is_equal(scan_3) - assert scan_1.get_hashable() != scan_3.get_hashable() - assert hash(scan_1) != hash(scan_3) - - -def test_fused_scan_is_equal() -> None: - parquet_options = ParquetOptions(prefetch_file_metadata=True) - base_scan = _make_parquet_scan( - ["file_1.parquet", "file_2.parquet"], parquet_options=parquet_options - ) - other_base_scan = _make_parquet_scan( - ["other.parquet"], parquet_options=parquet_options - ) - cached_info = [ - CachedParquetInfo(path="file_1.parquet", size=100, file_metadata=None), # type: ignore[arg-type] - CachedParquetInfo(path="file_2.parquet", size=200, file_metadata=None), # type: ignore[arg-type] - ] - - scan_1 = FusedScan( - base_scan.schema, - base_scan, - ["file_1.parquet", "file_2.parquet"], - parquet_options, - None, - ) - scan_2 = FusedScan( - base_scan.schema, - base_scan, - ["file_1.parquet", "file_2.parquet"], - parquet_options, - cached_info, - ) - scan_3 = FusedScan( - other_base_scan.schema, - other_base_scan, - ["other.parquet"], - parquet_options, - None, - ) - - assert scan_1.is_equal(scan_2) - assert scan_1.get_hashable() == scan_2.get_hashable() - assert hash(scan_1) == hash(scan_2) - - assert not scan_1.is_equal(scan_3) - assert scan_1.get_hashable() != scan_3.get_hashable() - assert hash(scan_1) != hash(scan_3) - - -def test_streaming_scan_is_equal() -> None: - parquet_options = ParquetOptions(prefetch_file_metadata=True) - base_scan = _make_parquet_scan(["file.parquet"], parquet_options=parquet_options) - cached_info = [ - CachedParquetInfo(path="file.parquet", size=100, file_metadata=None) # type: ignore[arg-type] - ] - - split_1 = SplitScan( - base_scan.schema, - base_scan, - ["file.parquet"], - 0, - 2, - parquet_options, - None, - ) - split_2 = SplitScan( - base_scan.schema, - base_scan, - ["file.parquet"], - 0, - 2, - parquet_options, - cached_info, - ) - split_3 = SplitScan( - base_scan.schema, - base_scan, - ["file.parquet"], - 1, - 2, - parquet_options, - None, - ) - - streaming_1 = StreamingScan([split_1], base_scan, "split") - streaming_2 = StreamingScan([split_2], base_scan, "split") - streaming_3 = StreamingScan([split_3], base_scan, "split") - - assert streaming_1.is_equal(streaming_2) - assert streaming_1.get_hashable() == streaming_2.get_hashable() - assert hash(streaming_1) == hash(streaming_2) - - assert not streaming_1.is_equal(streaming_3) - assert streaming_1.get_hashable() != streaming_3.get_hashable() - assert hash(streaming_1) != hash(streaming_3) diff --git a/python/cudf_polars/tests/test_scan.py b/python/cudf_polars/tests/test_scan.py index 0ae0d6317816..abd9b518e7f4 100644 --- a/python/cudf_polars/tests/test_scan.py +++ b/python/cudf_polars/tests/test_scan.py @@ -17,7 +17,7 @@ import polars as pl from cudf_polars.containers import DataType -from cudf_polars.dsl.ir import CachedParquetInfo, IRExecutionContext, Scan +from cudf_polars.dsl.ir import IRExecutionContext, Scan from cudf_polars.testing.asserts import ( assert_gpu_result_equal, assert_ir_translation_raises, @@ -855,47 +855,3 @@ def test_scan_parquet_is_between_literal_dtype_mismatch_22622( ) assert_gpu_result_equal(q, engine=engine) - - -def test_scan_is_equal() -> None: - parquet_options = ParquetOptions(prefetch_file_metadata=True) - kwargs = { - "schema": {"a": DataType(pl.Int64())}, - "typ": "parquet", - "reader_options": {}, - "cloud_options": {}, - "paths": ["file.parquet"], - "with_columns": None, - "skip_rows": 0, - "n_rows": -1, - "row_index": None, - "include_file_paths": None, - "predicate": None, - "parquet_options": parquet_options, - "cached_parquet_info": None, - } - - scan_1 = Scan(**kwargs) # type: ignore[arg-type] - scan_2 = Scan(**kwargs) # type: ignore[arg-type] - - assert scan_1.is_equal(scan_2) - assert hash(scan_1) == hash(scan_2) - - kwargs3 = { - **kwargs, - "cached_parquet_info": [ - CachedParquetInfo(path="file.parquet", size=100, file_metadata=None) # type: ignore[arg-type] - ], - } - - scan_3 = Scan(**kwargs3) # type: ignore[arg-type] - assert scan_1.is_equal(scan_3) - assert hash(scan_1) == hash(scan_3) - - kwargs4 = { - **kwargs, - "paths": ["file2.parquet"], - } - scan_4 = Scan(**kwargs4) # type: ignore[arg-type] - assert not scan_1.is_equal(scan_4) - assert hash(scan_1) != hash(scan_4) From 36cb18113f7b4488227d379af2ce9e7b010c39ce Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 30 Jun 2026 13:34:50 -0700 Subject: [PATCH 32/51] move --- python/cudf_polars/cudf_polars/dsl/ir.py | 159 +--------------- .../cudf_polars/cudf_polars/dsl/utils/io.py | 174 ++++++++++++++++++ python/cudf_polars/cudf_polars/engine/core.py | 2 +- .../cudf_polars/cudf_polars/streaming/io.py | 2 +- .../cudf_polars/tests/streaming/test_scan.py | 2 +- 5 files changed, 179 insertions(+), 160 deletions(-) create mode 100644 python/cudf_polars/cudf_polars/dsl/utils/io.py diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 1789945cb106..34d993dff2dc 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -14,7 +14,6 @@ from __future__ import annotations import asyncio -import concurrent.futures import contextlib import contextvars import functools @@ -49,7 +48,6 @@ from cudf_polars.dsl.nodebase import Node from cudf_polars.dsl.to_ast import _DECIMAL_IDS, to_ast, to_parquet_filter from cudf_polars.dsl.tracing import log_do_evaluate, nvtx_annotate_cudf_polars -from cudf_polars.dsl.traversal import traversal from cudf_polars.dsl.utils.reshape import broadcast from cudf_polars.dsl.utils.windows import ( offsets_to_windows, @@ -67,6 +65,7 @@ ) if TYPE_CHECKING: + import concurrent.futures from collections.abc import Callable, Generator, Hashable, Iterable, Sequence from typing import Literal, Self @@ -75,7 +74,7 @@ from rmm.pylibrmm.stream import Stream from cudf_polars.containers.dataframe import NamedColumn - from cudf_polars.streaming.base import StatsCollector + from cudf_polars.dsl.utils.io import CachedParquetInfo from cudf_polars.typing import CSECache, ClosedInterval, Schema, Slice as Zlice from cudf_polars.utils.config import ParquetOptions from cudf_polars.utils.timer import Timer @@ -112,34 +111,6 @@ ] -@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 - - @dataclass(frozen=True) class IRExecutionContext: """ @@ -215,132 +186,6 @@ def stream_ordered_after(self, *dfs: DataFrame) -> Generator[Stream, None, None] yield result_stream -@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: - 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): - # 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": - for path in node.paths: - all_paths.add(path) - - 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 - - _BINOPS = { plc.binaryop.BinaryOperator.EQUAL, plc.binaryop.BinaryOperator.NOT_EQUAL, 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..f058a01ee8e3 --- /dev/null +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -0,0 +1,174 @@ +# 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 + +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: + 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): + # 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": + for path in node.paths: + all_paths.add(path) + + 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 diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index ad661e7397f5..b2058b9fc3ec 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -28,9 +28,9 @@ from cudf_polars.dsl.ir import ( IRExecutionContext, Scan, - prefetch_parquet_file_metadata_for_ir, ) from cudf_polars.dsl.traversal import traversal +from cudf_polars.dsl.utils.io import prefetch_parquet_file_metadata_for_ir from cudf_polars.dsl.utils.replace import replace from cudf_polars.streaming.actor_graph.collectives import ReserveOpIDs from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index a2233982f001..54ef50780ddf 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -23,9 +23,9 @@ Empty, Scan, Sink, - _prefetch_parquet_footers_for_paths, ) from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars +from cudf_polars.dsl.utils.io import _prefetch_parquet_footers_for_paths from cudf_polars.streaming.base import ( IOPartitionFlavor, IOPartitionPlan, diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 6c1cef391fcb..1819806d4c4e 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -16,8 +16,8 @@ Empty, IRExecutionContext, Scan, - prefetch_parquet_file_metadata_for_ir, ) +from cudf_polars.dsl.utils.io import prefetch_parquet_file_metadata_for_ir from cudf_polars.engine.options import StreamingOptions from cudf_polars.streaming.base import IOPartitionFlavor, IOPartitionPlan from cudf_polars.streaming.io import ( From ef8a87e491b421fad2474af01b5867a86168b1d0 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 30 Jun 2026 13:51:39 -0700 Subject: [PATCH 33/51] One more --- python/cudf_polars/cudf_polars/streaming/io.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 54ef50780ddf..152b62d2c982 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -465,10 +465,9 @@ def with_prefetched_parquet_metadata( ) -> Self: """Create a new FusedScan node, with prefetched parquet metadata set.""" cached_parquet_info = [cached_parquet_info_map[path] for path in node.paths] - if node.paths != [info.path for info in cached_parquet_info]: - raise AssertionError( - f"Paths do not match cached parquet info. {node.paths} != {[info.path for info in cached_parquet_info]}" - ) + Scan._validate_cached_parquet_info( + node.paths, node.parquet_options, cached_parquet_info + ) return cls( node.schema, node.base_scan, From 479036eb38e1dcb52d629d935564349e5b668bbe Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 30 Jun 2026 13:53:55 -0700 Subject: [PATCH 34/51] fixup --- python/cudf_polars/cudf_polars/streaming/io.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 152b62d2c982..50683e850681 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -25,7 +25,6 @@ Sink, ) from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars -from cudf_polars.dsl.utils.io import _prefetch_parquet_footers_for_paths from cudf_polars.streaming.base import ( IOPartitionFlavor, IOPartitionPlan, @@ -1032,6 +1031,8 @@ class ParquetMetadata: @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 From c4a6fe016c5ab8ca08115725a6031189365699b3 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 30 Jun 2026 13:54:16 -0700 Subject: [PATCH 35/51] fixup --- python/cudf_polars/cudf_polars/streaming/io.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 50683e850681..4bd91a7faae4 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -274,7 +274,6 @@ def with_prefetched_parquet_metadata( new_base = Scan.with_prefetched_parquet_metadata( node.base_scan, cached_parquet_info ) - # assert new_base.paths == [info.path for info in cached_parquet_info] return cls( node.schema, new_base, From 44f0eb708d9d4d21d51acc337a800c1c024afb1d Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 1 Jul 2026 06:01:18 -0700 Subject: [PATCH 36/51] cleanup --- python/cudf_polars/cudf_polars/dsl/ir.py | 41 ++++++++++--------- .../cudf_polars/cudf_polars/streaming/io.py | 26 +++++------- .../cudf_polars/tests/streaming/test_stats.py | 4 -- 3 files changed, 32 insertions(+), 39 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 34d993dff2dc..49c4abd929e6 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -523,7 +523,7 @@ def __init__( self.parquet_options = parquet_options self.cached_parquet_info = cached_parquet_info - # Scan._validate_cached_parquet_info(self.paths, self.parquet_options, self.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 @@ -628,14 +628,9 @@ def with_prefetched_parquet_metadata( @staticmethod def _validate_cached_parquet_info( paths: list[str], - parquet_options: ParquetOptions, cached_parquet_info: list[CachedParquetInfo] | None, ) -> None: - if parquet_options.prefetch_file_metadata and not cached_parquet_info: - raise AssertionError( - "Prefetching is enabled, but no cached parquet info was provided." - ) - elif cached_parquet_info is not None and paths != [ + if cached_parquet_info is not None and paths != [ info.path for info in cached_parquet_info ]: raise AssertionError( @@ -708,10 +703,13 @@ def _get_parquet_row_count_from_metadata( # Zero-width parquet files lose their row count when read through # pylibcudf. See https://github.com/rapidsai/cudf/issues/21428 if parquet_options.prefetch_file_metadata: - Scan._validate_cached_parquet_info( - paths, parquet_options, cached_parquet_info - ) - parquet_metadatas = [info.file_metadata for info in cached_parquet_info] # type: ignore[union-attr] + 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] num_rows = sum(metadata.num_rows for metadata in parquet_metadatas) else: meta = plc.io.parquet_metadata.read_parquet_metadata( @@ -857,16 +855,19 @@ def read_csv_header( ) elif typ == "parquet": if parquet_options.prefetch_file_metadata: - Scan._validate_cached_parquet_info( - paths, parquet_options, cached_parquet_info - ) - source_info = plc.io.SourceInfo( - [ + 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) - for info in cached_parquet_info # type: ignore[union-attr] - ] - ) - parquet_metadatas = [info.file_metadata for info in cached_parquet_info] # type: ignore[union-attr] + ) + parquet_metadatas.append(info.file_metadata) + source_info = plc.io.SourceInfo(filepath_sources) else: parquet_metadatas = None source_info = plc.io.SourceInfo(paths) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 4bd91a7faae4..b87b7c57a8a4 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -251,7 +251,7 @@ def __init__( def with_prefetched_parquet_metadata( cls, node: SplitScan, - cached_parquet_info: list[CachedParquetInfo], + cached_parquet_info_map: dict[str, CachedParquetInfo], ) -> Self: """ Create a new SplitScan node, with prefetched parquet metadata set. @@ -263,20 +263,21 @@ def with_prefetched_parquet_metadata( ---------- node The SplitScan node to create a new node from. - cached_parquet_info - The cached parquet metadata to set on the new node. This will be a - length-1 list, matching the length-1 ``path`` for the base scan node. + cached_parquet_info_map + A dictionary mapping file paths to cached parquet metadata. This should contain + all the file paths, including those from the base scan (which has been split + into multiple SplitScan nodes). Returns ------- The new SplitScan node. """ - new_base = Scan.with_prefetched_parquet_metadata( - node.base_scan, cached_parquet_info - ) + # cached_parquet_info_map *might* not contain all the paths for the base scan, + # if, e.g., this worker has only been assigned a subset of the paths. + cached_parquet_info = [cached_parquet_info_map[path] for path in node.paths] return cls( node.schema, - new_base, + node.base_scan, node.paths, node.split_index, node.total_splits, @@ -463,9 +464,7 @@ def with_prefetched_parquet_metadata( ) -> Self: """Create a new FusedScan node, with prefetched parquet metadata set.""" cached_parquet_info = [cached_parquet_info_map[path] for path in node.paths] - Scan._validate_cached_parquet_info( - node.paths, node.parquet_options, cached_parquet_info - ) + Scan._validate_cached_parquet_info(node.paths, cached_parquet_info) return cls( node.schema, node.base_scan, @@ -694,13 +693,10 @@ def with_prefetched_parquet_metadata( if node.scan_type == "split": new_scans = [] for scan in node.scans: - new_parquet_info = [ - cached_parquet_info_map[path] for path in scan.paths - ] # SplitScan should be generic / overload based on type. new_scan = SplitScan.with_prefetched_parquet_metadata( scan, # type: ignore[arg-type] - new_parquet_info, + cached_parquet_info_map, ) assert new_scan.cached_parquet_info is not None assert new_scan.paths == [ diff --git a/python/cudf_polars/tests/streaming/test_stats.py b/python/cudf_polars/tests/streaming/test_stats.py index 507974fd2646..22c7a1795ee1 100644 --- a/python/cudf_polars/tests/streaming/test_stats.py +++ b/python/cudf_polars/tests/streaming/test_stats.py @@ -290,10 +290,6 @@ def test_parquet_source_info_omits_footers_when_paths_are_sampled( assert info.cached_parquet_info is None -# TODO: Add a test that calls prefetch_parquet_file_metadata_for_ir -# and ensure that we don't refetch the metadata from stats - - def test_parquet_metadata_reads_footers( tmp_path: pathlib.Path, df_and_schema: tuple[pl.DataFrame, Schema], From fbdfa0bf0713b26ad6f8b1a955ed547f4932cd40 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 1 Jul 2026 09:27:51 -0700 Subject: [PATCH 37/51] Fix fast_count --- python/cudf_polars/cudf_polars/streaming/select.py | 7 +++++-- python/cudf_polars/tests/streaming/test_scan.py | 14 ++++++++++++++ 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/select.py b/python/cudf_polars/cudf_polars/streaming/select.py index df946e653851..a4bb5f2ad0bf 100644 --- a/python/cudf_polars/cudf_polars/streaming/select.py +++ b/python/cudf_polars/cudf_polars/streaming/select.py @@ -4,6 +4,7 @@ from __future__ import annotations +import dataclasses from collections import defaultdict from typing import TYPE_CHECKING @@ -438,8 +439,10 @@ def _( scan_child.paths, scan_child.skip_rows, scan_child.n_rows, - scan_child.parquet_options, - context=None, + dataclasses.replace( + scan_child.parquet_options, prefetch_file_metadata=False + ), + None, ) dtype = ir.exprs[0].value.dtype diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 1819806d4c4e..05bd4fe4d886 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -127,6 +127,20 @@ def test_prefetch_parquet_file_metadata_no_parquet_scans() -> 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 # --------------------------------------------------------------------------- From 95430fded59a0cad9139b8576c6f80cc96b5bb9a Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 1 Jul 2026 09:37:21 -0700 Subject: [PATCH 38/51] hash bug fix --- python/cudf_polars/cudf_polars/engine/core.py | 9 +++++++-- python/cudf_polars/tests/streaming/test_scan.py | 16 ++++++++++++++++ 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index b2058b9fc3ec..cd70b8175cd1 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -717,15 +717,20 @@ def evaluate_on_rank( node, [cached_parquet_info_map[path] for path in node.paths] ) replacements[node] = new_node - partition_info[new_node] = partition_info[node] elif isinstance(node, StreamingScan): new_node = StreamingScan.with_prefetched_parquet_metadata( node, cached_parquet_info_map ) replacements[node] = new_node - partition_info[new_node] = partition_info[node] + old_ir = ir ir = replace([ir], replacements)[0] + partition_info = { + new_node: partition_info[old_node] + for old_node, new_node in zip( + traversal([old_ir]), traversal([ir]), strict=True + ) + } with ReserveOpIDs(ir, config_options) as collective_id_map: return execute_ir_on_rank( diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 05bd4fe4d886..87ddd9c7d286 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -414,3 +414,19 @@ def test_split_scan_do_evaluate_missing_prefetch_metadata() -> None: [], 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) From d01be7492816a2274dd85345a1163c5fc0a7e961 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 1 Jul 2026 11:45:17 -0700 Subject: [PATCH 39/51] Truncate error message --- python/cudf_polars/cudf_polars/dsl/ir.py | 6 +++- python/cudf_polars/tests/test_select.py | 45 ++---------------------- 2 files changed, 8 insertions(+), 43 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 49c4abd929e6..b13c5dbbd97b 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 dataclasses import dataclass, field @@ -633,8 +634,11 @@ def _validate_cached_parquet_info( 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. {paths} != {[info.path for info in cached_parquet_info]}" + f"Paths do not match cached parquet info. Missing paths: {missing}" ) def get_hashable(self) -> Hashable: diff --git a/python/cudf_polars/tests/test_select.py b/python/cudf_polars/tests/test_select.py index 6365d34115af..beb95e882f4e 100644 --- a/python/cudf_polars/tests/test_select.py +++ b/python/cudf_polars/tests/test_select.py @@ -8,7 +8,7 @@ import polars as pl -from cudf_polars.dsl.ir import IRExecutionContext, Scan +from cudf_polars.dsl.ir import Scan from cudf_polars.testing.asserts import ( assert_gpu_result_equal, assert_ir_translation_raises, @@ -159,15 +159,6 @@ def parquet_fast_count_df() -> pl.DataFrame: return pl.DataFrame({"a": range(PARQUET_FAST_COUNT_ROWS)}) -@pytest.fixture -def prefetch_engine() -> pl.GPUEngine: - return pl.GPUEngine( - executor="in-memory", - raise_on_fail=True, - parquet_options={"prefetch_file_metadata": True}, - ) - - @pytest.fixture( params=[ pytest.param({"skip_rows": 0, "n_rows": None}, id="all_rows"), @@ -185,48 +176,18 @@ def parquet_scan_row_bounds(request) -> dict[str, int | None]: return request.param -def test_select_fast_count_parquet_prefetch_metadata( - tmp_path, - parquet_fast_count_df: pl.DataFrame, - prefetch_engine: pl.GPUEngine, - parquet_scan_row_bounds: dict[str, int | None], -) -> None: - skip_rows = parquet_scan_row_bounds["skip_rows"] - assert skip_rows is not None - n_rows = parquet_scan_row_bounds["n_rows"] - - file = tmp_path / "data.parquet" - parquet_fast_count_df.write_parquet(file) - - if skip_rows == 0 and n_rows is None: - q = pl.scan_parquet(file) - elif skip_rows == 0: - q = pl.scan_parquet(file, n_rows=n_rows) - elif n_rows is None: - q = pl.scan_parquet(file).slice(skip_rows) - else: - q = pl.scan_parquet(file).slice(skip_rows, n_rows) - - q = q.select(pl.len()) - assert_gpu_result_equal(q, engine=prefetch_engine) - - def test_get_parquet_row_count_from_metadata_missing_prefetch() -> None: paths = ["/some/missing/file.parquet"] parquet_options = ParquetOptions(prefetch_file_metadata=True) - context = IRExecutionContext() with pytest.raises( AssertionError, - match=( - r"Parquet file metadata was not prefetched for paths: " - r"\['/some/missing/file\.parquet'\]\." - ), + 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, - context=context, + cached_parquet_info=[], ) From 4e7f4dc5fe16212ec6c737cd87b540d038c16073 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 1 Jul 2026 13:20:44 -0700 Subject: [PATCH 40/51] Coverage, simplification --- python/cudf_polars/cudf_polars/dsl/ir.py | 40 +++++++++---------- .../cudf_polars/cudf_polars/dsl/utils/io.py | 5 +-- python/cudf_polars/cudf_polars/engine/core.py | 12 +++--- .../cudf_polars/tests/streaming/test_scan.py | 21 ++++++++++ python/cudf_polars/tests/test_select.py | 11 ++++- 5 files changed, 59 insertions(+), 30 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index b13c5dbbd97b..87c27af9573f 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -605,26 +605,26 @@ def __init__( "Reading only parquet metadata to produce row index." ) - @classmethod - def with_prefetched_parquet_metadata( - cls, scan: Scan, cached_parquet_info: list[CachedParquetInfo] - ) -> Self: - """Create a new scan node, with prefetched parquet metadata set.""" - return cls( - scan.schema, - scan.typ, - scan.reader_options, - scan.cloud_options, - scan.paths, - scan.with_columns, - scan.skip_rows, - scan.n_rows, - scan.row_index, - scan.include_file_paths, - scan.predicate, - scan.parquet_options, - cached_parquet_info, - ) + # @classmethod + # def with_prefetched_parquet_metadata( + # cls, scan: Scan, cached_parquet_info: list[CachedParquetInfo] + # ) -> Self: + # """Create a new scan node, with prefetched parquet metadata set.""" + # return cls( + # scan.schema, + # scan.typ, + # scan.reader_options, + # scan.cloud_options, + # scan.paths, + # scan.with_columns, + # scan.skip_rows, + # scan.n_rows, + # scan.row_index, + # scan.include_file_paths, + # scan.predicate, + # scan.parquet_options, + # cached_parquet_info, + # ) @staticmethod def _validate_cached_parquet_info( diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index f058a01ee8e3..8450d026b20d 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -137,9 +137,8 @@ def prefetch_parquet_file_metadata_for_ir( for scan in node.scans: for path in scan.paths: all_paths.add(path) - elif isinstance(node, Scan) and node.typ == "parquet": - for path in node.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: diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index cd70b8175cd1..403c8890e7ad 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -712,17 +712,17 @@ def evaluate_on_rank( new_node: Scan | StreamingScan for node in traversal([ir]): - if isinstance(node, Scan): - new_node = Scan.with_prefetched_parquet_metadata( - node, [cached_parquet_info_map[path] for path in node.paths] - ) - replacements[node] = new_node - elif isinstance(node, StreamingScan): + if isinstance(node, StreamingScan): new_node = StreamingScan.with_prefetched_parquet_metadata( node, cached_parquet_info_map ) replacements[node] = new_node + elif isinstance(node, Scan) and node.typ == "parquet": # pragma: no cover + raise RuntimeError( + "Unexpected parquet 'Scan' node in lowered IR graph." + ) + old_ir = ir ir = replace([ir], replacements)[0] partition_info = { diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 87ddd9c7d286..0b7e262f7fe1 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -354,6 +354,27 @@ def test_scan_missing_prefetch_metadata_raises() -> None: ["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", diff --git a/python/cudf_polars/tests/test_select.py b/python/cudf_polars/tests/test_select.py index beb95e882f4e..435bec5d2031 100644 --- a/python/cudf_polars/tests/test_select.py +++ b/python/cudf_polars/tests/test_select.py @@ -176,10 +176,19 @@ def parquet_scan_row_bounds(request) -> dict[str, int | None]: return request.param -def test_get_parquet_row_count_from_metadata_missing_prefetch() -> None: +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."), From 6f62dd240c1ec6238746d851fbaa8c78363a7626 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 1 Jul 2026 13:24:16 -0700 Subject: [PATCH 41/51] kvikio pragmas --- python/cudf_polars/cudf_polars/dsl/utils/io.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index 8450d026b20d..74b53a54a75b 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -72,13 +72,15 @@ def _prefetch_parquet_footers_for_paths(paths: list[str]) -> list[CachedParquetI # For now, we'll just use kvikio to explicitly get the size. sizes = [] - try: + 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): + 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 From 848b48d6b5f1df221edf3e36b9110c795c2f6157 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 1 Jul 2026 13:26:57 -0700 Subject: [PATCH 42/51] coverage --- python/cudf_polars/tests/streaming/test_scan.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 0b7e262f7fe1..41cc450ad3f2 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -94,7 +94,7 @@ def test_scan_parquet_use_rapidsmpf_native(tmp_path, df, streaming_engine_factor @pytest.mark.parametrize( - "target_partition_size_and_n_files", [(1_000, 1), (1_000, 2), (1_000_000, 3)] + "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, From 2734181661faef4059efde6e489ca7d5062d9f1e Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Thu, 2 Jul 2026 09:57:45 -0700 Subject: [PATCH 43/51] Workaround Projection(Cache(StreamingScan)) issue --- .../cudf_polars/dsl/utils/replace.py | 18 ++- .../cudf_polars/cudf_polars/streaming/io.py | 19 ++- .../cudf_polars/tests/streaming/test_scan.py | 116 +++++++++++++++++- 3 files changed, 145 insertions(+), 8 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/replace.py b/python/cudf_polars/cudf_polars/dsl/utils/replace.py index e5006f14fe20..eedb409ed307 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/replace.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/replace.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 """Utilities for replacing nodes in a DAG.""" @@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Generic -from cudf_polars.dsl.traversal import CachingVisitor, reuse_if_unchanged +from cudf_polars.dsl.traversal import CachingVisitor from cudf_polars.typing import NodeT, TypedDict if TYPE_CHECKING: @@ -39,9 +39,19 @@ def _replace( # in translate.py, which also skips code coverage # See the TODO there for more details. try: - return fn.state["replacements"][node] + r = fn.state["replacements"][node] except KeyError: - return reuse_if_unchanged(node, fn) + # replacement must propagate when children are rebuilt, + # even if child __eq__ intentionally ignores children (e.g. Cache). + # So we use identity, not equality like reuse_if_unchanged, to decide unchanged. + new_children = [fn(c) for c in node.children] + if all( + new is old for new, old in zip(new_children, node.children, strict=True) + ): + return node + return node.reconstruct(new_children) + else: + return r def replace( diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index b87b7c57a8a4..ce45fa87d099 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -287,14 +287,22 @@ def with_prefetched_parquet_metadata( def get_hashable(self) -> Hashable: """Hashable representation of the node.""" + if self.cached_parquet_info is not None: + cached_parquet_info = tuple( + (info.path, info.size) for info in self.cached_parquet_info + ) + else: + cached_parquet_info = () + return ( - type(self), + type(self), # type: ignore[arg-type] tuple(self.schema.items()), self.base_scan.get_hashable(), tuple(self.paths), self.split_index, self.total_splits, self.parquet_options, + *cached_parquet_info, ) @classmethod @@ -475,12 +483,19 @@ def with_prefetched_parquet_metadata( def get_hashable(self) -> Hashable: """Hashable representation of the node.""" + if self.cached_parquet_info is not None: + cached_parquet_info = tuple( + (info.path, info.size) for info in self.cached_parquet_info + ) + else: + cached_parquet_info = () return ( - type(self), + type(self), # type: ignore[arg-type] tuple(self.schema.items()), self.base_scan.get_hashable(), tuple(self.paths), self.parquet_options, + *cached_parquet_info, ) @classmethod diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 41cc450ad3f2..cb4cae976002 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 @@ -17,7 +17,10 @@ IRExecutionContext, Scan, ) -from cudf_polars.dsl.utils.io import prefetch_parquet_file_metadata_for_ir +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 IOPartitionFlavor, IOPartitionPlan from cudf_polars.streaming.io import ( @@ -39,6 +42,8 @@ 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 @@ -451,3 +456,110 @@ def test_prefetch_file_metadata_join( 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=10), + ) + split_diff_metadata = SplitScan( + base.schema, + base, + base.paths, + 0, + 2, + base.parquet_options, + _make_cached_parquet_info(base.paths, size=11), + ) + + a = StreamingScan([split], base, "split") + b = StreamingScan([split_same], base, "split") + c = StreamingScan([split_diff_metadata], base, "split") + + assert a == b + assert hash(a) == hash(b) + assert a != c From e8f853434a82ba4f13c1c02b32dc116e2d70dfb9 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Thu, 2 Jul 2026 11:32:57 -0700 Subject: [PATCH 44/51] coverage --- python/cudf_polars/cudf_polars/dsl/ir.py | 29 +++++------------------- 1 file changed, 6 insertions(+), 23 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 71470750fb25..73c485c7619a 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -823,27 +823,6 @@ def __init__( "Reading only parquet metadata to produce row index." ) - # @classmethod - # def with_prefetched_parquet_metadata( - # cls, scan: Scan, cached_parquet_info: list[CachedParquetInfo] - # ) -> Self: - # """Create a new scan node, with prefetched parquet metadata set.""" - # return cls( - # scan.schema, - # scan.typ, - # scan.reader_options, - # scan.cloud_options, - # scan.paths, - # scan.with_columns, - # scan.skip_rows, - # scan.n_rows, - # scan.row_index, - # scan.include_file_paths, - # scan.predicate, - # scan.parquet_options, - # cached_parquet_info, - # ) - @staticmethod def _validate_cached_parquet_info( paths: list[str], @@ -931,8 +910,12 @@ def _get_parquet_row_count_from_metadata( ) Scan._validate_cached_parquet_info(paths, cached_parquet_info) - parquet_metadatas = [info.file_metadata for info in cached_parquet_info] - num_rows = sum(metadata.num_rows for metadata in parquet_metadatas) + 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) From bd0245047c984e10ae5dd06e6efdbce08dac4f4e Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 6 Jul 2026 06:31:42 -0700 Subject: [PATCH 45/51] revert replace --- python/cudf_polars/cudf_polars/dsl/utils/replace.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/replace.py b/python/cudf_polars/cudf_polars/dsl/utils/replace.py index eedb409ed307..6f5716aada4d 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/replace.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/replace.py @@ -39,7 +39,7 @@ def _replace( # in translate.py, which also skips code coverage # See the TODO there for more details. try: - r = fn.state["replacements"][node] + return fn.state["replacements"][node] except KeyError: # replacement must propagate when children are rebuilt, # even if child __eq__ intentionally ignores children (e.g. Cache). @@ -50,8 +50,6 @@ def _replace( ): return node return node.reconstruct(new_children) - else: - return r def replace( From c8e77a0eedf680c06939596090fa05bf1ee7b93a Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 6 Jul 2026 09:23:52 -0700 Subject: [PATCH 46/51] Update hashing, equality of prefetched metadata We treat this as an optimization, so exclude prefetched metadata from hashing & equality of Scan-type nodes. This makes it safe to mutate the Scan nodes after theyre' created by lowering, which simplifies updating the IR graph after metadata have been prefetched. --- python/cudf_polars/cudf_polars/dsl/ir.py | 3 +- .../cudf_polars/cudf_polars/dsl/utils/io.py | 36 ++++- .../cudf_polars/dsl/utils/replace.py | 14 +- python/cudf_polars/cudf_polars/engine/core.py | 38 +---- .../cudf_polars/cudf_polars/streaming/io.py | 133 +----------------- .../cudf_polars/tests/streaming/test_scan.py | 48 ++++++- 6 files changed, 93 insertions(+), 179 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index e3d200f7827c..615f3ff9e18f 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -666,7 +666,6 @@ class Scan(IR): "include_file_paths", "predicate", "parquet_options", - "cached_parquet_info", ) _n_non_child_args = 12 typ: str @@ -711,7 +710,7 @@ def __init__( include_file_paths: str | None, predicate: expr.NamedExpr | None, parquet_options: ParquetOptions, - cached_parquet_info: list[CachedParquetInfo] | None, + cached_parquet_info: list[CachedParquetInfo] | None = None, ): self.schema = schema self.typ = typ diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index 74b53a54a75b..f8af29cc6570 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -13,7 +13,7 @@ from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars from cudf_polars.dsl.traversal import traversal -from cudf_polars.streaming.io import Scan +from cudf_polars.streaming.io import FusedScan, Scan, SplitScan, StreamingScan if TYPE_CHECKING: from cudf_polars.dsl.ir import IR @@ -173,3 +173,37 @@ def prefetch_parquet_file_metadata_for_ir( for info in future.result(): cached_parquet_info[info.path] = info return cached_parquet_info + + +def _attach_cached_parquet_info( + node: SplitScan | FusedScan, + cached_parquet_info_map: dict[str, CachedParquetInfo], +) -> None: + cached = [cached_parquet_info_map[path] for path in node.paths] + Scan._validate_cached_parquet_info(node.paths, cached) + node.cached_parquet_info = cached + node._non_child_args = (*node._non_child_args[:-1], cached) + + +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: + _attach_cached_parquet_info(scan, cached_parquet_info_map) + elif isinstance(node, SplitScan | FusedScan): + _attach_cached_parquet_info(node, cached_parquet_info_map) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/replace.py b/python/cudf_polars/cudf_polars/dsl/utils/replace.py index 6f5716aada4d..e5006f14fe20 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/replace.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/replace.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 """Utilities for replacing nodes in a DAG.""" @@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Generic -from cudf_polars.dsl.traversal import CachingVisitor +from cudf_polars.dsl.traversal import CachingVisitor, reuse_if_unchanged from cudf_polars.typing import NodeT, TypedDict if TYPE_CHECKING: @@ -41,15 +41,7 @@ def _replace( try: return fn.state["replacements"][node] except KeyError: - # replacement must propagate when children are rebuilt, - # even if child __eq__ intentionally ignores children (e.g. Cache). - # So we use identity, not equality like reuse_if_unchanged, to decide unchanged. - new_children = [fn(c) for c in node.children] - if all( - new is old for new, old in zip(new_children, node.children, strict=True) - ): - return node - return node.reconstruct(new_children) + return reuse_if_unchanged(node, fn) def replace( diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 403c8890e7ad..bc82c3bc2ec5 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -25,13 +25,11 @@ from rapidsmpf.streaming.core.actor import run_actor_network from cudf_polars.containers import DataFrame -from cudf_polars.dsl.ir import ( - IRExecutionContext, - Scan, +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.dsl.traversal import traversal -from cudf_polars.dsl.utils.io import prefetch_parquet_file_metadata_for_ir -from cudf_polars.dsl.utils.replace import replace 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 @@ -684,8 +682,6 @@ def evaluate_on_rank( metadata Collected channel metadata. """ - from cudf_polars.streaming.io import StreamingScan - stats = allgather_stats(comm, ctx.br(), ir, config_options, py_executor) ir, partition_info = lower_ir_graph( ir, config_options, stats, rank=comm.rank, nranks=comm.nranks @@ -706,31 +702,7 @@ def evaluate_on_rank( ir_context.py_executor, stats=stats, ) - # We'll replace scan nodes with variants that have the prefetched metadata set. - # We also update partition_info to point to the new nodes. - replacements: dict[IR, IR] = {} - - new_node: Scan | StreamingScan - for node in traversal([ir]): - if isinstance(node, StreamingScan): - new_node = StreamingScan.with_prefetched_parquet_metadata( - node, cached_parquet_info_map - ) - replacements[node] = new_node - - elif isinstance(node, Scan) and node.typ == "parquet": # pragma: no cover - raise RuntimeError( - "Unexpected parquet 'Scan' node in lowered IR graph." - ) - - old_ir = ir - ir = replace([ir], replacements)[0] - partition_info = { - new_node: partition_info[old_node] - for old_node, new_node in zip( - traversal([old_ir]), traversal([ir]), strict=True - ) - } + attach_cached_parquet_metadata(ir, cached_parquet_info_map) with ReserveOpIDs(ir, config_options) as collective_id_map: return execute_ir_on_rank( diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 700f432da3d5..058943124f6c 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -208,7 +208,6 @@ class SplitScan(IR): "split_index", "total_splits", "parquet_options", - "cached_parquet_info", ) _n_non_child_args = 13 base_scan: Scan @@ -231,7 +230,7 @@ def __init__( split_index: int, total_splits: int, parquet_options: ParquetOptions, - cached_parquet_info: list[CachedParquetInfo] | None, + cached_parquet_info: list[CachedParquetInfo] | None = None, ): self.schema = schema self.base_scan = base_scan @@ -262,62 +261,17 @@ def __init__( f"Unhandled Scan type for file splitting: {base_scan.typ}" ) - @classmethod - def with_prefetched_parquet_metadata( - cls, - node: SplitScan, - cached_parquet_info_map: dict[str, CachedParquetInfo], - ) -> Self: - """ - Create a new SplitScan node, with prefetched parquet metadata set. - - Because SplitScan is a single-file scan, each composed Scan nodes will - use the same cached parquet metadata. - - Parameters - ---------- - node - The SplitScan node to create a new node from. - cached_parquet_info_map - A dictionary mapping file paths to cached parquet metadata. This should contain - all the file paths, including those from the base scan (which has been split - into multiple SplitScan nodes). - - Returns - ------- - The new SplitScan node. - """ - # cached_parquet_info_map *might* not contain all the paths for the base scan, - # if, e.g., this worker has only been assigned a subset of the paths. - cached_parquet_info = [cached_parquet_info_map[path] for path in node.paths] - return cls( - node.schema, - node.base_scan, - node.paths, - node.split_index, - node.total_splits, - node.parquet_options, - cached_parquet_info, - ) - def get_hashable(self) -> Hashable: """Hashable representation of the node.""" - if self.cached_parquet_info is not None: - cached_parquet_info = tuple( - (info.path, info.size) for info in self.cached_parquet_info - ) - else: - cached_parquet_info = () - + # cached_parquet_info is deliberately not included in the hash data. return ( - type(self), # type: ignore[arg-type] + type(self), tuple(self.schema.items()), self.base_scan.get_hashable(), tuple(self.paths), self.split_index, self.total_splits, self.parquet_options, - *cached_parquet_info, ) @classmethod @@ -438,7 +392,6 @@ class FusedScan(IR): "base_scan", "paths", "parquet_options", - "cached_parquet_info", ) _n_non_child_args = 11 base_scan: Scan @@ -456,7 +409,7 @@ def __init__( base_scan: Scan, paths: list[str], parquet_options: ParquetOptions, - cached_parquet_info: list[CachedParquetInfo] | None, + cached_parquet_info: list[CachedParquetInfo] | None = None, ): self.schema = schema self.base_scan = base_scan @@ -479,38 +432,15 @@ def __init__( ) self.children = () - @classmethod - def with_prefetched_parquet_metadata( - cls, - node: FusedScan, - cached_parquet_info_map: dict[str, CachedParquetInfo], - ) -> Self: - """Create a new FusedScan node, with prefetched parquet metadata set.""" - cached_parquet_info = [cached_parquet_info_map[path] for path in node.paths] - Scan._validate_cached_parquet_info(node.paths, cached_parquet_info) - return cls( - node.schema, - node.base_scan, - node.paths, - node.parquet_options, - cached_parquet_info, - ) - def get_hashable(self) -> Hashable: """Hashable representation of the node.""" - if self.cached_parquet_info is not None: - cached_parquet_info = tuple( - (info.path, info.size) for info in self.cached_parquet_info - ) - else: - cached_parquet_info = () + # cached_parquet_info is deliberately not included in the hash data. return ( - type(self), # type: ignore[arg-type] + type(self), tuple(self.schema.items()), self.base_scan.get_hashable(), tuple(self.paths), self.parquet_options, - *cached_parquet_info, ) @classmethod @@ -714,54 +644,6 @@ def __init__( self._non_child_args = (scans, base_scan, scan_type) self.children = () - @classmethod - def with_prefetched_parquet_metadata( - cls, - node: StreamingScan, - cached_parquet_info_map: dict[str, CachedParquetInfo], - ) -> Self: - """ - Create a new StreamingScan node, with prefetched parquet metadata set. - - Parameters - ---------- - node: StreamingScan - The StreamingScan node to create a new node from. - cached_parquet_info_map - The cached parquet metadata to set on the new node. - - Returns - ------- - Self: The new StreamingScan node. - """ - new_scans: list[SplitScan | FusedScan] = [] - if node.scan_type == "split": - new_scans = [] - for scan in node.scans: - # SplitScan should be generic / overload based on type. - new_scan = SplitScan.with_prefetched_parquet_metadata( - scan, # type: ignore[arg-type] - cached_parquet_info_map, - ) - assert new_scan.cached_parquet_info is not None - assert new_scan.paths == [ - info.path for info in new_scan.cached_parquet_info - ] - new_scans.append(new_scan) - else: - new_scans = [ - FusedScan.with_prefetched_parquet_metadata( - scan, # type: ignore[arg-type] - cached_parquet_info_map, - ) - for scan in node.scans - ] - for scan in new_scans: - assert scan.cached_parquet_info is not None - assert scan.paths == [info.path for info in scan.cached_parquet_info] - - return cls(new_scans, node.base_scan, node.scan_type) # type: ignore[arg-type] - @classmethod def for_split_files( cls, @@ -783,9 +665,6 @@ def for_split_files( splits_created = 0 for path in local_paths: while sindex < plan.factor and splits_created < local_count: - # TODO: We should replace base_scan first, and then ensure cached_parquet_info - # is set properly from the start. Then we can remove the with_prefetched_parquet_metadata - # alternate constructor. scans.append( SplitScan( base_scan.schema, diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 8673a0511df5..13e88ead7731 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -550,27 +550,65 @@ def test_streaming_scan_identity_equality() -> None: 0, 2, base.parquet_options, - _make_cached_parquet_info(base.paths, size=10), + _make_cached_parquet_info(base.paths, size=11), ) - split_diff_metadata = SplitScan( + split_diff = SplitScan( base.schema, base, base.paths, - 0, + 1, 2, base.parquet_options, - _make_cached_parquet_info(base.paths, size=11), + _make_cached_parquet_info(base.paths, size=10), ) a = StreamingScan([split], base, "split") b = StreamingScan([split_same], base, "split") - c = StreamingScan([split_diff_metadata], 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 From d22d7cc94b8ddfde8a4239d12b2b2ea7083b6d18 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 6 Jul 2026 10:59:35 -0700 Subject: [PATCH 47/51] Coverage --- python/cudf_polars/cudf_polars/dsl/utils/io.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index f8af29cc6570..2cd62d906a78 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -205,5 +205,8 @@ def attach_cached_parquet_metadata( if isinstance(node, StreamingScan): for scan in node.scans: _attach_cached_parquet_info(scan, cached_parquet_info_map) - elif isinstance(node, SplitScan | FusedScan): - _attach_cached_parquet_info(node, cached_parquet_info_map) + elif isinstance(node, SplitScan | FusedScan): # pragma: no cover + # This should be called on a lowered IR graph. All SplitScan and + # FusedScan nodes should be wrapped in a StreamingScan node. + msg = "Unexpected 'SplitScan' or 'FusedScan' node in lowered IR graph." + raise TypeError(msg) From 90659c7c8ba687148f82939904a5d72884aae556 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 6 Jul 2026 11:02:31 -0700 Subject: [PATCH 48/51] Coverage --- python/cudf_polars/cudf_polars/dsl/utils/io.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index 2cd62d906a78..900d1772f88b 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -13,11 +13,12 @@ from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars from cudf_polars.dsl.traversal import traversal -from cudf_polars.streaming.io import FusedScan, Scan, SplitScan, StreamingScan +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 + from cudf_polars.streaming.io import FusedScan, SplitScan @dataclass(frozen=True) @@ -205,8 +206,3 @@ def attach_cached_parquet_metadata( if isinstance(node, StreamingScan): for scan in node.scans: _attach_cached_parquet_info(scan, cached_parquet_info_map) - elif isinstance(node, SplitScan | FusedScan): # pragma: no cover - # This should be called on a lowered IR graph. All SplitScan and - # FusedScan nodes should be wrapped in a StreamingScan node. - msg = "Unexpected 'SplitScan' or 'FusedScan' node in lowered IR graph." - raise TypeError(msg) From 677b6f2e8b3a81418aff3351bd525e799a5ec55e Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 6 Jul 2026 14:28:08 -0700 Subject: [PATCH 49/51] Link to blocker --- python/cudf_polars/cudf_polars/dsl/ir.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 615f3ff9e18f..4d89bcad5afe 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -572,7 +572,8 @@ def do_evaluate( def _parquet_physical_types( paths: list[str], columns: list[str] | None ) -> dict[str, plc.DataType]: - # This may not be able use prefetched metadata, since we don't (currently) have a Schema. + # 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() From 0b7898bb3c6b6098b8d0c492c2a8db712ef96dcb Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 6 Jul 2026 14:29:40 -0700 Subject: [PATCH 50/51] Inline --- python/cudf_polars/cudf_polars/dsl/utils/io.py | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index 900d1772f88b..cc49522d498a 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -18,7 +18,6 @@ if TYPE_CHECKING: from cudf_polars.dsl.ir import IR from cudf_polars.streaming.base import StatsCollector - from cudf_polars.streaming.io import FusedScan, SplitScan @dataclass(frozen=True) @@ -176,16 +175,6 @@ def prefetch_parquet_file_metadata_for_ir( return cached_parquet_info -def _attach_cached_parquet_info( - node: SplitScan | FusedScan, - cached_parquet_info_map: dict[str, CachedParquetInfo], -) -> None: - cached = [cached_parquet_info_map[path] for path in node.paths] - Scan._validate_cached_parquet_info(node.paths, cached) - node.cached_parquet_info = cached - node._non_child_args = (*node._non_child_args[:-1], cached) - - def attach_cached_parquet_metadata( root: IR, cached_parquet_info_map: dict[str, CachedParquetInfo], @@ -205,4 +194,7 @@ def attach_cached_parquet_metadata( for node in traversal([root]): if isinstance(node, StreamingScan): for scan in node.scans: - _attach_cached_parquet_info(scan, cached_parquet_info_map) + 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) From ef5bb7a867cbdfd7b7d980c96f31a684f0e424f8 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 6 Jul 2026 14:30:41 -0700 Subject: [PATCH 51/51] Remove uninformative comments --- python/cudf_polars/cudf_polars/dsl/ir.py | 1 - python/cudf_polars/cudf_polars/streaming/io.py | 2 -- 2 files changed, 3 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 4d89bcad5afe..68b6da4a1f16 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -845,7 +845,6 @@ def get_hashable(self) -> Hashable: The options dictionaries are serialised for hashing purposes as json strings. """ - # cached_parquet_info is deliberately not included in the hash data. schema_hash = tuple(self.schema.items()) return ( type(self), diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 058943124f6c..b3b4438812c5 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -263,7 +263,6 @@ def __init__( def get_hashable(self) -> Hashable: """Hashable representation of the node.""" - # cached_parquet_info is deliberately not included in the hash data. return ( type(self), tuple(self.schema.items()), @@ -434,7 +433,6 @@ def __init__( def get_hashable(self) -> Hashable: """Hashable representation of the node.""" - # cached_parquet_info is deliberately not included in the hash data. return ( type(self), tuple(self.schema.items()),