mirror of
https://github.com/langchain-ai/datafusion.git
synced 2026-08-27 02:21:16 -04:00
Albert Skalt
34dad2ccee
- closes https://github.com/apache/datafusion/issues/19796 This patch aims to implement a fast-path for the ExecutionPlan::with_new_children function for some plans, moving closer to a physical plan re-use implementation and improving planning performance. If the passed children properties are the same as in self, we do not actually recompute self's properties (which could be costly if projection mapping is required). Instead, we just replace the children and re-use self's properties as-is. To be able to compare two different properties -- ExecutionPlan::properties(...) signature is modified and now returns `&Arc<PlanProperties>`. If `children` properties are the same in `with_new_children` -- we clone our properties arc and then a parent plan will consider our properties as unchanged, doing the same. - Return `&Arc<PlanProperties>` from `ExecutionPlan::properties(...)` instead of a reference. - Implement `with_new_children` fast-path if there is no children properties changes for all major plans. Note: currently, `reset_plan_states` does not allow to re-use plan in general: it is not supported for dynamic filters and recursive queries features, as in this case state reset should update pointers in the children plans. --------- Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
821 lines
28 KiB
Rust
821 lines
28 KiB
Rust
// Licensed to the Apache Software Foundation (ASF) under one
|
|
// or more contributor license agreements. See the NOTICE file
|
|
// distributed with this work for additional information
|
|
// regarding copyright ownership. The ASF licenses this file
|
|
// to you under the Apache License, Version 2.0 (the
|
|
// "License"); you may not use this file except in compliance
|
|
// with the License. You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing,
|
|
// software distributed under the License is distributed on an
|
|
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
// KIND, either express or implied. See the License for the
|
|
// specific language governing permissions and limitations
|
|
// under the License.
|
|
|
|
//! # TABLESAMPLE Example
|
|
//!
|
|
//! This example demonstrates implementing SQL `TABLESAMPLE` support using
|
|
//! DataFusion's extensibility APIs.
|
|
//!
|
|
//! This is a working `TABLESAMPLE` implementation that can serve as a starting
|
|
//! point for your own projects. It also works as a template for adding other
|
|
//! custom SQL operators, covering the full pipeline from parsing to execution.
|
|
//!
|
|
//! It shows how to:
|
|
//!
|
|
//! 1. **Parse** TABLESAMPLE syntax via a custom [`RelationPlanner`]
|
|
//! 2. **Plan** sampling as a custom logical node ([`TableSamplePlanNode`])
|
|
//! 3. **Execute** sampling via a custom physical operator ([`SampleExec`])
|
|
//!
|
|
//! ## Supported Syntax
|
|
//!
|
|
//! ```sql
|
|
//! -- Bernoulli sampling (each row has N% chance of selection)
|
|
//! SELECT * FROM table TABLESAMPLE BERNOULLI(10 PERCENT)
|
|
//!
|
|
//! -- Fractional sampling (0.0 to 1.0)
|
|
//! SELECT * FROM table TABLESAMPLE (0.1)
|
|
//!
|
|
//! -- Row count limit
|
|
//! SELECT * FROM table TABLESAMPLE (100 ROWS)
|
|
//!
|
|
//! -- Reproducible sampling with a seed
|
|
//! SELECT * FROM table TABLESAMPLE (10 PERCENT) REPEATABLE(42)
|
|
//! ```
|
|
//!
|
|
//! ## Architecture
|
|
//!
|
|
//! ```text
|
|
//! ┌─────────────────────────────────────────────────────────────────┐
|
|
//! │ SQL Query │
|
|
//! │ SELECT * FROM t TABLESAMPLE BERNOULLI(10 PERCENT) REPEATABLE(1)│
|
|
//! └─────────────────────────────────────────────────────────────────┘
|
|
//! │
|
|
//! ▼
|
|
//! ┌─────────────────────────────────────────────────────────────────┐
|
|
//! │ TableSamplePlanner │
|
|
//! │ (RelationPlanner: parses TABLESAMPLE, creates logical node) │
|
|
//! └─────────────────────────────────────────────────────────────────┘
|
|
//! │
|
|
//! ▼
|
|
//! ┌─────────────────────────────────────────────────────────────────┐
|
|
//! │ TableSamplePlanNode │
|
|
//! │ (UserDefinedLogicalNode: stores sampling params) │
|
|
//! └─────────────────────────────────────────────────────────────────┘
|
|
//! │
|
|
//! ▼
|
|
//! ┌─────────────────────────────────────────────────────────────────┐
|
|
//! │ TableSampleExtensionPlanner │
|
|
//! │ (ExtensionPlanner: creates physical execution plan) │
|
|
//! └─────────────────────────────────────────────────────────────────┘
|
|
//! │
|
|
//! ▼
|
|
//! ┌─────────────────────────────────────────────────────────────────┐
|
|
//! │ SampleExec │
|
|
//! │ (ExecutionPlan: performs actual row sampling at runtime) │
|
|
//! └─────────────────────────────────────────────────────────────────┘
|
|
//! ```
|
|
|
|
use std::{
|
|
any::Any,
|
|
fmt::{self, Debug, Formatter},
|
|
hash::{Hash, Hasher},
|
|
pin::Pin,
|
|
sync::Arc,
|
|
task::{Context, Poll},
|
|
};
|
|
|
|
use arrow::datatypes::{Float64Type, Int64Type};
|
|
use arrow::{
|
|
array::{ArrayRef, Int32Array, RecordBatch, StringArray, UInt32Array},
|
|
compute,
|
|
};
|
|
use arrow_schema::SchemaRef;
|
|
use futures::{
|
|
ready,
|
|
stream::{Stream, StreamExt},
|
|
};
|
|
use rand::{Rng, SeedableRng, rngs::StdRng};
|
|
use tonic::async_trait;
|
|
|
|
use datafusion::optimizer::simplify_expressions::simplify_literal::parse_literal;
|
|
use datafusion::{
|
|
execution::{
|
|
RecordBatchStream, SendableRecordBatchStream, SessionState, SessionStateBuilder,
|
|
TaskContext, context::QueryPlanner,
|
|
},
|
|
physical_expr::EquivalenceProperties,
|
|
physical_plan::{
|
|
DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
|
|
metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RecordOutput},
|
|
},
|
|
physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner},
|
|
prelude::*,
|
|
};
|
|
use datafusion_common::{
|
|
DFSchemaRef, DataFusionError, Result, Statistics, internal_err, not_impl_err,
|
|
plan_datafusion_err, plan_err,
|
|
};
|
|
use datafusion_expr::{
|
|
UserDefinedLogicalNode, UserDefinedLogicalNodeCore,
|
|
logical_plan::{Extension, LogicalPlan, LogicalPlanBuilder},
|
|
planner::{
|
|
PlannedRelation, RelationPlanner, RelationPlannerContext, RelationPlanning,
|
|
},
|
|
};
|
|
use datafusion_sql::sqlparser::ast::{
|
|
self, TableFactor, TableSampleMethod, TableSampleUnit,
|
|
};
|
|
use insta::assert_snapshot;
|
|
|
|
// ============================================================================
|
|
// Example Entry Point
|
|
// ============================================================================
|
|
|
|
/// Runs the TABLESAMPLE examples demonstrating various sampling techniques.
|
|
pub async fn table_sample() -> Result<()> {
|
|
// Build session with custom query planner for physical planning
|
|
let state = SessionStateBuilder::new()
|
|
.with_default_features()
|
|
.with_query_planner(Arc::new(TableSampleQueryPlanner))
|
|
.build();
|
|
|
|
let ctx = SessionContext::new_with_state(state);
|
|
|
|
// Register custom relation planner for logical planning
|
|
ctx.register_relation_planner(Arc::new(TableSamplePlanner))?;
|
|
register_sample_data(&ctx)?;
|
|
|
|
println!("TABLESAMPLE Example");
|
|
println!("===================\n");
|
|
|
|
run_examples(&ctx).await
|
|
}
|
|
|
|
async fn run_examples(ctx: &SessionContext) -> Result<()> {
|
|
// Example 1: Baseline - full table scan
|
|
let results = run_example(
|
|
ctx,
|
|
"Example 1: Full table (baseline)",
|
|
"SELECT * FROM sample_data",
|
|
)
|
|
.await?;
|
|
assert_snapshot!(results, @r"
|
|
+---------+---------+
|
|
| column1 | column2 |
|
|
+---------+---------+
|
|
| 1 | row_1 |
|
|
| 2 | row_2 |
|
|
| 3 | row_3 |
|
|
| 4 | row_4 |
|
|
| 5 | row_5 |
|
|
| 6 | row_6 |
|
|
| 7 | row_7 |
|
|
| 8 | row_8 |
|
|
| 9 | row_9 |
|
|
| 10 | row_10 |
|
|
+---------+---------+
|
|
");
|
|
|
|
// Example 2: Percentage-based Bernoulli sampling
|
|
// REPEATABLE(seed) ensures deterministic results for snapshot testing
|
|
let results = run_example(
|
|
ctx,
|
|
"Example 2: BERNOULLI percentage sampling",
|
|
"SELECT * FROM sample_data TABLESAMPLE BERNOULLI(30 PERCENT) REPEATABLE(123)",
|
|
)
|
|
.await?;
|
|
assert_snapshot!(results, @r"
|
|
+---------+---------+
|
|
| column1 | column2 |
|
|
+---------+---------+
|
|
| 1 | row_1 |
|
|
| 2 | row_2 |
|
|
| 7 | row_7 |
|
|
| 8 | row_8 |
|
|
+---------+---------+
|
|
");
|
|
|
|
// Example 3: Fractional sampling (0.0 to 1.0)
|
|
// REPEATABLE(seed) ensures deterministic results for snapshot testing
|
|
let results = run_example(
|
|
ctx,
|
|
"Example 3: Fractional sampling",
|
|
"SELECT * FROM sample_data TABLESAMPLE (0.5) REPEATABLE(456)",
|
|
)
|
|
.await?;
|
|
assert_snapshot!(results, @r"
|
|
+---------+---------+
|
|
| column1 | column2 |
|
|
+---------+---------+
|
|
| 2 | row_2 |
|
|
| 4 | row_4 |
|
|
| 8 | row_8 |
|
|
+---------+---------+
|
|
");
|
|
|
|
// Example 4: Row count limit (deterministic, no seed needed)
|
|
let results = run_example(
|
|
ctx,
|
|
"Example 4: Row count limit",
|
|
"SELECT * FROM sample_data TABLESAMPLE (3 ROWS)",
|
|
)
|
|
.await?;
|
|
assert_snapshot!(results, @r"
|
|
+---------+---------+
|
|
| column1 | column2 |
|
|
+---------+---------+
|
|
| 1 | row_1 |
|
|
| 2 | row_2 |
|
|
| 3 | row_3 |
|
|
+---------+---------+
|
|
");
|
|
|
|
// Example 5: Sampling combined with filtering
|
|
let results = run_example(
|
|
ctx,
|
|
"Example 5: Sampling with WHERE clause",
|
|
"SELECT * FROM sample_data TABLESAMPLE (5 ROWS) WHERE column1 > 2",
|
|
)
|
|
.await?;
|
|
assert_snapshot!(results, @r"
|
|
+---------+---------+
|
|
| column1 | column2 |
|
|
+---------+---------+
|
|
| 3 | row_3 |
|
|
| 4 | row_4 |
|
|
| 5 | row_5 |
|
|
+---------+---------+
|
|
");
|
|
|
|
// Example 6: Sampling in JOIN queries
|
|
// REPEATABLE(seed) ensures deterministic results for snapshot testing
|
|
let results = run_example(
|
|
ctx,
|
|
"Example 6: Sampling in JOINs",
|
|
r#"SELECT t1.column1, t2.column1, t1.column2, t2.column2
|
|
FROM sample_data t1 TABLESAMPLE (0.7) REPEATABLE(789)
|
|
JOIN sample_data t2 TABLESAMPLE (0.7) REPEATABLE(123)
|
|
ON t1.column1 = t2.column1"#,
|
|
)
|
|
.await?;
|
|
assert_snapshot!(results, @r"
|
|
+---------+---------+---------+---------+
|
|
| column1 | column1 | column2 | column2 |
|
|
+---------+---------+---------+---------+
|
|
| 2 | 2 | row_2 | row_2 |
|
|
| 5 | 5 | row_5 | row_5 |
|
|
| 7 | 7 | row_7 | row_7 |
|
|
| 8 | 8 | row_8 | row_8 |
|
|
| 10 | 10 | row_10 | row_10 |
|
|
+---------+---------+---------+---------+
|
|
");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Helper to run a single example query and capture results.
|
|
async fn run_example(ctx: &SessionContext, title: &str, sql: &str) -> Result<String> {
|
|
println!("{title}:\n{sql}\n");
|
|
let df = ctx.sql(sql).await?;
|
|
println!("{}\n", df.logical_plan().display_indent());
|
|
|
|
let batches = df.collect().await?;
|
|
let results = arrow::util::pretty::pretty_format_batches(&batches)?.to_string();
|
|
println!("{results}\n");
|
|
|
|
Ok(results)
|
|
}
|
|
|
|
/// Register test data: 10 rows with column1=1..10 and column2="row_1".."row_10"
|
|
fn register_sample_data(ctx: &SessionContext) -> Result<()> {
|
|
let column1: ArrayRef = Arc::new(Int32Array::from((1..=10).collect::<Vec<i32>>()));
|
|
let column2: ArrayRef = Arc::new(StringArray::from(
|
|
(1..=10).map(|i| format!("row_{i}")).collect::<Vec<_>>(),
|
|
));
|
|
let batch =
|
|
RecordBatch::try_from_iter(vec![("column1", column1), ("column2", column2)])?;
|
|
ctx.register_batch("sample_data", batch)?;
|
|
Ok(())
|
|
}
|
|
|
|
// ============================================================================
|
|
// Logical Planning: TableSamplePlanner + TableSamplePlanNode
|
|
// ============================================================================
|
|
|
|
/// Relation planner that intercepts `TABLESAMPLE` clauses in SQL and creates
|
|
/// [`TableSamplePlanNode`] logical nodes.
|
|
#[derive(Debug)]
|
|
struct TableSamplePlanner;
|
|
|
|
impl RelationPlanner for TableSamplePlanner {
|
|
fn plan_relation(
|
|
&self,
|
|
relation: TableFactor,
|
|
context: &mut dyn RelationPlannerContext,
|
|
) -> Result<RelationPlanning> {
|
|
// Only handle Table relations with TABLESAMPLE clause
|
|
let TableFactor::Table {
|
|
sample: Some(sample),
|
|
alias,
|
|
name,
|
|
args,
|
|
with_hints,
|
|
version,
|
|
with_ordinality,
|
|
partitions,
|
|
json_path,
|
|
index_hints,
|
|
} = relation
|
|
else {
|
|
return Ok(RelationPlanning::Original(Box::new(relation)));
|
|
};
|
|
|
|
// Extract sample spec (handles both before/after alias positions)
|
|
let sample = match sample {
|
|
ast::TableSampleKind::BeforeTableAlias(s)
|
|
| ast::TableSampleKind::AfterTableAlias(s) => s,
|
|
};
|
|
|
|
// Validate sampling method
|
|
if let Some(method) = &sample.name
|
|
&& *method != TableSampleMethod::Bernoulli
|
|
&& *method != TableSampleMethod::Row
|
|
{
|
|
return not_impl_err!(
|
|
"Sampling method {} is not supported (only BERNOULLI and ROW)",
|
|
method
|
|
);
|
|
}
|
|
|
|
// Offset sampling (ClickHouse-style) not supported
|
|
if sample.offset.is_some() {
|
|
return not_impl_err!(
|
|
"TABLESAMPLE with OFFSET is not supported (requires total row count)"
|
|
);
|
|
}
|
|
|
|
// Parse optional REPEATABLE seed
|
|
let seed = sample
|
|
.seed
|
|
.map(|s| {
|
|
s.value.to_string().parse::<u64>().map_err(|_| {
|
|
plan_datafusion_err!("REPEATABLE seed must be an integer")
|
|
})
|
|
})
|
|
.transpose()?;
|
|
|
|
// Plan the underlying table without the sample clause
|
|
let base_relation = TableFactor::Table {
|
|
sample: None,
|
|
alias: alias.clone(),
|
|
name,
|
|
args,
|
|
with_hints,
|
|
version,
|
|
with_ordinality,
|
|
partitions,
|
|
json_path,
|
|
index_hints,
|
|
};
|
|
let input = context.plan(base_relation)?;
|
|
|
|
// Handle bucket sampling (Hive-style: TABLESAMPLE(BUCKET x OUT OF y))
|
|
if let Some(bucket) = sample.bucket {
|
|
if bucket.on.is_some() {
|
|
return not_impl_err!(
|
|
"TABLESAMPLE BUCKET with ON clause requires CLUSTERED BY table"
|
|
);
|
|
}
|
|
let bucket_num: u64 =
|
|
bucket.bucket.to_string().parse().map_err(|_| {
|
|
plan_datafusion_err!("bucket number must be an integer")
|
|
})?;
|
|
let total: u64 =
|
|
bucket.total.to_string().parse().map_err(|_| {
|
|
plan_datafusion_err!("bucket total must be an integer")
|
|
})?;
|
|
|
|
let fraction = bucket_num as f64 / total as f64;
|
|
let plan = TableSamplePlanNode::new(input, fraction, seed).into_plan();
|
|
return Ok(RelationPlanning::Planned(Box::new(PlannedRelation::new(
|
|
plan, alias,
|
|
))));
|
|
}
|
|
|
|
// Handle quantity-based sampling
|
|
let Some(quantity) = sample.quantity else {
|
|
return plan_err!(
|
|
"TABLESAMPLE requires a quantity (percentage, fraction, or row count)"
|
|
);
|
|
};
|
|
let quantity_value_expr = context.sql_to_expr(quantity.value, input.schema())?;
|
|
|
|
match quantity.unit {
|
|
// TABLESAMPLE (N ROWS) - exact row limit
|
|
Some(TableSampleUnit::Rows) => {
|
|
let rows: i64 = parse_literal::<Int64Type>(&quantity_value_expr)?;
|
|
if rows < 0 {
|
|
return plan_err!("row count must be non-negative, got {}", rows);
|
|
}
|
|
let plan = LogicalPlanBuilder::from(input)
|
|
.limit(0, Some(rows as usize))?
|
|
.build()?;
|
|
Ok(RelationPlanning::Planned(Box::new(PlannedRelation::new(
|
|
plan, alias,
|
|
))))
|
|
}
|
|
|
|
// TABLESAMPLE (N PERCENT) - percentage sampling
|
|
Some(TableSampleUnit::Percent) => {
|
|
let percent: f64 = parse_literal::<Float64Type>(&quantity_value_expr)?;
|
|
let fraction = percent / 100.0;
|
|
let plan = TableSamplePlanNode::new(input, fraction, seed).into_plan();
|
|
Ok(RelationPlanning::Planned(Box::new(PlannedRelation::new(
|
|
plan, alias,
|
|
))))
|
|
}
|
|
|
|
// TABLESAMPLE (N) - fraction if <1.0, row limit if >=1.0
|
|
None => {
|
|
let value = parse_literal::<Float64Type>(&quantity_value_expr)?;
|
|
if value < 0.0 {
|
|
return plan_err!("sample value must be non-negative, got {}", value);
|
|
}
|
|
let plan = if value >= 1.0 {
|
|
// Interpret as row limit
|
|
LogicalPlanBuilder::from(input)
|
|
.limit(0, Some(value as usize))?
|
|
.build()?
|
|
} else {
|
|
// Interpret as fraction
|
|
TableSamplePlanNode::new(input, value, seed).into_plan()
|
|
};
|
|
Ok(RelationPlanning::Planned(Box::new(PlannedRelation::new(
|
|
plan, alias,
|
|
))))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Custom logical plan node representing a TABLESAMPLE operation.
|
|
///
|
|
/// Stores sampling parameters (bounds, seed) and wraps the input plan.
|
|
/// Gets converted to [`SampleExec`] during physical planning.
|
|
#[derive(Debug, Clone, Hash, Eq, PartialEq, PartialOrd)]
|
|
struct TableSamplePlanNode {
|
|
input: LogicalPlan,
|
|
lower_bound: HashableF64,
|
|
upper_bound: HashableF64,
|
|
seed: u64,
|
|
}
|
|
|
|
impl TableSamplePlanNode {
|
|
/// Create a new sampling node with the given fraction (0.0 to 1.0).
|
|
fn new(input: LogicalPlan, fraction: f64, seed: Option<u64>) -> Self {
|
|
Self {
|
|
input,
|
|
lower_bound: HashableF64(0.0),
|
|
upper_bound: HashableF64(fraction),
|
|
seed: seed.unwrap_or_else(rand::random),
|
|
}
|
|
}
|
|
|
|
/// Wrap this node in a LogicalPlan::Extension.
|
|
fn into_plan(self) -> LogicalPlan {
|
|
LogicalPlan::Extension(Extension {
|
|
node: Arc::new(self),
|
|
})
|
|
}
|
|
}
|
|
|
|
impl UserDefinedLogicalNodeCore for TableSamplePlanNode {
|
|
fn name(&self) -> &str {
|
|
"TableSample"
|
|
}
|
|
|
|
fn inputs(&self) -> Vec<&LogicalPlan> {
|
|
vec![&self.input]
|
|
}
|
|
|
|
fn schema(&self) -> &DFSchemaRef {
|
|
self.input.schema()
|
|
}
|
|
|
|
fn expressions(&self) -> Vec<Expr> {
|
|
vec![]
|
|
}
|
|
|
|
fn fmt_for_explain(&self, f: &mut Formatter) -> fmt::Result {
|
|
write!(
|
|
f,
|
|
"Sample: bounds=[{}, {}], seed={}",
|
|
self.lower_bound.0, self.upper_bound.0, self.seed
|
|
)
|
|
}
|
|
|
|
fn with_exprs_and_inputs(
|
|
&self,
|
|
_exprs: Vec<Expr>,
|
|
mut inputs: Vec<LogicalPlan>,
|
|
) -> Result<Self> {
|
|
Ok(Self {
|
|
input: inputs.swap_remove(0),
|
|
lower_bound: self.lower_bound,
|
|
upper_bound: self.upper_bound,
|
|
seed: self.seed,
|
|
})
|
|
}
|
|
}
|
|
|
|
/// Wrapper for f64 that implements Hash and Eq (required for LogicalPlan).
|
|
#[derive(Debug, Clone, Copy, PartialOrd)]
|
|
struct HashableF64(f64);
|
|
|
|
impl PartialEq for HashableF64 {
|
|
fn eq(&self, other: &Self) -> bool {
|
|
self.0.to_bits() == other.0.to_bits()
|
|
}
|
|
}
|
|
|
|
impl Eq for HashableF64 {}
|
|
|
|
impl Hash for HashableF64 {
|
|
fn hash<H: Hasher>(&self, state: &mut H) {
|
|
self.0.to_bits().hash(state);
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Physical Planning: TableSampleQueryPlanner + TableSampleExtensionPlanner
|
|
// ============================================================================
|
|
|
|
/// Custom query planner that registers [`TableSampleExtensionPlanner`] to
|
|
/// convert [`TableSamplePlanNode`] into [`SampleExec`].
|
|
#[derive(Debug)]
|
|
struct TableSampleQueryPlanner;
|
|
|
|
#[async_trait]
|
|
impl QueryPlanner for TableSampleQueryPlanner {
|
|
async fn create_physical_plan(
|
|
&self,
|
|
logical_plan: &LogicalPlan,
|
|
session_state: &SessionState,
|
|
) -> Result<Arc<dyn ExecutionPlan>> {
|
|
let planner = DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(
|
|
TableSampleExtensionPlanner,
|
|
)]);
|
|
planner
|
|
.create_physical_plan(logical_plan, session_state)
|
|
.await
|
|
}
|
|
}
|
|
|
|
/// Extension planner that converts [`TableSamplePlanNode`] to [`SampleExec`].
|
|
struct TableSampleExtensionPlanner;
|
|
|
|
#[async_trait]
|
|
impl ExtensionPlanner for TableSampleExtensionPlanner {
|
|
async fn plan_extension(
|
|
&self,
|
|
_planner: &dyn PhysicalPlanner,
|
|
node: &dyn UserDefinedLogicalNode,
|
|
_logical_inputs: &[&LogicalPlan],
|
|
physical_inputs: &[Arc<dyn ExecutionPlan>],
|
|
_session_state: &SessionState,
|
|
) -> Result<Option<Arc<dyn ExecutionPlan>>> {
|
|
let Some(sample_node) = node.as_any().downcast_ref::<TableSamplePlanNode>()
|
|
else {
|
|
return Ok(None);
|
|
};
|
|
|
|
let exec = SampleExec::try_new(
|
|
Arc::clone(&physical_inputs[0]),
|
|
sample_node.lower_bound.0,
|
|
sample_node.upper_bound.0,
|
|
sample_node.seed,
|
|
)?;
|
|
Ok(Some(Arc::new(exec)))
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Physical Execution: SampleExec + BernoulliSampler
|
|
// ============================================================================
|
|
|
|
/// Physical execution plan that samples rows from its input using Bernoulli sampling.
|
|
///
|
|
/// Each row is independently selected with probability `(upper_bound - lower_bound)`
|
|
/// and appears at most once.
|
|
#[derive(Debug, Clone)]
|
|
pub struct SampleExec {
|
|
input: Arc<dyn ExecutionPlan>,
|
|
lower_bound: f64,
|
|
upper_bound: f64,
|
|
seed: u64,
|
|
metrics: ExecutionPlanMetricsSet,
|
|
cache: Arc<PlanProperties>,
|
|
}
|
|
|
|
impl SampleExec {
|
|
/// Create a new SampleExec with Bernoulli sampling (without replacement).
|
|
///
|
|
/// # Arguments
|
|
/// * `input` - The input execution plan
|
|
/// * `lower_bound` - Lower bound of sampling range (typically 0.0)
|
|
/// * `upper_bound` - Upper bound of sampling range (0.0 to 1.0)
|
|
/// * `seed` - Random seed for reproducible sampling
|
|
pub fn try_new(
|
|
input: Arc<dyn ExecutionPlan>,
|
|
lower_bound: f64,
|
|
upper_bound: f64,
|
|
seed: u64,
|
|
) -> Result<Self> {
|
|
if lower_bound < 0.0 || upper_bound > 1.0 || lower_bound > upper_bound {
|
|
return internal_err!(
|
|
"Sampling bounds must satisfy 0.0 <= lower <= upper <= 1.0, got [{}, {}]",
|
|
lower_bound,
|
|
upper_bound
|
|
);
|
|
}
|
|
|
|
let cache = PlanProperties::new(
|
|
EquivalenceProperties::new(input.schema()),
|
|
input.properties().partitioning.clone(),
|
|
input.properties().emission_type,
|
|
input.properties().boundedness,
|
|
);
|
|
|
|
Ok(Self {
|
|
input,
|
|
lower_bound,
|
|
upper_bound,
|
|
seed,
|
|
metrics: ExecutionPlanMetricsSet::new(),
|
|
cache: Arc::new(cache),
|
|
})
|
|
}
|
|
|
|
/// Create a sampler for the given partition.
|
|
fn create_sampler(&self, partition: usize) -> BernoulliSampler {
|
|
let seed = self.seed.wrapping_add(partition as u64);
|
|
BernoulliSampler::new(self.lower_bound, self.upper_bound, seed)
|
|
}
|
|
}
|
|
|
|
impl DisplayAs for SampleExec {
|
|
fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result {
|
|
write!(
|
|
f,
|
|
"SampleExec: bounds=[{}, {}], seed={}",
|
|
self.lower_bound, self.upper_bound, self.seed
|
|
)
|
|
}
|
|
}
|
|
|
|
impl ExecutionPlan for SampleExec {
|
|
fn name(&self) -> &'static str {
|
|
"SampleExec"
|
|
}
|
|
|
|
fn as_any(&self) -> &dyn Any {
|
|
self
|
|
}
|
|
|
|
fn properties(&self) -> &Arc<PlanProperties> {
|
|
&self.cache
|
|
}
|
|
|
|
fn maintains_input_order(&self) -> Vec<bool> {
|
|
// Sampling preserves row order (rows are filtered, not reordered)
|
|
vec![true]
|
|
}
|
|
|
|
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
|
|
vec![&self.input]
|
|
}
|
|
|
|
fn with_new_children(
|
|
self: Arc<Self>,
|
|
mut children: Vec<Arc<dyn ExecutionPlan>>,
|
|
) -> Result<Arc<dyn ExecutionPlan>> {
|
|
Ok(Arc::new(Self::try_new(
|
|
children.swap_remove(0),
|
|
self.lower_bound,
|
|
self.upper_bound,
|
|
self.seed,
|
|
)?))
|
|
}
|
|
|
|
fn execute(
|
|
&self,
|
|
partition: usize,
|
|
context: Arc<TaskContext>,
|
|
) -> Result<SendableRecordBatchStream> {
|
|
Ok(Box::pin(SampleStream {
|
|
input: self.input.execute(partition, context)?,
|
|
sampler: self.create_sampler(partition),
|
|
metrics: BaselineMetrics::new(&self.metrics, partition),
|
|
}))
|
|
}
|
|
|
|
fn metrics(&self) -> Option<MetricsSet> {
|
|
Some(self.metrics.clone_inner())
|
|
}
|
|
|
|
fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
|
|
let mut stats = self.input.partition_statistics(partition)?;
|
|
let ratio = self.upper_bound - self.lower_bound;
|
|
|
|
// Scale statistics by sampling ratio (inexact due to randomness)
|
|
stats.num_rows = stats
|
|
.num_rows
|
|
.map(|n| (n as f64 * ratio) as usize)
|
|
.to_inexact();
|
|
stats.total_byte_size = stats
|
|
.total_byte_size
|
|
.map(|n| (n as f64 * ratio) as usize)
|
|
.to_inexact();
|
|
|
|
Ok(stats)
|
|
}
|
|
}
|
|
|
|
/// Bernoulli sampler: includes each row with probability `(upper - lower)`.
|
|
/// This is sampling **without replacement** - each row appears at most once.
|
|
struct BernoulliSampler {
|
|
lower_bound: f64,
|
|
upper_bound: f64,
|
|
rng: StdRng,
|
|
}
|
|
|
|
impl BernoulliSampler {
|
|
fn new(lower_bound: f64, upper_bound: f64, seed: u64) -> Self {
|
|
Self {
|
|
lower_bound,
|
|
upper_bound,
|
|
rng: StdRng::seed_from_u64(seed),
|
|
}
|
|
}
|
|
|
|
fn sample(&mut self, batch: &RecordBatch) -> Result<RecordBatch> {
|
|
let range = self.upper_bound - self.lower_bound;
|
|
if range <= 0.0 {
|
|
return Ok(RecordBatch::new_empty(batch.schema()));
|
|
}
|
|
|
|
// Select rows where random value falls in [lower, upper)
|
|
let indices: Vec<u32> = (0..batch.num_rows())
|
|
.filter(|_| {
|
|
let r: f64 = self.rng.random();
|
|
r >= self.lower_bound && r < self.upper_bound
|
|
})
|
|
.map(|i| i as u32)
|
|
.collect();
|
|
|
|
if indices.is_empty() {
|
|
return Ok(RecordBatch::new_empty(batch.schema()));
|
|
}
|
|
|
|
compute::take_record_batch(batch, &UInt32Array::from(indices))
|
|
.map_err(DataFusionError::from)
|
|
}
|
|
}
|
|
|
|
/// Stream adapter that applies sampling to each batch.
|
|
struct SampleStream {
|
|
input: SendableRecordBatchStream,
|
|
sampler: BernoulliSampler,
|
|
metrics: BaselineMetrics,
|
|
}
|
|
|
|
impl Stream for SampleStream {
|
|
type Item = Result<RecordBatch>;
|
|
|
|
fn poll_next(
|
|
mut self: Pin<&mut Self>,
|
|
cx: &mut Context<'_>,
|
|
) -> Poll<Option<Self::Item>> {
|
|
match ready!(self.input.poll_next_unpin(cx)) {
|
|
Some(Ok(batch)) => {
|
|
let elapsed = self.metrics.elapsed_compute().clone();
|
|
let _timer = elapsed.timer();
|
|
let result = self.sampler.sample(&batch);
|
|
Poll::Ready(Some(result.record_output(&self.metrics)))
|
|
}
|
|
Some(Err(e)) => Poll::Ready(Some(Err(e))),
|
|
None => Poll::Ready(None),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl RecordBatchStream for SampleStream {
|
|
fn schema(&self) -> SchemaRef {
|
|
self.input.schema()
|
|
}
|
|
}
|