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-23649-3b1ce16be6beb3377c59ec36fd9f8ae62998c4c0 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 4184b0778e0c691ed490c8ab0fa2b0f6f157722c Author: Tim Saucer <[email protected]> AuthorDate: Mon Jul 20 08:12:04 2026 -0400 refactor: pass SubqueryContext explicitly through planner traits (#23649) ## Which issue does this PR close? - Alternative to #22340 ## Rationale for this change When we turn logical plans into physical plans, we have a work around to get the context for scalar subqueries. Specifically we clone the session and mutate the `execution_props` where the subquery context currently sits. This is even called out in the code with comments: ```rust // Ideally, the subquery state would live in a dedicated planning // context rather than in `ExecutionProps`. It's here because // `create_physical_expr` only receives `&ExecutionProps`. ``` The real reason for this is that we want to expose the `QueryPlanner` via FFI and currently it depends upon downcasts from `Session` to `SessionState`. That does not work for the FFI crate where we have a different `Session` implementation. Since we cannot do the downcast, we cannot use the work around in the current code. This PR is designed to plumb through the required context properly rather than rely on manipulating the execution properties. ## What changes are included in this PR? The major change here is to remove `subquery_indexes` and `subquery_results` from `ExecutionProps` and put them into their own struct `SubqueryContext`. Everything else is just refactoring and plumbing to align around this split of the data. ## Are these changes tested? - Existing unit tests all pass. - Added additional unit test to cover user defined expressions. ## Are there any user-facing changes? Yes, this adds a single parameter into `create_physical_expr`, the subquery context. The migration guide demonstrates how to add a default context in for users, if necessary. --- .../examples/dataframe/cache_factory.rs | 2 + .../examples/query_planning/expr_api.rs | 12 +- .../examples/query_planning/pruning.rs | 9 +- .../examples/relation_planner/table_sample.rs | 2 + datafusion/catalog-listing/src/helpers.rs | 8 +- datafusion/catalog-listing/src/table.rs | 2 + datafusion/catalog/src/memory/table.rs | 25 +- datafusion/catalog/src/streaming.rs | 9 +- datafusion/core/src/execution/context/mod.rs | 2 + datafusion/core/src/execution/session_state.rs | 8 +- datafusion/core/src/physical_planner.rs | 319 +++++++++++++++++---- datafusion/core/src/test_util/parquet.rs | 9 +- datafusion/core/tests/parquet/page_pruning.rs | 9 +- .../core/tests/user_defined/user_defined_plan.rs | 2 + datafusion/datasource/src/file_scan_config/mod.rs | 15 +- datafusion/datasource/src/projection.rs | 14 +- datafusion/expr/src/execution_props.rs | 149 +--------- datafusion/expr/src/lib.rs | 1 + datafusion/expr/src/physical_planning_context.rs | 211 ++++++++++++++ datafusion/functions-nested/src/array_any_match.rs | 3 + datafusion/functions-nested/src/lambda_utils.rs | 2 + .../src/simplify_expressions/expr_simplifier.rs | 15 +- datafusion/optimizer/src/utils.rs | 10 +- datafusion/physical-expr/src/aggregate.rs | 36 ++- datafusion/physical-expr/src/analysis.rs | 28 +- datafusion/physical-expr/src/physical_expr.rs | 35 ++- datafusion/physical-expr/src/planner.rs | 258 +++++++++++++---- datafusion/physical-expr/src/scalar_subquery.rs | 2 +- datafusion/physical-plan/src/scalar_subquery.rs | 6 +- datafusion/proto/src/physical_plan/mod.rs | 2 +- .../proto/tests/cases/roundtrip_physical_plan.rs | 2 +- docs/source/library-user-guide/upgrading/55.0.0.md | 95 ++++++ 32 files changed, 993 insertions(+), 309 deletions(-) diff --git a/datafusion-examples/examples/dataframe/cache_factory.rs b/datafusion-examples/examples/dataframe/cache_factory.rs index a92c3dc4ce..dd145c715f 100644 --- a/datafusion-examples/examples/dataframe/cache_factory.rs +++ b/datafusion-examples/examples/dataframe/cache_factory.rs @@ -29,6 +29,7 @@ use datafusion::error::Result; use datafusion::execution::context::QueryPlanner; use datafusion::execution::session_state::CacheFactory; use datafusion::execution::{SessionState, SessionStateBuilder}; +use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext; use datafusion::logical_expr::{ Extension, LogicalPlan, UserDefinedLogicalNode, UserDefinedLogicalNodeCore, }; @@ -146,6 +147,7 @@ impl ExtensionPlanner for CacheNodePlanner { logical_inputs: &[&LogicalPlan], physical_inputs: &[Arc<dyn ExecutionPlan>], session_state: &SessionState, + _planning_ctx: &PhysicalPlanningContext, ) -> Result<Option<Arc<dyn ExecutionPlan>>> { if let Some(cache_node) = node.as_any().downcast_ref::<CacheNode>() { assert_eq!(logical_inputs.len(), 1, "Inconsistent number of inputs"); diff --git a/datafusion-examples/examples/query_planning/expr_api.rs b/datafusion-examples/examples/query_planning/expr_api.rs index c087019c68..dd5145def3 100644 --- a/datafusion-examples/examples/query_planning/expr_api.rs +++ b/datafusion-examples/examples/query_planning/expr_api.rs @@ -33,6 +33,7 @@ use datafusion::functions_aggregate::first_last::first_value_udaf; use datafusion::logical_expr::execution_props::ExecutionProps; use datafusion::logical_expr::expr::BinaryExpr; use datafusion::logical_expr::interval_arithmetic::Interval; +use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext; use datafusion::logical_expr::simplify::SimplifyContext; use datafusion::logical_expr::{ColumnarValue, ExprFunctionExt, ExprSchemable, Operator}; use datafusion::optimizer::analyzer::type_coercion::TypeCoercionRewriter; @@ -541,8 +542,12 @@ fn type_coercion_demo() -> Result<()> { // Evaluation with an expression that has not been type coerced cannot succeed. let props = ExecutionProps::default(); - let physical_expr = - datafusion::physical_expr::create_physical_expr(&expr, &df_schema, &props)?; + let physical_expr = datafusion::physical_expr::create_physical_expr( + &expr, + &df_schema, + &props, + &PhysicalPlanningContext::default(), + )?; let e = physical_expr.evaluate(&batch).unwrap_err(); assert!( e.find_root() @@ -566,6 +571,7 @@ fn type_coercion_demo() -> Result<()> { &coerced_expr, &df_schema, &props, + &PhysicalPlanningContext::default(), )?; assert!(physical_expr.evaluate(&batch).is_ok()); @@ -578,6 +584,7 @@ fn type_coercion_demo() -> Result<()> { &coerced_expr, &df_schema, &props, + &PhysicalPlanningContext::default(), )?; assert!(physical_expr.evaluate(&batch).is_ok()); @@ -606,6 +613,7 @@ fn type_coercion_demo() -> Result<()> { &coerced_expr, &df_schema, &props, + &PhysicalPlanningContext::default(), )?; assert!(physical_expr.evaluate(&batch).is_ok()); diff --git a/datafusion-examples/examples/query_planning/pruning.rs b/datafusion-examples/examples/query_planning/pruning.rs index 7fdc4a7952..df26aa57b6 100644 --- a/datafusion-examples/examples/query_planning/pruning.rs +++ b/datafusion-examples/examples/query_planning/pruning.rs @@ -26,6 +26,7 @@ use datafusion::common::pruning::PruningStatistics; use datafusion::common::{DFSchema, ScalarValue}; use datafusion::error::Result; use datafusion::execution::context::ExecutionProps; +use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext; use datafusion::physical_expr::create_physical_expr; use datafusion::physical_optimizer::pruning::PruningPredicate; use datafusion::prelude::*; @@ -194,7 +195,13 @@ impl PruningStatistics for MyCatalog { fn create_pruning_predicate(expr: Expr, schema: &SchemaRef) -> PruningPredicate { let df_schema = DFSchema::try_from(Arc::clone(schema)).unwrap(); let props = ExecutionProps::new(); - let physical_expr = create_physical_expr(&expr, &df_schema, &props).unwrap(); + let physical_expr = create_physical_expr( + &expr, + &df_schema, + &props, + &PhysicalPlanningContext::default(), + ) + .unwrap(); PruningPredicate::try_new(physical_expr, Arc::clone(schema)).unwrap() } diff --git a/datafusion-examples/examples/relation_planner/table_sample.rs b/datafusion-examples/examples/relation_planner/table_sample.rs index 2c696d92d7..b0ccff8d10 100644 --- a/datafusion-examples/examples/relation_planner/table_sample.rs +++ b/datafusion-examples/examples/relation_planner/table_sample.rs @@ -119,6 +119,7 @@ use datafusion_common::{ DFSchemaRef, DataFusionError, Result, Statistics, internal_err, not_impl_err, plan_datafusion_err, plan_err, }; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{ UserDefinedLogicalNode, UserDefinedLogicalNodeCore, logical_plan::{Extension, LogicalPlan, LogicalPlanBuilder}, @@ -587,6 +588,7 @@ impl ExtensionPlanner for TableSampleExtensionPlanner { _logical_inputs: &[&LogicalPlan], physical_inputs: &[Arc<dyn ExecutionPlan>], _session_state: &SessionState, + _planning_ctx: &PhysicalPlanningContext, ) -> Result<Option<Arc<dyn ExecutionPlan>>> { let Some(sample_node) = node.as_any().downcast_ref::<TableSamplePlanNode>() else { diff --git a/datafusion/catalog-listing/src/helpers.rs b/datafusion/catalog-listing/src/helpers.rs index 31f00b62ef..098f3d51ef 100644 --- a/datafusion/catalog-listing/src/helpers.rs +++ b/datafusion/catalog-listing/src/helpers.rs @@ -34,6 +34,7 @@ use arrow::{ record_batch::RecordBatch, }; use datafusion_expr::execution_props::ExecutionProps; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use futures::stream::FuturesUnordered; use futures::{StreamExt, TryStreamExt, stream::BoxStream}; use log::{debug, trace}; @@ -328,7 +329,12 @@ pub fn filter_partitioned_file( let filter = utils::conjunction(filters.iter().cloned()).unwrap_or_else(|| lit(true)); let props = ExecutionProps::new(); - let expr = create_physical_expr(&filter, df_schema, &props)?; + let expr = create_physical_expr( + &filter, + df_schema, + &props, + &PhysicalPlanningContext::default(), + )?; // Since we're only operating on a single file, our batch and resulting "array" holds only one // value indicating if the input file matches the provided filters diff --git a/datafusion/catalog-listing/src/table.rs b/datafusion/catalog-listing/src/table.rs index fabc7ca2a0..b3328cc063 100644 --- a/datafusion/catalog-listing/src/table.rs +++ b/datafusion/catalog-listing/src/table.rs @@ -44,6 +44,7 @@ use datafusion_execution::cache::cache_manager::{ }; use datafusion_expr::dml::InsertOp; use datafusion_expr::execution_props::ExecutionProps; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{ Expr, Partitioning as LogicalPartitioning, TableProviderFilterPushDown, TableType, }; @@ -614,6 +615,7 @@ impl TableProvider for ListingTable { output_partitioning, &df_schema, state.execution_props(), + &PhysicalPlanningContext::default(), )? } }; diff --git a/datafusion/catalog/src/memory/table.rs b/datafusion/catalog/src/memory/table.rs index 075e462f4f..5d07133799 100644 --- a/datafusion/catalog/src/memory/table.rs +++ b/datafusion/catalog/src/memory/table.rs @@ -36,6 +36,7 @@ use datafusion_datasource::memory::{MemSink, MemorySourceConfig}; use datafusion_datasource::sink::DataSinkExec; use datafusion_datasource::source::DataSourceExec; use datafusion_expr::dml::InsertOp; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{Expr, SortExpr, TableType}; use datafusion_physical_expr::{ LexOrdering, PhysicalExpr, create_physical_expr, create_physical_sort_exprs, @@ -209,8 +210,12 @@ impl TableProvider for MemTable { let eqp = state.execution_props(); let mut file_sort_order = vec![]; for sort_exprs in sort_order.iter() { - let physical_exprs = - create_physical_sort_exprs(sort_exprs, &df_schema, eqp)?; + let physical_exprs = create_physical_sort_exprs( + sort_exprs, + &df_schema, + eqp, + &PhysicalPlanningContext::default(), + )?; file_sort_order.extend(LexOrdering::new(physical_exprs)); } source = source.try_with_sort_information(file_sort_order)?; @@ -356,8 +361,12 @@ impl TableProvider for MemTable { let physical_assignments: HashMap<String, Arc<dyn PhysicalExpr>> = assignments .iter() .map(|(name, expr)| { - let physical_expr = - create_physical_expr(expr, &df_schema, state.execution_props())?; + let physical_expr = create_physical_expr( + expr, + &df_schema, + state.execution_props(), + &PhysicalPlanningContext::default(), + )?; Ok((name.clone(), physical_expr)) }) .collect::<Result<_>>()?; @@ -470,8 +479,12 @@ fn evaluate_filters_to_mask( let mut combined_mask: Option<BooleanArray> = None; for filter_expr in filters { - let physical_expr = - create_physical_expr(filter_expr, df_schema, execution_props)?; + let physical_expr = create_physical_expr( + filter_expr, + df_schema, + execution_props, + &PhysicalPlanningContext::default(), + )?; let result = physical_expr.evaluate(batch)?; let array = result.into_array(batch.num_rows())?; diff --git a/datafusion/catalog/src/streaming.rs b/datafusion/catalog/src/streaming.rs index 5bfecef1fb..50f05355aa 100644 --- a/datafusion/catalog/src/streaming.rs +++ b/datafusion/catalog/src/streaming.rs @@ -22,6 +22,7 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use async_trait::async_trait; use datafusion_common::{DFSchema, Result, plan_err}; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{Expr, SortExpr, TableType}; use datafusion_physical_expr::equivalence::project_ordering; use datafusion_physical_expr::projection::ProjectionMapping; @@ -131,8 +132,12 @@ impl TableProvider for StreamingTable { let df_schema = DFSchema::try_from(Arc::clone(&self.schema))?; let eqp = state.execution_props(); - let original_sort_exprs = - create_physical_sort_exprs(&self.sort_order, &df_schema, eqp)?; + let original_sort_exprs = create_physical_sort_exprs( + &self.sort_order, + &df_schema, + eqp, + &PhysicalPlanningContext::default(), + )?; if let Some(p) = projection { // When performing a projection, the output columns will not match diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs index 08c7463e21..281cb4dd79 100644 --- a/datafusion/core/src/execution/context/mod.rs +++ b/datafusion/core/src/execution/context/mod.rs @@ -2383,6 +2383,7 @@ mod tests { use arrow_schema::FieldRef; use datafusion_common::DataFusionError; use datafusion_common::datatype::DataTypeExt; + use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use std::error::Error; use std::path::PathBuf; @@ -2851,6 +2852,7 @@ mod tests { _expr: &Expr, _input_dfschema: &DFSchema, _session_state: &SessionState, + _planning_ctx: &PhysicalPlanningContext, ) -> Result<Arc<dyn PhysicalExpr>> { unimplemented!() } diff --git a/datafusion/core/src/execution/session_state.rs b/datafusion/core/src/execution/session_state.rs index f7117c89ef..ff4ec20cfc 100644 --- a/datafusion/core/src/execution/session_state.rs +++ b/datafusion/core/src/execution/session_state.rs @@ -55,6 +55,7 @@ use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_expr::TableSource; use datafusion_expr::execution_props::ExecutionProps; use datafusion_expr::expr_rewriter::FunctionRewrite; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::planner::ExprPlanner; #[cfg(feature = "sql")] use datafusion_expr::planner::{RelationPlanner, TypePlanner}; @@ -799,7 +800,12 @@ impl SessionState { .transform_up(|expr| rewrite.rewrite(expr, df_schema, config_options))? .data; } - create_physical_expr(&expr, df_schema, self.execution_props()) + create_physical_expr( + &expr, + df_schema, + self.execution_props(), + &PhysicalPlanningContext::default(), + ) } /// Return the session ID diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index b6d28e7b21..aef8036c74 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -80,7 +80,6 @@ use datafusion_common::{ use datafusion_datasource::file_groups::FileGroup; use datafusion_datasource::memory::MemorySourceConfig; use datafusion_expr::dml::{CopyTo, InsertOp}; -use datafusion_expr::execution_props::{ScalarSubqueryResults, SubqueryIndex}; use datafusion_expr::expr::{ Alias, GroupingSet, NullTreatment, WindowFunction, WindowFunctionParams, physical_name, @@ -88,6 +87,9 @@ use datafusion_expr::expr::{ use datafusion_expr::expr_rewriter::unnormalize_cols; use datafusion_expr::logical_plan::Subquery; use datafusion_expr::logical_plan::builder::wrap_projection_for_join_if_necessary; +use datafusion_expr::physical_planning_context::{ + PhysicalPlanningContext, ScalarSubqueryResults, SubqueryIndex, +}; use datafusion_expr::utils::{expr_to_columns, split_conjunction}; use datafusion_expr::{ Analyze, BinaryExpr, DescribeTable, DmlStatement, Explain, ExplainFormat, Extension, @@ -135,11 +137,19 @@ pub trait PhysicalPlanner: Send + Sync { /// `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>>; } @@ -156,6 +166,12 @@ pub trait ExtensionPlanner { /// 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, @@ -163,6 +179,7 @@ pub trait ExtensionPlanner { logical_inputs: &[&LogicalPlan], physical_inputs: &[Arc<dyn ExecutionPlan>], session_state: &SessionState, + planning_ctx: &PhysicalPlanningContext, ) -> Result<Option<Arc<dyn ExecutionPlan>>>; /// Create a physical plan for a [`LogicalPlan::TableScan`]. @@ -202,6 +219,7 @@ pub trait ExtensionPlanner { /// _logical_inputs: &[&LogicalPlan], /// _physical_inputs: &[Arc<dyn ExecutionPlan>], /// _session_state: &SessionState, + /// _planning_ctx: &PhysicalPlanningContext, /// ) -> Result<Option<Arc<dyn ExecutionPlan>>> { /// Ok(None) /// } @@ -211,6 +229,7 @@ pub trait ExtensionPlanner { /// _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>() { @@ -235,6 +254,7 @@ pub trait ExtensionPlanner { _planner: &dyn PhysicalPlanner, _scan: &TableScan, _session_state: &SessionState, + _planning_ctx: &PhysicalPlanningContext, ) -> Result<Option<Arc<dyn ExecutionPlan>>> { Ok(None) } @@ -295,8 +315,14 @@ impl PhysicalPlanner for DefaultPhysicalPlanner { expr: &Expr, input_dfschema: &DFSchema, session_state: &SessionState, + planning_ctx: &PhysicalPlanningContext, ) -> Result<Arc<dyn PhysicalExpr>> { - create_physical_expr(expr, input_dfschema, session_state.execution_props()) + create_physical_expr( + expr, + input_dfschema, + session_state.execution_props(), + planning_ctx, + ) } } @@ -417,9 +443,9 @@ impl DefaultPhysicalPlanner { /// collected, planned as separate physical plans, and each assigned an /// index in a shared [`ScalarSubqueryResults`] container that will hold its /// result at execution time. The index map and shared results container are - /// registered in [`ExecutionProps`] so that [`create_physical_expr`] can - /// convert `Expr::ScalarSubquery` into [`ScalarSubqueryExpr`] nodes that - /// read from that container. + /// stored in a [`PhysicalPlanningContext`] and passed explicitly to + /// [`create_physical_expr`] so it can convert `Expr::ScalarSubquery` into + /// [`ScalarSubqueryExpr`] nodes that read from that container. /// /// The resulting physical plan is wrapped in a [`ScalarSubqueryExec`] node /// that executes those subquery plans before any data flows through the @@ -458,27 +484,25 @@ impl DefaultPhysicalPlanner { if links.is_empty() { return self - .create_initial_plan_inner(logical_plan, session_state) + .create_initial_plan_inner( + logical_plan, + session_state, + &PhysicalPlanningContext::default(), + ) .await; } - // Create the shared `ScalarSubqueryResults` container and register - // it in `ExecutionProps` so that `create_physical_expr` can resolve - // `Expr::ScalarSubquery` into `ScalarSubqueryExpr` nodes. We clone - // the `SessionState` so these are available throughout physical - // planning without mutating the caller's state. - // - // Ideally, the subquery state would live in a dedicated planning - // context rather than in `ExecutionProps`. It's here because - // `create_physical_expr` only receives `&ExecutionProps`. + // Build a `PhysicalPlanningContext` that carries the index map and + // shared results container into calls that create physical expressions. + // The context is threaded explicitly through physical planning rather + // than being stashed in `ExecutionProps`, so the planner does not need + // a mutable `SessionState` and each recursively planned subtree receives + // the correct context. let results = ScalarSubqueryResults::new(links.len()); - let mut owned = session_state.clone(); - owned.execution_props_mut().subquery_indexes = index_map; - owned.execution_props_mut().subquery_results = results.clone(); - let session_state = Cow::Owned(owned); + let planning_ctx = PhysicalPlanningContext::new(index_map, results.clone()); let plan = self - .create_initial_plan_inner(logical_plan, &session_state) + .create_initial_plan_inner(logical_plan, session_state, &planning_ctx) .await?; Ok(Arc::new(ScalarSubqueryExec::new(plan, links, results))) }) @@ -490,6 +514,7 @@ impl DefaultPhysicalPlanner { &self, logical_plan: &LogicalPlan, session_state: &SessionState, + planning_ctx: &PhysicalPlanningContext, ) -> Result<Arc<dyn ExecutionPlan>> { // DFS the tree to flatten it into a Vec. // This will allow us to build the Physical Plan from the leaves up @@ -540,9 +565,9 @@ impl DefaultPhysicalPlanner { let max_concurrency = planning_concurrency.min(flat_tree_leaf_indices.len()); // Spawning tasks which will traverse leaf up to the root. - let tasks = flat_tree_leaf_indices - .into_iter() - .map(|index| self.task_helper(index, Arc::clone(&flat_tree), session_state)); + let tasks = flat_tree_leaf_indices.into_iter().map(|index| { + self.task_helper(index, Arc::clone(&flat_tree), session_state, planning_ctx) + }); let mut outputs = futures::stream::iter(tasks) .buffer_unordered(max_concurrency) .try_collect::<Vec<_>>() @@ -570,6 +595,7 @@ impl DefaultPhysicalPlanner { leaf_starter_index: usize, flat_tree: Arc<Vec<LogicalNode<'a>>>, session_state: &'a SessionState, + planning_ctx: &'a PhysicalPlanningContext, ) -> Result<Option<Arc<dyn ExecutionPlan>>> { // We always start with a leaf, so can ignore status and pass empty children let mut node = flat_tree.get(leaf_starter_index).ok_or_else(|| { @@ -581,6 +607,7 @@ impl DefaultPhysicalPlanner { .map_logical_node_to_physical( node.node, session_state, + planning_ctx, ChildrenContainer::None, ) .await?; @@ -598,6 +625,7 @@ impl DefaultPhysicalPlanner { .map_logical_node_to_physical( node.node, session_state, + planning_ctx, ChildrenContainer::One(plan), ) .await?; @@ -634,7 +662,12 @@ impl DefaultPhysicalPlanner { let children = children.into_iter().map(|epc| epc.plan).collect(); let children = ChildrenContainer::Multiple(children); plan = self - .map_logical_node_to_physical(node.node, session_state, children) + .map_logical_node_to_physical( + node.node, + session_state, + planning_ctx, + children, + ) .await?; } } @@ -649,6 +682,7 @@ impl DefaultPhysicalPlanner { &self, node: &LogicalPlan, session_state: &SessionState, + planning_ctx: &PhysicalPlanningContext, children: ChildrenContainer, ) -> Result<Arc<dyn ExecutionPlan>> { let execution_props = session_state.execution_props(); @@ -687,8 +721,9 @@ impl DefaultPhysicalPlanner { break; } - maybe_plan = - planner.plan_table_scan(self, scan, session_state).await?; + maybe_plan = planner + .plan_table_scan(self, scan, session_state, planning_ctx) + .await?; } let plan = match maybe_plan { @@ -714,7 +749,12 @@ impl DefaultPhysicalPlanner { .map(|row| { row.iter() .map(|expr| { - create_physical_expr(expr, schema, execution_props) + create_physical_expr( + expr, + schema, + execution_props, + planning_ctx, + ) }) .collect::<Result<Vec<Arc<dyn PhysicalExpr>>>>() }) @@ -973,7 +1013,14 @@ impl DefaultPhysicalPlanner { let logical_schema = node.schema(); let window_expr = window_expr .iter() - .map(|e| create_window_expr(e, logical_schema, execution_props)) + .map(|e| { + create_window_expr( + e, + logical_schema, + execution_props, + planning_ctx, + ) + }) .collect::<Result<Vec<_>>>()?; let can_repartition = session_state.config().target_partitions() > 1 @@ -1079,6 +1126,7 @@ impl DefaultPhysicalPlanner { logical_input_schema, &physical_input_schema, execution_props, + planning_ctx, )?; let agg_filter = aggr_expr @@ -1089,6 +1137,7 @@ impl DefaultPhysicalPlanner { logical_input_schema, &physical_input_schema, execution_props, + planning_ctx, ) .build() .map(lowered_aggregate_to_tuple) @@ -1186,6 +1235,7 @@ impl DefaultPhysicalPlanner { LogicalPlan::Projection(Projection { input, expr, .. }) => self .create_project_physical_exec_with_props( execution_props, + planning_ctx, children.one()?, input, expr, @@ -1195,8 +1245,12 @@ impl DefaultPhysicalPlanner { }) => { let physical_input = children.one()?; let input_dfschema = input.schema(); - let runtime_expr = - create_physical_expr(predicate, input_dfschema, execution_props)?; + let runtime_expr = create_physical_expr( + predicate, + input_dfschema, + execution_props, + planning_ctx, + )?; let input_schema = input.schema(); let filter = match self.try_plan_async_exprs( @@ -1258,6 +1312,7 @@ impl DefaultPhysicalPlanner { partitioning_scheme, input_dfschema, execution_props, + planning_ctx, )?; Arc::new(RepartitionExec::try_new( physical_input, @@ -1269,8 +1324,12 @@ impl DefaultPhysicalPlanner { }) => { let physical_input = children.one()?; let input_dfschema = input.as_ref().schema(); - let sort_exprs = - create_physical_sort_exprs(expr, input_dfschema, execution_props)?; + let sort_exprs = create_physical_sort_exprs( + expr, + input_dfschema, + execution_props, + planning_ctx, + )?; let Some(ordering) = LexOrdering::new(sort_exprs) else { return internal_err!( "SortExec requires at least one sort expression" @@ -1399,6 +1458,7 @@ impl DefaultPhysicalPlanner { LogicalPlan::Projection(Projection { input, expr, .. }), ) => self.create_project_physical_exec_with_props( execution_props, + planning_ctx, physical_left, input, expr, @@ -1412,6 +1472,7 @@ impl DefaultPhysicalPlanner { LogicalPlan::Projection(Projection { input, expr, .. }), ) => self.create_project_physical_exec_with_props( execution_props, + planning_ctx, physical_right, input, expr, @@ -1479,9 +1540,18 @@ impl DefaultPhysicalPlanner { let join_on = keys .iter() .map(|(l, r)| { - let l = create_physical_expr(l, left_df_schema, execution_props)?; - let r = - create_physical_expr(r, right_df_schema, execution_props)?; + let l = create_physical_expr( + l, + left_df_schema, + execution_props, + planning_ctx, + )?; + let r = create_physical_expr( + r, + right_df_schema, + execution_props, + planning_ctx, + )?; Ok((l, r)) }) .collect::<Result<join_utils::JoinOn>>()?; @@ -1582,6 +1652,7 @@ impl DefaultPhysicalPlanner { expr, &filter_df_schema, execution_props, + planning_ctx, )?; let column_indices = join_utils::JoinFilter::build_column_indices( left_field_indices, @@ -1702,11 +1773,13 @@ impl DefaultPhysicalPlanner { lhs_logical, left_df_schema, execution_props, + planning_ctx, )?; let on_right = create_physical_expr( rhs_logical, right_df_schema, execution_props, + planning_ctx, )?; Arc::new(PiecewiseMergeJoinExec::try_new( @@ -1778,6 +1851,7 @@ impl DefaultPhysicalPlanner { if let Some((input, expr)) = new_project { self.create_project_physical_exec_with_props( execution_props, + planning_ctx, join, input, expr, @@ -1820,6 +1894,7 @@ impl DefaultPhysicalPlanner { &logical_input, &children, session_state, + planning_ctx, ) .await?; } @@ -1882,6 +1957,7 @@ impl DefaultPhysicalPlanner { input_dfschema: &DFSchema, input_schema: &Schema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result<PhysicalGroupBy> { if group_expr.len() == 1 { match &group_expr[0] { @@ -1891,6 +1967,7 @@ impl DefaultPhysicalPlanner { input_dfschema, input_schema, execution_props, + planning_ctx, ) } Expr::GroupingSet(GroupingSet::Cube(exprs)) => create_cube_physical_expr( @@ -1898,6 +1975,7 @@ impl DefaultPhysicalPlanner { input_dfschema, input_schema, execution_props, + planning_ctx, ), Expr::GroupingSet(GroupingSet::Rollup(exprs)) => { create_rollup_physical_expr( @@ -1905,10 +1983,16 @@ impl DefaultPhysicalPlanner { input_dfschema, input_schema, execution_props, + planning_ctx, ) } expr => Ok(PhysicalGroupBy::new_single(vec![tuple_err(( - create_physical_expr(expr, input_dfschema, execution_props), + create_physical_expr( + expr, + input_dfschema, + execution_props, + planning_ctx, + ), physical_name(expr), ))?])), } @@ -1922,7 +2006,12 @@ impl DefaultPhysicalPlanner { .iter() .map(|e| { tuple_err(( - create_physical_expr(e, input_dfschema, execution_props), + create_physical_expr( + e, + input_dfschema, + execution_props, + planning_ctx, + ), physical_name(e), )) }) @@ -1947,6 +2036,7 @@ fn merge_grouping_set_physical_expr( input_dfschema: &DFSchema, input_schema: &Schema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result<PhysicalGroupBy> { let num_groups = grouping_sets.len(); let mut all_exprs: Vec<Expr> = vec![]; @@ -1961,6 +2051,7 @@ fn merge_grouping_set_physical_expr( expr, input_dfschema, execution_props, + planning_ctx, )?); null_exprs.push(get_null_physical_expr_pair( @@ -1968,6 +2059,7 @@ fn merge_grouping_set_physical_expr( input_dfschema, input_schema, execution_props, + planning_ctx, )?); } } @@ -1998,6 +2090,7 @@ fn create_cube_physical_expr( input_dfschema: &DFSchema, input_schema: &Schema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result<PhysicalGroupBy> { let num_of_exprs = exprs.len(); let num_groups = num_of_exprs * num_of_exprs; @@ -2013,12 +2106,14 @@ fn create_cube_physical_expr( input_dfschema, input_schema, execution_props, + planning_ctx, )?); all_exprs.push(get_physical_expr_pair( expr, input_dfschema, execution_props, + planning_ctx, )?) } @@ -2044,6 +2139,7 @@ fn create_rollup_physical_expr( input_dfschema: &DFSchema, input_schema: &Schema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result<PhysicalGroupBy> { let num_of_exprs = exprs.len(); @@ -2060,12 +2156,14 @@ fn create_rollup_physical_expr( input_dfschema, input_schema, execution_props, + planning_ctx, )?); all_exprs.push(get_physical_expr_pair( expr, input_dfschema, execution_props, + planning_ctx, )?) } @@ -2092,8 +2190,10 @@ fn get_null_physical_expr_pair( input_dfschema: &DFSchema, input_schema: &Schema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result<(Arc<dyn PhysicalExpr>, String)> { - let physical_expr = create_physical_expr(expr, input_dfschema, execution_props)?; + let physical_expr = + create_physical_expr(expr, input_dfschema, execution_props, planning_ctx)?; let physical_name = physical_name(&expr.clone())?; let data_type = physical_expr.data_type(input_schema)?; @@ -2163,8 +2263,10 @@ fn get_physical_expr_pair( expr: &Expr, input_dfschema: &DFSchema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result<(Arc<dyn PhysicalExpr>, String)> { - let physical_expr = create_physical_expr(expr, input_dfschema, execution_props)?; + let physical_expr = + create_physical_expr(expr, input_dfschema, execution_props, planning_ctx)?; let physical_name = physical_name(expr)?; Ok((physical_expr, physical_name)) } @@ -2417,11 +2519,14 @@ pub fn is_window_frame_bound_valid(window_frame: &WindowFrame) -> bool { } /// Create a window expression with a name from a logical expression +/// +/// See [`create_physical_expr`] for details on the `planning_ctx` argument. pub fn create_window_expr_with_name( e: &Expr, name: impl Into<String>, logical_schema: &DFSchema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result<Arc<dyn WindowExpr>> { let name = name.into(); let physical_schema = Arc::clone(logical_schema.inner()); @@ -2440,12 +2545,24 @@ pub fn create_window_expr_with_name( filter, }, } = window_fun.as_ref(); - let physical_args = - create_physical_exprs(args, logical_schema, execution_props)?; - let partition_by = - create_physical_exprs(partition_by, logical_schema, execution_props)?; - let order_by = - create_physical_sort_exprs(order_by, logical_schema, execution_props)?; + let physical_args = create_physical_exprs( + args, + logical_schema, + execution_props, + planning_ctx, + )?; + let partition_by = create_physical_exprs( + partition_by, + logical_schema, + execution_props, + planning_ctx, + )?; + let order_by = create_physical_sort_exprs( + order_by, + logical_schema, + execution_props, + planning_ctx, + )?; if !is_window_frame_bound_valid(window_frame) { return plan_err!( @@ -2460,7 +2577,9 @@ pub fn create_window_expr_with_name( == NullTreatment::IgnoreNulls; let physical_filter = filter .as_ref() - .map(|f| create_physical_expr(f, logical_schema, execution_props)) + .map(|f| { + create_physical_expr(f, logical_schema, execution_props, planning_ctx) + }) .transpose()?; windows::create_window_expr( @@ -2481,10 +2600,13 @@ pub fn create_window_expr_with_name( } /// Create a window expression from a logical expression or an alias +/// +/// See [`create_physical_expr`] for details on the `planning_ctx` argument. pub fn create_window_expr( e: &Expr, logical_schema: &DFSchema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result<Arc<dyn WindowExpr>> { // unpack aliased logical expressions, e.g. "sum(col) over () as total" let (name, e) = match e { @@ -2494,7 +2616,7 @@ pub fn create_window_expr( ), _ => (e.schema_name().to_string(), e.clone()), }; - create_window_expr_with_name(&e, name, logical_schema, execution_props) + create_window_expr_with_name(&e, name, logical_schema, execution_props, planning_ctx) } type AggregateExprWithOptionalArgs = ( @@ -2515,11 +2637,13 @@ pub fn create_aggregate_expr_with_name_and_maybe_filter( physical_input_schema: &Schema, execution_props: &ExecutionProps, ) -> Result<AggregateExprWithOptionalArgs> { + let planning_ctx = PhysicalPlanningContext::default(); let mut builder = LoweredAggregateBuilder::new( e, logical_input_schema, physical_input_schema, execution_props, + &planning_ctx, ) .with_human_display(human_display); @@ -2554,11 +2678,13 @@ pub fn create_aggregate_expr_and_maybe_filter( _ => (None, String::default(), e.clone()), }; + let planning_ctx = PhysicalPlanningContext::default(); let mut builder = LoweredAggregateBuilder::new( &e, logical_input_schema, physical_input_schema, execution_props, + &planning_ctx, ) .with_human_display(human_display); @@ -2953,6 +3079,7 @@ impl DefaultPhysicalPlanner { fn create_project_physical_exec_with_props( &self, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, input_exec: Arc<dyn ExecutionPlan>, input: &Arc<LogicalPlan>, expr: &[Expr], @@ -2990,8 +3117,12 @@ impl DefaultPhysicalPlanner { physical_name(e) }; - let physical_expr = - create_physical_expr(e, input_logical_schema, execution_props); + let physical_expr = create_physical_expr( + e, + input_logical_schema, + execution_props, + planning_ctx, + ); tuple_err((physical_expr, physical_name)) }) @@ -3248,6 +3379,7 @@ mod tests { Accumulator, AggregateUDF, AggregateUDFImpl, ExprFunctionExt, LogicalPlanBuilder, Partitioning as LogicalPartitioning, RangePartitioning, Signature, TableSource, UserDefinedLogicalNodeCore, Volatility, WindowFunctionDefinition, col, lit, + scalar_subquery, }; use datafusion_functions_aggregate::count::{count_all, count_udaf}; use datafusion_functions_aggregate::expr_fn::sum; @@ -3357,8 +3489,12 @@ mod tests { )) .alias_with_metadata("window_alias", Some(metadata)); - let window_expr = - create_window_expr(&expr, &logical_schema, &ExecutionProps::new())?; + let window_expr = create_window_expr( + &expr, + &logical_schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + )?; assert_eq!(window_expr.name(), "window_alias"); Ok(()) @@ -3613,6 +3749,7 @@ mod tests { logical_input_schema, physical_input_schema, session_state.execution_props(), + &PhysicalPlanningContext::default(), ); insta::assert_debug_snapshot!(cube, @r#" @@ -3744,6 +3881,7 @@ mod tests { logical_input_schema, physical_input_schema, session_state.execution_props(), + &PhysicalPlanningContext::default(), ); insta::assert_debug_snapshot!(rollup, @r#" @@ -3848,6 +3986,7 @@ mod tests { &col("a").not(), &dfschema, &make_session_state(), + &PhysicalPlanningContext::default(), )?; let expected = expressions::not(expressions::col("a", &schema)?)?; @@ -3875,6 +4014,29 @@ mod tests { Ok(()) } + #[tokio::test] + async fn scalar_subquery_in_extension_expr_plans() -> Result<()> { + let subquery = LogicalPlanBuilder::empty(true) + .project(vec![lit(42_i32)])? + .build()?; + let logical_plan = LogicalPlan::Extension(Extension { + node: Arc::new(NoOpExtensionNode { + expressions: vec![scalar_subquery(Arc::new(subquery))], + ..Default::default() + }), + }); + let planner = DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new( + ExpressionExtensionPlanner, + )]); + + let plan = planner + .create_physical_plan(&logical_plan, &make_session_state()) + .await?; + + assert_contains!(format!("{plan:?}"), "ScalarSubqueryExec"); + Ok(()) + } + #[tokio::test] async fn error_during_extension_planning() { let session_state = make_session_state(); @@ -4396,6 +4558,7 @@ mod tests { _logical_inputs: &[&LogicalPlan], _physical_inputs: &[Arc<dyn ExecutionPlan>], _session_state: &SessionState, + _planning_ctx: &PhysicalPlanningContext, ) -> Result<Option<Arc<dyn ExecutionPlan>>> { internal_err!("BOOM") } @@ -4404,6 +4567,7 @@ mod tests { #[derive(PartialEq, Eq, Hash)] struct NoOpExtensionNode { schema: DFSchemaRef, + expressions: Vec<Expr>, } impl Default for NoOpExtensionNode { @@ -4416,6 +4580,7 @@ mod tests { ) .unwrap(), ), + expressions: vec![], } } } @@ -4448,7 +4613,7 @@ mod tests { } fn expressions(&self) -> Vec<Expr> { - vec![] + self.expressions.clone() } fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { @@ -4457,10 +4622,13 @@ mod tests { fn with_exprs_and_inputs( &self, - _exprs: Vec<Expr>, + exprs: Vec<Expr>, _inputs: Vec<LogicalPlan>, ) -> Result<Self> { - unimplemented!("NoOp"); + Ok(Self { + schema: Arc::clone(&self.schema), + expressions: exprs, + }) } fn supports_limit_pushdown(&self) -> bool { @@ -4522,9 +4690,13 @@ mod tests { fn with_new_children( self: Arc<Self>, - _children: Vec<Arc<dyn ExecutionPlan>>, + children: Vec<Arc<dyn ExecutionPlan>>, ) -> Result<Arc<dyn ExecutionPlan>> { - unimplemented!("NoOpExecutionPlan::with_new_children"); + if children.is_empty() { + Ok(self) + } else { + exec_err!("NoOpExecutionPlan does not support children") + } } fn execute( @@ -4536,6 +4708,33 @@ mod tests { } } + struct ExpressionExtensionPlanner; + + #[async_trait] + impl ExtensionPlanner for ExpressionExtensionPlanner { + 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>>> { + for expr in node.expressions() { + planner.create_physical_expr( + &expr, + node.schema(), + session_state, + planning_ctx, + )?; + } + Ok(Some(Arc::new(NoOpExecutionPlan::new(Arc::clone( + node.schema().inner(), + ))))) + } + } + // Produces an execution plan where the schema is mismatched from // the logical plan node. struct BadExtensionPlanner {} @@ -4550,6 +4749,7 @@ mod tests { _logical_inputs: &[&LogicalPlan], _physical_inputs: &[Arc<dyn ExecutionPlan>], _session_state: &SessionState, + _planning_ctx: &PhysicalPlanningContext, ) -> Result<Option<Arc<dyn ExecutionPlan>>> { Ok(Some(Arc::new(NoOpExecutionPlan::new(SchemaRef::new( Schema::new(vec![Field::new("b", DataType::Int32, false)]), @@ -4995,9 +5195,8 @@ digraph { } #[tokio::test] - // When schemas match, planning proceeds past the schema_satisfied_by check. - // It then panics on unimplemented error in NoOpExecutionPlan. - #[should_panic(expected = "NoOpExecutionPlan")] + // When schemas match, planning proceeds past the schema_satisfied_by check + // and succeeds. async fn test_aggregate_schema_check_passes() { let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)])); @@ -5184,6 +5383,7 @@ digraph { _logical_inputs: &[&LogicalPlan], _physical_inputs: &[Arc<dyn ExecutionPlan>], _session_state: &SessionState, + _planning_ctx: &PhysicalPlanningContext, ) -> Result<Option<Arc<dyn ExecutionPlan>>> { Ok(None) } @@ -5193,6 +5393,7 @@ digraph { _planner: &dyn PhysicalPlanner, scan: &TableScan, _session_state: &SessionState, + _planning_ctx: &PhysicalPlanningContext, ) -> Result<Option<Arc<dyn ExecutionPlan>>> { if scan.source.is::<MockTableSource>() { Ok(Some(Arc::new(EmptyExec::new(Arc::clone( diff --git a/datafusion/core/src/test_util/parquet.rs b/datafusion/core/src/test_util/parquet.rs index c534954213..d1018f3fb0 100644 --- a/datafusion/core/src/test_util/parquet.rs +++ b/datafusion/core/src/test_util/parquet.rs @@ -29,6 +29,7 @@ use crate::datasource::object_store::ObjectStoreUrl; use crate::datasource::physical_plan::ParquetSource; use crate::error::Result; use crate::logical_expr::execution_props::ExecutionProps; +use crate::logical_expr::physical_planning_context::PhysicalPlanningContext; use crate::logical_expr::simplify::SimplifyContext; use crate::optimizer::simplify_expressions::ExprSimplifier; use crate::physical_expr::create_physical_expr; @@ -172,8 +173,12 @@ impl TestParquetFile { if let Some(filter) = maybe_filter { let simplifier = ExprSimplifier::new(context); let filter = simplifier.coerce(filter, &df_schema).unwrap(); - let physical_filter_expr = - create_physical_expr(&filter, &df_schema, &ExecutionProps::default())?; + let physical_filter_expr = create_physical_expr( + &filter, + &df_schema, + &ExecutionProps::default(), + &PhysicalPlanningContext::default(), + )?; let source = Arc::new( ParquetSource::new(Arc::clone(&self.schema)) diff --git a/datafusion/core/tests/parquet/page_pruning.rs b/datafusion/core/tests/parquet/page_pruning.rs index a41803191a..372a7a601d 100644 --- a/datafusion/core/tests/parquet/page_pruning.rs +++ b/datafusion/core/tests/parquet/page_pruning.rs @@ -38,6 +38,7 @@ use datafusion_expr::{Expr, col, lit}; use datafusion_physical_expr::create_physical_expr; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use futures::StreamExt; use object_store::ObjectMeta; use object_store::path::Path; @@ -74,7 +75,13 @@ async fn get_parquet_exec( let df_schema = schema.clone().to_dfschema().unwrap(); let execution_props = ExecutionProps::new(); - let predicate = create_physical_expr(&filter, &df_schema, &execution_props).unwrap(); + let predicate = create_physical_expr( + &filter, + &df_schema, + &execution_props, + &PhysicalPlanningContext::default(), + ) + .unwrap(); let source = Arc::new( ParquetSource::new(schema.clone()) diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs b/datafusion/core/tests/user_defined/user_defined_plan.rs index 75738bcfe1..354a1b3110 100644 --- a/datafusion/core/tests/user_defined/user_defined_plan.rs +++ b/datafusion/core/tests/user_defined/user_defined_plan.rs @@ -99,6 +99,7 @@ use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; use async_trait::async_trait; use datafusion_common::cast::as_string_view_array; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use futures::{Stream, StreamExt}; /// Execute the specified sql and return the resulting record batches @@ -630,6 +631,7 @@ impl ExtensionPlanner for TopKPlanner { logical_inputs: &[&LogicalPlan], physical_inputs: &[Arc<dyn ExecutionPlan>], _session_state: &SessionState, + _planning_ctx: &PhysicalPlanningContext, ) -> Result<Option<Arc<dyn ExecutionPlan>>> { Ok( if let Some(topk_node) = node.as_any().downcast_ref::<TopKPlanNode>() { diff --git a/datafusion/datasource/src/file_scan_config/mod.rs b/datafusion/datasource/src/file_scan_config/mod.rs index c0b47a9a52..962df06302 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/datafusion/datasource/src/file_scan_config/mod.rs @@ -1664,6 +1664,7 @@ mod tests { use chrono::TimeZone; use datafusion_common::DFSchema; use datafusion_expr::execution_props::ExecutionProps; + use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use object_store::{ObjectMeta, path::Path}; struct File { @@ -1877,6 +1878,7 @@ mod tests { &expr, &DFSchema::try_from(Arc::clone(&table_schema))?, &ExecutionProps::default(), + &PhysicalPlanningContext::default(), ) }) .collect::<Result<Vec<_>>>()?, @@ -2243,7 +2245,10 @@ mod tests { #[test] fn test_split_groups_by_statistics_with_target_partitions() -> Result<()> { use datafusion_common::DFSchema; - use datafusion_expr::{col, execution_props::ExecutionProps}; + use datafusion_expr::{ + col, execution_props::ExecutionProps, + physical_planning_context::PhysicalPlanningContext, + }; let schema = Arc::new(Schema::new(vec![Field::new( "value", @@ -2257,7 +2262,13 @@ mod tests { let sort_expr = [col("value").sort(true, false)]; let sort_ordering = sort_expr .map(|expr| { - create_physical_sort_expr(&expr, &df_schema, &exec_props).unwrap() + create_physical_sort_expr( + &expr, + &df_schema, + &exec_props, + &PhysicalPlanningContext::default(), + ) + .unwrap() }) .into(); diff --git a/datafusion/datasource/src/projection.rs b/datafusion/datasource/src/projection.rs index de822ae602..3cf4f29a77 100644 --- a/datafusion/datasource/src/projection.rs +++ b/datafusion/datasource/src/projection.rs @@ -298,7 +298,10 @@ mod test { use arrow::datatypes as arrow_schema; use arrow::datatypes::{DataType, Field, SchemaRef}; use datafusion_common::{DFSchema, ScalarValue, config::ConfigOptions}; - use datafusion_expr::{Expr, ScalarUDF, col, execution_props::ExecutionProps}; + use datafusion_expr::{ + Expr, ScalarUDF, col, execution_props::ExecutionProps, + physical_planning_context::PhysicalPlanningContext, + }; use datafusion_functions::core::input_file_name::InputFileNameFunc; use datafusion_physical_expr::{ ScalarFunctionExpr, create_physical_exprs, projection::ProjectionExpr, @@ -325,8 +328,13 @@ mod test { schema: &SchemaRef, ) -> ProjectionExprs { let df_schema = DFSchema::try_from(Arc::clone(schema)).unwrap(); - let physical_exprs = - create_physical_exprs(exprs, &df_schema, &ExecutionProps::default()).unwrap(); + let physical_exprs = create_physical_exprs( + exprs, + &df_schema, + &ExecutionProps::default(), + &PhysicalPlanningContext::default(), + ) + .unwrap(); let projection_exprs = physical_exprs .into_iter() .enumerate() diff --git a/datafusion/expr/src/execution_props.rs b/datafusion/expr/src/execution_props.rs index 649f74ed39..9910918c6e 100644 --- a/datafusion/expr/src/execution_props.rs +++ b/datafusion/expr/src/execution_props.rs @@ -18,14 +18,10 @@ use crate::var_provider::{VarProvider, VarType}; use chrono::{DateTime, Utc}; use datafusion_common::HashMap; -use datafusion_common::ScalarValue; use datafusion_common::TableReference; use datafusion_common::alias::AliasGenerator; use datafusion_common::config::ConfigOptions; -use datafusion_common::{Result, internal_err}; -use std::fmt; -use std::hash::{Hash, Hasher}; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; /// Holds properties and scratch state used while optimizing a [`LogicalPlan`] /// and translating it into an executable physical plan, such as the statement @@ -64,12 +60,6 @@ pub struct ExecutionProps { pub config_options: Option<Arc<ConfigOptions>>, /// Providers for scalar variables pub var_providers: Option<HashMap<VarType, Arc<dyn VarProvider + Send + Sync>>>, - /// Maps each logical `Subquery` to its index in `subquery_results`. - /// Populated by the physical planner before calling `create_physical_expr`. - pub subquery_indexes: HashMap<crate::logical_plan::Subquery, SubqueryIndex>, - /// Shared results container for uncorrelated scalar subquery values. - /// Populated at execution time by `ScalarSubqueryExec`. - pub subquery_results: ScalarSubqueryResults, /// Maps each lambda variable name to its lambda qualifier generated /// during physical planning. Populated by the physical planner for /// each lambda before calling `create_physical_expr`. @@ -90,8 +80,6 @@ impl ExecutionProps { alias_generator: Arc::new(AliasGenerator::new()), config_options: None, var_providers: None, - subquery_indexes: HashMap::new(), - subquery_results: ScalarSubqueryResults::default(), lambda_variable_qualifier: HashMap::new(), } } @@ -169,103 +157,6 @@ impl ExecutionProps { } } -/// Index of a scalar subquery within a [`ScalarSubqueryResults`] container. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] -pub struct SubqueryIndex(usize); - -impl SubqueryIndex { - /// Creates a new subquery index. - pub const fn new(index: usize) -> Self { - Self(index) - } - - /// Returns the underlying slot index. - pub const fn as_usize(self) -> usize { - self.0 - } -} - -/// Shared results container for uncorrelated scalar subqueries. -/// -/// Each entry corresponds to one scalar subquery, identified by its index. -/// Each slot is populated at execution time by `ScalarSubqueryExec`, read by -/// `ScalarSubqueryExpr` instances that share this container, and cleared when -/// the plan is reset for re-execution. -#[derive(Clone, Default)] -pub struct ScalarSubqueryResults { - slots: Arc<Vec<Mutex<Option<ScalarValue>>>>, -} - -impl ScalarSubqueryResults { - /// Creates a new shared results container with `n` empty slots. - pub fn new(n: usize) -> Self { - Self { - slots: Arc::new((0..n).map(|_| Mutex::new(None)).collect()), - } - } - - /// Returns the scalar value stored at `index`, if it has been populated. - pub fn get(&self, index: SubqueryIndex) -> Option<ScalarValue> { - let slot = self.slots.get(index.as_usize())?; - slot.lock().unwrap().clone() - } - - /// Stores `value` in the slot at `index`. - pub fn set(&self, index: SubqueryIndex, value: ScalarValue) -> Result<()> { - let Some(slot) = self.slots.get(index.as_usize()) else { - return internal_err!( - "ScalarSubqueryResults: result index {} is out of bounds", - index.as_usize() - ); - }; - - let mut slot = slot.lock().unwrap(); - if slot.is_some() { - return internal_err!( - "ScalarSubqueryResults: result for index {} was already populated", - index.as_usize() - ); - } - *slot = Some(value); - - Ok(()) - } - - /// Clears all populated results so the container can be reused. - pub fn clear(&self) { - for slot in self.slots.iter() { - *slot.lock().unwrap() = None; - } - } - - /// Returns true if `this` and `other` point to the same shared container. - pub fn ptr_eq(this: &Self, other: &Self) -> bool { - Arc::ptr_eq(&this.slots, &other.slots) - } -} - -impl fmt::Debug for ScalarSubqueryResults { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_list() - .entries(self.slots.iter().map(|slot| slot.lock().unwrap().clone())) - .finish() - } -} - -impl PartialEq for ScalarSubqueryResults { - fn eq(&self, other: &Self) -> bool { - Self::ptr_eq(self, other) - } -} - -impl Eq for ScalarSubqueryResults {} - -impl Hash for ScalarSubqueryResults { - fn hash<H: Hasher>(&self, state: &mut H) { - Arc::as_ptr(&self.slots).hash(state); - } -} - #[cfg(test)] mod test { use super::*; @@ -274,44 +165,8 @@ mod test { fn debug() { let props = ExecutionProps::new(); assert_eq!( - "ExecutionProps { query_execution_start_time: None, alias_generator: AliasGenerator { next_id: 1 }, config_options: None, var_providers: None, subquery_indexes: {}, subquery_results: [], lambda_variable_qualifier: {} }", + "ExecutionProps { query_execution_start_time: None, alias_generator: AliasGenerator { next_id: 1 }, config_options: None, var_providers: None, lambda_variable_qualifier: {} }", format!("{props:?}") ); } - - #[test] - fn scalar_subquery_results_set_and_get() -> Result<()> { - let results = ScalarSubqueryResults::new(1); - assert_eq!(results.get(SubqueryIndex::new(0)), None); - - results.set(SubqueryIndex::new(0), ScalarValue::Int32(Some(42)))?; - assert_eq!( - results.get(SubqueryIndex::new(0)), - Some(ScalarValue::Int32(Some(42))) - ); - assert!( - results - .set(SubqueryIndex::new(0), ScalarValue::Int32(Some(7))) - .is_err() - ); - - Ok(()) - } - - #[test] - fn scalar_subquery_results_clear() -> Result<()> { - let results = ScalarSubqueryResults::new(1); - results.set(SubqueryIndex::new(0), ScalarValue::Int32(Some(42)))?; - - results.clear(); - - assert_eq!(results.get(SubqueryIndex::new(0)), None); - results.set(SubqueryIndex::new(0), ScalarValue::Int32(Some(7)))?; - assert_eq!( - results.get(SubqueryIndex::new(0)), - Some(ScalarValue::Int32(Some(7))) - ); - - Ok(()) - } } diff --git a/datafusion/expr/src/lib.rs b/datafusion/expr/src/lib.rs index 43cb3fdc20..1033952642 100644 --- a/datafusion/expr/src/lib.rs +++ b/datafusion/expr/src/lib.rs @@ -55,6 +55,7 @@ pub mod expr_rewriter; pub mod expr_schema; pub mod extension_types; pub mod function; +pub mod physical_planning_context; pub mod select_expr; pub mod groups_accumulator { pub use datafusion_expr_common::groups_accumulator::*; diff --git a/datafusion/expr/src/physical_planning_context.rs b/datafusion/expr/src/physical_planning_context.rs new file mode 100644 index 0000000000..b1ba63e071 --- /dev/null +++ b/datafusion/expr/src/physical_planning_context.rs @@ -0,0 +1,211 @@ +// 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. + +use std::fmt; +use std::hash::{Hash, Hasher}; +use std::sync::{Arc, Mutex}; + +use datafusion_common::{HashMap, Result, ScalarValue, internal_err}; + +/// Context used while converting a logical plan subtree into a physical plan. +/// +/// Unlike [`ExecutionProps`](crate::execution_props::ExecutionProps), which +/// applies to the overall planning and execution of a query, this context can +/// differ between recursively planned subtrees. It currently carries the state +/// needed to create physical expressions for [`Expr::ScalarSubquery`] nodes +/// that read from a shared +/// [`ScalarSubqueryResults`] container. +/// +/// The physical planner builds this context from the set of uncorrelated scalar +/// subqueries it has scheduled for a subtree. It is then passed explicitly +/// through `create_physical_expr` so that function can find the slot index for +/// each [`Subquery`]. +/// +/// An empty [`PhysicalPlanningContext`] (the [`Default`]) is what every +/// non-physical-planner caller passes; if such a caller encounters a scalar +/// subquery, `create_physical_expr` returns a `not_impl_err`. +/// +/// [`Expr::ScalarSubquery`]: crate::Expr::ScalarSubquery +/// [`Subquery`]: crate::logical_plan::Subquery +#[derive(Clone, Debug, Default)] +pub struct PhysicalPlanningContext { + indexes: HashMap<crate::logical_plan::Subquery, SubqueryIndex>, + results: ScalarSubqueryResults, +} + +impl PhysicalPlanningContext { + /// Create a [`PhysicalPlanningContext`] from an index map and a shared + /// results container. The index map must use the same indices as slots in + /// `results`. + pub fn new( + indexes: HashMap<crate::logical_plan::Subquery, SubqueryIndex>, + results: ScalarSubqueryResults, + ) -> Self { + Self { indexes, results } + } + + /// Returns the slot index assigned to `subquery`, if any. + pub fn index_of( + &self, + subquery: &crate::logical_plan::Subquery, + ) -> Option<SubqueryIndex> { + self.indexes.get(subquery).copied() + } + + /// Returns the shared results container. + pub fn results(&self) -> &ScalarSubqueryResults { + &self.results + } +} + +/// Index of a scalar subquery within a [`ScalarSubqueryResults`] container. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct SubqueryIndex(usize); + +impl SubqueryIndex { + /// Creates a new subquery index. + pub const fn new(index: usize) -> Self { + Self(index) + } + + /// Returns the underlying slot index. + pub const fn as_usize(self) -> usize { + self.0 + } +} + +/// Shared results container for uncorrelated scalar subqueries. +/// +/// Each entry corresponds to one scalar subquery, identified by its index. +/// Each slot is populated at execution time by `ScalarSubqueryExec`, read by +/// `ScalarSubqueryExpr` instances that share this container, and cleared when +/// the plan is reset for re-execution. +#[derive(Clone, Default)] +pub struct ScalarSubqueryResults { + slots: Arc<Vec<Mutex<Option<ScalarValue>>>>, +} + +impl ScalarSubqueryResults { + /// Creates a new shared results container with `n` empty slots. + pub fn new(n: usize) -> Self { + Self { + slots: Arc::new((0..n).map(|_| Mutex::new(None)).collect()), + } + } + + /// Returns the scalar value stored at `index`, if it has been populated. + pub fn get(&self, index: SubqueryIndex) -> Option<ScalarValue> { + let slot = self.slots.get(index.as_usize())?; + slot.lock().unwrap().clone() + } + + /// Stores `value` in the slot at `index`. + pub fn set(&self, index: SubqueryIndex, value: ScalarValue) -> Result<()> { + let Some(slot) = self.slots.get(index.as_usize()) else { + return internal_err!( + "ScalarSubqueryResults: result index {} is out of bounds", + index.as_usize() + ); + }; + + let mut slot = slot.lock().unwrap(); + if slot.is_some() { + return internal_err!( + "ScalarSubqueryResults: result for index {} was already populated", + index.as_usize() + ); + } + *slot = Some(value); + + Ok(()) + } + + /// Clears all populated results so the container can be reused. + pub fn clear(&self) { + for slot in self.slots.iter() { + *slot.lock().unwrap() = None; + } + } + + /// Returns true if `this` and `other` point to the same shared container. + pub fn ptr_eq(this: &Self, other: &Self) -> bool { + Arc::ptr_eq(&this.slots, &other.slots) + } +} + +impl fmt::Debug for ScalarSubqueryResults { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_list() + .entries(self.slots.iter().map(|slot| slot.lock().unwrap().clone())) + .finish() + } +} + +impl PartialEq for ScalarSubqueryResults { + fn eq(&self, other: &Self) -> bool { + Self::ptr_eq(self, other) + } +} + +impl Eq for ScalarSubqueryResults {} + +impl Hash for ScalarSubqueryResults { + fn hash<H: Hasher>(&self, state: &mut H) { + Arc::as_ptr(&self.slots).hash(state); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn scalar_subquery_results_set_and_get() -> Result<()> { + let results = ScalarSubqueryResults::new(1); + assert_eq!(results.get(SubqueryIndex::new(0)), None); + + results.set(SubqueryIndex::new(0), ScalarValue::Int32(Some(42)))?; + assert_eq!( + results.get(SubqueryIndex::new(0)), + Some(ScalarValue::Int32(Some(42))) + ); + assert!( + results + .set(SubqueryIndex::new(0), ScalarValue::Int32(Some(7))) + .is_err() + ); + + Ok(()) + } + + #[test] + fn scalar_subquery_results_clear() -> Result<()> { + let results = ScalarSubqueryResults::new(1); + results.set(SubqueryIndex::new(0), ScalarValue::Int32(Some(42)))?; + + results.clear(); + + assert_eq!(results.get(SubqueryIndex::new(0)), None); + results.set(SubqueryIndex::new(0), ScalarValue::Int32(Some(7)))?; + assert_eq!( + results.get(SubqueryIndex::new(0)), + Some(ScalarValue::Int32(Some(7))) + ); + + Ok(()) + } +} diff --git a/datafusion/functions-nested/src/array_any_match.rs b/datafusion/functions-nested/src/array_any_match.rs index 8e6e67ed2a..b83c56e9e2 100644 --- a/datafusion/functions-nested/src/array_any_match.rs +++ b/datafusion/functions-nested/src/array_any_match.rs @@ -255,6 +255,7 @@ mod tests { execution_props::ExecutionProps, expr::{HigherOrderFunction, LambdaVariable}, lambda, lit, + physical_planning_context::PhysicalPlanningContext, }; use datafusion_physical_expr::create_physical_expr; @@ -290,6 +291,7 @@ mod tests { )), &schema, &ExecutionProps::new(), + &PhysicalPlanningContext::default(), )? .evaluate(&RecordBatch::try_new( Arc::clone(schema.inner()), @@ -323,6 +325,7 @@ mod tests { )), &schema, &ExecutionProps::new(), + &PhysicalPlanningContext::default(), )? .evaluate(&RecordBatch::try_new( Arc::clone(schema.inner()), diff --git a/datafusion/functions-nested/src/lambda_utils.rs b/datafusion/functions-nested/src/lambda_utils.rs index ce596d0f1b..927ca5a514 100644 --- a/datafusion/functions-nested/src/lambda_utils.rs +++ b/datafusion/functions-nested/src/lambda_utils.rs @@ -141,6 +141,7 @@ pub(crate) mod test_utils { execution_props::ExecutionProps, expr::{HigherOrderFunction, LambdaVariable}, lambda, + physical_planning_context::PhysicalPlanningContext, }; use datafusion_physical_expr::create_physical_expr; @@ -175,6 +176,7 @@ pub(crate) mod test_utils { )), &schema, &ExecutionProps::new(), + &PhysicalPlanningContext::default(), )? .evaluate(&RecordBatch::try_new( Arc::clone(schema.inner()), diff --git a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs index 39c8541b51..e4a22a3419 100644 --- a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs +++ b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs @@ -40,6 +40,7 @@ use datafusion_common::{ tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRewriter}, }; use datafusion_expr::expr::HigherOrderFunction; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{ BinaryExpr, Case, ColumnarValue, Expr, ExprSchemable, Like, Operator, Volatility, and, binary::BinaryTypeCoercer, lit, or, preimage::PreimageResult, @@ -707,11 +708,15 @@ impl ConstEvaluator { return ConstSimplifyResult::NotSimplified(s, m); } - let phys_expr = - match create_physical_expr(&expr, &DUMMY_DF_SCHEMA, &self.execution_props) { - Ok(e) => e, - Err(err) => return ConstSimplifyResult::SimplifyRuntimeError(err, expr), - }; + let phys_expr = match create_physical_expr( + &expr, + &DUMMY_DF_SCHEMA, + &self.execution_props, + &PhysicalPlanningContext::default(), + ) { + Ok(e) => e, + Err(err) => return ConstSimplifyResult::SimplifyRuntimeError(err, expr), + }; let metadata = phys_expr .return_field(DUMMY_BATCH.schema_ref()) .ok() diff --git a/datafusion/optimizer/src/utils.rs b/datafusion/optimizer/src/utils.rs index b29649e9ea..4ea1589cfa 100644 --- a/datafusion/optimizer/src/utils.rs +++ b/datafusion/optimizer/src/utils.rs @@ -29,6 +29,7 @@ use datafusion_common::{Column, DFSchema, Result, ScalarValue}; use datafusion_expr::execution_props::ExecutionProps; use datafusion_expr::expr::{Exists, InSubquery, SetComparison}; use datafusion_expr::expr_rewriter::replace_col; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{ColumnarValue, Expr, logical_plan::LogicalPlan}; use datafusion_physical_expr::create_physical_expr; use log::{debug, trace}; @@ -233,8 +234,13 @@ fn evaluate_expr_with_null_column<'a>( let replaced_predicate = replace_col(predicate, &join_cols_to_replace)?; let coerced_predicate = coerce(replaced_predicate, &input_schema)?; - create_physical_expr(&coerced_predicate, &input_schema, &execution_props)? - .evaluate(&input_batch) + create_physical_expr( + &coerced_predicate, + &input_schema, + &execution_props, + &PhysicalPlanningContext::default(), + )? + .evaluate(&input_batch) } fn coerce(expr: Expr, schema: &DFSchema) -> Result<Expr> { diff --git a/datafusion/physical-expr/src/aggregate.rs b/datafusion/physical-expr/src/aggregate.rs index b774658679..013779cf8c 100644 --- a/datafusion/physical-expr/src/aggregate.rs +++ b/datafusion/physical-expr/src/aggregate.rs @@ -51,6 +51,7 @@ use datafusion_expr::execution_props::ExecutionProps; use datafusion_expr::expr::{ AggregateFunction, AggregateFunctionParams, NullTreatment, physical_name, }; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{AggregateUDF, Expr, ReversedUDAF, SetMonotonicity}; use datafusion_expr_common::accumulator::Accumulator; use datafusion_expr_common::groups_accumulator::GroupsAccumulator; @@ -423,6 +424,7 @@ pub struct LoweredAggregateBuilder<'a> { logical_input_schema: &'a DFSchema, physical_input_schema: &'a Schema, execution_props: &'a ExecutionProps, + planning_ctx: &'a PhysicalPlanningContext, } impl<'a> LoweredAggregateBuilder<'a> { @@ -430,12 +432,17 @@ impl<'a> LoweredAggregateBuilder<'a> { /// /// `logical_input_schema` is used to resolve logical expressions such as /// columns, while `physical_input_schema` is the input schema used by the - /// physical aggregate expression. + /// physical aggregate expression. `planning_ctx` is used when creating + /// physical expressions that reference uncorrelated scalar subqueries. + /// Callers creating physical aggregates outside of physical planning should + /// pass `&PhysicalPlanningContext::default()`, in which case converting a + /// scalar-subquery expression returns a planning error. pub fn new( expr: &'a Expr, logical_input_schema: &'a DFSchema, physical_input_schema: &'a Schema, execution_props: &'a ExecutionProps, + planning_ctx: &'a PhysicalPlanningContext, ) -> Self { Self { expr, @@ -446,6 +453,7 @@ impl<'a> LoweredAggregateBuilder<'a> { logical_input_schema, physical_input_schema, execution_props, + planning_ctx, } } @@ -484,6 +492,7 @@ impl<'a> LoweredAggregateBuilder<'a> { logical_input_schema, physical_input_schema, execution_props, + planning_ctx, } = self; let (name, human_display, output_metadata, expr) = lower_aggregate_display( @@ -515,16 +524,29 @@ impl<'a> LoweredAggregateBuilder<'a> { physical_name(&expr)? }; - let physical_args = - create_physical_exprs(args, logical_input_schema, execution_props)?; + let physical_args = create_physical_exprs( + args, + logical_input_schema, + execution_props, + planning_ctx, + )?; let filter = filter .as_ref() .map(|filter| { - create_physical_expr(filter, logical_input_schema, execution_props) + create_physical_expr( + filter, + logical_input_schema, + execution_props, + planning_ctx, + ) }) .transpose()?; - let order_bys = - create_physical_sort_exprs(order_by, logical_input_schema, execution_props)?; + let order_bys = create_physical_sort_exprs( + order_by, + logical_input_schema, + execution_props, + planning_ctx, + )?; let ignore_nulls = null_treatment.unwrap_or(NullTreatment::RespectNulls) == NullTreatment::IgnoreNulls; @@ -1162,6 +1184,7 @@ mod tests { &logical_schema, &schema, &ExecutionProps::new(), + &PhysicalPlanningContext::default(), ) .build()?; @@ -1185,6 +1208,7 @@ mod tests { &logical_schema, &schema, &ExecutionProps::new(), + &PhysicalPlanningContext::default(), ) .with_human_display(expr.human_display().to_string()) .build()?; diff --git a/datafusion/physical-expr/src/analysis.rs b/datafusion/physical-expr/src/analysis.rs index 1dca36b75f..a00fc19ae9 100644 --- a/datafusion/physical-expr/src/analysis.rs +++ b/datafusion/physical-expr/src/analysis.rs @@ -350,6 +350,7 @@ mod tests { use datafusion_common::{DFSchema, ScalarValue, assert_contains, stats::Precision}; use datafusion_expr::{ Expr, col, execution_props::ExecutionProps, interval_arithmetic::Interval, lit, + physical_planning_context::PhysicalPlanningContext, }; use crate::{AnalysisContext, create_physical_expr, expressions::Column}; @@ -412,8 +413,13 @@ mod tests { for (expr, lower, upper) in test_cases { let boundaries = ExprBoundaries::try_new_unbounded(&schema).unwrap(); let df_schema = DFSchema::try_from(Arc::clone(&schema)).unwrap(); - let physical_expr = - create_physical_expr(&expr, &df_schema, &ExecutionProps::new()).unwrap(); + let physical_expr = create_physical_expr( + &expr, + &df_schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + ) + .unwrap(); let analysis_result = analyze( &physical_expr, AnalysisContext::new(boundaries), @@ -453,8 +459,13 @@ mod tests { for expr in test_cases { let boundaries = ExprBoundaries::try_new_unbounded(&schema).unwrap(); let df_schema = DFSchema::try_from(Arc::clone(&schema)).unwrap(); - let physical_expr = - create_physical_expr(&expr, &df_schema, &ExecutionProps::new()).unwrap(); + let physical_expr = create_physical_expr( + &expr, + &df_schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + ) + .unwrap(); let analysis_result = analyze( &physical_expr, AnalysisContext::new(boundaries), @@ -475,8 +486,13 @@ mod tests { let expected_error = "OR operator cannot yet propagate true intervals"; let boundaries = ExprBoundaries::try_new_unbounded(&schema).unwrap(); let df_schema = DFSchema::try_from(Arc::clone(&schema)).unwrap(); - let physical_expr = - create_physical_expr(&expr, &df_schema, &ExecutionProps::new()).unwrap(); + let physical_expr = create_physical_expr( + &expr, + &df_schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + ) + .unwrap(); let analysis_error = analyze( &physical_expr, AnalysisContext::new(boundaries), diff --git a/datafusion/physical-expr/src/physical_expr.rs b/datafusion/physical-expr/src/physical_expr.rs index cfc9866fc8..d45d0fe149 100644 --- a/datafusion/physical-expr/src/physical_expr.rs +++ b/datafusion/physical-expr/src/physical_expr.rs @@ -26,6 +26,7 @@ use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_common::{DFSchema, HashMap, ScalarValue, SplitPoint}; use datafusion_common::{Result, plan_err}; use datafusion_expr::execution_props::ExecutionProps; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{Expr, Partitioning as LogicalPartitioning, SortExpr}; use datafusion_expr_common::casts::try_cast_literal_to_type; @@ -190,47 +191,68 @@ pub fn create_lex_ordering( exprs, &df_schema, execution_props, + &PhysicalPlanningContext::default(), )?)); } Ok(all_sort_orders) } /// Create a physical sort expression from a logical expression +/// +/// See [`create_physical_expr`] for details on the `planning_ctx` argument. pub fn create_physical_sort_expr( e: &SortExpr, input_dfschema: &DFSchema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result<PhysicalSortExpr> { - create_physical_expr(&e.expr, input_dfschema, execution_props).map(|expr| { - let options = SortOptions::new(!e.asc, e.nulls_first); - PhysicalSortExpr::new(expr, options) - }) + create_physical_expr(&e.expr, input_dfschema, execution_props, planning_ctx).map( + |expr| { + let options = SortOptions::new(!e.asc, e.nulls_first); + PhysicalSortExpr::new(expr, options) + }, + ) } /// Create vector of physical sort expression from a vector of logical expression +/// +/// See [`create_physical_expr`] for details on the `planning_ctx` argument. pub fn create_physical_sort_exprs( exprs: &[SortExpr], input_dfschema: &DFSchema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result<Vec<PhysicalSortExpr>> { exprs .iter() - .map(|e| create_physical_sort_expr(e, input_dfschema, execution_props)) + .map(|e| { + create_physical_sort_expr(e, input_dfschema, execution_props, planning_ctx) + }) .collect() } /// Create physical partitioning from logical partitioning. +/// +/// See [`create_physical_expr`] for details on the `planning_ctx` argument. pub fn create_physical_partitioning( partitioning: &LogicalPartitioning, input_dfschema: &DFSchema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result<Partitioning> { match partitioning { LogicalPartitioning::RoundRobinBatch(n) => Ok(Partitioning::RoundRobinBatch(*n)), LogicalPartitioning::Hash(exprs, partition_count) => { let exprs = exprs .iter() - .map(|expr| create_physical_expr(expr, input_dfschema, execution_props)) + .map(|expr| { + create_physical_expr( + expr, + input_dfschema, + execution_props, + planning_ctx, + ) + }) .collect::<Result<Vec<_>>>()?; Ok(Partitioning::Hash(exprs, *partition_count)) } @@ -239,6 +261,7 @@ pub fn create_physical_partitioning( range.ordering(), input_dfschema, execution_props, + planning_ctx, )?; let Some(ordering) = LexOrdering::new(ordering) else { return plan_err!("Range partitioning requires non-empty ordering"); diff --git a/datafusion/physical-expr/src/planner.rs b/datafusion/physical-expr/src/planner.rs index d0d0508a10..3cdd64f7a7 100644 --- a/datafusion/physical-expr/src/planner.rs +++ b/datafusion/physical-expr/src/planner.rs @@ -37,6 +37,7 @@ use datafusion_expr::expr::{ Alias, Cast, HigherOrderFunction, InList, Lambda, LambdaVariable, Placeholder, ScalarFunction, }; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::var_provider::VarType; use datafusion_expr::var_provider::is_system_variables; use datafusion_expr::{ @@ -63,6 +64,7 @@ use datafusion_expr::{ /// # use datafusion_expr::{Expr, col, lit}; /// # use datafusion_physical_expr::create_physical_expr; /// # use datafusion_expr::execution_props::ExecutionProps; +/// # use datafusion_expr::physical_planning_context::PhysicalPlanningContext; /// // For a logical expression `a = 1`, we can create a physical expression /// let expr = col("a").eq(lit(1)); /// // To create a PhysicalExpr we need 1. a schema @@ -70,8 +72,11 @@ use datafusion_expr::{ /// let df_schema = DFSchema::try_from(schema).unwrap(); /// // 2. ExecutionProps /// let props = ExecutionProps::new(); -/// // We can now create a PhysicalExpr: -/// let physical_expr = create_physical_expr(&expr, &df_schema, &props).unwrap(); +/// // We can now create a PhysicalExpr. Expressions with no scalar +/// // subqueries use an empty `PhysicalPlanningContext`: +/// let physical_expr = +/// create_physical_expr(&expr, &df_schema, &props, &PhysicalPlanningContext::default()) +/// .unwrap(); /// ``` /// /// # Example: Executing a PhysicalExpr to obtain [ColumnarValue] @@ -83,12 +88,15 @@ use datafusion_expr::{ /// # use datafusion_expr::{Expr, col, lit, ColumnarValue}; /// # use datafusion_physical_expr::create_physical_expr; /// # use datafusion_expr::execution_props::ExecutionProps; +/// # use datafusion_expr::physical_planning_context::PhysicalPlanningContext; /// # let expr = col("a").eq(lit(1)); /// # let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]); /// # let df_schema = DFSchema::try_from(schema.clone()).unwrap(); /// # let props = ExecutionProps::new(); /// // Given a PhysicalExpr, for `a = 1` we can evaluate it against a RecordBatch like this: -/// let physical_expr = create_physical_expr(&expr, &df_schema, &props).unwrap(); +/// let physical_expr = +/// create_physical_expr(&expr, &df_schema, &props, &PhysicalPlanningContext::default()) +/// .unwrap(); /// // Input of [1,2,3] /// let input_batch = RecordBatch::try_from_iter(vec![ /// ("a", Arc::new(Int32Array::from(vec![1, 2, 3])) as _) @@ -111,11 +119,20 @@ use datafusion_expr::{ /// * `e` - The logical expression /// * `input_dfschema` - The DataFusion schema for the input, used to resolve `Column` references /// to qualified or unqualified fields by name. +/// * `execution_props` - Per-execution properties such as the query start time. +/// * `planning_ctx` - The [`PhysicalPlanningContext`] used to resolve +/// `Expr::ScalarSubquery` nodes. The physical planner threads the subquery +/// index map and shared results container from its `ScalarSubqueryExec` +/// construction into calls to `create_physical_expr`. Callers creating +/// physical expressions outside of physical planning should pass +/// `&PhysicalPlanningContext::default()`; converting a scalar subquery then returns a +/// planning error. #[cfg_attr(feature = "recursive_protection", recursive::recursive)] pub fn create_physical_expr( e: &Expr, input_dfschema: &DFSchema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result<Arc<dyn PhysicalExpr>> { let input_schema = input_dfschema.as_arrow(); @@ -131,7 +148,12 @@ pub fn create_physical_expr( new_metadata, ))) } else { - Ok(create_physical_expr(expr, input_dfschema, execution_props)?) + Ok(create_physical_expr( + expr, + input_dfschema, + execution_props, + planning_ctx, + )?) } } Expr::Column(c) => { @@ -167,12 +189,22 @@ pub fn create_physical_expr( Operator::IsNotDistinctFrom, lit(true), ); - create_physical_expr(&binary_op, input_dfschema, execution_props) + create_physical_expr( + &binary_op, + input_dfschema, + execution_props, + planning_ctx, + ) } Expr::IsNotTrue(expr) => { let binary_op = binary_expr(expr.as_ref().clone(), Operator::IsDistinctFrom, lit(true)); - create_physical_expr(&binary_op, input_dfschema, execution_props) + create_physical_expr( + &binary_op, + input_dfschema, + execution_props, + planning_ctx, + ) } Expr::IsFalse(expr) => { let binary_op = binary_expr( @@ -180,12 +212,22 @@ pub fn create_physical_expr( Operator::IsNotDistinctFrom, lit(false), ); - create_physical_expr(&binary_op, input_dfschema, execution_props) + create_physical_expr( + &binary_op, + input_dfschema, + execution_props, + planning_ctx, + ) } Expr::IsNotFalse(expr) => { let binary_op = binary_expr(expr.as_ref().clone(), Operator::IsDistinctFrom, lit(false)); - create_physical_expr(&binary_op, input_dfschema, execution_props) + create_physical_expr( + &binary_op, + input_dfschema, + execution_props, + planning_ctx, + ) } Expr::IsUnknown(expr) => { let binary_op = binary_expr( @@ -193,7 +235,12 @@ pub fn create_physical_expr( Operator::IsNotDistinctFrom, Expr::Literal(ScalarValue::Boolean(None), None), ); - create_physical_expr(&binary_op, input_dfschema, execution_props) + create_physical_expr( + &binary_op, + input_dfschema, + execution_props, + planning_ctx, + ) } Expr::IsNotUnknown(expr) => { let binary_op = binary_expr( @@ -201,12 +248,27 @@ pub fn create_physical_expr( Operator::IsDistinctFrom, Expr::Literal(ScalarValue::Boolean(None), None), ); - create_physical_expr(&binary_op, input_dfschema, execution_props) + create_physical_expr( + &binary_op, + input_dfschema, + execution_props, + planning_ctx, + ) } Expr::BinaryExpr(BinaryExpr { left, op, right }) => { // Create physical expressions for left and right operands - let lhs = create_physical_expr(left, input_dfschema, execution_props)?; - let rhs = create_physical_expr(right, input_dfschema, execution_props)?; + let lhs = create_physical_expr( + left, + input_dfschema, + execution_props, + planning_ctx, + )?; + let rhs = create_physical_expr( + right, + input_dfschema, + execution_props, + planning_ctx, + )?; // Note that the logical planner is responsible // for type coercion on the arguments (e.g. if one // argument was originally Int32 and one was @@ -229,10 +291,18 @@ pub fn create_physical_expr( "LIKE does not support escape_char other than the backslash (\\)" ); } - let physical_expr = - create_physical_expr(expr, input_dfschema, execution_props)?; - let physical_pattern = - create_physical_expr(pattern, input_dfschema, execution_props)?; + let physical_expr = create_physical_expr( + expr, + input_dfschema, + execution_props, + planning_ctx, + )?; + let physical_pattern = create_physical_expr( + pattern, + input_dfschema, + execution_props, + planning_ctx, + )?; like( *negated, *case_insensitive, @@ -251,10 +321,18 @@ pub fn create_physical_expr( if escape_char.is_some() { return exec_err!("SIMILAR TO does not support escape_char yet"); } - let physical_expr = - create_physical_expr(expr, input_dfschema, execution_props)?; - let physical_pattern = - create_physical_expr(pattern, input_dfschema, execution_props)?; + let physical_expr = create_physical_expr( + expr, + input_dfschema, + execution_props, + planning_ctx, + )?; + let physical_pattern = create_physical_expr( + pattern, + input_dfschema, + execution_props, + planning_ctx, + )?; similar_to(*negated, *case_insensitive, physical_expr, physical_pattern) } Expr::Case(case) => { @@ -263,6 +341,7 @@ pub fn create_physical_expr( e.as_ref(), input_dfschema, execution_props, + planning_ctx, )?) } else { None @@ -272,10 +351,18 @@ pub fn create_physical_expr( .iter() .map(|(w, t)| (w.as_ref(), t.as_ref())) .unzip(); - let when_expr = - create_physical_exprs(when_expr, input_dfschema, execution_props)?; - let then_expr = - create_physical_exprs(then_expr, input_dfschema, execution_props)?; + let when_expr = create_physical_exprs( + when_expr, + input_dfschema, + execution_props, + planning_ctx, + )?; + let then_expr = create_physical_exprs( + then_expr, + input_dfschema, + execution_props, + planning_ctx, + )?; let when_then_expr: Vec<(Arc<dyn PhysicalExpr>, Arc<dyn PhysicalExpr>)> = when_expr .iter() @@ -288,6 +375,7 @@ pub fn create_physical_expr( e.as_ref(), input_dfschema, execution_props, + planning_ctx, )?) } else { None @@ -295,7 +383,7 @@ pub fn create_physical_expr( Ok(expressions::case(expr, when_then_expr, else_expr)?) } Expr::Cast(Cast { expr, field }) => expressions::cast_with_target_field( - create_physical_expr(expr, input_dfschema, execution_props)?, + create_physical_expr(expr, input_dfschema, execution_props, planning_ctx)?, input_schema, Arc::clone(field), None, @@ -314,31 +402,45 @@ pub fn create_physical_expr( } expressions::try_cast( - create_physical_expr(expr, input_dfschema, execution_props)?, + create_physical_expr( + expr, + input_dfschema, + execution_props, + planning_ctx, + )?, input_schema, field.data_type().clone(), ) } - Expr::Not(expr) => { - expressions::not(create_physical_expr(expr, input_dfschema, execution_props)?) - } + Expr::Not(expr) => expressions::not(create_physical_expr( + expr, + input_dfschema, + execution_props, + planning_ctx, + )?), Expr::Negative(expr) => expressions::negative( - create_physical_expr(expr, input_dfschema, execution_props)?, + create_physical_expr(expr, input_dfschema, execution_props, planning_ctx)?, input_schema, ), Expr::IsNull(expr) => expressions::is_null(create_physical_expr( expr, input_dfschema, execution_props, + planning_ctx, )?), Expr::IsNotNull(expr) => expressions::is_not_null(create_physical_expr( expr, input_dfschema, execution_props, + planning_ctx, )?), Expr::ScalarFunction(ScalarFunction { func, args }) => { - let physical_args = - create_physical_exprs(args, input_dfschema, execution_props)?; + let physical_args = create_physical_exprs( + args, + input_dfschema, + execution_props, + planning_ctx, + )?; let config_options = match execution_props.config_options.as_ref() { Some(config_options) => Arc::clone(config_options), None => Arc::new(ConfigOptions::default()), @@ -357,9 +459,20 @@ pub fn create_physical_expr( low, high, }) => { - let value_expr = create_physical_expr(expr, input_dfschema, execution_props)?; - let low_expr = create_physical_expr(low, input_dfschema, execution_props)?; - let high_expr = create_physical_expr(high, input_dfschema, execution_props)?; + let value_expr = create_physical_expr( + expr, + input_dfschema, + execution_props, + planning_ctx, + )?; + let low_expr = + create_physical_expr(low, input_dfschema, execution_props, planning_ctx)?; + let high_expr = create_physical_expr( + high, + input_dfschema, + execution_props, + planning_ctx, + )?; // rewrite the between into the two binary operators let binary_expr = binary( @@ -394,17 +507,25 @@ pub fn create_physical_expr( Ok(expressions::lit(ScalarValue::Boolean(None))) } _ => { - let value_expr = - create_physical_expr(expr, input_dfschema, execution_props)?; + let value_expr = create_physical_expr( + expr, + input_dfschema, + execution_props, + planning_ctx, + )?; - let list_exprs = - create_physical_exprs(list, input_dfschema, execution_props)?; + let list_exprs = create_physical_exprs( + list, + input_dfschema, + execution_props, + planning_ctx, + )?; expressions::in_list(value_expr, list_exprs, negated, input_schema) } }, Expr::ScalarSubquery(sq) => { - match execution_props.subquery_indexes.get(sq) { - Some(&index) => { + match planning_ctx.index_of(sq) { + Some(index) => { let schema = sq.subquery.schema(); if schema.fields().len() != 1 { return plan_err!( @@ -418,7 +539,7 @@ pub fn create_physical_expr( dt, nullable, index, - execution_props.subquery_results.clone(), + planning_ctx.results().clone(), ))) } None => { @@ -495,9 +616,19 @@ pub fn create_physical_expr( .clone() .with_qualified_lambda_variables(&qualifier, &lambda.params); - create_physical_expr(arg, &lambda_schema, &execution_props) + create_physical_expr( + arg, + &lambda_schema, + &execution_props, + planning_ctx, + ) } - _ => create_physical_expr(arg, input_dfschema, execution_props), + _ => create_physical_expr( + arg, + input_dfschema, + execution_props, + planning_ctx, + ), }) .collect::<Result<_>>()?; @@ -515,7 +646,7 @@ pub fn create_physical_expr( } Expr::Lambda(Lambda { params, body }) => expressions::lambda( params, - create_physical_expr(body, input_dfschema, execution_props)?, + create_physical_expr(body, input_dfschema, execution_props, planning_ctx)?, ), Expr::LambdaVariable(LambdaVariable { name, @@ -572,17 +703,22 @@ pub fn create_physical_expr( } /// Create vector of Physical Expression from a vector of logical expression +/// +/// See [`create_physical_expr`] for details on the `planning_ctx` argument. pub fn create_physical_exprs<'a, I>( exprs: I, input_dfschema: &DFSchema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result<Vec<Arc<dyn PhysicalExpr>>> where I: IntoIterator<Item = &'a Expr>, { exprs .into_iter() - .map(|expr| create_physical_expr(expr, input_dfschema, execution_props)) + .map(|expr| { + create_physical_expr(expr, input_dfschema, execution_props, planning_ctx) + }) .collect() } @@ -591,7 +727,13 @@ pub fn logical2physical(expr: &Expr, schema: &Schema) -> Arc<dyn PhysicalExpr> { // TODO this makes a deep copy of the Schema. Should take SchemaRef instead and avoid deep copy let df_schema = schema.clone().to_dfschema().unwrap(); let execution_props = ExecutionProps::new(); - create_physical_expr(expr, &df_schema, &execution_props).unwrap() + create_physical_expr( + expr, + &df_schema, + &execution_props, + &PhysicalPlanningContext::default(), + ) + .unwrap() } #[cfg(test)] @@ -608,7 +750,12 @@ mod tests { fn lower_cast_expr(expr: &Expr, schema: &Schema) -> Result<Arc<dyn PhysicalExpr>> { let df_schema = DFSchema::try_from(schema.clone())?; - create_physical_expr(expr, &df_schema, &ExecutionProps::new()) + create_physical_expr( + expr, + &df_schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + ) } fn as_planner_cast(physical: &Arc<dyn PhysicalExpr>) -> &expressions::CastExpr { @@ -623,7 +770,12 @@ mod tests { let schema = Schema::new(vec![Field::new("letter", DataType::Utf8, false)]); let df_schema = DFSchema::try_from_qualified_schema("data", &schema)?; - let p = create_physical_expr(&expr, &df_schema, &ExecutionProps::new())?; + let p = create_physical_expr( + &expr, + &df_schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + )?; let batch = RecordBatch::try_new( Arc::new(schema), @@ -728,8 +880,12 @@ mod tests { let df_schema = DFSchema::try_from(schema)?; // This should not stack overflow - let _physical_expr = - create_physical_expr(&expr, &df_schema, &ExecutionProps::new())?; + let _physical_expr = create_physical_expr( + &expr, + &df_schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + )?; Ok(()) } diff --git a/datafusion/physical-expr/src/scalar_subquery.rs b/datafusion/physical-expr/src/scalar_subquery.rs index f7270f8054..473b52a5cb 100644 --- a/datafusion/physical-expr/src/scalar_subquery.rs +++ b/datafusion/physical-expr/src/scalar_subquery.rs @@ -24,7 +24,7 @@ use std::sync::Arc; use arrow::datatypes::{DataType, Field, FieldRef, Schema}; use arrow::record_batch::RecordBatch; use datafusion_common::{Result, internal_datafusion_err}; -use datafusion_expr::execution_props::{ScalarSubqueryResults, SubqueryIndex}; +use datafusion_expr::physical_planning_context::{ScalarSubqueryResults, SubqueryIndex}; use datafusion_expr_common::columnar_value::ColumnarValue; use datafusion_expr_common::sort_properties::{ExprProperties, SortProperties}; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; diff --git a/datafusion/physical-plan/src/scalar_subquery.rs b/datafusion/physical-plan/src/scalar_subquery.rs index 74de1f11bf..2e04b5456b 100644 --- a/datafusion/physical-plan/src/scalar_subquery.rs +++ b/datafusion/physical-plan/src/scalar_subquery.rs @@ -29,7 +29,7 @@ use std::sync::Arc; use datafusion_common::{Result, ScalarValue, Statistics, exec_err, internal_err}; use datafusion_execution::TaskContext; -use datafusion_expr::execution_props::{ScalarSubqueryResults, SubqueryIndex}; +use datafusion_expr::physical_planning_context::{ScalarSubqueryResults, SubqueryIndex}; use crate::execution_plan::{CardinalityEffect, ExecutionPlan, PlanProperties}; use crate::joins::utils::{OnceAsync, OnceFut}; @@ -202,9 +202,9 @@ impl ExecutionPlan for ScalarSubqueryExec { ) -> Result<SendableRecordBatchStream> { let subqueries = self.subqueries.clone(); let results = self.results.clone(); - let subquery_ctx = Arc::clone(&context); + let planning_ctx = Arc::clone(&context); let mut subquery_future = self.subquery_future.try_once(move || { - Ok(async move { execute_subqueries(subqueries, results, subquery_ctx).await }) + Ok(async move { execute_subqueries(subqueries, results, planning_ctx).await }) })?; let input = Arc::clone(&self.input); let schema = self.schema(); diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index e5f8aa072d..4f72668813 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -53,7 +53,7 @@ use datafusion_datasource_parquet::source::ParquetSource; #[cfg(feature = "parquet")] use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{FunctionRegistry, TaskContext}; -use datafusion_expr::execution_props::{ScalarSubqueryResults, SubqueryIndex}; +use datafusion_expr::physical_planning_context::{ScalarSubqueryResults, SubqueryIndex}; use datafusion_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF, WindowUDF}; use datafusion_functions_table::generate_series::{ Empty, GenSeriesArgs, GenerateSeriesTable, GenericSeriesState, TimestampValue, diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index cbd6fd912a..d87efcf986 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -111,7 +111,7 @@ use datafusion_expr::{ Accumulator, AccumulatorFactoryFunction, AggregateUDF, ColumnarValue, HigherOrderUDF, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, SimpleAggregateUDF, WindowFrame, WindowFrameBound, WindowUDF, - execution_props::{ScalarSubqueryResults, SubqueryIndex}, + physical_planning_context::{ScalarSubqueryResults, SubqueryIndex}, }; use datafusion_functions_aggregate::approx_percentile_cont::approx_percentile_cont_udaf; use datafusion_functions_aggregate::array_agg::array_agg_udaf; 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 28ceb27d63..67f90c649e 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -681,3 +681,98 @@ months are ignored, as in PostgreSQL. The result keeps the input time's unit unit is truncated -- so `time(s) + interval '1 nanosecond'` is a no-op. See [PR #23279](https://github.com/apache/datafusion/pull/23279) for details. + +### Scalar-subquery state moved to an explicit `PhysicalPlanningContext` + +The `subquery_indexes` and `subquery_results` public fields on +`datafusion_expr::execution_props::ExecutionProps` have been removed. They were +added in `54.0.0` as the channel through which the physical planner passed +uncorrelated scalar-subquery state to functions that create physical +`Arc<dyn PhysicalExpr>` values from logical `Expr` values. + +That state is now carried by a dedicated +`datafusion_expr::physical_planning_context::PhysicalPlanningContext` passed explicitly +through functions and planner traits. Unlike `ExecutionProps`, which applies +throughout the planning of an entire query, this context is scoped to the +logical plan subtree currently being converted. This removes the need for the +physical planner to clone and mutate a `SessionState`, is a prerequisite for +letting the planner take `&dyn Session`, and lets `ExtensionPlanner` +implementations create physical +expressions containing scalar subqueries against the same subquery state as the +rest of the plan. + +The following functions take a new trailing +`planning_ctx: &PhysicalPlanningContext` parameter: + +- `datafusion_physical_expr::create_physical_expr` / `create_physical_exprs` +- `datafusion_physical_expr::create_physical_sort_expr` / + `create_physical_sort_exprs` / `create_physical_partitioning` +- `datafusion::physical_planner::create_window_expr` / + `create_window_expr_with_name` +- `datafusion_physical_expr::aggregate::LoweredAggregateBuilder::new` + +The planner traits changed accordingly: + +- `PhysicalPlanner::create_physical_expr` takes + `planning_ctx: &PhysicalPlanningContext` +- `ExtensionPlanner::plan_extension` and `plan_table_scan` receive + `planning_ctx: &PhysicalPlanningContext` and should forward it to + `PhysicalPlanner::create_physical_expr` when creating physical expressions + +Convenience methods such as `SessionContext::create_physical_expr` and +`SessionState::create_physical_expr` are unchanged. + +**Who is affected:** + +- Code calling the functions above: pass + `&PhysicalPlanningContext::default()` unless you are creating physical + expressions as part of a physical plan that contains uncorrelated scalar + subqueries. +- Custom `PhysicalPlanner` or `ExtensionPlanner` implementations: add the new + parameter and forward it. +- Code that read or wrote `execution_props.subquery_indexes` / + `execution_props.subquery_results`: build a `PhysicalPlanningContext` instead. + +**Migration guide:** + +When creating a physical expression outside of physical planning, pass an empty +context: + +```rust,ignore +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; +use datafusion_physical_expr::create_physical_expr; + +// Before +let phys = create_physical_expr(&expr, &schema, &props)?; + +// After +let phys = create_physical_expr( + &expr, + &schema, + &props, + &PhysicalPlanningContext::default(), +)?; +``` + +For `ExtensionPlanner` implementations, accept and forward the context: + +```rust,ignore +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, // 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)?; + } + // ... +} +``` + +See [PR #23649](https://github.com/apache/datafusion/pull/23649) for details. --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
