feat: TABLESAMPLE SYSTEM end-to-end + row-group / row 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 is independent of (and lands cleanly without) the API proposal in Query-aware statistics requests via ScanArgs / ScanResult (RFC for #21624) #21996 — they're orthogonal building blocks.
• Closes Feature: Support Sample #16533 for the SYSTEM method (BERNOULLI / ROW counts / BUCKET are rejected at planning time with a pointer to RelationPlanner for custom semantics).

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, and no SQL surface at all. That works for one-off code but is awkward for:

• TABLESAMPLE SQL — most users want a sample by writing SELECT … FROM t TABLESAMPLE SYSTEM (10), not by reaching into ParquetSource builders.
• Ad-hoc data exploration — "give me 1% of this dataset" via the CLI / DataFrame API.
• Layered helpers that want to compute approximate stats over a bounded slice of data without scanning everything (e.g. an optimizer-fed sampled-stats helper — see the linked POC).
• EXPLAIN ANALYZE-driven debug runs against a representative slice instead of the full table.

Two prior in-core attempts at TABLESAMPLE were closed without merging on semantics-fragmentation grounds: #16505 (Spark-style row-level + Poisson) and #16325 (random()-rewrite, which also hit a correctness bug from random() getting pushed into the parquet executor — fixed separately in #16545). The discussion in #13563 settled on "use the RelationPlanner extension API" as the way forward, and #17843 (merged Dec 9 2025) shipped that API specifically with TABLESAMPLE as the canonical motivating example. The blog post Extending SQL in DataFusion: from ->> to TABLESAMPLE documents that pattern.

This PR adds two layers in one go:

1. Low-level builders on ParquetSource — opt-in, additive, no behavior change for existing scans.
2. End-to-end SQL — a built-in RelationPlanner (auto-registered) that lifts TABLESAMPLE SYSTEM(p%) [REPEATABLE(n)] to a Sample extension node; a SamplePushdown optimizer rule pushes it into ParquetSource via the cube-root hybrid across files / row groups / rows. Both SessionStateBuilder::with_default_features() and DefaultPhysicalPlanner::default() register the necessary glue, so datafusion-cli and any default SessionContext get TABLESAMPLE SYSTEM out of the box. Downstream consumers can register a RelationPlanner ahead of the built-in to add other flavors (BERNOULLI, ROW, BUCKET); the existing relation_planner/table_sample.rs example shows that pattern.

What changes are included in this PR?

Layer 1 — ParquetSource sampling builders

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_name, row_group_count, fraction) — re-runs match.
• 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_name, row_group_index, fraction, cluster_size).

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.

Layer 2 — TABLESAMPLE SYSTEM SQL → cube-root pushdown

• Logical: Sample UserDefinedLogicalNodeCore extension node in datafusion-expr (logical_plan/sample.rs). Schema-preserving; validates fraction ∈ (0, 1]. Encodes SampleMethod::System only — by design, since BERNOULLI / BUCKET / row-count semantics differ across DBs.
• SQL surface: new datafusion_sql::sample::TableSampleSystemPlanner RelationPlanner, public, auto-registered via SessionStateDefaults::default_relation_planners() from SessionStateBuilder::with_default_features(). Handles TABLESAMPLE SYSTEM(p%) [REPEATABLE(n)]; rejects BERNOULLI, ROW count, BUCKET ... OUT OF ..., OFFSET with errors that explicitly point users at registering a custom RelationPlanner first. register_relation_planner prepends, so a user-supplied planner that returns Planned for the same syntax wins; returning Original falls through to the built-in.
• Physical placeholder: SampleExec in datafusion-physical-plan. Errors at execute (it is a marker node, not an executor — the SamplePushdown rule is expected to remove it). Implements filter/sort pushdown passthrough so unrelated optimizer rules see straight through it.
• Built-in ExtensionPlanner: new public SamplePhysicalPlanner in datafusion::physical_planner, pre-registered in DefaultPhysicalPlanner::default(). Lowers Sample → SampleExec. Callers using with_extension_planners(...) replace the defaults and need to re-add SamplePhysicalPlanner if they want sampling.
• Pushdown contract: new try_push_sample method on ExecutionPlan and FileSource, returning Absorbed { inner } / Passthrough / Unsupported { reason }. Passthrough overrides on filter, projection, coalesce, repartition, and non-fetch sort.
• ParquetSource::try_push_sample runs the cube-root hybrid: q = p^(1/3) applied at all three levels so the IO win at small fractions doesn't concentrate at one granularity. Returns keep_files: Option<Vec<usize>> that the rule uses to rebuild FileScanConfig.file_groups.
• SamplePushdown optimizer rule (between PushdownSort and EnsureCooperative) walks top-down. On Absorbed it replaces SampleExec with the rebuilt source; on Passthrough it pushes through the single-child node and recurses; on Unsupported it errors at planning time. There is intentionally no generic post-scan SampleExec yet — that's a follow-up.
• EXPLAIN visibility: ParquetSource::fmt_extra surfaces sample_row_group_fraction and sample_row_fraction when set, so EXPLAIN reflects the pushdown.

