You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/05/26 19:05:21 UTC
(datafusion) 01/01: Add minimal APIs / hooks for granular statistics collection in TableProvider implementations (#22300)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-22300-7d862d64078158ef4e7b456ea7c1baa3f7144f1f
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 633595d0d4017364944d150907df5a9651d6a042
Author: Adrian Garcia Badaracco <17...@users.noreply.github.com>
AuthorDate: Tue May 26 14:05:09 2026 -0500
Add minimal APIs / hooks for granular statistics collection in TableProvider implementations (#22300)
## Which issue does this PR close?
- Related to #21996
- Related to #21624
This is **not** a replacement for #21996 — it is a minimal subset of it,
carved out so the feature can be discussed/merged in smaller pieces.
## Rationale for this change
#21996 ("Query-aware statistics requests via ScanArgs / ScanResult") is
a full vertical slice: new statistics types, request threading optimizer
→ planner → provider, a built-in `RequestStatistics` optimizer rule, and
a consumer integration (`FilePruner` / `ListingTable`).
This PR extracts **only the framework hooks** — just enough that the
rest can be implemented *entirely outside* of DataFusion. A third party
can write their own optimizer rule to derive statistics requests, and
their own `TableProvider` to consume them, without DataFusion shipping
any rule or consumer of its own.
In stock DataFusion nothing observable changes: no rule populates the
new field, and the built-in providers ignore it.
## What changes are included in this PR?
Five small, independently-reviewable commits:
1. **`refactor: add TableScanBuilder, deprecate TableScan::try_new`** —
`TableScan::try_new` takes five positional args and bare `TableScan { ..
}` literals are fragile to field additions. Introduce `TableScanBuilder`
(with `From<TableScan>`), move schema derivation into `build()`,
deprecate `try_new` (delegates to the builder), migrate all in-tree
callers. Pure refactor.
2. **`feat: add StatisticsRequest type`** — a new public
`StatisticsRequest` enum in `datafusion-expr-common::statistics`
(Min/Max/NullCount/DistinctCount/Sum/ByteSize per column, plus
RowCount/TotalByteSize). Nothing consumes it yet.
3. **`feat: add TableScan::statistics_requests field`** — an advisory
`BTreeSet<StatisticsRequest>` on `TableScan`. A **set** so
request-deriving optimizer rules stay idempotent under fixpoint
iteration (a rule `insert`s its requests; re-running is a no-op and
composes with other rules — no per-rule dedup). Empty by default;
DataFusion's own rules never populate it.
4. **`feat: thread statistics requests into ScanArgs`** — `ScanArgs`
gains `statistics_requests`; the physical planner threads
`TableScan::statistics_requests` into it so the request reaches
`TableProvider::scan_with_args`.
5. **`test: e2e statistics-request flow via a custom optimizer rule`** —
an integration test playing both external roles.
Deliberately **left out** vs #21996, since this PR is request-side only:
the built-in `RequestStatistics` optimizer rule, the `FilePruner` /
`ListingTable` consumer integration, the
`PartitionedFile::satisfied_stats` per-file response field, and the
response-side types (`StatisticsValue` / `SatisfiedStatistics`). Those
belong with whatever actually wires the response side.
## Are these changes tested?
Yes:
- `datafusion/core/tests/user_defined/statistics_requests.rs`: an
end-to-end integration test where a custom `OptimizerRule` annotates
`TableScan` and a custom `TableProvider` asserts the requests reach
`scan_with_args` — plus a test that without such a rule the provider
sees an empty request list.
- All existing `datafusion-expr` / `datafusion-optimizer` /
`datafusion-proto` tests pass against the `TableScanBuilder` refactor.
## Are there any user-facing changes?
Yes — this needs the `api change` label:
- New public type `StatisticsRequest` (re-exported via
`datafusion_expr::statistics`).
- New `TableScanBuilder`; `TableScan::try_new` is **deprecated** (still
works, delegates to the builder).
- `TableScan` gains a new public field `statistics_requests` — this
breaks exhaustive `TableScan { .. }` struct literals downstream (the
recommended fix is `TableScanBuilder`).
- `ScanArgs` gains `with_statistics_requests` / `statistics_requests`.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Claude Opus 4.7 (1M context) <no...@anthropic.com>
Co-authored-by: Andrew Lamb <an...@nerdnetworks.org>
---
datafusion/catalog/src/table.rs | 28 +++
datafusion/core/src/physical_planner.rs | 6 +-
datafusion/core/tests/user_defined/mod.rs | 4 +
.../core/tests/user_defined/statistics_requests.rs | 221 +++++++++++++++++++++
datafusion/expr-common/src/statistics.rs | 44 ++++
datafusion/expr/src/logical_plan/builder.rs | 11 +-
datafusion/expr/src/logical_plan/mod.rs | 4 +-
datafusion/expr/src/logical_plan/plan.rs | 114 ++++++++++-
datafusion/expr/src/logical_plan/tree_node.rs | 2 +
.../optimizer/src/optimize_projections/mod.rs | 20 +-
datafusion/optimizer/src/push_down_filter.rs | 1 +
datafusion/proto/src/logical_plan/mod.rs | 15 +-
12 files changed, 432 insertions(+), 38 deletions(-)
diff --git a/datafusion/catalog/src/table.rs b/datafusion/catalog/src/table.rs
index 5d1391bed1..c6468fd5ad 100644
--- a/datafusion/catalog/src/table.rs
+++ b/datafusion/catalog/src/table.rs
@@ -26,6 +26,7 @@ use async_trait::async_trait;
use datafusion_common::{Constraints, Statistics, not_impl_err};
use datafusion_common::{Result, internal_err};
use datafusion_expr::Expr;
+use datafusion_expr::statistics::StatisticsRequest;
use datafusion_expr::dml::InsertOp;
use datafusion_expr::{
@@ -406,6 +407,7 @@ pub struct ScanArgs<'a> {
filters: Option<&'a [Expr]>,
projection: Option<&'a [usize]>,
limit: Option<usize>,
+ statistics_requests: &'a [StatisticsRequest],
}
impl<'a> ScanArgs<'a> {
@@ -467,6 +469,32 @@ impl<'a> ScanArgs<'a> {
pub fn limit(&self) -> Option<usize> {
self.limit
}
+
+ /// Specifies the statistics the caller may use when optimizing the query.
+ ///
+ /// This is intended to allow the `TableProvider` to cheaply provide
+ /// statistics that may help, such as those it has in an in-memory catalog
+ /// or from some other metadata source.
+ ///
+ /// `TableProvider`s read these via [`Self::statistics_requests()`]; anything
+ /// a `TableProvider` cannot answer cheaply it simply ignores. DataFusion's
+ /// own `TableProvider`s ignore this field — it exists so a request can be
+ /// threaded from a custom optimizer rule (which annotates
+ /// `TableScan::statistics_requests`) through to a custom `TableProvider`.
+ pub fn with_statistics_requests(
+ mut self,
+ statistics_requests: &'a [StatisticsRequest],
+ ) -> Self {
+ self.statistics_requests = statistics_requests;
+ self
+ }
+
+ /// Get the statistics requests for the scan. Empty if none were set.
+ ///
+ /// See [`Self::with_statistics_requests`] for more details
+ pub fn statistics_requests(&self) -> &'a [StatisticsRequest] {
+ self.statistics_requests
+ }
}
/// Result of a table scan operation from [`TableProvider::scan_with_args`].
diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs
index ee97309c27..a00d07a09f 100644
--- a/datafusion/core/src/physical_planner.rs
+++ b/datafusion/core/src/physical_planner.rs
@@ -647,6 +647,7 @@ impl DefaultPhysicalPlanner {
filters,
fetch,
projected_schema,
+ statistics_requests,
..
} = scan;
@@ -656,10 +657,13 @@ impl DefaultPhysicalPlanner {
// referred to in the query
let filters = unnormalize_cols(filters.iter().cloned());
let filters_vec = filters.into_iter().collect::<Vec<_>>();
+ let stats_requests =
+ statistics_requests.iter().cloned().collect::<Vec<_>>();
let opts = ScanArgs::default()
.with_projection(projection.as_deref())
.with_filters(Some(&filters_vec))
- .with_limit(*fetch);
+ .with_limit(*fetch)
+ .with_statistics_requests(&stats_requests);
let res = source.scan_with_args(session_state, opts).await?;
Arc::clone(res.plan())
} else {
diff --git a/datafusion/core/tests/user_defined/mod.rs b/datafusion/core/tests/user_defined/mod.rs
index bc9949f5d6..4dad3ec457 100644
--- a/datafusion/core/tests/user_defined/mod.rs
+++ b/datafusion/core/tests/user_defined/mod.rs
@@ -41,3 +41,7 @@ mod relation_planner;
/// Tests for insert operations
mod insert_operation;
+
+/// Tests for `StatisticsRequest`s flowing from a custom optimizer rule
+/// through the physical planner into a custom `TableProvider`.
+mod statistics_requests;
diff --git a/datafusion/core/tests/user_defined/statistics_requests.rs b/datafusion/core/tests/user_defined/statistics_requests.rs
new file mode 100644
index 0000000000..64c6676c23
--- /dev/null
+++ b/datafusion/core/tests/user_defined/statistics_requests.rs
@@ -0,0 +1,221 @@
+// 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.
+
+//! End-to-end test that a *custom* optimizer rule can annotate a
+//! `TableScan` with `StatisticsRequest`s and have them reach a *custom*
+//! `TableProvider`'s `scan_with_args`.
+//!
+//! DataFusion ships no rule that populates `TableScan::statistics_requests`
+//! and no provider that consumes `ScanArgs::statistics_requests`. This test
+//! plays both roles, demonstrating that the request-side hooks are
+//! sufficient to build the whole feature outside of DataFusion.
+
+use std::sync::{Arc, Mutex};
+
+use arrow::array::{Int64Array, RecordBatch};
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use async_trait::async_trait;
+use datafusion::catalog::{ScanArgs, ScanResult, Session, TableProvider};
+use datafusion::common::tree_node::Transformed;
+use datafusion::common::{Column, Result};
+use datafusion::datasource::TableType;
+use datafusion::datasource::memory::MemorySourceConfig;
+use datafusion::execution::context::SessionContext;
+use datafusion::execution::session_state::SessionStateBuilder;
+use datafusion::logical_expr::statistics::StatisticsRequest;
+use datafusion::logical_expr::{Expr, LogicalPlan};
+use datafusion::optimizer::{ApplyOrder, OptimizerConfig, OptimizerRule};
+use datafusion::physical_plan::ExecutionPlan;
+
+/// A custom optimizer rule that annotates every `TableScan` with a
+/// `RowCount` request plus a `Min` request for each of its columns.
+///
+/// This stands in for whatever request-derivation logic an external
+/// implementer would write (e.g. Min/Max for sort keys, DistinctCount for
+/// join keys). Here it is intentionally trivial and deterministic.
+#[derive(Debug)]
+struct RequestColumnStatistics;
+
+impl OptimizerRule for RequestColumnStatistics {
+ fn name(&self) -> &str {
+ "test_request_column_statistics"
+ }
+
+ fn apply_order(&self) -> Option<ApplyOrder> {
+ Some(ApplyOrder::TopDown)
+ }
+
+ fn supports_rewrite(&self) -> bool {
+ true
+ }
+
+ fn rewrite(
+ &self,
+ plan: LogicalPlan,
+ _config: &dyn OptimizerConfig,
+ ) -> Result<Transformed<LogicalPlan>> {
+ let LogicalPlan::TableScan(mut scan) = plan else {
+ return Ok(Transformed::no(plan));
+ };
+ // Insert into the scan's existing request set. `BTreeSet::insert`
+ // reports whether the value was new, so the rule is idempotent — and
+ // composes with other rules' requests for free: re-inserting an
+ // existing request is a no-op, and we report `Transformed::yes` only
+ // when something was actually added, so the optimizer reaches a
+ // fixpoint without a manual "already visited" guard.
+ let mut changed = scan.statistics_requests.insert(StatisticsRequest::RowCount);
+ for field in scan.projected_schema.fields() {
+ let req =
+ StatisticsRequest::Min(Arc::new(Column::new_unqualified(field.name())));
+ changed |= scan.statistics_requests.insert(req);
+ }
+ Ok(if changed {
+ Transformed::yes(LogicalPlan::TableScan(scan))
+ } else {
+ Transformed::no(LogicalPlan::TableScan(scan))
+ })
+ }
+}
+
+/// A `TableProvider` that records the `statistics_requests` it was asked
+/// for, so the test can assert what reached it.
+#[derive(Debug)]
+struct RecordingTable {
+ schema: SchemaRef,
+ batch: RecordBatch,
+ last_requests: Arc<Mutex<Vec<StatisticsRequest>>>,
+}
+
+#[async_trait]
+impl TableProvider for RecordingTable {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+
+ fn table_type(&self) -> TableType {
+ TableType::Base
+ }
+
+ async fn scan(
+ &self,
+ _state: &dyn Session,
+ projection: Option<&Vec<usize>>,
+ _filters: &[Expr],
+ _limit: Option<usize>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(MemorySourceConfig::try_new_exec(
+ &[vec![self.batch.clone()]],
+ Arc::clone(&self.schema),
+ projection.cloned(),
+ )?)
+ }
+
+ async fn scan_with_args<'a>(
+ &self,
+ state: &dyn Session,
+ args: ScanArgs<'a>,
+ ) -> Result<ScanResult> {
+ // Record what reached us, then delegate to `scan`.
+ *self.last_requests.lock().unwrap() = args.statistics_requests().to_vec();
+ let plan = self
+ .scan(
+ state,
+ args.projection().map(|p| p.to_vec()).as_ref(),
+ args.filters().unwrap_or(&[]),
+ args.limit(),
+ )
+ .await?;
+ Ok(ScanResult::new(plan))
+ }
+}
+
+fn make_table() -> (Arc<RecordingTable>, Arc<Mutex<Vec<StatisticsRequest>>>) {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int64, false),
+ Field::new("b", DataType::Int64, false),
+ ]));
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(Int64Array::from(vec![1, 2, 3])),
+ Arc::new(Int64Array::from(vec![10, 20, 30])),
+ ],
+ )
+ .unwrap();
+ let last_requests = Arc::new(Mutex::new(Vec::new()));
+ let provider = Arc::new(RecordingTable {
+ schema,
+ batch,
+ last_requests: Arc::clone(&last_requests),
+ });
+ (provider, last_requests)
+}
+
+#[tokio::test]
+async fn custom_rule_requests_reach_custom_provider() -> Result<()> {
+ let (provider, last_requests) = make_table();
+
+ let state = SessionStateBuilder::new()
+ .with_default_features()
+ .with_optimizer_rule(Arc::new(RequestColumnStatistics))
+ .build();
+ let ctx = SessionContext::new_with_state(state);
+ ctx.register_table("t", provider)?;
+
+ ctx.sql("SELECT a, b FROM t").await?.collect().await?;
+
+ let got = last_requests.lock().unwrap().clone();
+ assert_eq!(
+ got.len(),
+ 3,
+ "expected RowCount + Min(a) + Min(b), got {got:?}"
+ );
+ assert!(
+ got.contains(&StatisticsRequest::RowCount),
+ "expected RowCount, got {got:?}"
+ );
+ assert!(
+ got.contains(&StatisticsRequest::Min(Arc::new(Column::new_unqualified(
+ "a"
+ )))),
+ "expected Min(a), got {got:?}"
+ );
+ assert!(
+ got.contains(&StatisticsRequest::Min(Arc::new(Column::new_unqualified(
+ "b"
+ )))),
+ "expected Min(b), got {got:?}"
+ );
+ Ok(())
+}
+
+#[tokio::test]
+async fn no_requests_without_a_rule() -> Result<()> {
+ // Without a rule populating `TableScan::statistics_requests`, the
+ // provider sees an empty request list — stock DataFusion behavior.
+ let (provider, last_requests) = make_table();
+ let ctx = SessionContext::new();
+ ctx.register_table("t", provider)?;
+
+ ctx.sql("SELECT a, b FROM t").await?.collect().await?;
+
+ assert!(
+ last_requests.lock().unwrap().is_empty(),
+ "expected no requests without a custom rule"
+ );
+ Ok(())
+}
diff --git a/datafusion/expr-common/src/statistics.rs b/datafusion/expr-common/src/statistics.rs
index c94c181615..034358b043 100644
--- a/datafusion/expr-common/src/statistics.rs
+++ b/datafusion/expr-common/src/statistics.rs
@@ -1694,3 +1694,47 @@ mod tests {
all_ops.into_iter().collect()
}
}
+
+use std::sync::Arc;
+
+use datafusion_common::Column;
+
+/// A statistic a caller would like a provider to supply, if it can do so
+/// cheaply.
+///
+/// A small, query-aware extension to the existing `Statistics` model: instead
+/// of "give me everything you have for every column", a caller can ask for a
+/// specific list of stats by name. `StatisticsRequest` is just that vocabulary
+/// — DataFusion itself does not populate or consume it. It exists so a request
+/// can be threaded from a `TableScan` (see `TableScan::statistics_requests`)
+/// through `ScanArgs::statistics_requests` to a `TableProvider`, which is enough
+/// for a query-aware statistics feature to be implemented outside of DataFusion.
+///
+/// Each variant maps onto a field of [`datafusion_common::Statistics`] /
+/// [`datafusion_common::ColumnStatistics`], so a provider that already
+/// populates one can answer the request trivially.
+///
+/// The per-column variants hold an `Arc<Column>` rather than an owned
+/// [`Column`] (which carries owned strings) so cloning a request — and the
+/// `BTreeSet<StatisticsRequest>` stored on `TableScan`, which is cloned with
+/// the plan during optimization — stays cheap.
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
+pub enum StatisticsRequest {
+ /// Smallest non-null value of `column`.
+ Min(Arc<Column>),
+ /// Largest non-null value of `column`.
+ Max(Arc<Column>),
+ /// Number of NULLs in `column`.
+ NullCount(Arc<Column>),
+ /// Number of distinct values in `column` (exact or estimated).
+ DistinctCount(Arc<Column>),
+ /// Sum of values in `column` (numerics, widened per
+ /// `ColumnStatistics::sum_value`).
+ Sum(Arc<Column>),
+ /// Encoded/output byte size of `column`.
+ ByteSize(Arc<Column>),
+ /// Number of rows in the container (table / file).
+ RowCount,
+ /// Total byte size of the container's output.
+ TotalByteSize,
+}
diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs
index 8c03374578..7bc705e0f4 100644
--- a/datafusion/expr/src/logical_plan/builder.rs
+++ b/datafusion/expr/src/logical_plan/builder.rs
@@ -33,8 +33,8 @@ use crate::expr_rewriter::{
use crate::logical_plan::{
Aggregate, Analyze, Distinct, DistinctOn, EmptyRelation, Explain, Filter, Join,
JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, Prepare,
- Projection, Repartition, Sort, SubqueryAlias, TableScan, Union, Unnest, Values,
- Window,
+ Projection, Repartition, Sort, SubqueryAlias, TableScanBuilder, Union, Unnest,
+ Values, Window,
};
use crate::select_expr::SelectExpr;
use crate::utils::{
@@ -515,8 +515,11 @@ impl LogicalPlanBuilder {
filters: Vec<Expr>,
fetch: Option<usize>,
) -> Result<Self> {
- let table_scan =
- TableScan::try_new(table_name, table_source, projection, filters, fetch)?;
+ let table_scan = TableScanBuilder::new(table_name, table_source)
+ .with_projection(projection)
+ .with_filters(filters)
+ .with_fetch(fetch)
+ .build()?;
// Inline TableScan
if table_scan.filters.is_empty()
diff --git a/datafusion/expr/src/logical_plan/mod.rs b/datafusion/expr/src/logical_plan/mod.rs
index c2b01868c9..5087b25178 100644
--- a/datafusion/expr/src/logical_plan/mod.rs
+++ b/datafusion/expr/src/logical_plan/mod.rs
@@ -42,8 +42,8 @@ pub use plan::{
EmptyRelation, Explain, ExplainOption, Extension, FetchType, Filter, Join,
JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, Projection,
RecursiveQuery, Repartition, SkipType, Sort, StringifiedPlan, Subquery,
- SubqueryAlias, TableScan, ToStringifiedPlan, Union, Unnest, Values, Window,
- projection_schema,
+ SubqueryAlias, TableScan, TableScanBuilder, ToStringifiedPlan, Union, Unnest, Values,
+ Window, projection_schema,
};
pub use statement::{
Deallocate, Execute, Prepare, ResetVariable, SetVariable, Statement,
diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs
index 2f1061c438..e7e03bcac5 100644
--- a/datafusion/expr/src/logical_plan/plan.rs
+++ b/datafusion/expr/src/logical_plan/plan.rs
@@ -18,7 +18,7 @@
//! Logical plan types
use std::cmp::Ordering;
-use std::collections::{HashMap, HashSet};
+use std::collections::{BTreeSet, HashMap, HashSet};
use std::fmt::{self, Debug, Display, Formatter};
use std::hash::{Hash, Hasher};
use std::sync::{Arc, LazyLock};
@@ -50,6 +50,7 @@ use crate::{
WindowFunctionDefinition, build_join_schema, expr_vec_fmt, requalify_sides_if_needed,
};
+use crate::statistics::StatisticsRequest;
use arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef};
use datafusion_common::cse::{NormalizeEq, Normalizeable};
use datafusion_common::format::ExplainFormat;
@@ -2793,6 +2794,12 @@ pub struct TableScan {
pub filters: Vec<Expr>,
/// Optional number of rows to read
pub fetch: Option<usize>,
+ /// Statistics the planner would like the provider to answer for this
+ /// scan, typically attached by a custom optimizer rule from the
+ /// surrounding plan (e.g. Min/Max for sort keys).
+ ///
+ /// A [`BTreeSet`], not a `Vec` to keep the resulting plan deterministic.
+ pub statistics_requests: BTreeSet<StatisticsRequest>,
}
impl Debug for TableScan {
@@ -2867,6 +2874,7 @@ impl Hash for TableScan {
impl TableScan {
/// Initialize TableScan with appropriate schema from the given
/// arguments.
+ #[deprecated(since = "54.0.0", note = "use `TableScanBuilder` instead")]
pub fn try_new(
table_name: impl Into<TableReference>,
table_source: Arc<dyn TableSource>,
@@ -2874,14 +2882,92 @@ impl TableScan {
filters: Vec<Expr>,
fetch: Option<usize>,
) -> Result<Self> {
- let table_name = table_name.into();
+ TableScanBuilder::new(table_name, table_source)
+ .with_projection(projection)
+ .with_filters(filters)
+ .with_fetch(fetch)
+ .build()
+ }
+}
+
+/// Builder for [`TableScan`].
+///
+/// Prefer this over constructing a [`TableScan`] directly: it derives the
+/// `projected_schema` from the source schema and projection, and is resilient
+/// to new fields being added to [`TableScan`]. An existing scan can be turned
+/// back into a builder with `TableScanBuilder::from(scan)`, tweaked, and
+/// rebuilt with [`TableScanBuilder::build`].
+pub struct TableScanBuilder {
+ table_name: TableReference,
+ source: Arc<dyn TableSource>,
+ projection: Option<Vec<usize>>,
+ filters: Vec<Expr>,
+ fetch: Option<usize>,
+ statistics_requests: BTreeSet<StatisticsRequest>,
+}
+
+impl TableScanBuilder {
+ /// Create a new builder for a scan of `source` named `table_name`.
+ pub fn new(
+ table_name: impl Into<TableReference>,
+ source: Arc<dyn TableSource>,
+ ) -> Self {
+ Self {
+ table_name: table_name.into(),
+ source,
+ projection: None,
+ filters: vec![],
+ fetch: None,
+ statistics_requests: BTreeSet::new(),
+ }
+ }
+
+ /// Set the column projection (indices into the source schema).
+ pub fn with_projection(mut self, projection: Option<Vec<usize>>) -> Self {
+ self.projection = projection;
+ self
+ }
+
+ /// Set the filter expressions offered to the table provider.
+ pub fn with_filters(mut self, filters: Vec<Expr>) -> Self {
+ self.filters = filters;
+ self
+ }
+
+ /// Set the maximum number of rows to read.
+ pub fn with_fetch(mut self, fetch: Option<usize>) -> Self {
+ self.fetch = fetch;
+ self
+ }
+
+ /// Set the statistics requests for the scan. See
+ /// [`TableScan::statistics_requests`].
+ pub fn with_statistics_requests(
+ mut self,
+ statistics_requests: BTreeSet<StatisticsRequest>,
+ ) -> Self {
+ self.statistics_requests = statistics_requests;
+ self
+ }
+
+ /// Build the [`TableScan`], deriving its `projected_schema` from the
+ /// source schema and projection.
+ pub fn build(self) -> Result<TableScan> {
+ let TableScanBuilder {
+ table_name,
+ source,
+ projection,
+ filters,
+ fetch,
+ statistics_requests,
+ } = self;
if table_name.table().is_empty() {
return plan_err!("table_name cannot be empty");
}
- let schema = table_source.schema();
+ let schema = source.schema();
let func_dependencies = FunctionalDependencies::new_from_constraints(
- table_source.constraints(),
+ source.constraints(),
schema.fields.len(),
);
let projected_schema = projection
@@ -2907,17 +2993,31 @@ impl TableScan {
})?;
let projected_schema = Arc::new(projected_schema);
- Ok(Self {
+ Ok(TableScan {
table_name,
- source: table_source,
+ source,
projection,
projected_schema,
filters,
fetch,
+ statistics_requests,
})
}
}
+impl From<TableScan> for TableScanBuilder {
+ fn from(scan: TableScan) -> Self {
+ Self {
+ table_name: scan.table_name,
+ source: scan.source,
+ projection: scan.projection,
+ filters: scan.filters,
+ fetch: scan.fetch,
+ statistics_requests: scan.statistics_requests,
+ }
+ }
+}
+
// Repartition the plan based on a partitioning scheme.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
pub struct Repartition {
@@ -5152,6 +5252,7 @@ mod tests {
projected_schema: Arc::clone(&schema),
filters: vec![],
fetch: None,
+ statistics_requests: BTreeSet::new(),
}));
let col = schema.field_names()[0].clone();
@@ -5182,6 +5283,7 @@ mod tests {
projected_schema: Arc::clone(&unique_schema),
filters: vec![],
fetch: None,
+ statistics_requests: BTreeSet::new(),
}));
let col = schema.field_names()[0].clone();
diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs
index 1f58de37e9..801caddcd0 100644
--- a/datafusion/expr/src/logical_plan/tree_node.rs
+++ b/datafusion/expr/src/logical_plan/tree_node.rs
@@ -615,6 +615,7 @@ impl LogicalPlan {
projected_schema,
filters,
fetch,
+ statistics_requests,
}) => filters.map_elements(f)?.update_data(|filters| {
LogicalPlan::TableScan(TableScan {
table_name,
@@ -623,6 +624,7 @@ impl LogicalPlan {
projected_schema,
filters,
fetch,
+ statistics_requests,
})
}),
LogicalPlan::Distinct(Distinct::On(DistinctOn {
diff --git a/datafusion/optimizer/src/optimize_projections/mod.rs b/datafusion/optimizer/src/optimize_projections/mod.rs
index 59109a822b..b9f22a3f9e 100644
--- a/datafusion/optimizer/src/optimize_projections/mod.rs
+++ b/datafusion/optimizer/src/optimize_projections/mod.rs
@@ -29,8 +29,8 @@ use datafusion_common::{
};
use datafusion_expr::expr::Alias;
use datafusion_expr::{
- Aggregate, Distinct, EmptyRelation, Expr, Projection, TableScan, Unnest, Window,
- logical_plan::LogicalPlan,
+ Aggregate, Distinct, EmptyRelation, Expr, Projection, TableScanBuilder, Unnest,
+ Window, logical_plan::LogicalPlan,
};
use crate::optimize_projections::required_indices::RequiredIndices;
@@ -269,23 +269,15 @@ fn optimize_projections(
.transform_data(|plan| optimize_subqueries(plan, config));
}
LogicalPlan::TableScan(table_scan) => {
- let TableScan {
- table_name,
- source,
- projection,
- filters,
- fetch,
- projected_schema: _,
- } = table_scan;
-
// Get indices referred to in the original (schema with all fields)
// given projected indices.
- let projection = match &projection {
+ let projection = match &table_scan.projection {
Some(projection) => indices.into_mapped_indices(|idx| projection[idx]),
None => indices.into_inner(),
};
- let new_scan =
- TableScan::try_new(table_name, source, Some(projection), filters, fetch)?;
+ let new_scan = TableScanBuilder::from(table_scan)
+ .with_projection(Some(projection))
+ .build()?;
return Transformed::yes(LogicalPlan::TableScan(new_scan))
.transform_data(|plan| optimize_subqueries(plan, config));
diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs
index 9c2ac07ff0..54878d2f54 100644
--- a/datafusion/optimizer/src/push_down_filter.rs
+++ b/datafusion/optimizer/src/push_down_filter.rs
@@ -3119,6 +3119,7 @@ mod tests {
projection,
source: Arc::new(test_provider),
fetch: None,
+ statistics_requests: std::collections::BTreeSet::new(),
});
Ok(LogicalPlanBuilder::from(table_scan))
diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs
index 542cae890d..e378532667 100644
--- a/datafusion/proto/src/logical_plan/mod.rs
+++ b/datafusion/proto/src/logical_plan/mod.rs
@@ -40,7 +40,7 @@ use datafusion_catalog::empty::EmptyTable;
use datafusion_common::file_options::file_type::FileType;
use datafusion_common::format::ExplainFormat;
use datafusion_common::{
- NullEquality, Result, TableReference, ToDFSchema, assert_or_internal_err, context,
+ NullEquality, Result, TableReference, assert_or_internal_err, context,
internal_datafusion_err, internal_err, not_impl_err, plan_err,
};
use datafusion_datasource::file_format::FileFormat;
@@ -66,7 +66,8 @@ use datafusion_expr::{
logical_plan::{
Aggregate, CreateCatalog, CreateCatalogSchema, CreateExternalTable, CreateView,
DdlStatement, Distinct, EmptyRelation, Extension, Join, Prepare, Projection,
- Repartition, Sort, SubqueryAlias, TableScan, Values, Window, builder::project,
+ Repartition, Sort, SubqueryAlias, TableScan, TableScanBuilder, Values, Window,
+ builder::project,
},
};
@@ -371,15 +372,7 @@ fn from_table_source(
target: Arc<dyn TableSource>,
extension_codec: &dyn LogicalExtensionCodec,
) -> Result<LogicalPlanNode> {
- let projected_schema = target.schema().to_dfschema_ref()?;
- let r = LogicalPlan::TableScan(TableScan {
- table_name,
- source: target,
- projection: None,
- projected_schema,
- filters: vec![],
- fetch: None,
- });
+ let r = LogicalPlan::TableScan(TableScanBuilder::new(table_name, target).build()?);
LogicalPlanNode::try_from_logical_plan(&r, extension_codec)
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org