feat(parquet): row-group and row-range sampling on ParquetSource

[adriangb]

Which issue does this PR close?

• Relates to datafusion.execution.collect_statistics on wide tables #21624 (datafusion.execution.collect_statistics on wide tables).

This PR extracts the first layer of #22000 — the opt-in ParquetSource sampling primitives — as a self-contained change. The TABLESAMPLE SQL surface, Sample logical/physical nodes, and SamplePushdown rule are deliberately not included; they will land in follow-ups.

Rationale for this change

DataFusion has the machinery for fine-grained parquet sampling (ParquetAccessPlan with Skip / Scan / Selection(RowSelection)) but no public way to ask for a sample without constructing the access plan by hand and stuffing it into PartitionedFile.extensions. That works for one-off code but is awkward for ad-hoc data exploration, layered helpers that want to compute approximate stats over a bounded slice, and EXPLAIN ANALYZE-driven debug runs against a representative slice.

This PR adds the lowest layer: opt-in builders on ParquetSource that translate fractions into a ParquetAccessPlan lazily inside the opener (after the footer is loaded, so we sample by real row-group index). It is additive and has no behavior change for existing scans. The SQL surface in #22000 is built on top of these primitives.

What changes are included in this PR?

ParquetSource::new(schema)
    .with_row_group_sampling(0.1)   // keep ~10% of row groups per file
    .with_row_fraction(0.05)        // within each kept row group, keep ~5% of rows
    .with_row_cluster_size(8192);   // controls window granularity (default 32 768)

with_row_group_sampling(fraction):

• Selection is deferred until the opener has loaded the parquet footer, so we sample by real row-group index.
• Deterministic per (file_index, row_group_count, fraction) — re-runs match. The opener passes the execution partition_index as the stable file_index, so sampling is reproducible across environments without depending on object-store paths.
• Always keeps at least one row group (target = max(1, ceil(N * fraction))).
• No-op when fraction >= 1.0.

with_row_fraction(fraction):

• Translates the per-row-group target into K contiguous windows spread evenly across the row group, each placed at a random offset within its stride. Window count = ceil(target / cluster_size).
• Materializes a RowSelection per kept row group; the parquet reader uses the page index to read only the data pages covering the selected rows. This gives "page-level" IO savings without requiring per-column page alignment (which doesn't exist in parquet).
• Falls back gracefully when the page index is missing — the reader still returns the right rows, the IO win just disappears.
• Deterministic per (file_index, row_group_index, fraction, cluster_size).
• Window-layout math is extracted into a dedicated build_row_window_selectors function and fuzz-tested across thousands of configurations to guarantee no overlap, in-bounds positions, and full coverage.

The two layers compose: row_group_fraction = 0.1 × row_fraction = 0.1 reads ~1% of the rows from ~10% of the row groups, with windows spread out so the sample isn't clustered at one end of each row group.

Internals

• New ParquetSampling struct re-exported at the crate root.
• Plumbed through ParquetMorselizer → PreparedParquetOpen.
• Two free functions invoked from prune_row_groups right after create_initial_plan.
• New dep: rand with the small_rng feature (already in workspace Cargo.toml).

Differences vs. the original commit in #22000

Two pieces of review feedback on the parent PR are folded in here:

• apply_*_sampling keys on a stable file_index: usize rather than file_name: &str, addressing feat: TABLESAMPLE SYSTEM end-to-end + row-group / row sampling on ParquetSource #22000 (comment). The opener passes the execution partition_index. This removes the on-disk-path dependency from the seed inputs while still decorrelating files in different partitions.
• The row-window layout math is extracted into build_row_window_selectors and fuzz-tested (feat: TABLESAMPLE SYSTEM end-to-end + row-group / row sampling on ParquetSource #22000 (comment)). Fuzzing surfaced an overlap bug at target_rows ≈ total_rows where window_size = ceil(target / n_windows) could exceed stride = total_rows / n_windows; the fix caps window_size at stride.

Are these changes tested?

12 tests in datafusion-datasource-parquet:

Row-group sampling (sampling::tests):

• row_group_sampling_keeps_target_count — ceil(N * fraction) math.
• row_group_sampling_is_deterministic — same inputs → same selection.
• row_group_sampling_differs_per_file_index — different file_index → different sample.
• row_group_sampling_no_op_when_fraction_is_one — fraction ≥ 1.0 keeps everything.
• row_group_sampling_target_at_least_one — fraction = 0.001 over 100 row groups still keeps 1.
• row_group_sampling_no_op_when_unset — None is a no-op.

Row-window selection (sampling::tests):

• row_window_selection_basic_layout — hand-checked anchor case.
• row_window_selection_returns_none_on_invalid_input — degenerate inputs (zero row group, zero target, zero cluster) return None.
• row_window_selection_full_target_no_overlap — the previously-buggy target_rows == total_rows case.
• row_window_selection_fuzz_invariants — 5 000 randomized (total_rows, target_rows, cluster_size, seed) configurations, asserting full coverage, in-bounds positions, and no overlap.
• row_window_selection_fuzz_determinism — 1 000 iterations verifying identical seeds produce identical layouts.

End-to-end (opener::test):

• row_group_sampling_end_to_end — writes a 4-row-group parquet to InMemory, scans with fraction = 0.5, asserts exactly 6 rows out (2 row groups × 3 rows).
• row_fraction_end_to_end — writes a 100-row single-row-group parquet, scans with row_fraction = 0.1 and cluster_size = 4, asserts the result is in the expected range.

cargo build, cargo fmt --all, and cargo clippy -p datafusion-datasource-parquet --all-targets --all-features -- -D warnings are clean.

Are there any user-facing changes?

• New public Rust API on ParquetSource: with_row_group_sampling, with_row_fraction, with_row_cluster_size, sampling(), plus the ParquetSampling struct.
• Existing scans / queries that don't opt in are unaffected.
• No SQL surface yet — that lands in a follow-up to feat: TABLESAMPLE SYSTEM end-to-end + row-group / row sampling on ParquetSource #22000.

🤖 Generated with Claude Code