How the override path works

SessionStateBuilder::register_relation_planner inserts at the front of the planner chain. The first planner to return Planned wins; returning Original falls through. So a downstream user who wants BERNOULLI (or anything else) can:

let mut state = SessionStateBuilder::new_with_default_features().build();
ctx.register_relation_planner(Arc::new(MyTableSamplePlanner));

MyTableSamplePlanner runs first; if it doesn't handle the input it returns Original and the built-in TableSampleSystemPlanner gets the next look. This is the same composition the existing relation_planner/table_sample.rs example demonstrates.

Internals

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

Are these changes tested?

Layer 1 — ParquetSource sampling builders. 7 tests in datafusion-datasource-parquet::opener::test:

• 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 — different file_name, 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_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 (1, 16] range.

Layer 2 — pushdown contract. 4 unit tests on the Sample logical node + 7 unit tests on ParquetSource::try_push_sample directly:

• try_push_sample_system_full_is_noop — fraction >= 1.0 doesn't configure any sampling.
• try_push_sample_system_configures_cube_root_on_source — both row_group_fraction and row_fraction are set to cbrt(p).
• try_push_sample_system_drops_files_for_multi_file_scan — file-level keep_files is the right size and sorted/unique.
• try_push_sample_system_keeps_at_least_one_file — never drops to zero.
• try_push_sample_system_skips_file_drop_for_single_file — single-file scans return keep_files = None.
• try_push_sample_system_repeatable_seed_is_deterministic — same seed → same selection; different seeds → different selection.
• try_push_sample_system_target_clamped_to_num_files — when target == num_files, returns None to skip the rebuild.

Layer 2 — end-to-end SQL. New datafusion/sqllogictest/test_files/tablesample.slt runs against a default SessionContext, exercising the same path datafusion-cli users get:

• SYSTEM(100) returns every row.
• SYSTEM(50) REPEATABLE(42) is deterministic — re-runs match. Asserts the cube-root math gives 813 rows for a single-file / single-row-group parquet.
• EXPLAIN confirms SampleExec is gone and ParquetSource advertises its sampling config (sample_row_group_fraction=0.7937, sample_row_fraction=0.7937).
• Each rejection path (BERNOULLI, ROW count, fraction out of range) produces the right error.

cargo build --workspace, cargo fmt --all, and cargo clippy --workspace --exclude datafusion-benchmarks --all-targets -- -D warnings are clean.

Are there any user-facing changes?

