Add parquet-sampling configuration options (rapidsai#19423)

Closes rapidsai#19389

Adds `max_footer_samples` and `max_row_group_samples` configuration options to control metadata/row-group sampling. Although these configuration options are only *used* by the streaming executor, it felt more natural to add these to `ParquetOptions` (since they are definitely Parquet specific).

Authors:
  - Richard (Rick) Zamora (https://github.com/rjzamora)

Approvers:
  - Tom Augspurger (https://github.com/TomAugspurger)
  - Matthew Murray (https://github.com/Matt711)

URL: rapidsai#19423

Author: rjzamora

python/cudf_polars/cudf_polars/experimental/io.py

@@ -116,7 +116,7 @@ def from_scan(ir: Scan, config_options: ConfigOptions) -> ScanPartitionPlan:
             )
 
             blocksize: int = config_options.executor.target_partition_size
-            column_stats = _extract_scan_stats(ir)
+            column_stats = _extract_scan_stats(ir, config_options)
             column_sizes: list[int] = []
             for name, cs in column_stats.items():
                 storage_size = cs.source_info.storage_size(name)
@@ -593,13 +593,13 @@ class ParquetMetadata:
     ----------
     paths
         Parquet-dataset paths.
-    max_file_samples
-        Maximum number of files to sample for metadata.
+    max_footer_samples
+        Maximum number of file footers to sample metadata from.
     """
 
     __slots__ = (
         "column_names",
-        "max_file_samples",
+        "max_footer_samples",
         "mean_size_per_file",
         "num_row_groups_per_file",
         "paths",
@@ -609,8 +609,8 @@ class ParquetMetadata:
 
     paths: tuple[str, ...]
     """Parquet-dataset paths."""
-    max_file_samples: int
-    """Maximum number of files to sample for metadata."""
+    max_footer_samples: int
+    """Maximum number of file footers to sample metadata from."""
     row_count: ColumnStat[int]
     """Total row-count estimate."""
     num_row_groups_per_file: tuple[int, ...]
@@ -622,15 +622,17 @@ class ParquetMetadata:
     sample_paths: tuple[str, ...]
     """Sampled file paths."""
 
-    def __init__(self, paths: tuple[str, ...], max_file_samples: int):
+    def __init__(self, paths: tuple[str, ...], max_footer_samples: int):
         self.paths = paths
-        self.max_file_samples = max_file_samples
+        self.max_footer_samples = max_footer_samples
         self.row_count = ColumnStat[int]()
         self.num_row_groups_per_file = ()
         self.mean_size_per_file = {}
         self.column_names = ()
-        stride = max(1, int(len(paths) / max_file_samples)) if max_file_samples else 1
-        self.sample_paths = paths[: stride * max_file_samples : stride]
+        stride = (
+            max(1, int(len(paths) / max_footer_samples)) if max_footer_samples else 1
+        )
+        self.sample_paths = paths[: stride * max_footer_samples : stride]
 
         if not self.sample_paths:
             # No paths to sample from
@@ -685,29 +687,29 @@ class ParquetSourceInfo(DataSourceInfo):
     ----------
     paths
         Parquet-dataset paths.
-    max_file_samples
-        Maximum number of files to sample metadata from.
-    max_rg_samples
+    max_footer_samples
+        Maximum number of file footers to sample metadata from.
+    max_row_group_samples
         Maximum number of row-groups to sample data from.
     """
 
     def __init__(
         self,
         paths: tuple[str, ...],
-        max_file_samples: int,
-        max_rg_samples: int,
+        max_footer_samples: int,
+        max_row_group_samples: int,
     ):
         self.paths = paths
-        self.max_file_samples = max_file_samples
-        self.max_rg_samples = max_rg_samples
+        self.max_footer_samples = max_footer_samples
+        self.max_row_group_samples = max_row_group_samples
         # Helper attributes
         self._key_columns: set[str] = set()  # Used to fuse lazy row-group sampling
         self._unique_stats: dict[str, UniqueStats] = {}
 
     @functools.cached_property
     def metadata(self) -> ParquetMetadata:
         """Return Parquet metadata."""
-        return ParquetMetadata(self.paths, self.max_file_samples)
+        return ParquetMetadata(self.paths, self.max_footer_samples)
 
     @property
     def row_count(self) -> ColumnStat[int]:
@@ -717,7 +719,7 @@ def row_count(self) -> ColumnStat[int]:
     def _sample_row_groups(self) -> None:
         """Estimate unique-value statistics from a row-group sample."""
         sample_paths = self.metadata.sample_paths
-        if not sample_paths or self.max_rg_samples < 1:
+        if not sample_paths or self.max_row_group_samples < 1:
             # No row-groups to sample from
             return
 
@@ -742,14 +744,14 @@ def _sample_row_groups(self) -> None:
             for rg_id in range(num_rgs):
                 n += 1
                 samples[path].append(rg_id)
-                if n == self.max_rg_samples:
+                if n == self.max_row_group_samples:
                     break
-            if n == self.max_rg_samples:
+            if n == self.max_row_group_samples:
                 break
 
-        exact = sampled_file_count == len(self.paths) and self.max_rg_samples >= sum(
-            num_row_groups_per_file
-        )
+        exact = sampled_file_count == len(
+            self.paths
+        ) and self.max_row_group_samples >= sum(num_row_groups_per_file)
 
         options = plc.io.parquet.ParquetReaderOptions.builder(
             plc.io.SourceInfo(list(samples))
@@ -809,22 +811,23 @@ def add_unique_stats_column(self, column: str) -> None:
 @functools.cache
 def _sample_pq_stats(
     paths: tuple[str, ...],
-    max_file_samples: int,
-    max_rg_samples: int,
+    max_footer_samples: int,
+    max_row_group_samples: int,
 ) -> ParquetSourceInfo:
     """Return Parquet datasource information."""
-    return ParquetSourceInfo(paths, max_file_samples, max_rg_samples)
+    return ParquetSourceInfo(paths, max_footer_samples, max_row_group_samples)
 
 
 def _extract_scan_stats(
-    ir: Scan, *, max_file_samples: int = 3, max_rg_samples: int = 1
+    ir: Scan,
+    config_options: ConfigOptions,
 ) -> dict[str, ColumnStats]:
     """Extract base ColumnStats for a Scan node."""
     if ir.typ == "parquet":
-        # TODO: Add max_file_samples and max_rg_samples
-        # to the ConfigOption system.
         source_info = _sample_pq_stats(
-            tuple(ir.paths), max_file_samples, max_rg_samples
+            tuple(ir.paths),
+            config_options.parquet_options.max_footer_samples,
+            config_options.parquet_options.max_row_group_samples,
         )
         return {
             name: ColumnStats(

python/cudf_polars/cudf_polars/utils/config.py

@@ -160,6 +160,18 @@ class ParquetOptions:
     pass_read_limit
         Limit on the amount of memory used for reading and decompressing data
         or 0 if there is no limit.
+    max_footer_samples
+        Maximum number of file footers to sample for metadata. This
+        option is currently used by the streaming executor to gather
+        datasource statistics before generating a physical plan. Set to
+        0 to avoid metadata sampling. Default is 3.
+    max_row_group_samples
+        Maximum number of row-groups to sample for unique-value statistics.
+        This option may be used by the streaming executor to optimize
+        the physical plan. Default is 1.
+
+        Set to 0 to avoid row-group sampling. Note that row-group sampling
+        will also be skipped if ``max_footer_samples`` is 0.
     """
 
     _env_prefix = "CUDF_POLARS__PARQUET_OPTIONS"
@@ -179,6 +191,16 @@ class ParquetOptions:
             f"{_env_prefix}__PASS_READ_LIMIT", int, default=0
         )
     )
+    max_footer_samples: int = dataclasses.field(
+        default_factory=_make_default_factory(
+            f"{_env_prefix}__MAX_FOOTER_SAMPLES", int, default=3
+        )
+    )
+    max_row_group_samples: int = dataclasses.field(
+        default_factory=_make_default_factory(
+            f"{_env_prefix}__MAX_ROW_GROUP_SAMPLES", int, default=1
+        )
+    )
 
     def __post_init__(self) -> None:  # noqa: D105
         if not isinstance(self.chunked, bool):
@@ -187,6 +209,10 @@ def __post_init__(self) -> None:  # noqa: D105
             raise TypeError("chunk_read_limit must be an int")
         if not isinstance(self.pass_read_limit, int):
             raise TypeError("pass_read_limit must be an int")
+        if not isinstance(self.max_footer_samples, int):
+            raise TypeError("max_footer_samples must be an int")
+        if not isinstance(self.max_row_group_samples, int):
+            raise TypeError("max_row_group_samples must be an int")
 
 
 def default_blocksize(scheduler: str) -> int:

python/cudf_polars/tests/experimental/test_scan.py

@@ -88,15 +88,15 @@ def test_split_scan_predicate(tmp_path, df, mask):
 
 @pytest.mark.parametrize("n_files", [1, 3])
 @pytest.mark.parametrize("row_group_size", [None, 10_000])
-@pytest.mark.parametrize("max_file_samples", [3, 0])
-@pytest.mark.parametrize("max_rg_samples", [1, 0])
+@pytest.mark.parametrize("max_footer_samples", [3, 0])
+@pytest.mark.parametrize("max_row_group_samples", [1, 0])
 def test_source_statistics(
     tmp_path,
     df,
     n_files,
     row_group_size,
-    max_file_samples,
-    max_rg_samples,
+    max_footer_samples,
+    max_row_group_samples,
 ):
     from cudf_polars.experimental.io import (
         _clear_source_info_cache,
@@ -119,26 +119,26 @@ def test_source_statistics(
             "target_partition_size": 10_000,
             "scheduler": DEFAULT_SCHEDULER,
         },
+        parquet_options={
+            "max_footer_samples": max_footer_samples,
+            "max_row_group_samples": max_row_group_samples,
+        },
     )
     ir = Translator(q._ldf.visit(), engine).translate_ir()
-    column_stats = _extract_scan_stats(
-        ir,
-        max_file_samples=max_file_samples,
-        max_rg_samples=max_rg_samples,
-    )
+    column_stats = _extract_scan_stats(ir, ConfigOptions.from_polars_engine(engine))
 
     # Source info is the same for all columns
     source_info = column_stats["x"].source_info
     assert source_info is column_stats["y"].source_info
     assert source_info is column_stats["z"].source_info
-    if max_file_samples:
+    if max_footer_samples:
         assert source_info.row_count.value == df.height
         assert source_info.row_count.exact
     else:
         assert source_info.row_count.value is None
 
     # Storage stats should be available
-    if max_file_samples:
+    if max_footer_samples:
         assert source_info.storage_size("x").value > 0
         assert source_info.storage_size("y").value > 0
     else:
@@ -153,21 +153,21 @@ def test_source_statistics(
     # source._unique_stats should be empty
     assert set(source_info._unique_stats) == set()
 
-    if max_file_samples and max_rg_samples:
+    if max_footer_samples and max_row_group_samples:
         assert source_info.unique_stats("x").count.value == df.height
         assert source_info.unique_stats("x").fraction.value == 1.0
     else:
         assert source_info.unique_stats("x").count.value is None
         assert source_info.unique_stats("x").fraction.value is None
 
     # source_info._unique_stats should only contain 'x'
-    if max_file_samples and max_rg_samples:
+    if max_footer_samples and max_row_group_samples:
         assert set(source_info._unique_stats) == {"x"}
     else:
         assert set(source_info._unique_stats) == set()
 
     # Check add_unique_stats_column behavior
-    if max_file_samples and max_rg_samples:
+    if max_footer_samples and max_row_group_samples:
         # Can add a "bad"/missing key column
         source_info.add_unique_stats_column("foo")
         assert set(source_info._unique_stats) == {"x"}
@@ -198,7 +198,7 @@ def test_source_statistics_csv(tmp_path, df):
         },
     )
     ir = Translator(q._ldf.visit(), engine).translate_ir()
-    column_stats = _extract_scan_stats(ir)
+    column_stats = _extract_scan_stats(ir, ConfigOptions.from_polars_engine(engine))
 
     # Source info should be empty for CSV
     source_info = column_stats["x"].source_info

python/cudf_polars/tests/test_config.py

@@ -301,13 +301,17 @@ def test_parquet_options_from_env(monkeypatch: pytest.MonkeyPatch) -> None:
         m.setenv("CUDF_POLARS__PARQUET_OPTIONS__CHUNKED", "0")
         m.setenv("CUDF_POLARS__PARQUET_OPTIONS__CHUNK_READ_LIMIT", "100")
         m.setenv("CUDF_POLARS__PARQUET_OPTIONS__PASS_READ_LIMIT", "200")
+        m.setenv("CUDF_POLARS__PARQUET_OPTIONS__MAX_FOOTER_SAMPLES", "0")
+        m.setenv("CUDF_POLARS__PARQUET_OPTIONS__MAX_ROW_GROUP_SAMPLES", "0")
 
         # Test default
         engine = pl.GPUEngine()
         config = ConfigOptions.from_polars_engine(engine)
         assert config.parquet_options.chunked is False
         assert config.parquet_options.chunk_read_limit == 100
         assert config.parquet_options.pass_read_limit == 200
+        assert config.parquet_options.max_footer_samples == 0
+        assert config.parquet_options.max_row_group_samples == 0
 
     with monkeypatch.context() as m:
         m.setenv("CUDF_POLARS__PARQUET_OPTIONS__CHUNKED", "foo")
@@ -393,7 +397,16 @@ def test_cardinality_factor_compat() -> None:
         )
 
 
-@pytest.mark.parametrize("option", ["chunked", "chunk_read_limit", "pass_read_limit"])
+@pytest.mark.parametrize(
+    "option",
+    [
+        "chunked",
+        "chunk_read_limit",
+        "pass_read_limit",
+        "max_footer_samples",
+        "max_row_group_samples",
+    ],
+)
 def test_validate_parquet_options(option: str) -> None:
     with pytest.raises(TypeError, match=f"{option} must be"):
         ConfigOptions.from_polars_engine(