From 5e2f3d22b8345e2eb3dc687dfb53be66df6c7602 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Wed, 26 Aug 2026 23:49:56 +0000 Subject: [PATCH 1/2] Reuse sampled parquet footers from statistics collection --- .../cudf_polars/cudf_polars/streaming/io.py | 20 +++---- .../cudf_polars/tests/streaming/test_scan.py | 55 +++++++++++++++++++ .../cudf_polars/tests/streaming/test_stats.py | 14 +++-- 3 files changed, 74 insertions(+), 15 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 4e061e6a3cb2..8fabeaa76f5f 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -870,7 +870,7 @@ class ParquetMetadata: sample_paths: tuple[str, ...] """Sampled file paths.""" cached_parquet_info: list[CachedParquetInfo] | None - """Cached parquet info for the sampled paths. Only set if all files were sampled.""" + """Cached parquet info for the sampled paths.""" @nvtx_annotate_cudf_polars(message="ParquetMetadata") def __init__(self, paths: tuple[str, ...], max_footer_samples: int): @@ -904,10 +904,10 @@ def __init__(self, paths: tuple[str, ...], max_footer_samples: int): ) sample_footers = [info.file_metadata for info in sample_parquet_info] + self.cached_parquet_info = 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 @@ -1080,15 +1080,13 @@ def from_paths( else max(footer_mean, decoded_floor) ) - cached_parquet_info: list[CachedParquetInfo] | None - if ( - metadata.sampled_file_count == metadata.total_file_count - and metadata.cached_parquet_info is not None - ): - cached_parquet_info = list(metadata.cached_parquet_info) - else: - cached_parquet_info = None - return cls(row_count, per_file_means, cached_parquet_info=cached_parquet_info) + return cls( + row_count, + per_file_means, + cached_parquet_info=list(metadata.cached_parquet_info) + if metadata.cached_parquet_info is not None + else None, + ) 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_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index a44af628ea1c..268c17d44134 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -126,6 +126,61 @@ def test_prefetch_parquet_file_metadata_no_parquet_scans() -> None: assert result == {} +def test_prefetch_skips_paths_cached_by_stats_collection( + tmp_path, + df: pl.DataFrame, + monkeypatch: pytest.MonkeyPatch, + parquet_stats_executor, +) -> None: + import cudf_polars.dsl.utils.io as io_module + from cudf_polars.streaming.io import _clear_source_info_cache + + _clear_source_info_cache() + n_files = 5 + max_footer_samples = 2 + make_partitioned_source(df, tmp_path, "parquet", n_files=n_files) + paths = sorted(str(p) for p in tmp_path.glob("*.parquet")) + + engine = pl.GPUEngine( + raise_on_fail=True, + executor="streaming", + parquet_options={"max_footer_samples": max_footer_samples}, + ) + q = pl.scan_parquet(tmp_path) + from cudf_polars import Translator + + ir = Translator(q._ldf.visit(), engine).translate_ir() + config = ConfigOptions.from_polars_engine(engine) + stats = collect_statistics(ir, config, parquet_stats_executor) + + source = stats.scan_stats[ir] + assert source.cached_parquet_info is not None + sampled_paths = {info.path for info in source.cached_parquet_info} + assert len(sampled_paths) == max_footer_samples + + fetched_paths: list[str] = [] + real_prefetch = io_module._prefetch_parquet_footers_for_paths + + def recording_prefetch(paths_arg: list[str]) -> list: + fetched_paths.extend(paths_arg) + return real_prefetch(paths_arg) + + monkeypatch.setattr( + io_module, "_prefetch_parquet_footers_for_paths", recording_prefetch + ) + + scan = _make_parquet_scan(paths) + fused = FusedScan(scan.schema, scan, tuple(paths), scan.parquet_options, None) + streaming_scan = StreamingScan([fused], scan, "fused") + + result = prefetch_parquet_file_metadata_for_ir( + streaming_scan, py_executor=None, stats=stats + ) + + assert set(result) == set(paths) + assert set(fetched_paths) == set(paths) - sampled_paths + + def test_prefetch_parquet_file_metadata_remote_only(tmp_path, df) -> None: make_partitioned_source(df, tmp_path, "parquet", n_files=1) local_path = str(next(tmp_path.glob("*.parquet"))) diff --git a/python/cudf_polars/tests/streaming/test_stats.py b/python/cudf_polars/tests/streaming/test_stats.py index 22c7a1795ee1..ad926dd1830c 100644 --- a/python/cudf_polars/tests/streaming/test_stats.py +++ b/python/cudf_polars/tests/streaming/test_stats.py @@ -166,6 +166,7 @@ def __init__(self, paths: tuple[str, ...], max_footer_samples: int) -> None: self.max_footer_samples = max_footer_samples self.sampled_file_count = 1 self.total_file_count = len(paths) + self.cached_parquet_info = None sampled_cols: list[str] = [] @@ -271,23 +272,28 @@ def test_parquet_source_info_stores_footers_when_all_files_sampled( ) -def test_parquet_source_info_omits_footers_when_paths_are_sampled( +def test_parquet_source_info_stores_sampled_footers_when_partially_sampled( tmp_path: pathlib.Path, df_and_schema: tuple[pl.DataFrame, Schema], ) -> None: _clear_source_info_cache() df, schema = df_and_schema - make_partitioned_source(df, tmp_path, "parquet", n_files=5) + n_files = 5 + max_footer_samples = 2 + make_partitioned_source(df, tmp_path, "parquet", n_files=n_files) paths = tuple(str(p) for p in sorted(tmp_path.iterdir())) info = _build_parquet_source( paths, frozenset(df.columns), tuple(schema.items()), - max_footer_samples=2, + max_footer_samples=max_footer_samples, max_row_group_samples=0, ) - assert info.cached_parquet_info is None + assert info.cached_parquet_info is not None + assert len(info.cached_parquet_info) == max_footer_samples + cached_paths = {cached.path for cached in info.cached_parquet_info} + assert cached_paths <= set(paths) def test_parquet_metadata_reads_footers( From 4bff73d421e39b783aabf343e711a20a6ce7b0f9 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 27 Aug 2026 01:21:49 +0000 Subject: [PATCH 2/2] address review --- .../cudf_polars/cudf_polars/streaming/io.py | 15 +++++++------- .../cudf_polars/tests/streaming/test_scan.py | 2 +- .../cudf_polars/tests/streaming/test_stats.py | 20 +++++++++++++++++++ 3 files changed, 28 insertions(+), 9 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 8fabeaa76f5f..6e964d02208c 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -1038,9 +1038,14 @@ def from_paths( file_count = len(paths) per_file_means: dict[str, int] = {} + cached_parquet_info = ( + list(metadata.cached_parquet_info) + if metadata.cached_parquet_info is not None + else None + ) if not (file_count and row_count and needed_cols): - return cls(row_count, {}) + return cls(row_count, {}, cached_parquet_info=cached_parquet_info) rows_per_file = max(1, row_count // file_count) schema_map = dict(schema) @@ -1080,13 +1085,7 @@ def from_paths( else max(footer_mean, decoded_floor) ) - return cls( - row_count, - per_file_means, - cached_parquet_info=list(metadata.cached_parquet_info) - if metadata.cached_parquet_info is not None - else 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_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 268c17d44134..d4852ee71155 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -170,7 +170,7 @@ def recording_prefetch(paths_arg: list[str]) -> list: ) scan = _make_parquet_scan(paths) - fused = FusedScan(scan.schema, scan, tuple(paths), scan.parquet_options, None) + fused = FusedScan(scan.schema, scan, paths, scan.parquet_options, None) streaming_scan = StreamingScan([fused], scan, "fused") result = prefetch_parquet_file_metadata_for_ir( diff --git a/python/cudf_polars/tests/streaming/test_stats.py b/python/cudf_polars/tests/streaming/test_stats.py index ad926dd1830c..0eb12d72cd1e 100644 --- a/python/cudf_polars/tests/streaming/test_stats.py +++ b/python/cudf_polars/tests/streaming/test_stats.py @@ -296,6 +296,26 @@ def test_parquet_source_info_stores_sampled_footers_when_partially_sampled( assert cached_paths <= set(paths) +def test_parquet_source_info_preserves_footers_on_empty_needed_cols( + tmp_path: pathlib.Path, + df_and_schema: tuple[pl.DataFrame, Schema], +) -> None: + _clear_source_info_cache() + df, schema = df_and_schema + make_partitioned_source(df, tmp_path, "parquet", n_files=2) + paths = tuple(str(p) for p in sorted(tmp_path.iterdir())) + info = _build_parquet_source( + paths, + frozenset(), + tuple(schema.items()), + max_footer_samples=10, + max_row_group_samples=0, + ) + + assert info.cached_parquet_info is not None + assert len(info.cached_parquet_info) == len(paths) + + def test_parquet_metadata_reads_footers( tmp_path: pathlib.Path, df_and_schema: tuple[pl.DataFrame, Schema],