• New SQL surface (out of the box): SELECT … FROM t TABLESAMPLE SYSTEM (p) [REPEATABLE (seed)] works against any plan that bottoms out at a ParquetSource on a default SessionContext. datafusion-cli picks this up with no flag changes. Other forms error at planning time with messages that point at registering a custom RelationPlanner first.
• New public Rust API on ParquetSource: with_row_group_sampling, with_row_fraction, with_row_cluster_size, sampling(), plus the ParquetSampling struct.
• New public Rust API on datafusion-expr: Sample extension node + SampleMethod::System + sample_plan helper.
• New public Rust API on datafusion-sql: sample::TableSampleSystemPlanner (the built-in RelationPlanner).
• New trait methods on ExecutionPlan and FileSource: try_push_sample (default Unsupported), and the SamplePushdownResult / SampleSpec / SampleMethod / FileSourceSampleResult types in datafusion-physical-plan and datafusion-datasource.
• New public SamplePhysicalPlanner in datafusion::physical_planner. Pre-registered in DefaultPhysicalPlanner::default(). Callers using with_extension_planners(...) replace the defaults and need to re-add it.
• New SamplePushdown optimizer rule in the default physical pipeline. No-op when no SampleExec is in the plan; errors at planning time if a SampleExec exists and can't be pushed (no generic post-scan executor yet).
• New default RelationPlanner registered by SessionStateBuilder::with_default_features(). A user-supplied RelationPlanner registered via register_relation_planner runs first and can override or supplement built-in semantics.
• Existing scans / queries that don't use TABLESAMPLE are unaffected.

[adriangb]

I've looked into the history of #16533 a bit.

In particular the decision in #16505 to add hooks (#17843) instead of integrating into core directly. The example in #17843 does post-scan sampling, once the IO cost has been paid. This PR attempts to do much more efficient IO level sampling (files, row groups, pages). In order for this to work we need to add APIs in ParquetSource, ExecutionPlan, etc. At that point I'm not sure an "external extension" approach is worth it. It seems that it's simpler to just bake the feature into DataFusion. The implementation in this PR splits the difference: it uses the RelationPlanner and extension logical plan nodes, but bundles it with SessionStateBuilder::with_default_features(). Personally I would like to at least add Sample to enum LogicalPlan, I don't find the WASM argument very strong but maybe that's just me.

@geoffreyclaude @alamb wdyt?

[geoffreyclaude]

@geoffreyclaude @alamb wdyt?

Thanks for the ping @adriangb!

My take is that the 'RelationPlanner' direction still makes sense here. The point was not that TABLESAMPLE should never be shipped by DataFusion, but that the SQL semantics should remain overridable.

So I think your split makes a lot of sense: put the parquet and physical pushdown in core, but leave the SQL semantics to RelationPlanner though the TableSampleSystemPlanner you introduced. Since custom planners are inserted first (by design), downstream users can still override the built-in SYSTEM semantics if they need to.

On adding Sample to enum LogicalPlan, I'm not fundamentally against it, but it should probably be done in a dedicated follow-up PR. Maybe that follow-up PR should also wire it up more completely with a DataFrame API, generic SampleExec for when pushdown isn't supported, sampled-stats as discussed in #21624... On its own, I don't see the advantage of Sample over UserDefinedLogicalNodeCore, especially given the amount of new boilerplate code it will introduce.

[adriangb]

Okay makes sense then, we're on the same page 😄. Since you've worked on this before would you like to review this PR?

[adriangb]

I'm also happy to split it up:

• PR 1: add methods to ParquetSource, etc.
• PR 2: add ExecutionPlan trait methods, SampleExec physical plan and pushdown optimizer rule
• PR 3: add TableSampleSystemPlanner and wire it all up

Up to you...

[adriangb]

I've cleaned up the history into stacked commits in case that's helpful.

[geoffreyclaude]

I've cleaned up the history into stacked commits in case that's helpful.

That's super useful, thanks! I'll give it a look tomorrow.

[alamb]

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, and no SQL surface at all. That works for one-off code but is awkward for:

My personal rationale was that all the different SQL systems did sampling differently -- so any particular choice for sampling is probably fine but I wasn't at all sure there would be enough commonality across implementations to put it into DataFusion core

ALso, what is the problem with constructing the access plan by hand? For this type of low level access pattern (particular sampling methods etc) it seems like low level construction is just the escape valve that is needed (super fine grained control)

I am very wary of complicating the built in Parquet reader any more -- it is already very complicated with lots of behaviros (and new ones getting added all ghe time, for example the sortedness ones from @zhuqi-lucas and @xudong963 )

So adding APIs to make it easier to extend / modify plans makes sense to me, but hard coding more sampling into the core is much less clear to me

[alamb]

Did you try and implement whatever sampling you need with just the existing APIs? Aka no change to the core? If that wasn't possible, what was missing in the API?

[adriangb]

