Skip to content

Commit 88326d4

Browse files
adriangbclaude
andcommitted
move TABLESAMPLE SQL surface out of core, into RelationPlanner
Background: PR apache#17843 (merged Dec 9 2025) added a `RelationPlanner` extension API specifically motivated by TABLESAMPLE, and the blog post "Extending SQL in DataFusion: from ->> to TABLESAMPLE" canonizes that as the way to add SQL syntax extensions. Two prior in-core attempts (apache#16505 Spark-style, apache#16325 random()-rewrite) closed without merging on semantics-fragmentation grounds. alamb's stated reason for keeping the surface out of core — sampling semantics differ widely across DBs, plus WASM bloat — applies directly to the parsing we had baked into `datafusion/sql/src/relation/mod.rs`. This commit reshapes the PR around that: - Revert the inline TABLESAMPLE parsing in `relation/mod.rs`. - Drop the SQL planning unit tests that exercised the inline parsing. - Replace the hacky inline `Sample` fallback in `DefaultPhysicalPlanner` with a public `SamplePhysicalPlanner` `ExtensionPlanner` that callers register alongside their `RelationPlanner`. Same effect, cleaner shape. - Move the e2e fixture from a sqllogictest `.slt` (which can't register a `RelationPlanner`) to a Rust integration test under `datafusion/core/tests/parquet/tablesample.rs`. The test uses a small inline `RelationPlanner` to emit the core `Sample` node, proving the cube-root pushdown works end to end against parquet. What stays in core (the actual contribution): - `Sample` logical extension node and `SampleExec` placeholder. - `SampleSpec` / `SamplePushdownResult` / `FileSourceSampleResult` types and the `try_push_sample` trait method on `ExecutionPlan` and `FileSource`. - Per-node Passthrough overrides (filter, projection, coalesce, repartition, non-fetch sort). - `ParquetSource::try_push_sample` cube-root absorption + `DataSourceExec::try_push_sample` rebuilding `file_groups`. - `SamplePushdown` optimizer rule registered in the default pipeline. These are the building blocks any SQL surface (the existing example under `datafusion-examples/examples/relation_planner/table_sample.rs` included) can use to push sampling into the source instead of filtering per-batch with `random() < p`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent d769eb6 commit 88326d4

6 files changed

Lines changed: 387 additions & 292 deletions

File tree

datafusion/core/src/physical_planner.rs

Lines changed: 71 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -237,6 +237,77 @@ pub trait ExtensionPlanner {
237237
}
238238
}
239239

240+
/// [`ExtensionPlanner`] that lowers the [`Sample`] logical extension
241+
/// node into a [`SampleExec`] physical node.
242+
///
243+
/// The pushdown machinery (cube-root absorption into `ParquetSource`,
244+
/// the `SamplePushdown` optimizer rule, per-node `Passthrough`
245+
/// overrides) is wired into the default optimizer pipeline, so once
246+
/// a `Sample` reaches the physical planner it will be pushed into
247+
/// the source — but it has to *get* there first. Register this
248+
/// planner on a [`SessionStateBuilder`] / [`DefaultPhysicalPlanner`]
249+
/// alongside whichever [`RelationPlanner`] (or other front-end) you
250+
/// use to emit the `Sample` logical node:
251+
///
252+
/// ```rust,ignore
253+
/// use std::sync::Arc;
254+
/// use datafusion::physical_planner::{DefaultPhysicalPlanner, SamplePhysicalPlanner};
255+
///
256+
/// let planner = DefaultPhysicalPlanner::with_extension_planners(vec![
257+
/// Arc::new(SamplePhysicalPlanner),
258+
/// ]);
259+
/// ```
260+
///
261+
/// `Sample` is intentionally not built into [`DefaultPhysicalPlanner`]
262+
/// by default: that would force callers who use a different SQL
263+
/// surface (or who don't want any TABLESAMPLE story) to carry the
264+
/// physical glue. The planner is small enough to opt into.
265+
///
266+
/// [`Sample`]: datafusion_expr::logical_plan::sample::Sample
267+
/// [`SampleExec`]: datafusion_physical_plan::sample::SampleExec
268+
/// [`RelationPlanner`]: datafusion_expr::planner::RelationPlanner
269+
/// [`SessionStateBuilder`]: crate::execution::session_state::SessionStateBuilder
270+
#[derive(Debug, Default)]
271+
pub struct SamplePhysicalPlanner;
272+
273+
#[async_trait]
274+
impl ExtensionPlanner for SamplePhysicalPlanner {
275+
async fn plan_extension(
276+
&self,
277+
_planner: &dyn PhysicalPlanner,
278+
node: &dyn UserDefinedLogicalNode,
279+
_logical_inputs: &[&LogicalPlan],
280+
physical_inputs: &[Arc<dyn ExecutionPlan>],
281+
_session_state: &SessionState,
282+
) -> Result<Option<Arc<dyn ExecutionPlan>>> {
283+
let Some(sample) = node
284+
.as_any()
285+
.downcast_ref::<datafusion_expr::logical_plan::sample::Sample>()
286+
else {
287+
return Ok(None);
288+
};
289+
if physical_inputs.len() != 1 {
290+
return plan_err!(
291+
"Sample expects exactly one input; got {}",
292+
physical_inputs.len()
293+
);
294+
}
295+
let method = match sample.method {
296+
datafusion_expr::logical_plan::sample::SampleMethod::System => {
297+
datafusion_physical_plan::sample_pushdown::SampleMethod::System
298+
}
299+
};
300+
Ok(Some(Arc::new(
301+
datafusion_physical_plan::sample::SampleExec::new(
302+
Arc::clone(&physical_inputs[0]),
303+
method,
304+
sample.fraction,
305+
sample.seed,
306+
),
307+
)))
308+
}
309+
}
310+
240311
/// Default single node physical query planner that converts a
241312
/// `LogicalPlan` to an `ExecutionPlan` suitable for execution.
242313
///
@@ -1810,35 +1881,6 @@ impl DefaultPhysicalPlanner {
18101881
.await?;
18111882
}
18121883

1813-
// Built-in fallback: the `Sample` extension node is part of
1814-
// core SQL (TABLESAMPLE), so plan it without requiring a
1815-
// user-installed ExtensionPlanner. User planners run first,
1816-
// so a custom override still wins.
1817-
if maybe_plan.is_none()
1818-
&& let Some(sample) = node
1819-
.as_any()
1820-
.downcast_ref::<datafusion_expr::logical_plan::sample::Sample>(
1821-
)
1822-
{
1823-
if children.len() != 1 {
1824-
return plan_err!(
1825-
"Sample expects exactly one input; got {}",
1826-
children.len()
1827-
);
1828-
}
1829-
let method = match sample.method {
1830-
datafusion_expr::logical_plan::sample::SampleMethod::System =>
1831-
datafusion_physical_plan::sample_pushdown::SampleMethod::System,
1832-
};
1833-
maybe_plan =
1834-
Some(Arc::new(datafusion_physical_plan::sample::SampleExec::new(
1835-
Arc::clone(&children[0]),
1836-
method,
1837-
sample.fraction,
1838-
sample.seed,
1839-
)) as Arc<dyn ExecutionPlan>);
1840-
}
1841-
18421884
let plan = match maybe_plan {
18431885
Some(v) => Ok(v),
18441886
_ => plan_err!(

datafusion/core/tests/parquet/mod.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ mod page_pruning;
5757
mod row_group_pruning;
5858
mod schema;
5959
mod schema_coercion;
60+
mod tablesample;
6061
mod utils;
6162

6263
#[cfg(test)]

0 commit comments

Comments
 (0)