diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 4e061e6a3cb2..6e964d02208c 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 @@ -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,14 +1085,6 @@ 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) def column_storage_size(self, column: str) -> int | None: diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index a44af628ea1c..d4852ee71155 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, 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..0eb12d72cd1e 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,48 @@ 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 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_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 None + assert info.cached_parquet_info is not None + assert len(info.cached_parquet_info) == len(paths) def test_parquet_metadata_reads_footers(