Thanks for reviewing @alamb!

Did you try and implement whatever sampling you need with just the existing APIs? Aka no change to the core? If that wasn't possible, what was missing in the API?

Yes, it won't work sadly. With current APIs you could maybe sample at the file level, but not at the row group or page level (which requires reading parquet footers for all files, parsing them, etc). The way I implemented things here that is lazily deferred to file opening so there is no wasted work.

Also requiring doing rust bits to make it work precludes this working with e.g. datafusion-cli. I think it is a useful primitive to have in datafusion-cli (e.g. for data exploration) and can be a strong building block for work like #21624

On the ExecutionPlan level unless we add some APIs there would be a lot of tight coupling (mainly via downcast matching on both ends) between the custom optimizer rule, the custom physical plan and the data source. I.e. the optimizer rule needs to downcast match every plan in the path between the SampleExec and ParquetSource and understand them. It is incompatible with any other custom plans, custom data sources, etc (i.e. composition is broken). The way I've set it up in this PR if e.g. Vortex wanted to implement table sampling they could do that in a self contained way inside of their file source implementation quite easily.

[adriangb]

The actual changes to opener.rs are quite small, just 2 blocks of (IMO) easily understandable code to update the access plan: 3d0dc4a.

[adriangb]

I am very wary of complicating the built in Parquet reader any more -- it is already very complicated with lots of behaviros (and new ones getting added all ghe time, for example the sortedness ones from @zhuqi-lucas and @xudong963 )

I agree it is a complex piece of software but I think we can continue to add the right abstractions and simplifications (like you recently did with the moralization work 😄 ). Ultimately the file reader is going to be a key piece of a data toolkit like DataFusion so it's unsurprising (to me) that it holds a lot of the complexity.

[alamb]

I am very wary of complicating the built in Parquet reader any more -- it is already very complicated with lots of behaviros (and new ones getting added all ghe time, for example the sortedness ones from @zhuqi-lucas and @xudong963 )

I agree it is a complex piece of software but I think we can continue to add the right abstractions and simplifications (like you recently did with the moralization work 😄 ). Ultimately the file reader is going to be a key piece of a data toolkit like DataFusion so it's unsurprising (to me) that it holds a lot of the complexity.

yeah -- maybe I am over sensitive as I feel like as soon as we are able to refactor away some of the complexity then it get all complicated again 😆

[adriangb]

I am very wary of complicating the built in Parquet reader any more -- it is already very complicated with lots of behaviros (and new ones getting added all ghe time, for example the sortedness ones from @zhuqi-lucas and @xudong963 )

I agree it is a complex piece of software but I think we can continue to add the right abstractions and simplifications (like you recently did with the moralization work 😄 ). Ultimately the file reader is going to be a key piece of a data toolkit like DataFusion so it's unsurprising (to me) that it holds a lot of the complexity.

yeah -- maybe I am over sensitive as I feel like as soon as we are able to refactor away some of the complexity then it get all complicated again 😆

No you are right: it is a big risk that this code turns into feature spaghetti. It's just not one I think we can necessarily avoid. We should be cautious about introducing complexity and push back (like you have here) but if this is the right place to put it and we can factor it into a shape that only adds complexity, not multiplies or exponentiates it, then maybe we just need to deal with it over time.

[adriangb]

@alamb I've restructured some more and pulled out new code into sampling.rs so that the changes to opener.rs (aside from tests, imports, etc.) is ~ 10 LOC. I hope this helps 🙇🏻

2026412#diff-bbd611d7b35d7f17633eebbf32a07dc9e394f20135754ed949751e8030049e38

[adriangb]

Regarding how other systems implement this: it seems most columnar, analytical DBs are approximate for the same reasons we are:

https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-syntax-qry-select-sampling

https://docs.cloud.google.com/bigquery/docs/table-sampling

https://duckdb.org/docs/current/sql/samples

Similar for Spark and Trino.

I think it's also worth highlighting that DuckDB has the behavior built in.

In other words @alamb: if we merged 2026412 that's at least enough for someone like me to make it work in my system / in future DataFusion use cases. But I think it's valuable to take it one step further and expose this to users of datafusion-cli, etc.