diff --git a/Cargo.lock b/Cargo.lock index ed63f60b41519..9ed8b14a2b0e3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2453,6 +2453,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 dd145c715f3c6..ffbce298b4f17 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], - session_state: &SessionState, + session_state: &dyn Session, _planning_ctx: &PhysicalPlanningContext, ) -> Result>> { if let Some(cache_node) = node.as_any().downcast_ref::() { @@ -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> { 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 b0ccff8d10d8c..c019e136ccd8b 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> { 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], - _session_state: &SessionState, + _session_state: &dyn Session, _planning_ctx: &PhysicalPlanningContext, ) -> Result>> { let Some(sample_node) = node.as_any().downcast_ref::() diff --git a/datafusion/core/src/datasource/listing_table_factory.rs b/datafusion/core/src/datasource/listing_table_factory.rs index 68f6743189447..cbecd2a8da49e 100644 --- a/datafusion/core/src/datasource/listing_table_factory.rs +++ b/datafusion/core/src/datasource/listing_table_factory.rs @@ -838,7 +838,10 @@ mod tests { use datafusion_execution::config::SessionConfig; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_plan::ExecutionPlan; - use datafusion_session::{CatalogProviderList, EmptyCatalogProviderList}; + use datafusion_session::{ + CatalogProviderList, EmptyCatalogProviderList, QueryPlanner, + UnsupportedQueryPlanner, + }; use std::any::Any; use std::collections::HashMap; @@ -857,6 +860,9 @@ mod tests { fn catalog_list(&self) -> Arc { Arc::new(EmptyCatalogProviderList) } + fn query_planner(&self) -> Arc { + Arc::new(UnsupportedQueryPlanner) + } async fn create_physical_plan( &self, _logical_plan: &datafusion_expr::LogicalPlan, diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs index 281cb4dd79d4d..d4791d2b293d2 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}; @@ -2187,16 +2186,9 @@ impl From 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>; -} +pub use datafusion_session::{QueryPlanner, UnsupportedQueryPlanner}; /// Interface for handling `CREATE FUNCTION` statements and interacting with /// [SessionState] to create and register functions ([`ScalarUDF`], @@ -2396,6 +2388,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; @@ -2842,7 +2835,7 @@ mod tests { async fn create_physical_plan( &self, _logical_plan: &LogicalPlan, - _session_state: &SessionState, + _session_state: &dyn Session, ) -> Result> { not_impl_err!("query not supported") } @@ -2851,7 +2844,7 @@ mod tests { &self, _expr: &Expr, _input_dfschema: &DFSchema, - _session_state: &SessionState, + _session_state: &dyn Session, _planning_ctx: &PhysicalPlanningContext, ) -> Result> { unimplemented!() @@ -2866,7 +2859,7 @@ mod tests { async fn create_physical_plan( &self, logical_plan: &LogicalPlan, - session_state: &SessionState, + session_state: &dyn Session, ) -> Result> { 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 dfdbb1617efde..2efeb67338e4a 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,22 @@ impl Session for SessionState { Arc::clone(self.catalog_list()) } + fn query_planner(&self) -> Arc { + Arc::clone(SessionState::query_planner(self)) + } + + fn optimize(&self, plan: &LogicalPlan) -> datafusion_common::Result { + SessionState::optimize(self, plan) + } + + fn physical_optimizers(&self) -> &[Arc] { + 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 +2336,7 @@ impl QueryPlanner for DefaultQueryPlanner { async fn create_physical_plan( &self, logical_plan: &LogicalPlan, - session_state: &SessionState, + session_state: &dyn Session, ) -> datafusion_common::Result> { let planner = DefaultPhysicalPlanner::default(); planner diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index aef8036c749a8..1f441f8421fe3 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, }; @@ -57,6 +55,7 @@ use crate::physical_plan::{ displayable, windows, }; use crate::schema_equivalence::schema_satisfied_by; +use datafusion_session::{PhysicalOptimizerContext, Session}; use arrow::array::{RecordBatch, builder::StringBuilder}; use arrow::compute::SortOptions; @@ -103,7 +102,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 +109,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::PhysicalOptimizerRule; use async_trait::async_trait; use datafusion_physical_plan::async_func::{AsyncFuncExec, AsyncMapper}; @@ -120,162 +119,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>; +// 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>; +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], - session_state: &SessionState, - planning_ctx: &PhysicalPlanningContext, - ) -> Result>>; +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], - /// _session_state: &SessionState, - /// _planning_ctx: &PhysicalPlanningContext, - /// ) -> Result>> { - /// Ok(None) - /// } - /// - /// async fn plan_table_scan( - /// &self, - /// _planner: &dyn PhysicalPlanner, - /// scan: &TableScan, - /// _session_state: &SessionState, - /// _planning_ctx: &PhysicalPlanningContext, - /// ) -> Result>> { - /// // Check if this is your custom table source - /// if scan.source.is::() { - /// // 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>> { - 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 +157,7 @@ impl PhysicalPlanner for DefaultPhysicalPlanner { async fn create_physical_plan( &self, logical_plan: &LogicalPlan, - session_state: &SessionState, + session_state: &dyn Session, ) -> Result> { if let Some(plan) = self .handle_explain_or_analyze(logical_plan, session_state) @@ -314,7 +182,7 @@ impl PhysicalPlanner for DefaultPhysicalPlanner { &self, expr: &Expr, input_dfschema: &DFSchema, - session_state: &SessionState, + session_state: &dyn Session, planning_ctx: &PhysicalPlanningContext, ) -> Result> { create_physical_expr( @@ -461,7 +329,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>> { Box::pin(async move { // When `enable_physical_uncorrelated_scalar_subquery` is disabled, the @@ -513,7 +381,7 @@ impl DefaultPhysicalPlanner { async fn create_initial_plan_inner( &self, logical_plan: &LogicalPlan, - session_state: &SessionState, + session_state: &dyn Session, planning_ctx: &PhysicalPlanningContext, ) -> Result> { // DFS the tree to flatten it into a Vec. @@ -594,7 +462,7 @@ impl DefaultPhysicalPlanner { &'a self, leaf_starter_index: usize, flat_tree: Arc>>, - session_state: &'a SessionState, + session_state: &'a dyn Session, planning_ctx: &'a PhysicalPlanningContext, ) -> Result>> { // We always start with a leaf, so can ignore status and pass empty children @@ -681,7 +549,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> { @@ -2710,7 +2578,7 @@ impl DefaultPhysicalPlanner { async fn handle_explain_or_analyze( &self, logical_plan: &LogicalPlan, - session_state: &SessionState, + session_state: &dyn Session, ) -> Result>> { let execution_plan = match logical_plan { LogicalPlan::Explain(e) => self.handle_explain(e, session_state).await?, @@ -2724,7 +2592,7 @@ impl DefaultPhysicalPlanner { async fn handle_explain( &self, e: &Explain, - session_state: &SessionState, + session_state: &dyn Session, ) -> Result> { use PlanType::*; let mut stringified_plans = vec![]; @@ -2914,7 +2782,7 @@ impl DefaultPhysicalPlanner { async fn handle_analyze( &self, a: &Analyze, - session_state: &SessionState, + session_state: &dyn Session, ) -> Result> { let input = self.create_physical_plan(&a.input, session_state).await?; let schema = Arc::clone(a.schema.inner()); @@ -2922,7 +2790,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 @@ -2950,7 +2818,7 @@ impl DefaultPhysicalPlanner { pub fn optimize_physical_plan( &self, plan: Arc, - session_state: &SessionState, + session_state: &dyn Session, mut observer: F, ) -> Result> where @@ -2971,10 +2839,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)) })?; @@ -3054,7 +2925,7 @@ impl DefaultPhysicalPlanner { async fn plan_scalar_subqueries( &self, subqueries: Vec, - session_state: &SessionState, + session_state: &dyn Session, ) -> Result<(Vec, DFHashMap)> { let mut links = Vec::with_capacity(subqueries.len()); let mut index_map = DFHashMap::with_capacity(subqueries.len()); @@ -3350,6 +3221,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; @@ -3361,11 +3233,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, @@ -3375,16 +3250,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, + } + + impl PhysicalOptimizerRule for ContextCheckingRule { + fn optimize( + &self, + plan: Arc, + _config: &ConfigOptions, + ) -> Result> { + Ok(plan) + } + + fn optimize_with_context( + &self, + plan: Arc, + context: &dyn PhysicalOptimizerContext, + ) -> Result> { + 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, + } + + #[async_trait] + impl QueryPlanner for TestQueryPlanner { + async fn create_physical_plan( + &self, + logical_plan: &LogicalPlan, + session: &dyn Session, + ) -> Result> { + self.invoked.store(true, AtomicOrdering::Relaxed); + DefaultPhysicalPlanner::default() + .create_physical_plan(logical_plan, session) + .await + } + } + + struct TestSession { + inner: SessionState, + query_planner: Arc, + } + + #[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 { + Arc::clone(self.inner.catalog_list()) + } + + fn query_planner(&self) -> Arc { + Arc::clone(&self.query_planner) + } + + fn optimize(&self, plan: &LogicalPlan) -> Result { + self.inner.optimize(plan) + } + + fn physical_optimizers(&self) -> &[Arc] { + 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> { + 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> { + Session::create_physical_expr(&self.inner, expr, df_schema) + } + + fn scalar_functions(&self) -> &HashMap> { + Session::scalar_functions(&self.inner) + } + + fn higher_order_functions(&self) -> &HashMap> { + Session::higher_order_functions(&self.inner) + } + + fn aggregate_functions(&self) -> &HashMap> { + Session::aggregate_functions(&self.inner) + } + + fn window_functions(&self) -> &HashMap> { + Session::window_functions(&self.inner) + } + + fn extension_type_registry(&self) -> &ExtensionTypeRegistryRef { + Session::extension_type_registry(&self.inner) + } + + fn runtime_env(&self) -> &Arc { + 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 { + self.inner.task_ctx() + } + } fn make_session_state() -> SessionState { let runtime = Arc::new(RuntimeEnv::default()); @@ -3407,6 +3437,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::().is_none()); + + let logical_plan = LogicalPlanBuilder::empty(false).build()?; + let physical_plan = session.create_physical_plan(&logical_plan).await?; + assert!(physical_plan.is::()); + assert!(query_planner_invoked.load(AtomicOrdering::Relaxed)); + assert!(invoked.load(AtomicOrdering::Relaxed)); + Ok(()) + } + async fn aggregate_explain(logical_plan: &LogicalPlan) -> Result { let physical_plan = plan(logical_plan).await?; Ok(displayable(physical_plan.as_ref()).indent(true).to_string()) @@ -4028,9 +4087,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::().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"); @@ -4557,7 +4623,7 @@ mod tests { _node: &dyn UserDefinedLogicalNode, _logical_inputs: &[&LogicalPlan], _physical_inputs: &[Arc], - _session_state: &SessionState, + _session_state: &dyn Session, _planning_ctx: &PhysicalPlanningContext, ) -> Result>> { internal_err!("BOOM") @@ -4718,7 +4784,7 @@ mod tests { node: &dyn UserDefinedLogicalNode, _logical_inputs: &[&LogicalPlan], _physical_inputs: &[Arc], - session_state: &SessionState, + session_state: &dyn Session, planning_ctx: &PhysicalPlanningContext, ) -> Result>> { for expr in node.expressions() { @@ -4748,7 +4814,7 @@ mod tests { _node: &dyn UserDefinedLogicalNode, _logical_inputs: &[&LogicalPlan], _physical_inputs: &[Arc], - _session_state: &SessionState, + _session_state: &dyn Session, _planning_ctx: &PhysicalPlanningContext, ) -> Result>> { Ok(Some(Arc::new(NoOpExecutionPlan::new(SchemaRef::new( @@ -5382,7 +5448,7 @@ digraph { _node: &dyn UserDefinedLogicalNode, _logical_inputs: &[&LogicalPlan], _physical_inputs: &[Arc], - _session_state: &SessionState, + _session_state: &dyn Session, _planning_ctx: &PhysicalPlanningContext, ) -> Result>> { Ok(None) @@ -5392,7 +5458,7 @@ digraph { &self, _planner: &dyn PhysicalPlanner, scan: &TableScan, - _session_state: &SessionState, + _session_state: &dyn Session, _planning_ctx: &PhysicalPlanningContext, ) -> Result>> { if scan.source.is::() { diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs b/datafusion/core/tests/user_defined/user_defined_plan.rs index 354a1b3110250..99363ba500b81 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> { // 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], - _session_state: &SessionState, + _session_state: &dyn Session, _planning_ctx: &PhysicalPlanningContext, ) -> Result>> { Ok( diff --git a/datafusion/datasource-arrow/src/file_format.rs b/datafusion/datasource-arrow/src/file_format.rs index c50ad98dfca0b..ad8bb46bde7e3 100644 --- a/datafusion/datasource-arrow/src/file_format.rs +++ b/datafusion/datasource-arrow/src/file_format.rs @@ -548,7 +548,10 @@ mod tests { AggregateUDF, Expr, HigherOrderUDF, LogicalPlan, ScalarUDF, WindowUDF, }; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; - use datafusion_session::{CatalogProviderList, EmptyCatalogProviderList}; + use datafusion_session::{ + CatalogProviderList, EmptyCatalogProviderList, QueryPlanner, + UnsupportedQueryPlanner, + }; use object_store::{chunked::ChunkedStore, memory::InMemory}; struct MockSession { @@ -579,6 +582,10 @@ mod tests { Arc::new(EmptyCatalogProviderList) } + fn query_planner(&self) -> Arc { + Arc::new(UnsupportedQueryPlanner) + } + async fn create_physical_plan( &self, _logical_plan: &LogicalPlan, diff --git a/datafusion/datasource/src/url.rs b/datafusion/datasource/src/url.rs index 9ac4c5f50d1f7..cc21e3a561067 100644 --- a/datafusion/datasource/src/url.rs +++ b/datafusion/datasource/src/url.rs @@ -523,7 +523,10 @@ mod tests { }; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_plan::ExecutionPlan; - use datafusion_session::{CatalogProviderList, EmptyCatalogProviderList}; + use datafusion_session::{ + CatalogProviderList, EmptyCatalogProviderList, QueryPlanner, + UnsupportedQueryPlanner, + }; use object_store::{ CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, PutMultipartOptions, PutPayload, @@ -1196,6 +1199,10 @@ mod tests { Arc::new(EmptyCatalogProviderList) } + fn query_planner(&self) -> Arc { + Arc::new(UnsupportedQueryPlanner) + } + async fn create_physical_plan( &self, _logical_plan: &LogicalPlan, diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs index 0c2d9fdeee819..831fe8f9bc98c 100644 --- a/datafusion/ffi/src/session/mod.rs +++ b/datafusion/ffi/src/session/mod.rs @@ -42,7 +42,9 @@ use datafusion_proto::logical_plan::LogicalExtensionCodec; use datafusion_proto::logical_plan::from_proto::parse_expr; use datafusion_proto::logical_plan::to_proto::serialize_expr; use datafusion_proto::protobuf::LogicalExprNode; -use datafusion_session::{CatalogProviderList, Session}; +use datafusion_session::{ + CatalogProviderList, QueryPlanner, Session, UnsupportedQueryPlanner, +}; use prost::Message; use stabby::str::Str as SStr; @@ -573,6 +575,10 @@ impl Session for ForeignSession { Arc::clone(&self.catalog_list) } + fn query_planner(&self) -> Arc { + Arc::new(UnsupportedQueryPlanner) + } + async fn create_physical_plan( &self, logical_plan: &LogicalPlan, @@ -730,6 +736,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 38c8a7c37211f..cb03303ac3c3f 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 0f81512b61c8e..2841afecf6ce3 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, - config: &ConfigOptions, - ) -> Result>; - - /// 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, - context: &dyn PhysicalOptimizerContext, - ) -> Result> { - 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 3b9ed7dacf1a8..6f7cfb7792c73 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 0000000000000..2b9154ba13a96 --- /dev/null +++ b/datafusion/session/src/physical_optimizer.rs @@ -0,0 +1,73 @@ +// 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 + } +} + +/// Transforms one [`ExecutionPlan`] into another that computes the same results, +/// but may do so more efficiently. +pub trait PhysicalOptimizerRule: Debug + std::any::Any { + /// Rewrite `plan` to an optimized form. + /// + /// Rules that need the statistics registry should override + /// [`Self::optimize_with_context`] instead. + fn optimize( + &self, + plan: Arc, + config: &ConfigOptions, + ) -> Result>; + + /// Rewrite `plan` with access to extended context. + fn optimize_with_context( + &self, + plan: Arc, + context: &dyn PhysicalOptimizerContext, + ) -> Result> { + self.optimize(plan, context.config_options()) + } + + /// A human-readable name for this optimizer rule. + fn name(&self) -> &str; + + /// Whether the physical planner should validate that this rule preserves + /// the plan schema. + fn schema_check(&self) -> bool; +} diff --git a/datafusion/session/src/planner.rs b/datafusion/session/src/planner.rs new file mode 100644 index 0000000000000..acb0aa4c629d6 --- /dev/null +++ b/datafusion/session/src/planner.rs @@ -0,0 +1,114 @@ +// 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 { + /// Create an [`ExecutionPlan`] from a [`LogicalPlan`]. + async fn create_physical_plan( + &self, + logical_plan: &LogicalPlan, + session: &dyn Session, + ) -> Result>; +} + +/// 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> { + 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>; + + /// Create a physical expression from a logical expression. + /// + /// `planning_ctx` resolves scalar subqueries for the plan subtree being + /// converted. Callers outside plan conversion should pass + /// `&PhysicalPlanningContext::default()`. + fn create_physical_expr( + &self, + expr: &Expr, + input_dfschema: &DFSchema, + session: &dyn Session, + planning_ctx: &PhysicalPlanningContext, + ) -> Result>; +} + +/// Plans user-defined logical nodes and table sources. +#[async_trait] +pub trait ExtensionPlanner { + /// Create a physical plan for a [`UserDefinedLogicalNode`]. + /// + /// Return `Ok(None)` when this planner does not support `node`. + async fn plan_extension( + &self, + planner: &dyn PhysicalPlanner, + node: &dyn UserDefinedLogicalNode, + logical_inputs: &[&LogicalPlan], + physical_inputs: &[Arc], + session: &dyn Session, + planning_ctx: &PhysicalPlanningContext, + ) -> Result>>; + + /// Create a physical plan for a [`TableScan`]. + /// + /// Return `Ok(None)` when this planner does not support `scan`. + async fn plan_table_scan( + &self, + _planner: &dyn PhysicalPlanner, + _scan: &TableScan, + _session: &dyn Session, + _planning_ctx: &PhysicalPlanningContext, + ) -> Result>> { + Ok(None) + } +} diff --git a/datafusion/session/src/session.rs b/datafusion/session/src/session.rs index cdac3f4bebc9e..0d585bce3edfb 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}; + /// 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,30 @@ pub trait Session: Send + Sync { self.config().options() } + /// Return the query planner for this session. + fn query_planner(&self) -> Arc; + + /// Optimize a logical plan. + /// + /// The default implementation returns the plan unchanged. Sessions that use + /// the default query planning implementation should override this method. + fn optimize(&self, plan: &LogicalPlan) -> Result { + Ok(plan.clone()) + } + + /// Return the physical optimizer rules for this session. + /// + /// The default implementation returns no rules. Sessions that use the + /// default physical planner should override this method. + fn physical_optimizers(&self) -> &[Arc] { + &[] + } + + /// 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 57da7b7dac248..40b00e05a43ce 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -763,13 +763,13 @@ async fn plan_extension( node: &dyn UserDefinedLogicalNode, logical_inputs: &[&LogicalPlan], physical_inputs: &[Arc], - session_state: &SessionState, + session: &dyn Session, planning_ctx: &PhysicalPlanningContext, // new parameter ) -> Result>> { 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)?; } // ... } @@ -777,20 +777,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: @@ -811,7 +823,32 @@ fn catalog_list(&self) -> Arc { } ``` -See [PR #23703](https://github.com/apache/datafusion/pull/23703) for details. +`Session` now also requires a `query_planner` method. Implementations that do +not expose a query planner can return the new `UnsupportedQueryPlanner`: + +```rust +use std::sync::Arc; +use datafusion_session::{QueryPlanner, UnsupportedQueryPlanner}; + +fn query_planner(&self) -> Arc { + Arc::new(UnsupportedQueryPlanner) +} +``` + +The `optimize`, `physical_optimizers`, and `statistics_registry` methods have +default implementations that leave logical plans unchanged, return no physical +optimizer rules, and return no statistics registry, respectively. A custom +session that uses `DefaultQueryPlanner` or `DefaultPhysicalPlanner` should +override these methods to expose its planning and optimization behavior. + +`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. ### `MSRV` updated to 1.94.0