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/07/31 10:23:04 UTC
(datafusion) branch gh-readonly-queue/main/pr-23842-aa9b7bb63f14211688ffb1c1610700ca54975338 created (now dbcb5c0f72)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-23842-aa9b7bb63f14211688ffb1c1610700ca54975338
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at dbcb5c0f72 refactor: move planning APIs to session crate (#23842)
This branch includes the following new commits:
new dbcb5c0f72 refactor: move planning APIs to session crate (#23842)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: refactor: move planning APIs to session crate (#23842)
Posted by gi...@apache.org.
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-23842-aa9b7bb63f14211688ffb1c1610700ca54975338
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48
Author: Tim Saucer <ti...@gmail.com>
AuthorDate: Fri Jul 31 06:22:43 2026 -0400
refactor: move planning APIs to session crate (#23842)
## Which issue does this PR close?
- Addresses a portion of
https://github.com/apache/datafusion/issues/23678
- Follow on to https://github.com/apache/datafusion/pull/23703
## Rationale for this change
This PR unlocks using a query planner across the FFI boundary.
## What changes are included in this PR?
Moves these traits to the `datafusion-session` crate:
- `QueryPlanner`
- `PhysicalPlanner`
- `ExtensionPlanner`
- `PhysicalOptimizerRule`
- `PhysicalOptimizerContext`
Additionally changed the method signatures from taking `&SessionState`
to `&dyn Session`.
Adds these methods on `Session` trait:
- `fn query_planner(&self) -> Arc<dyn QueryPlanner + Send + Sync>`
- `fn optimize(&self, plan: &LogicalPlan) -> Result<LogicalPlan>`
- `fn physical_optimizers(&self) -> &[Arc<dyn PhysicalOptimizerRule +
Send + Sync>]`
- `fn statistics_registry(&self) -> Option<&StatisticsRegistry>`
## Are these changes tested?
Unit tests are added for the new methods.
## Are there any user-facing changes?
Yes, users must implement new methods for `query_planner` for their
custom `Session`. This is not expected to impact many users, since it is
most common to use the existing `SessionState`.
---------
Co-authored-by: Claude Opus 4.8 (1M context) <no...@anthropic.com>
---
Cargo.lock | 1 +
.../examples/dataframe/cache_factory.rs | 5 +-
.../examples/relation_planner/table_sample.rs | 9 +-
datafusion/core/src/execution/context/mod.rs | 19 +-
datafusion/core/src/execution/session_state.rs | 25 +-
datafusion/core/src/physical_planner.rs | 413 ++++++++++++---------
.../core/tests/user_defined/user_defined_plan.rs | 7 +-
datafusion/ffi/src/session/mod.rs | 10 +
datafusion/physical-optimizer/Cargo.toml | 1 +
datafusion/physical-optimizer/src/optimizer.rs | 64 +---
datafusion/session/src/lib.rs | 10 +
datafusion/session/src/physical_optimizer.rs | 84 +++++
datafusion/session/src/planner.rs | 198 ++++++++++
datafusion/session/src/session.rs | 51 +++
docs/source/library-user-guide/upgrading/55.0.0.md | 77 +++-
15 files changed, 698 insertions(+), 276 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index 9cd5d12d5b..619fee7603 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2463,6 +2463,7 @@ dependencies = [
"datafusion-physical-expr-common",
"datafusion-physical-plan",
"datafusion-pruning",
+ "datafusion-session",
"insta",
"itertools 0.15.0",
"recursive",
diff --git a/datafusion-examples/examples/dataframe/cache_factory.rs b/datafusion-examples/examples/dataframe/cache_factory.rs
index dd145c715f..ffbce298b4 100644
--- a/datafusion-examples/examples/dataframe/cache_factory.rs
+++ b/datafusion-examples/examples/dataframe/cache_factory.rs
@@ -23,6 +23,7 @@ use std::sync::{Arc, RwLock};
use arrow::array::RecordBatch;
use async_trait::async_trait;
+use datafusion::catalog::Session;
use datafusion::catalog::memory::MemorySourceConfig;
use datafusion::common::DFSchemaRef;
use datafusion::error::Result;
@@ -146,7 +147,7 @@ impl ExtensionPlanner for CacheNodePlanner {
node: &dyn UserDefinedLogicalNode,
logical_inputs: &[&LogicalPlan],
physical_inputs: &[Arc<dyn ExecutionPlan>],
- session_state: &SessionState,
+ session_state: &dyn Session,
_planning_ctx: &PhysicalPlanningContext,
) -> Result<Option<Arc<dyn ExecutionPlan>>> {
if let Some(cache_node) = node.as_any().downcast_ref::<CacheNode>() {
@@ -200,7 +201,7 @@ impl QueryPlanner for CacheNodeQueryPlanner {
async fn create_physical_plan(
&self,
logical_plan: &LogicalPlan,
- session_state: &SessionState,
+ session_state: &dyn Session,
) -> Result<Arc<dyn ExecutionPlan>> {
let physical_planner =
DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(
diff --git a/datafusion-examples/examples/relation_planner/table_sample.rs b/datafusion-examples/examples/relation_planner/table_sample.rs
index b0ccff8d10..c019e136cc 100644
--- a/datafusion-examples/examples/relation_planner/table_sample.rs
+++ b/datafusion-examples/examples/relation_planner/table_sample.rs
@@ -102,9 +102,10 @@ use tonic::async_trait;
use datafusion::optimizer::simplify_expressions::simplify_literal::parse_literal;
use datafusion::{
+ catalog::Session,
execution::{
- RecordBatchStream, SendableRecordBatchStream, SessionState, SessionStateBuilder,
- TaskContext, context::QueryPlanner,
+ RecordBatchStream, SendableRecordBatchStream, SessionStateBuilder, TaskContext,
+ context::QueryPlanner,
},
physical_expr::EquivalenceProperties,
physical_plan::{
@@ -565,7 +566,7 @@ impl QueryPlanner for TableSampleQueryPlanner {
async fn create_physical_plan(
&self,
logical_plan: &LogicalPlan,
- session_state: &SessionState,
+ session_state: &dyn Session,
) -> Result<Arc<dyn ExecutionPlan>> {
let planner = DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(
TableSampleExtensionPlanner,
@@ -587,7 +588,7 @@ impl ExtensionPlanner for TableSampleExtensionPlanner {
node: &dyn UserDefinedLogicalNode,
_logical_inputs: &[&LogicalPlan],
physical_inputs: &[Arc<dyn ExecutionPlan>],
- _session_state: &SessionState,
+ _session_state: &dyn Session,
_planning_ctx: &PhysicalPlanningContext,
) -> Result<Option<Arc<dyn ExecutionPlan>>> {
let Some(sample_node) = node.as_any().downcast_ref::<TableSamplePlanNode>()
diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs
index cd30193e30..5b287f103a 100644
--- a/datafusion/core/src/execution/context/mod.rs
+++ b/datafusion/core/src/execution/context/mod.rs
@@ -17,7 +17,6 @@
//! [`SessionContext`] API for registering data sources and executing queries
-use std::any::Any;
use std::collections::HashSet;
use std::fmt::Debug;
use std::sync::{Arc, Weak};
@@ -2181,16 +2180,9 @@ impl From<SessionContext> for SessionStateBuilder {
}
}
+// Re-export from this module for backwards compatibility.
/// A planner used to add extensions to DataFusion logical and physical plans.
-#[async_trait]
-pub trait QueryPlanner: Any + Debug {
- /// Given a [`LogicalPlan`], create an [`ExecutionPlan`] suitable for execution
- async fn create_physical_plan(
- &self,
- logical_plan: &LogicalPlan,
- session_state: &SessionState,
- ) -> Result<Arc<dyn ExecutionPlan>>;
-}
+pub use datafusion_session::{QueryPlanner, UnsupportedQueryPlanner};
/// Interface for handling `CREATE FUNCTION` statements and interacting with
/// [SessionState] to create and register functions ([`ScalarUDF`],
@@ -2390,6 +2382,7 @@ mod tests {
use crate::physical_planner::PhysicalPlanner;
use async_trait::async_trait;
use datafusion_expr::planner::TypePlanner;
+ use datafusion_session::Session;
use sqlparser::ast;
use tempfile::TempDir;
@@ -2836,7 +2829,7 @@ mod tests {
async fn create_physical_plan(
&self,
_logical_plan: &LogicalPlan,
- _session_state: &SessionState,
+ _session_state: &dyn Session,
) -> Result<Arc<dyn ExecutionPlan>> {
not_impl_err!("query not supported")
}
@@ -2845,7 +2838,7 @@ mod tests {
&self,
_expr: &Expr,
_input_dfschema: &DFSchema,
- _session_state: &SessionState,
+ _session_state: &dyn Session,
_planning_ctx: &PhysicalPlanningContext,
) -> Result<Arc<dyn PhysicalExpr>> {
unimplemented!()
@@ -2860,7 +2853,7 @@ mod tests {
async fn create_physical_plan(
&self,
logical_plan: &LogicalPlan,
- session_state: &SessionState,
+ session_state: &dyn Session,
) -> Result<Arc<dyn ExecutionPlan>> {
let physical_planner = MyPhysicalPlanner {};
physical_planner
diff --git a/datafusion/core/src/execution/session_state.rs b/datafusion/core/src/execution/session_state.rs
index dfdbb1617e..bfd38faacf 100644
--- a/datafusion/core/src/execution/session_state.rs
+++ b/datafusion/core/src/execution/session_state.rs
@@ -73,12 +73,10 @@ use datafusion_optimizer::{
};
use datafusion_physical_expr::create_physical_expr;
use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
-use datafusion_physical_optimizer::PhysicalOptimizerContext;
-use datafusion_physical_optimizer::PhysicalOptimizerRule;
use datafusion_physical_optimizer::optimizer::PhysicalOptimizer;
use datafusion_physical_plan::ExecutionPlan;
use datafusion_physical_plan::operator_statistics::StatisticsRegistry;
-use datafusion_session::Session;
+use datafusion_session::{PhysicalOptimizerContext, PhysicalOptimizerRule, Session};
#[cfg(feature = "sql")]
use datafusion_sql::{
parser::{DFParserBuilder, Statement},
@@ -274,6 +272,25 @@ impl Session for SessionState {
Arc::clone(self.catalog_list())
}
+ fn query_planner(&self) -> Arc<dyn QueryPlanner + Send + Sync> {
+ // Disambiguate: `SessionState` has an inherent `query_planner` (returning
+ // `&Arc<...>`) with the same name as this trait method. The qualified path
+ // calls the inherent one; a bare `self.query_planner()` would recurse.
+ Arc::clone(SessionState::query_planner(self))
+ }
+
+ fn optimize(&self, plan: &LogicalPlan) -> datafusion_common::Result<LogicalPlan> {
+ SessionState::optimize(self, plan)
+ }
+
+ fn physical_optimizers(&self) -> &[Arc<dyn PhysicalOptimizerRule + Send + Sync>] {
+ SessionState::physical_optimizers(self)
+ }
+
+ fn statistics_registry(&self) -> Option<&StatisticsRegistry> {
+ SessionState::statistics_registry(self)
+ }
+
async fn create_physical_plan(
&self,
logical_plan: &LogicalPlan,
@@ -2322,7 +2339,7 @@ impl QueryPlanner for DefaultQueryPlanner {
async fn create_physical_plan(
&self,
logical_plan: &LogicalPlan,
- session_state: &SessionState,
+ session_state: &dyn Session,
) -> datafusion_common::Result<Arc<dyn ExecutionPlan>> {
let planner = DefaultPhysicalPlanner::default();
planner
diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs
index 4e914556b4..222511b901 100644
--- a/datafusion/core/src/physical_planner.rs
+++ b/datafusion/core/src/physical_planner.rs
@@ -26,14 +26,12 @@ use crate::datasource::listing::ListingTableUrl;
use crate::datasource::physical_plan::{FileOutputMode, FileSinkConfig};
use crate::datasource::{DefaultTableSource, source_as_provider};
use crate::error::{DataFusionError, Result};
-use crate::execution::context::{ExecutionProps, SessionState};
+use crate::execution::context::ExecutionProps;
use crate::logical_expr::utils::generate_sort_key;
use crate::logical_expr::{
Aggregate, EmptyRelation, Join, Projection, Sort, TableScan, Unnest, Values, Window,
};
-use crate::logical_expr::{
- Expr, LogicalPlan, PlanType, Repartition, UserDefinedLogicalNode,
-};
+use crate::logical_expr::{Expr, LogicalPlan, PlanType, Repartition};
use crate::physical_expr::{
create_physical_expr, create_physical_exprs, create_physical_partitioning,
};
@@ -103,7 +101,6 @@ use datafusion_physical_expr::expressions::Literal;
use datafusion_physical_expr::{
LexOrdering, PhysicalSortExpr, create_physical_sort_exprs,
};
-use datafusion_physical_optimizer::PhysicalOptimizerRule;
use datafusion_physical_plan::empty::EmptyExec;
use datafusion_physical_plan::execution_plan::InvariantLevel;
use datafusion_physical_plan::joins::PiecewiseMergeJoinExec;
@@ -111,6 +108,7 @@ use datafusion_physical_plan::placeholder_row::PlaceholderRowExec;
use datafusion_physical_plan::recursive_query::RecursiveQueryExec;
use datafusion_physical_plan::scalar_subquery::{ScalarSubqueryExec, ScalarSubqueryLink};
use datafusion_physical_plan::unnest::ListUnnest;
+use datafusion_session::{PhysicalOptimizerContext, PhysicalOptimizerRule, Session};
use async_trait::async_trait;
use datafusion_physical_plan::async_func::{AsyncFuncExec, AsyncMapper};
@@ -120,162 +118,31 @@ use itertools::{Itertools, multiunzip};
use log::debug;
use tokio::sync::Mutex;
-/// Physical query planner that converts a `LogicalPlan` to an
-/// `ExecutionPlan` suitable for execution.
-#[async_trait]
-pub trait PhysicalPlanner: Send + Sync {
- /// Create a physical plan from a logical plan
- async fn create_physical_plan(
- &self,
- logical_plan: &LogicalPlan,
- session_state: &SessionState,
- ) -> Result<Arc<dyn ExecutionPlan>>;
+// Re-export from this module for backwards compatibility.
+pub use datafusion_session::{ExtensionPlanner, PhysicalPlanner};
- /// Create a physical expression from a logical expression
- /// suitable for evaluation
- ///
- /// `expr`: the expression to convert
- ///
- /// `input_dfschema`: the logical plan schema for evaluating `expr`
- ///
- /// `planning_ctx`: the [`PhysicalPlanningContext`] used to resolve
- /// `Expr::ScalarSubquery` nodes. During physical planning the planner
- /// threads the context of the plan currently being converted to a physical
- /// plan (for example into [`ExtensionPlanner::plan_extension`], which
- /// should forward it here). Callers creating physical expressions outside
- /// of a plan should pass `&PhysicalPlanningContext::default()`.
- fn create_physical_expr(
- &self,
- expr: &Expr,
- input_dfschema: &DFSchema,
- session_state: &SessionState,
- planning_ctx: &PhysicalPlanningContext,
- ) -> Result<Arc<dyn PhysicalExpr>>;
+struct SessionOptimizerContext<'a> {
+ session: &'a dyn Session,
}
-/// This trait exposes the ability to plan an [`ExecutionPlan`] out of a [`LogicalPlan`].
-#[async_trait]
-pub trait ExtensionPlanner {
- /// Create a physical plan for a [`UserDefinedLogicalNode`].
- ///
- /// `input_dfschema`: the logical plan schema for the inputs to this node
- ///
- /// Returns an error when the planner knows how to plan the concrete
- /// implementation of `node` but errors while doing so.
- ///
- /// Returns `None` when the planner does not know how to plan the
- /// `node` and wants to delegate the planning to another
- /// [`ExtensionPlanner`].
- ///
- /// `planning_ctx` is the [`PhysicalPlanningContext`] of the plan subtree
- /// currently being converted to a physical plan. Forward it to
- /// [`PhysicalPlanner::create_physical_expr`] when creating this node's
- /// physical expressions so that scalar subqueries resolve against the same
- /// subquery state as the rest of the plan.
- async fn plan_extension(
- &self,
- planner: &dyn PhysicalPlanner,
- node: &dyn UserDefinedLogicalNode,
- logical_inputs: &[&LogicalPlan],
- physical_inputs: &[Arc<dyn ExecutionPlan>],
- session_state: &SessionState,
- planning_ctx: &PhysicalPlanningContext,
- ) -> Result<Option<Arc<dyn ExecutionPlan>>>;
+impl PhysicalOptimizerContext for SessionOptimizerContext<'_> {
+ fn config_options(&self) -> &datafusion_common::config::ConfigOptions {
+ self.session.config_options()
+ }
- /// Create a physical plan for a [`LogicalPlan::TableScan`].
- ///
- /// This is useful for planning valid [`TableSource`]s that are not [`TableProvider`]s.
- ///
- /// Returns:
- /// * `Ok(Some(plan))` if the planner knows how to plan the `scan`
- /// * `Ok(None)` if the planner does not know how to plan the `scan` and wants to delegate the planning to another [`ExtensionPlanner`]
- /// * `Err` if the planner knows how to plan the `scan` but errors while doing so
- ///
- /// # Example
- ///
- /// ```rust,ignore
- /// use std::sync::Arc;
- /// use datafusion::physical_plan::ExecutionPlan;
- /// use datafusion::logical_expr::TableScan;
- /// use datafusion::execution::context::SessionState;
- /// use datafusion::error::Result;
- /// use datafusion_physical_planner::{ExtensionPlanner, PhysicalPlanner};
- /// use async_trait::async_trait;
- ///
- /// // Your custom table source type
- /// struct MyCustomTableSource { /* ... */ }
- ///
- /// // Your custom execution plan
- /// struct MyCustomExec { /* ... */ }
- ///
- /// struct MyExtensionPlanner;
- ///
- /// #[async_trait]
- /// impl ExtensionPlanner for MyExtensionPlanner {
- /// async fn plan_extension(
- /// &self,
- /// _planner: &dyn PhysicalPlanner,
- /// _node: &dyn UserDefinedLogicalNode,
- /// _logical_inputs: &[&LogicalPlan],
- /// _physical_inputs: &[Arc<dyn ExecutionPlan>],
- /// _session_state: &SessionState,
- /// _planning_ctx: &PhysicalPlanningContext,
- /// ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
- /// Ok(None)
- /// }
- ///
- /// async fn plan_table_scan(
- /// &self,
- /// _planner: &dyn PhysicalPlanner,
- /// scan: &TableScan,
- /// _session_state: &SessionState,
- /// _planning_ctx: &PhysicalPlanningContext,
- /// ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
- /// // Check if this is your custom table source
- /// if scan.source.is::<MyCustomTableSource>() {
- /// // Create a custom execution plan for your table source
- /// let exec = MyCustomExec::new(
- /// scan.table_name.clone(),
- /// Arc::clone(scan.projected_schema.inner()),
- /// );
- /// Ok(Some(Arc::new(exec)))
- /// } else {
- /// // Return None to let other extension planners handle it
- /// Ok(None)
- /// }
- /// }
- /// }
- /// ```
- ///
- /// [`TableSource`]: datafusion_expr::TableSource
- /// [`TableProvider`]: datafusion_catalog::TableProvider
- async fn plan_table_scan(
+ fn statistics_registry(
&self,
- _planner: &dyn PhysicalPlanner,
- _scan: &TableScan,
- _session_state: &SessionState,
- _planning_ctx: &PhysicalPlanningContext,
- ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
- Ok(None)
+ ) -> Option<&datafusion_physical_plan::operator_statistics::StatisticsRegistry> {
+ self.session.statistics_registry()
}
}
/// Default single node physical query planner that converts a
/// `LogicalPlan` to an `ExecutionPlan` suitable for execution.
///
-/// This planner will first flatten the `LogicalPlan` tree via a
-/// depth first approach, which allows it to identify the leaves
-/// of the tree.
-///
-/// Tasks are spawned from these leaves and traverse back up the
-/// tree towards the root, converting each `LogicalPlan` node it
-/// reaches into their equivalent `ExecutionPlan` node. When these
-/// tasks reach a common node, they will terminate until the last
-/// task reaches the node which will then continue building up the
-/// tree.
-///
-/// Up to [`planning_concurrency`] tasks are buffered at once to
-/// execute concurrently.
+/// This planner first flattens the `LogicalPlan` tree with a depth-first
+/// traversal. It then builds the physical plan from the leaves to the root.
+/// Up to [`planning_concurrency`] tasks execute concurrently.
///
/// [`planning_concurrency`]: crate::config::ExecutionOptions::planning_concurrency
#[derive(Default)]
@@ -289,7 +156,7 @@ impl PhysicalPlanner for DefaultPhysicalPlanner {
async fn create_physical_plan(
&self,
logical_plan: &LogicalPlan,
- session_state: &SessionState,
+ session_state: &dyn Session,
) -> Result<Arc<dyn ExecutionPlan>> {
if let Some(plan) = self
.handle_explain_or_analyze(logical_plan, session_state)
@@ -314,7 +181,7 @@ impl PhysicalPlanner for DefaultPhysicalPlanner {
&self,
expr: &Expr,
input_dfschema: &DFSchema,
- session_state: &SessionState,
+ session_state: &dyn Session,
planning_ctx: &PhysicalPlanningContext,
) -> Result<Arc<dyn PhysicalExpr>> {
create_physical_expr(
@@ -461,7 +328,7 @@ impl DefaultPhysicalPlanner {
fn create_initial_plan<'a>(
&'a self,
logical_plan: &'a LogicalPlan,
- session_state: &'a SessionState,
+ session_state: &'a dyn Session,
) -> futures::future::BoxFuture<'a, Result<Arc<dyn ExecutionPlan>>> {
Box::pin(async move {
// When `enable_physical_uncorrelated_scalar_subquery` is disabled, the
@@ -513,7 +380,7 @@ impl DefaultPhysicalPlanner {
async fn create_initial_plan_inner(
&self,
logical_plan: &LogicalPlan,
- session_state: &SessionState,
+ session_state: &dyn Session,
planning_ctx: &PhysicalPlanningContext,
) -> Result<Arc<dyn ExecutionPlan>> {
// DFS the tree to flatten it into a Vec.
@@ -594,7 +461,7 @@ impl DefaultPhysicalPlanner {
&'a self,
leaf_starter_index: usize,
flat_tree: Arc<Vec<LogicalNode<'a>>>,
- session_state: &'a SessionState,
+ session_state: &'a dyn Session,
planning_ctx: &'a PhysicalPlanningContext,
) -> Result<Option<Arc<dyn ExecutionPlan>>> {
// We always start with a leaf, so can ignore status and pass empty children
@@ -681,7 +548,7 @@ impl DefaultPhysicalPlanner {
async fn map_logical_node_to_physical(
&self,
node: &LogicalPlan,
- session_state: &SessionState,
+ session_state: &dyn Session,
planning_ctx: &PhysicalPlanningContext,
children: ChildrenContainer,
) -> Result<Arc<dyn ExecutionPlan>> {
@@ -2715,7 +2582,7 @@ impl DefaultPhysicalPlanner {
async fn handle_explain_or_analyze(
&self,
logical_plan: &LogicalPlan,
- session_state: &SessionState,
+ session_state: &dyn Session,
) -> Result<Option<Arc<dyn ExecutionPlan>>> {
let execution_plan = match logical_plan {
LogicalPlan::Explain(e) => self.handle_explain(e, session_state).await?,
@@ -2729,7 +2596,7 @@ impl DefaultPhysicalPlanner {
async fn handle_explain(
&self,
e: &Explain,
- session_state: &SessionState,
+ session_state: &dyn Session,
) -> Result<Arc<dyn ExecutionPlan>> {
use PlanType::*;
let mut stringified_plans = vec![];
@@ -2919,7 +2786,7 @@ impl DefaultPhysicalPlanner {
async fn handle_analyze(
&self,
a: &Analyze,
- session_state: &SessionState,
+ session_state: &dyn Session,
) -> Result<Arc<dyn ExecutionPlan>> {
let input = self.create_physical_plan(&a.input, session_state).await?;
let schema = Arc::clone(a.schema.inner());
@@ -2927,7 +2794,7 @@ impl DefaultPhysicalPlanner {
// Statement-level overrides take precedence over the session config.
let analyze_level = a
.analyze_level
- .unwrap_or(session_state.config_options().explain.analyze_level);
+ .unwrap_or_else(|| session_state.config_options().explain.analyze_level);
let metric_types = analyze_level.included_types();
let analyze_categories = a.analyze_categories.clone().unwrap_or_else(|| {
session_state
@@ -2955,7 +2822,7 @@ impl DefaultPhysicalPlanner {
pub fn optimize_physical_plan<F>(
&self,
plan: Arc<dyn ExecutionPlan>,
- session_state: &SessionState,
+ session_state: &dyn Session,
mut observer: F,
) -> Result<Arc<dyn ExecutionPlan>>
where
@@ -2976,10 +2843,13 @@ impl DefaultPhysicalPlanner {
InvariantChecker(InvariantLevel::Always).check(&plan)?;
let mut new_plan = Arc::clone(&plan);
+ let optimizer_context = SessionOptimizerContext {
+ session: session_state,
+ };
for optimizer in optimizers {
let before_schema = new_plan.schema();
new_plan = optimizer
- .optimize_with_context(new_plan, session_state)
+ .optimize_with_context(new_plan, &optimizer_context)
.map_err(|e| {
DataFusionError::Context(optimizer.name().to_string(), Box::new(e))
})?;
@@ -3059,7 +2929,7 @@ impl DefaultPhysicalPlanner {
async fn plan_scalar_subqueries(
&self,
subqueries: Vec<Subquery>,
- session_state: &SessionState,
+ session_state: &dyn Session,
) -> Result<(Vec<ScalarSubqueryLink>, DFHashMap<Subquery, SubqueryIndex>)> {
let mut links = Vec::with_capacity(subqueries.len());
let mut index_map = DFHashMap::with_capacity(subqueries.len());
@@ -3355,6 +3225,7 @@ mod tests {
use std::fmt::{self, Debug};
use std::mem::size_of_val;
use std::ops::{BitAnd, Not};
+ use std::sync::atomic::{AtomicBool, Ordering as AtomicOrdering};
use super::*;
use crate::datasource::MemTable;
@@ -3366,11 +3237,14 @@ mod tests {
use crate::prelude::{SessionConfig, SessionContext};
use crate::test_util::{scan_empty, scan_empty_with_partitions};
+ use crate::execution::context::SessionState;
use crate::execution::session_state::SessionStateBuilder;
+ use crate::logical_expr::UserDefinedLogicalNode;
use arrow::array::{ArrayRef, DictionaryArray, Int32Array};
use arrow::datatypes::{DataType, Field, Int32Type};
use arrow_schema::{FieldRef, SchemaRef};
- use datafusion_common::config::ConfigOptions;
+ use datafusion_catalog::CatalogProviderList;
+ use datafusion_common::config::{ConfigOptions, TableOptions};
use datafusion_common::{
DFSchemaRef, ScalarValue, SplitPoint, TableReference, ToDFSchema as _,
assert_batches_eq, assert_contains,
@@ -3380,16 +3254,171 @@ mod tests {
use datafusion_expr::builder::subquery_alias;
use datafusion_expr::expr::AggregateFunctionParams;
use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
+ use datafusion_expr::registry::ExtensionTypeRegistryRef;
use datafusion_expr::{
- Accumulator, AggregateUDF, AggregateUDFImpl, ExprFunctionExt, LogicalPlanBuilder,
- Partitioning as LogicalPartitioning, RangePartitioning, Signature, TableSource,
- UserDefinedLogicalNodeCore, Volatility, WindowFunctionDefinition, col, lit,
- scalar_subquery,
+ Accumulator, AggregateUDF, AggregateUDFImpl, ExprFunctionExt, HigherOrderUDF,
+ LogicalPlanBuilder, Partitioning as LogicalPartitioning, RangePartitioning,
+ ScalarUDF, Signature, TableSource, UserDefinedLogicalNodeCore, Volatility,
+ WindowFunctionDefinition, WindowUDF, col, lit, scalar_subquery,
};
use datafusion_functions_aggregate::count::{count_all, count_udaf};
use datafusion_functions_aggregate::expr_fn::sum;
use datafusion_physical_expr::EquivalenceProperties;
use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
+ use datafusion_session::QueryPlanner;
+
+ #[derive(Debug)]
+ struct ContextCheckingRule {
+ invoked: Arc<AtomicBool>,
+ }
+
+ impl PhysicalOptimizerRule for ContextCheckingRule {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ _config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(plan)
+ }
+
+ fn optimize_with_context(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ context: &dyn PhysicalOptimizerContext,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ assert!(context.statistics_registry().is_some());
+ self.invoked.store(true, AtomicOrdering::Relaxed);
+ Ok(plan)
+ }
+
+ fn name(&self) -> &str {
+ "context_checking_rule"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+ }
+
+ #[derive(Debug)]
+ struct TestQueryPlanner {
+ invoked: Arc<AtomicBool>,
+ }
+
+ #[async_trait]
+ impl QueryPlanner for TestQueryPlanner {
+ async fn create_physical_plan(
+ &self,
+ logical_plan: &LogicalPlan,
+ session: &dyn Session,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.invoked.store(true, AtomicOrdering::Relaxed);
+ DefaultPhysicalPlanner::default()
+ .create_physical_plan(logical_plan, session)
+ .await
+ }
+ }
+
+ struct TestSession {
+ inner: SessionState,
+ query_planner: Arc<dyn QueryPlanner + Send + Sync>,
+ }
+
+ #[async_trait]
+ impl Session for TestSession {
+ fn session_id(&self) -> &str {
+ self.inner.session_id()
+ }
+
+ fn config(&self) -> &SessionConfig {
+ self.inner.config()
+ }
+
+ fn catalog_list(&self) -> Arc<dyn CatalogProviderList> {
+ Arc::clone(self.inner.catalog_list())
+ }
+
+ fn query_planner(&self) -> Arc<dyn QueryPlanner + Send + Sync> {
+ Arc::clone(&self.query_planner)
+ }
+
+ fn optimize(&self, plan: &LogicalPlan) -> Result<LogicalPlan> {
+ self.inner.optimize(plan)
+ }
+
+ fn physical_optimizers(&self) -> &[Arc<dyn PhysicalOptimizerRule + Send + Sync>] {
+ self.inner.physical_optimizers()
+ }
+
+ fn statistics_registry(
+ &self,
+ ) -> Option<&datafusion_physical_plan::operator_statistics::StatisticsRegistry>
+ {
+ self.inner.statistics_registry()
+ }
+
+ async fn create_physical_plan(
+ &self,
+ logical_plan: &LogicalPlan,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let logical_plan = self.optimize(logical_plan)?;
+ self.query_planner()
+ .create_physical_plan(&logical_plan, self)
+ .await
+ }
+
+ fn create_physical_expr(
+ &self,
+ expr: Expr,
+ df_schema: &DFSchema,
+ ) -> Result<Arc<dyn PhysicalExpr>> {
+ Session::create_physical_expr(&self.inner, expr, df_schema)
+ }
+
+ fn scalar_functions(&self) -> &HashMap<String, Arc<ScalarUDF>> {
+ Session::scalar_functions(&self.inner)
+ }
+
+ fn higher_order_functions(&self) -> &HashMap<String, Arc<HigherOrderUDF>> {
+ Session::higher_order_functions(&self.inner)
+ }
+
+ fn aggregate_functions(&self) -> &HashMap<String, Arc<AggregateUDF>> {
+ Session::aggregate_functions(&self.inner)
+ }
+
+ fn window_functions(&self) -> &HashMap<String, Arc<WindowUDF>> {
+ Session::window_functions(&self.inner)
+ }
+
+ fn extension_type_registry(&self) -> &ExtensionTypeRegistryRef {
+ Session::extension_type_registry(&self.inner)
+ }
+
+ fn runtime_env(&self) -> &Arc<RuntimeEnv> {
+ self.inner.runtime_env()
+ }
+
+ fn execution_props(&self) -> &ExecutionProps {
+ self.inner.execution_props()
+ }
+
+ fn as_any(&self) -> &dyn std::any::Any {
+ self
+ }
+
+ fn table_options(&self) -> &TableOptions {
+ self.inner.table_options()
+ }
+
+ fn table_options_mut(&mut self) -> &mut TableOptions {
+ self.inner.table_options_mut()
+ }
+
+ fn task_ctx(&self) -> Arc<TaskContext> {
+ self.inner.task_ctx()
+ }
+ }
fn make_session_state() -> SessionState {
let runtime = Arc::new(RuntimeEnv::default());
@@ -3412,6 +3441,35 @@ mod tests {
.await
}
+ #[tokio::test]
+ async fn plans_with_non_session_state_implementation() -> Result<()> {
+ let invoked = Arc::new(AtomicBool::new(false));
+ let inner = SessionStateBuilder::new()
+ .with_default_features()
+ .with_physical_optimizer_rules(vec![Arc::new(ContextCheckingRule {
+ invoked: Arc::clone(&invoked),
+ })])
+ .with_statistics_registry(
+ datafusion_physical_plan::operator_statistics::StatisticsRegistry::new(),
+ )
+ .build();
+ let query_planner_invoked = Arc::new(AtomicBool::new(false));
+ let session = TestSession {
+ inner,
+ query_planner: Arc::new(TestQueryPlanner {
+ invoked: Arc::clone(&query_planner_invoked),
+ }),
+ };
+ assert!(session.as_any().downcast_ref::<SessionState>().is_none());
+
+ let logical_plan = LogicalPlanBuilder::empty(false).build()?;
+ let physical_plan = session.create_physical_plan(&logical_plan).await?;
+ assert!(physical_plan.is::<EmptyExec>());
+ assert!(query_planner_invoked.load(AtomicOrdering::Relaxed));
+ assert!(invoked.load(AtomicOrdering::Relaxed));
+ Ok(())
+ }
+
async fn aggregate_explain(logical_plan: &LogicalPlan) -> Result<String> {
let physical_plan = plan(logical_plan).await?;
Ok(displayable(physical_plan.as_ref()).indent(true).to_string())
@@ -4033,9 +4091,16 @@ mod tests {
let planner = DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(
ExpressionExtensionPlanner,
)]);
+ let session = TestSession {
+ inner: make_session_state(),
+ query_planner: Arc::new(TestQueryPlanner {
+ invoked: Arc::new(AtomicBool::new(false)),
+ }),
+ };
+ assert!(session.as_any().downcast_ref::<SessionState>().is_none());
let plan = planner
- .create_physical_plan(&logical_plan, &make_session_state())
+ .create_physical_plan(&logical_plan, &session)
.await?;
assert_contains!(format!("{plan:?}"), "ScalarSubqueryExec");
@@ -4562,7 +4627,7 @@ mod tests {
_node: &dyn UserDefinedLogicalNode,
_logical_inputs: &[&LogicalPlan],
_physical_inputs: &[Arc<dyn ExecutionPlan>],
- _session_state: &SessionState,
+ _session_state: &dyn Session,
_planning_ctx: &PhysicalPlanningContext,
) -> Result<Option<Arc<dyn ExecutionPlan>>> {
internal_err!("BOOM")
@@ -4723,7 +4788,7 @@ mod tests {
node: &dyn UserDefinedLogicalNode,
_logical_inputs: &[&LogicalPlan],
_physical_inputs: &[Arc<dyn ExecutionPlan>],
- session_state: &SessionState,
+ session_state: &dyn Session,
planning_ctx: &PhysicalPlanningContext,
) -> Result<Option<Arc<dyn ExecutionPlan>>> {
for expr in node.expressions() {
@@ -4753,7 +4818,7 @@ mod tests {
_node: &dyn UserDefinedLogicalNode,
_logical_inputs: &[&LogicalPlan],
_physical_inputs: &[Arc<dyn ExecutionPlan>],
- _session_state: &SessionState,
+ _session_state: &dyn Session,
_planning_ctx: &PhysicalPlanningContext,
) -> Result<Option<Arc<dyn ExecutionPlan>>> {
Ok(Some(Arc::new(NoOpExecutionPlan::new(SchemaRef::new(
@@ -5387,7 +5452,7 @@ digraph {
_node: &dyn UserDefinedLogicalNode,
_logical_inputs: &[&LogicalPlan],
_physical_inputs: &[Arc<dyn ExecutionPlan>],
- _session_state: &SessionState,
+ _session_state: &dyn Session,
_planning_ctx: &PhysicalPlanningContext,
) -> Result<Option<Arc<dyn ExecutionPlan>>> {
Ok(None)
@@ -5397,7 +5462,7 @@ digraph {
&self,
_planner: &dyn PhysicalPlanner,
scan: &TableScan,
- _session_state: &SessionState,
+ _session_state: &dyn Session,
_planning_ctx: &PhysicalPlanningContext,
) -> Result<Option<Arc<dyn ExecutionPlan>>> {
if scan.source.is::<MockTableSource>() {
diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs b/datafusion/core/tests/user_defined/user_defined_plan.rs
index 354a1b3110..99363ba500 100644
--- a/datafusion/core/tests/user_defined/user_defined_plan.rs
+++ b/datafusion/core/tests/user_defined/user_defined_plan.rs
@@ -67,13 +67,14 @@ use arrow::{
array::Int64Array, datatypes::SchemaRef, record_batch::RecordBatch,
util::pretty::pretty_format_batches,
};
+use datafusion::catalog::Session;
use datafusion::execution::session_state::SessionStateBuilder;
use datafusion::{
common::cast::as_int64_array,
common::{DFSchemaRef, arrow_datafusion_err},
error::{DataFusionError, Result},
execution::{
- context::{QueryPlanner, SessionState, TaskContext},
+ context::{QueryPlanner, TaskContext},
runtime_env::RuntimeEnv,
},
logical_expr::{
@@ -467,7 +468,7 @@ impl QueryPlanner for TopKQueryPlanner {
async fn create_physical_plan(
&self,
logical_plan: &LogicalPlan,
- session_state: &SessionState,
+ session_state: &dyn Session,
) -> Result<Arc<dyn ExecutionPlan>> {
// Teach the default physical planner how to plan TopK nodes.
let physical_planner =
@@ -630,7 +631,7 @@ impl ExtensionPlanner for TopKPlanner {
node: &dyn UserDefinedLogicalNode,
logical_inputs: &[&LogicalPlan],
physical_inputs: &[Arc<dyn ExecutionPlan>],
- _session_state: &SessionState,
+ _session_state: &dyn Session,
_planning_ctx: &PhysicalPlanningContext,
) -> Result<Option<Arc<dyn ExecutionPlan>>> {
Ok(
diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs
index 0c2d9fdeee..519384379e 100644
--- a/datafusion/ffi/src/session/mod.rs
+++ b/datafusion/ffi/src/session/mod.rs
@@ -730,6 +730,16 @@ mod tests {
assert!(state.catalog_list().catalog("foreign_registered").is_some());
let logical_plan = LogicalPlan::default();
+ assert_eq!(foreign_session.optimize(&logical_plan)?, logical_plan);
+ assert!(foreign_session.physical_optimizers().is_empty());
+ assert!(foreign_session.statistics_registry().is_none());
+ let planner_error = foreign_session
+ .query_planner()
+ .create_physical_plan(&logical_plan, &foreign_session)
+ .await
+ .unwrap_err();
+ assert!(planner_error.to_string().contains("does not expose"));
+
let physical_plan = foreign_session.create_physical_plan(&logical_plan).await?;
assert_eq!(
format!("{physical_plan:?}"),
diff --git a/datafusion/physical-optimizer/Cargo.toml b/datafusion/physical-optimizer/Cargo.toml
index 38c8a7c372..cb03303ac3 100644
--- a/datafusion/physical-optimizer/Cargo.toml
+++ b/datafusion/physical-optimizer/Cargo.toml
@@ -50,6 +50,7 @@ datafusion-physical-expr = { workspace = true }
datafusion-physical-expr-common = { workspace = true }
datafusion-physical-plan = { workspace = true }
datafusion-pruning = { workspace = true }
+datafusion-session = { workspace = true }
itertools = { workspace = true }
recursive = { workspace = true, optional = true }
diff --git a/datafusion/physical-optimizer/src/optimizer.rs b/datafusion/physical-optimizer/src/optimizer.rs
index 0f81512b61..2841afecf6 100644
--- a/datafusion/physical-optimizer/src/optimizer.rs
+++ b/datafusion/physical-optimizer/src/optimizer.rs
@@ -39,29 +39,10 @@ use crate::hash_join_buffering::HashJoinBuffering;
use crate::limit_pushdown_past_window::LimitPushPastWindows;
use crate::pushdown_sort::PushdownSort;
use crate::window_topn::WindowTopN;
-use datafusion_common::Result;
use datafusion_common::config::ConfigOptions;
-use datafusion_physical_plan::ExecutionPlan;
-use datafusion_physical_plan::operator_statistics::StatisticsRegistry;
-/// Context available to physical optimizer rules.
-///
-/// This trait provides access to configuration options and optional statistics
-/// registry for enhanced statistics lookup. It allows optimizer rules to access
-/// extended context without changing the core [`PhysicalOptimizerRule::optimize`]
-/// signature.
-pub trait PhysicalOptimizerContext: Send + Sync {
- /// Returns the configuration options.
- fn config_options(&self) -> &ConfigOptions;
-
- /// Returns the statistics registry for enhanced statistics lookup.
- ///
- /// Returns `None` if no registry is configured, in which case rules
- /// should fall back to using `ExecutionPlan::partition_statistics()`.
- fn statistics_registry(&self) -> Option<&StatisticsRegistry> {
- None
- }
-}
+// Re-export from this module for backwards compatibility.
+pub use datafusion_session::{PhysicalOptimizerContext, PhysicalOptimizerRule};
/// Simple context wrapping [`ConfigOptions`] for backward compatibility.
///
@@ -85,47 +66,6 @@ impl PhysicalOptimizerContext for ConfigOnlyContext<'_> {
}
}
-/// `PhysicalOptimizerRule` transforms one ['ExecutionPlan'] into another which
-/// computes the same results, but in a potentially more efficient way.
-///
-/// Use [`SessionState::add_physical_optimizer_rule`] to register additional
-/// `PhysicalOptimizerRule`s.
-///
-/// [`SessionState::add_physical_optimizer_rule`]: https://docs.rs/datafusion/latest/datafusion/execution/session_state/struct.SessionState.html#method.add_physical_optimizer_rule
-pub trait PhysicalOptimizerRule: Debug + std::any::Any {
- /// Rewrite `plan` to an optimized form.
- ///
- /// This is the primary optimization method. For rules that need access to
- /// the statistics registry, override [`optimize_with_context`](Self::optimize_with_context) instead.
- fn optimize(
- &self,
- plan: Arc<dyn ExecutionPlan>,
- config: &ConfigOptions,
- ) -> Result<Arc<dyn ExecutionPlan>>;
-
- /// Rewrite `plan` with access to extended context (statistics registry, etc.).
- ///
- /// Override this method if you need access to the statistics registry for
- /// enhanced statistics lookup. The default implementation simply calls
- /// [`optimize`](Self::optimize) with the config options from the context.
- fn optimize_with_context(
- &self,
- plan: Arc<dyn ExecutionPlan>,
- context: &dyn PhysicalOptimizerContext,
- ) -> Result<Arc<dyn ExecutionPlan>> {
- self.optimize(plan, context.config_options())
- }
-
- /// A human readable name for this optimizer rule
- fn name(&self) -> &str;
-
- /// A flag to indicate whether the physical planner should validate that the rule will not
- /// change the schema of the plan after the rewriting.
- /// Some of the optimization rules might change the nullable properties of the schema
- /// and should disable the schema check.
- fn schema_check(&self) -> bool;
-}
-
/// A rule-based physical optimizer.
#[derive(Clone, Debug)]
pub struct PhysicalOptimizer {
diff --git a/datafusion/session/src/lib.rs b/datafusion/session/src/lib.rs
index 3b9ed7dacf..6f7cfb7792 100644
--- a/datafusion/session/src/lib.rs
+++ b/datafusion/session/src/lib.rs
@@ -32,6 +32,10 @@
//! * [`CatalogProviderList`], [`CatalogProvider`], and [`SchemaProvider`] -
//! Describe catalog hierarchies
//! * [`TableProvider`] - Provides data for query planning and execution
+//! * [`QueryPlanner`], [`PhysicalPlanner`], and [`ExtensionPlanner`] - Query and
+//! physical planning contracts
+//! * [`PhysicalOptimizerRule`] and [`PhysicalOptimizerContext`] - Physical
+//! optimization contracts
//! * [`SessionStore`] - Handles session persistence and retrieval
//!
//! The session system enables:
@@ -42,6 +46,8 @@
//! * Query state persistence
pub mod catalog;
+pub mod physical_optimizer;
+pub mod planner;
pub mod schema;
pub mod session;
pub mod table;
@@ -49,6 +55,10 @@ pub mod table;
pub use crate::catalog::{
CatalogProvider, CatalogProviderList, EmptyCatalogProviderList,
};
+pub use crate::physical_optimizer::{PhysicalOptimizerContext, PhysicalOptimizerRule};
+pub use crate::planner::{
+ ExtensionPlanner, PhysicalPlanner, QueryPlanner, UnsupportedQueryPlanner,
+};
pub use crate::schema::SchemaProvider;
pub use crate::session::{Session, SessionStore};
pub use crate::table::{
diff --git a/datafusion/session/src/physical_optimizer.rs b/datafusion/session/src/physical_optimizer.rs
new file mode 100644
index 0000000000..751a8e12d9
--- /dev/null
+++ b/datafusion/session/src/physical_optimizer.rs
@@ -0,0 +1,84 @@
+// 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.
+
+//! Physical optimizer interfaces.
+
+use std::fmt::Debug;
+use std::sync::Arc;
+
+use datafusion_common::Result;
+use datafusion_common::config::ConfigOptions;
+use datafusion_physical_plan::ExecutionPlan;
+use datafusion_physical_plan::operator_statistics::StatisticsRegistry;
+
+/// Context available to physical optimizer rules.
+///
+/// This trait provides access to configuration options and an optional statistics
+/// registry for enhanced statistics lookup.
+pub trait PhysicalOptimizerContext: Send + Sync {
+ /// Returns the configuration options.
+ fn config_options(&self) -> &ConfigOptions;
+
+ /// Returns the statistics registry for enhanced statistics lookup.
+ ///
+ /// Returns `None` if no registry is configured, in which case rules
+ /// should fall back to using [`ExecutionPlan::partition_statistics`].
+ fn statistics_registry(&self) -> Option<&StatisticsRegistry> {
+ None
+ }
+}
+
+/// `PhysicalOptimizerRule` transforms one [`ExecutionPlan`] into another which
+/// computes the same results, but in a potentially more efficient way.
+///
+/// Use [`SessionState::add_physical_optimizer_rule`] to register additional
+/// `PhysicalOptimizerRule`s.
+///
+/// [`SessionState::add_physical_optimizer_rule`]: https://docs.rs/datafusion/latest/datafusion/execution/session_state/struct.SessionState.html#method.add_physical_optimizer_rule
+pub trait PhysicalOptimizerRule: Debug + std::any::Any {
+ /// Rewrite `plan` to an optimized form.
+ ///
+ /// This is the primary optimization method. For rules that need access to
+ /// the statistics registry, override [`optimize_with_context`](Self::optimize_with_context) instead.
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>>;
+
+ /// Rewrite `plan` with access to extended context (statistics registry, etc.).
+ ///
+ /// Override this method if you need access to the statistics registry for
+ /// enhanced statistics lookup. The default implementation simply calls
+ /// [`optimize`](Self::optimize) with the config options from the context.
+ fn optimize_with_context(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ context: &dyn PhysicalOptimizerContext,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.optimize(plan, context.config_options())
+ }
+
+ /// A human readable name for this optimizer rule
+ fn name(&self) -> &str;
+
+ /// A flag to indicate whether the physical planner should validate that the rule will not
+ /// change the schema of the plan after the rewriting.
+ /// Some of the optimization rules might change the nullable properties of the schema
+ /// and should disable the schema check.
+ fn schema_check(&self) -> bool;
+}
diff --git a/datafusion/session/src/planner.rs b/datafusion/session/src/planner.rs
new file mode 100644
index 0000000000..37726009f0
--- /dev/null
+++ b/datafusion/session/src/planner.rs
@@ -0,0 +1,198 @@
+// 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.
+
+//! Query planner interfaces.
+
+use std::any::Any;
+use std::fmt::Debug;
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use datafusion_common::{DFSchema, Result, not_impl_err};
+use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
+use datafusion_expr::{Expr, LogicalPlan, TableScan, UserDefinedLogicalNode};
+use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr};
+
+use crate::Session;
+
+/// A planner that creates a physical plan for a query.
+#[async_trait]
+pub trait QueryPlanner: Any + Debug {
+ /// Given a [`LogicalPlan`], create an [`ExecutionPlan`] suitable for execution
+ async fn create_physical_plan(
+ &self,
+ logical_plan: &LogicalPlan,
+ session: &dyn Session,
+ ) -> Result<Arc<dyn ExecutionPlan>>;
+}
+
+/// A query planner that reports that planning is not implemented.
+///
+/// [`Session`] implementations that do not expose a query planner can return
+/// this planner explicitly.
+#[derive(Debug, Default)]
+pub struct UnsupportedQueryPlanner;
+
+#[async_trait]
+impl QueryPlanner for UnsupportedQueryPlanner {
+ async fn create_physical_plan(
+ &self,
+ _logical_plan: &LogicalPlan,
+ _session: &dyn Session,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ not_impl_err!("This session does not expose its query planner")
+ }
+}
+
+/// Physical query planner that converts a [`LogicalPlan`] to an
+/// [`ExecutionPlan`] suitable for execution.
+#[async_trait]
+pub trait PhysicalPlanner: Send + Sync {
+ /// Create a physical plan from a logical plan
+ async fn create_physical_plan(
+ &self,
+ logical_plan: &LogicalPlan,
+ session: &dyn Session,
+ ) -> Result<Arc<dyn ExecutionPlan>>;
+
+ /// Create a physical expression from a logical expression
+ /// suitable for evaluation
+ ///
+ /// `expr`: the expression to convert
+ ///
+ /// `input_dfschema`: the logical plan schema for evaluating `expr`
+ ///
+ /// `planning_ctx`: the [`PhysicalPlanningContext`] used to resolve
+ /// `Expr::ScalarSubquery` nodes. During physical planning the planner
+ /// threads the context of the plan currently being converted to a physical
+ /// plan (for example into [`ExtensionPlanner::plan_extension`], which
+ /// should forward it here). Callers creating physical expressions outside
+ /// of a plan should pass `&PhysicalPlanningContext::default()`.
+ fn create_physical_expr(
+ &self,
+ expr: &Expr,
+ input_dfschema: &DFSchema,
+ session: &dyn Session,
+ planning_ctx: &PhysicalPlanningContext,
+ ) -> Result<Arc<dyn PhysicalExpr>>;
+}
+
+/// This trait exposes the ability to plan an [`ExecutionPlan`] out of a [`LogicalPlan`].
+#[async_trait]
+pub trait ExtensionPlanner {
+ /// Create a physical plan for a [`UserDefinedLogicalNode`].
+ ///
+ /// `input_dfschema`: the logical plan schema for the inputs to this node
+ ///
+ /// Returns an error when the planner knows how to plan the concrete
+ /// implementation of `node` but errors while doing so.
+ ///
+ /// Returns `None` when the planner does not know how to plan the
+ /// `node` and wants to delegate the planning to another
+ /// [`ExtensionPlanner`].
+ ///
+ /// `planning_ctx` is the [`PhysicalPlanningContext`] of the plan subtree
+ /// currently being converted to a physical plan. Forward it to
+ /// [`PhysicalPlanner::create_physical_expr`] when creating this node's
+ /// physical expressions so that scalar subqueries resolve against the same
+ /// subquery state as the rest of the plan.
+ async fn plan_extension(
+ &self,
+ planner: &dyn PhysicalPlanner,
+ node: &dyn UserDefinedLogicalNode,
+ logical_inputs: &[&LogicalPlan],
+ physical_inputs: &[Arc<dyn ExecutionPlan>],
+ session: &dyn Session,
+ planning_ctx: &PhysicalPlanningContext,
+ ) -> Result<Option<Arc<dyn ExecutionPlan>>>;
+
+ /// Create a physical plan for a [`LogicalPlan::TableScan`].
+ ///
+ /// This is useful for planning valid [`TableSource`]s that are not `TableProvider`s.
+ ///
+ /// Returns:
+ /// * `Ok(Some(plan))` if the planner knows how to plan the `scan`
+ /// * `Ok(None)` if the planner does not know how to plan the `scan` and wants to delegate the planning to another [`ExtensionPlanner`]
+ /// * `Err` if the planner knows how to plan the `scan` but errors while doing so
+ ///
+ /// # Example
+ ///
+ /// ```rust,ignore
+ /// use std::sync::Arc;
+ /// use datafusion::physical_plan::ExecutionPlan;
+ /// use datafusion::logical_expr::TableScan;
+ /// use datafusion::catalog::Session;
+ /// use datafusion::error::Result;
+ /// use datafusion_session::{ExtensionPlanner, PhysicalPlanner};
+ /// use async_trait::async_trait;
+ ///
+ /// // Your custom table source type
+ /// struct MyCustomTableSource { /* ... */ }
+ ///
+ /// // Your custom execution plan
+ /// struct MyCustomExec { /* ... */ }
+ ///
+ /// struct MyExtensionPlanner;
+ ///
+ /// #[async_trait]
+ /// impl ExtensionPlanner for MyExtensionPlanner {
+ /// async fn plan_extension(
+ /// &self,
+ /// _planner: &dyn PhysicalPlanner,
+ /// _node: &dyn UserDefinedLogicalNode,
+ /// _logical_inputs: &[&LogicalPlan],
+ /// _physical_inputs: &[Arc<dyn ExecutionPlan>],
+ /// _session: &dyn Session,
+ /// _planning_ctx: &PhysicalPlanningContext,
+ /// ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
+ /// Ok(None)
+ /// }
+ ///
+ /// async fn plan_table_scan(
+ /// &self,
+ /// _planner: &dyn PhysicalPlanner,
+ /// scan: &TableScan,
+ /// _session: &dyn Session,
+ /// _planning_ctx: &PhysicalPlanningContext,
+ /// ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
+ /// // Check if this is your custom table source
+ /// if scan.source.is::<MyCustomTableSource>() {
+ /// // Create a custom execution plan for your table source
+ /// let exec = MyCustomExec::new(
+ /// scan.table_name.clone(),
+ /// Arc::clone(scan.projected_schema.inner()),
+ /// );
+ /// Ok(Some(Arc::new(exec)))
+ /// } else {
+ /// // Return None to let other extension planners handle it
+ /// Ok(None)
+ /// }
+ /// }
+ /// }
+ /// ```
+ ///
+ /// [`TableSource`]: datafusion_expr::TableSource
+ async fn plan_table_scan(
+ &self,
+ _planner: &dyn PhysicalPlanner,
+ _scan: &TableScan,
+ _session: &dyn Session,
+ _planning_ctx: &PhysicalPlanningContext,
+ ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
+ Ok(None)
+ }
+}
diff --git a/datafusion/session/src/session.rs b/datafusion/session/src/session.rs
index cdac3f4beb..f6143cc4a4 100644
--- a/datafusion/session/src/session.rs
+++ b/datafusion/session/src/session.rs
@@ -26,6 +26,7 @@ use datafusion_expr::registry::ExtensionTypeRegistryRef;
use datafusion_expr::{
AggregateUDF, Expr, HigherOrderUDF, LogicalPlan, ScalarUDF, WindowUDF,
};
+use datafusion_physical_plan::operator_statistics::StatisticsRegistry;
use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr};
use crate::CatalogProviderList;
@@ -34,6 +35,8 @@ use std::any::Any;
use std::collections::HashMap;
use std::sync::{Arc, Weak};
+use crate::{PhysicalOptimizerRule, QueryPlanner, UnsupportedQueryPlanner};
+
/// Interface for accessing [`SessionState`] from the catalog and data source.
///
/// This trait provides access to the information needed to plan and execute
@@ -89,6 +92,54 @@ pub trait Session: Send + Sync {
self.config().options()
}
+ /// Return the query planner for this session.
+ ///
+ /// # Warning
+ ///
+ /// The default implementation returns an [`UnsupportedQueryPlanner`], so
+ /// [`Session::create_physical_plan`] will fail. Sessions that support
+ /// physical planning should override this method (for example by returning
+ /// `SessionState::query_planner`).
+ fn query_planner(&self) -> Arc<dyn QueryPlanner + Send + Sync> {
+ Arc::new(UnsupportedQueryPlanner)
+ }
+
+ /// Optimize a logical plan.
+ ///
+ /// # Warning
+ ///
+ /// The default implementation returns the plan **unchanged**, applying no
+ /// logical optimizations whatsoever. This is almost never what you want:
+ /// without optimization, queries execute in their naive, unoptimized form
+ /// and may be dramatically slower or fail to run at all. The default exists
+ /// only so this crate need not depend on the optimizer; any real session
+ /// should override this method (for example by delegating to
+ /// `SessionState::optimize`).
+ fn optimize(&self, plan: &LogicalPlan) -> Result<LogicalPlan> {
+ Ok(plan.clone())
+ }
+
+ /// Return the physical optimizer rules for this session.
+ ///
+ /// # Warning
+ ///
+ /// The default implementation returns **no rules**. This is almost never
+ /// what you want: DataFusion relies on physical optimizer rules for
+ /// correctness-critical rewrites (such as inserting the repartitioning and
+ /// coalescing needed for parallel and multi-partition execution), so a
+ /// session with no rules will produce plans that are inefficient or that
+ /// fail to execute. The default exists only so this crate need not depend
+ /// on the optimizer; any real session should override this method (for
+ /// example by returning `SessionState::physical_optimizers`).
+ fn physical_optimizers(&self) -> &[Arc<dyn PhysicalOptimizerRule + Send + Sync>] {
+ &[]
+ }
+
+ /// Return the optional statistics registry used during physical optimization.
+ fn statistics_registry(&self) -> Option<&StatisticsRegistry> {
+ None
+ }
+
/// Creates a physical [`ExecutionPlan`] plan from a [`LogicalPlan`].
///
/// Note: this will optimize the provided plan first.
diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md
index 6097c8dc71..d97082e881 100644
--- a/docs/source/library-user-guide/upgrading/55.0.0.md
+++ b/docs/source/library-user-guide/upgrading/55.0.0.md
@@ -805,13 +805,13 @@ async fn plan_extension(
node: &dyn UserDefinedLogicalNode,
logical_inputs: &[&LogicalPlan],
physical_inputs: &[Arc<dyn ExecutionPlan>],
- session_state: &SessionState,
+ session: &dyn Session,
planning_ctx: &PhysicalPlanningContext, // new parameter
) -> Result<Option<Arc<dyn ExecutionPlan>>> {
for expr in node.expressions() {
// Forward the context so scalar subqueries in this node's
// expressions resolve against the plan's subquery state
- planner.create_physical_expr(&expr, node.schema(), session_state, planning_ctx)?;
+ planner.create_physical_expr(&expr, node.schema(), session, planning_ctx)?;
}
// ...
}
@@ -819,20 +819,32 @@ async fn plan_extension(
See [PR #23649](https://github.com/apache/datafusion/pull/23649) for details.
-### Catalog traits moved to `datafusion-session`
+### Catalog, planner, and optimizer contracts moved to `datafusion-session`
-The catalog contract traits now live in the `datafusion-session` crate so they
-can be reached without downcasting a `Session` to `SessionState` (in particular
-across the FFI boundary). The affected traits are `CatalogProviderList`,
-`CatalogProvider`, `SchemaProvider`, `TableProvider`, `TableProviderFactory`,
-and `TableFunctionImpl`. The related `TableFunction` struct also moved.
+The catalog, planner, and physical optimizer contract traits now live in the
+`datafusion-session` crate. This makes them available through `Session` without
+downcasting to `SessionState`, including across the FFI boundary.
-The `datafusion-catalog` crate re-exports all of them from their new location,
-so paths such as `datafusion::catalog::TableProvider` and
-`datafusion_catalog::CatalogProvider` continue to work unchanged. Most users do
-not need to do anything.
+The moved catalog traits are `CatalogProviderList`, `CatalogProvider`,
+`SchemaProvider`, `TableProvider`, `TableProviderFactory`, and
+`TableFunctionImpl`. The related `TableFunction` struct also moved. The
+`datafusion-catalog` crate re-exports these items from their new location, so
+paths such as `datafusion::catalog::TableProvider` and
+`datafusion_catalog::CatalogProvider` continue to work unchanged.
-### `Session` gains a required `catalog_list` method
+The moved planning and optimization traits are `QueryPlanner`,
+`PhysicalPlanner`, `ExtensionPlanner`, `PhysicalOptimizerRule`, and
+`PhysicalOptimizerContext`. Their previous paths also continue to work through
+re-exports:
+
+- `datafusion::execution::context::QueryPlanner`
+- `datafusion::physical_planner::{PhysicalPlanner, ExtensionPlanner}`
+- `datafusion_physical_optimizer::{PhysicalOptimizerRule, PhysicalOptimizerContext}`
+
+The session argument for methods on `QueryPlanner`, `PhysicalPlanner`, and
+`ExtensionPlanner` changed from `&SessionState` to `&dyn Session`. Custom planner
+implementations should update their signatures. Planner code should use methods
+on `Session` instead of downcasting it to `SessionState`.
The `Session` trait now requires a `catalog_list` method that returns the
catalogs registered with the session:
@@ -853,7 +865,44 @@ fn catalog_list(&self) -> Arc<dyn CatalogProviderList> {
}
```
-See [PR #23703](https://github.com/apache/datafusion/pull/23703) for details.
+`Session` gains a `query_planner` method alongside `optimize`,
+`physical_optimizers`, and `statistics_registry`. All four have default
+implementations, so existing `Session` implementations that do not perform
+physical planning require no changes: `query_planner` defaults to the new
+`UnsupportedQueryPlanner`, `optimize` returns the plan unchanged,
+`physical_optimizers` returns no rules, and `statistics_registry` returns
+`None`.
+
+A custom session that drives planning through `DefaultQueryPlanner` or
+`DefaultPhysicalPlanner` must override these methods to expose its planning and
+optimization behavior; the defaults will otherwise produce unoptimized plans or
+fail to plan at all. The simplest approach is to delegate to a `SessionState`:
+
+```rust
+use std::sync::Arc;
+use datafusion_session::{PhysicalOptimizerRule, QueryPlanner};
+
+fn query_planner(&self) -> Arc<dyn QueryPlanner + Send + Sync> {
+ self.inner.query_planner()
+}
+
+fn optimize(&self, plan: &LogicalPlan) -> Result<LogicalPlan> {
+ self.inner.optimize(plan)
+}
+
+fn physical_optimizers(&self) -> &[Arc<dyn PhysicalOptimizerRule + Send + Sync>] {
+ self.inner.physical_optimizers()
+}
+```
+
+`ForeignSession::create_physical_plan` continues to run the complete planning
+pipeline in the library that owns the session. `ForeignSession::query_planner`
+returns `UnsupportedQueryPlanner` until the query planner FFI interface is
+available. FFI wrappers for the individual planner and optimizer interfaces are
+not included in this release.
+
+See [PR #23703](https://github.com/apache/datafusion/pull/23703) for details on
+the catalog changes.
### Unused `async` removed from several public functions
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org