Reuse sampled parquet footers from statistics collection

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:

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")))

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(