diff --git a/datafusion-examples/examples/dataframe/cache_factory.rs b/datafusion-examples/examples/dataframe/cache_factory.rs index a92c3dc4ce26a..80c73d0867c7b 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::execution_props::SubqueryContext; use datafusion::logical_expr::{ Extension, LogicalPlan, UserDefinedLogicalNode, UserDefinedLogicalNodeCore, }; @@ -146,6 +147,7 @@ impl ExtensionPlanner for CacheNodePlanner { logical_inputs: &[&LogicalPlan], physical_inputs: &[Arc], session_state: &SessionState, + _subquery_ctx: &SubqueryContext, ) -> Result>> { if let Some(cache_node) = node.as_any().downcast_ref::() { 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 c087019c687c5..adaf0ed43194c 100644 --- a/datafusion-examples/examples/query_planning/expr_api.rs +++ b/datafusion-examples/examples/query_planning/expr_api.rs @@ -31,6 +31,7 @@ use datafusion::common::{ScalarValue, ToDFSchema}; use datafusion::error::Result; use datafusion::functions_aggregate::first_last::first_value_udaf; use datafusion::logical_expr::execution_props::ExecutionProps; +use datafusion::logical_expr::execution_props::SubqueryContext; use datafusion::logical_expr::expr::BinaryExpr; use datafusion::logical_expr::interval_arithmetic::Interval; use datafusion::logical_expr::simplify::SimplifyContext; @@ -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, + &SubqueryContext::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, + &SubqueryContext::default(), )?; assert!(physical_expr.evaluate(&batch).is_ok()); @@ -578,6 +584,7 @@ fn type_coercion_demo() -> Result<()> { &coerced_expr, &df_schema, &props, + &SubqueryContext::default(), )?; assert!(physical_expr.evaluate(&batch).is_ok()); @@ -606,6 +613,7 @@ fn type_coercion_demo() -> Result<()> { &coerced_expr, &df_schema, &props, + &SubqueryContext::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 7fdc4a7952d68..4ae793144cc92 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::execution_props::SubqueryContext; use datafusion::physical_expr::create_physical_expr; use datafusion::physical_optimizer::pruning::PruningPredicate; use datafusion::prelude::*; @@ -194,7 +195,9 @@ 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, &SubqueryContext::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 2c696d92d70b8..e553065b0ff8e 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::execution_props::SubqueryContext; use datafusion_expr::{ UserDefinedLogicalNode, UserDefinedLogicalNodeCore, logical_plan::{Extension, LogicalPlan, LogicalPlanBuilder}, @@ -587,6 +588,7 @@ impl ExtensionPlanner for TableSampleExtensionPlanner { _logical_inputs: &[&LogicalPlan], physical_inputs: &[Arc], _session_state: &SessionState, + _subquery_ctx: &SubqueryContext, ) -> Result>> { let Some(sample_node) = node.as_any().downcast_ref::() else { diff --git a/datafusion/catalog-listing/src/helpers.rs b/datafusion/catalog-listing/src/helpers.rs index 31f00b62ef236..bb060a51f11fb 100644 --- a/datafusion/catalog-listing/src/helpers.rs +++ b/datafusion/catalog-listing/src/helpers.rs @@ -33,7 +33,7 @@ use arrow::{ datatypes::{DataType, Field}, record_batch::RecordBatch, }; -use datafusion_expr::execution_props::ExecutionProps; +use datafusion_expr::execution_props::{ExecutionProps, SubqueryContext}; use futures::stream::FuturesUnordered; use futures::{StreamExt, TryStreamExt, stream::BoxStream}; use log::{debug, trace}; @@ -328,7 +328,8 @@ 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, &SubqueryContext::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 ce739e5019472..ad39064cf7ee1 100644 --- a/datafusion/catalog-listing/src/table.rs +++ b/datafusion/catalog-listing/src/table.rs @@ -43,7 +43,7 @@ use datafusion_execution::cache::cache_manager::{ CachedFileMetadata, FileStatisticsCache, SchemaFingerprint, TableScopedPath, }; use datafusion_expr::dml::InsertOp; -use datafusion_expr::execution_props::ExecutionProps; +use datafusion_expr::execution_props::{ExecutionProps, SubqueryContext}; use datafusion_expr::{ Expr, Partitioning as LogicalPartitioning, TableProviderFilterPushDown, TableType, }; @@ -614,6 +614,7 @@ impl TableProvider for ListingTable { output_partitioning, &df_schema, state.execution_props(), + &SubqueryContext::default(), )? } }; diff --git a/datafusion/catalog/src/memory/table.rs b/datafusion/catalog/src/memory/table.rs index 075e462f4fe2d..53572a4066f0a 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::execution_props::SubqueryContext; 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, + &SubqueryContext::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> = 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(), + &SubqueryContext::default(), + )?; Ok((name.clone(), physical_expr)) }) .collect::>()?; @@ -470,8 +479,12 @@ fn evaluate_filters_to_mask( let mut combined_mask: Option = 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, + &SubqueryContext::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 5bfecef1fb2ed..8a61243b5ac23 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::execution_props::SubqueryContext; 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, + &SubqueryContext::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 08c7463e211c6..2838d36c0d31c 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::execution_props::SubqueryContext; use std::error::Error; use std::path::PathBuf; @@ -2851,6 +2852,7 @@ mod tests { _expr: &Expr, _input_dfschema: &DFSchema, _session_state: &SessionState, + _subquery_ctx: &SubqueryContext, ) -> Result> { unimplemented!() } diff --git a/datafusion/core/src/execution/session_state.rs b/datafusion/core/src/execution/session_state.rs index 9ca16d5ba2dc2..f11feb98495a3 100644 --- a/datafusion/core/src/execution/session_state.rs +++ b/datafusion/core/src/execution/session_state.rs @@ -53,7 +53,7 @@ use datafusion_execution::config::SessionConfig; use datafusion_execution::runtime_env::RuntimeEnv; #[cfg(feature = "sql")] use datafusion_expr::TableSource; -use datafusion_expr::execution_props::ExecutionProps; +use datafusion_expr::execution_props::{ExecutionProps, SubqueryContext}; use datafusion_expr::expr_rewriter::FunctionRewrite; use datafusion_expr::planner::ExprPlanner; #[cfg(feature = "sql")] @@ -798,7 +798,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(), + &SubqueryContext::default(), + ) } /// Return the session ID diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index b6d28e7b21c79..4e857761cdde6 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -80,7 +80,9 @@ 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::execution_props::{ + ScalarSubqueryResults, SubqueryContext, SubqueryIndex, +}; use datafusion_expr::expr::{ Alias, GroupingSet, NullTreatment, WindowFunction, WindowFunctionParams, physical_name, @@ -135,11 +137,19 @@ pub trait PhysicalPlanner: Send + Sync { /// `expr`: the expression to convert /// /// `input_dfschema`: the logical plan schema for evaluating `expr` + /// + /// `subquery_ctx`: the [`SubqueryContext`] 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 `&SubqueryContext::default()`. fn create_physical_expr( &self, expr: &Expr, input_dfschema: &DFSchema, session_state: &SessionState, + subquery_ctx: &SubqueryContext, ) -> Result>; } @@ -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`]. + /// + /// `subquery_ctx` is the [`SubqueryContext`] of the plan 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], session_state: &SessionState, + subquery_ctx: &SubqueryContext, ) -> Result>>; /// Create a physical plan for a [`LogicalPlan::TableScan`]. @@ -202,6 +219,7 @@ pub trait ExtensionPlanner { /// _logical_inputs: &[&LogicalPlan], /// _physical_inputs: &[Arc], /// _session_state: &SessionState, + /// _subquery_ctx: &SubqueryContext, /// ) -> Result>> { /// Ok(None) /// } @@ -211,6 +229,7 @@ pub trait ExtensionPlanner { /// _planner: &dyn PhysicalPlanner, /// scan: &TableScan, /// _session_state: &SessionState, + /// _subquery_ctx: &SubqueryContext, /// ) -> Result>> { /// // Check if this is your custom table source /// if scan.source.is::() { @@ -235,6 +254,7 @@ pub trait ExtensionPlanner { _planner: &dyn PhysicalPlanner, _scan: &TableScan, _session_state: &SessionState, + _subquery_ctx: &SubqueryContext, ) -> Result>> { Ok(None) } @@ -295,8 +315,14 @@ impl PhysicalPlanner for DefaultPhysicalPlanner { expr: &Expr, input_dfschema: &DFSchema, session_state: &SessionState, + subquery_ctx: &SubqueryContext, ) -> Result> { - create_physical_expr(expr, input_dfschema, session_state.execution_props()) + create_physical_expr( + expr, + input_dfschema, + session_state.execution_props(), + subquery_ctx, + ) } } @@ -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, + &SubqueryContext::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 `SubqueryContext` 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 that the planner does not need + // a mutable `SessionState` and external callers of + // `create_physical_expr` are unaffected. 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 subquery_ctx = SubqueryContext::new(index_map, results.clone()); let plan = self - .create_initial_plan_inner(logical_plan, &session_state) + .create_initial_plan_inner(logical_plan, session_state, &subquery_ctx) .await?; Ok(Arc::new(ScalarSubqueryExec::new(plan, links, results))) }) @@ -490,6 +514,7 @@ impl DefaultPhysicalPlanner { &self, logical_plan: &LogicalPlan, session_state: &SessionState, + subquery_ctx: &SubqueryContext, ) -> Result> { // 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, subquery_ctx) + }); let mut outputs = futures::stream::iter(tasks) .buffer_unordered(max_concurrency) .try_collect::>() @@ -570,6 +595,7 @@ impl DefaultPhysicalPlanner { leaf_starter_index: usize, flat_tree: Arc>>, session_state: &'a SessionState, + subquery_ctx: &'a SubqueryContext, ) -> Result>> { // 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, + subquery_ctx, ChildrenContainer::None, ) .await?; @@ -598,6 +625,7 @@ impl DefaultPhysicalPlanner { .map_logical_node_to_physical( node.node, session_state, + subquery_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, + subquery_ctx, + children, + ) .await?; } } @@ -649,6 +682,7 @@ impl DefaultPhysicalPlanner { &self, node: &LogicalPlan, session_state: &SessionState, + subquery_ctx: &SubqueryContext, children: ChildrenContainer, ) -> Result> { 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, subquery_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, + subquery_ctx, + ) }) .collect::>>>() }) @@ -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, + subquery_ctx, + ) + }) .collect::>>()?; let can_repartition = session_state.config().target_partitions() > 1 @@ -1079,6 +1126,7 @@ impl DefaultPhysicalPlanner { logical_input_schema, &physical_input_schema, execution_props, + subquery_ctx, )?; let agg_filter = aggr_expr @@ -1089,6 +1137,7 @@ impl DefaultPhysicalPlanner { logical_input_schema, &physical_input_schema, execution_props, + subquery_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, + subquery_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, + subquery_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, + subquery_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, + subquery_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, + subquery_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, + subquery_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, + subquery_ctx, + )?; + let r = create_physical_expr( + r, + right_df_schema, + execution_props, + subquery_ctx, + )?; Ok((l, r)) }) .collect::>()?; @@ -1582,6 +1652,7 @@ impl DefaultPhysicalPlanner { expr, &filter_df_schema, execution_props, + subquery_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, + subquery_ctx, )?; let on_right = create_physical_expr( rhs_logical, right_df_schema, execution_props, + subquery_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, + subquery_ctx, join, input, expr, @@ -1820,6 +1894,7 @@ impl DefaultPhysicalPlanner { &logical_input, &children, session_state, + subquery_ctx, ) .await?; } @@ -1882,6 +1957,7 @@ impl DefaultPhysicalPlanner { input_dfschema: &DFSchema, input_schema: &Schema, execution_props: &ExecutionProps, + subquery_ctx: &SubqueryContext, ) -> Result { if group_expr.len() == 1 { match &group_expr[0] { @@ -1891,6 +1967,7 @@ impl DefaultPhysicalPlanner { input_dfschema, input_schema, execution_props, + subquery_ctx, ) } Expr::GroupingSet(GroupingSet::Cube(exprs)) => create_cube_physical_expr( @@ -1898,6 +1975,7 @@ impl DefaultPhysicalPlanner { input_dfschema, input_schema, execution_props, + subquery_ctx, ), Expr::GroupingSet(GroupingSet::Rollup(exprs)) => { create_rollup_physical_expr( @@ -1905,10 +1983,16 @@ impl DefaultPhysicalPlanner { input_dfschema, input_schema, execution_props, + subquery_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, + subquery_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, + subquery_ctx, + ), physical_name(e), )) }) @@ -1947,6 +2036,7 @@ fn merge_grouping_set_physical_expr( input_dfschema: &DFSchema, input_schema: &Schema, execution_props: &ExecutionProps, + subquery_ctx: &SubqueryContext, ) -> Result { let num_groups = grouping_sets.len(); let mut all_exprs: Vec = vec![]; @@ -1961,6 +2051,7 @@ fn merge_grouping_set_physical_expr( expr, input_dfschema, execution_props, + subquery_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, + subquery_ctx, )?); } } @@ -1998,6 +2090,7 @@ fn create_cube_physical_expr( input_dfschema: &DFSchema, input_schema: &Schema, execution_props: &ExecutionProps, + subquery_ctx: &SubqueryContext, ) -> Result { 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, + subquery_ctx, )?); all_exprs.push(get_physical_expr_pair( expr, input_dfschema, execution_props, + subquery_ctx, )?) } @@ -2044,6 +2139,7 @@ fn create_rollup_physical_expr( input_dfschema: &DFSchema, input_schema: &Schema, execution_props: &ExecutionProps, + subquery_ctx: &SubqueryContext, ) -> Result { let num_of_exprs = exprs.len(); @@ -2060,12 +2156,14 @@ fn create_rollup_physical_expr( input_dfschema, input_schema, execution_props, + subquery_ctx, )?); all_exprs.push(get_physical_expr_pair( expr, input_dfschema, execution_props, + subquery_ctx, )?) } @@ -2092,8 +2190,10 @@ fn get_null_physical_expr_pair( input_dfschema: &DFSchema, input_schema: &Schema, execution_props: &ExecutionProps, + subquery_ctx: &SubqueryContext, ) -> Result<(Arc, String)> { - let physical_expr = create_physical_expr(expr, input_dfschema, execution_props)?; + let physical_expr = + create_physical_expr(expr, input_dfschema, execution_props, subquery_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, + subquery_ctx: &SubqueryContext, ) -> Result<(Arc, String)> { - let physical_expr = create_physical_expr(expr, input_dfschema, execution_props)?; + let physical_expr = + create_physical_expr(expr, input_dfschema, execution_props, subquery_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 `subquery_ctx` argument. pub fn create_window_expr_with_name( e: &Expr, name: impl Into, logical_schema: &DFSchema, execution_props: &ExecutionProps, + subquery_ctx: &SubqueryContext, ) -> Result> { 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, + subquery_ctx, + )?; + let partition_by = create_physical_exprs( + partition_by, + logical_schema, + execution_props, + subquery_ctx, + )?; + let order_by = create_physical_sort_exprs( + order_by, + logical_schema, + execution_props, + subquery_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, subquery_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 `subquery_ctx` argument. pub fn create_window_expr( e: &Expr, logical_schema: &DFSchema, execution_props: &ExecutionProps, + subquery_ctx: &SubqueryContext, ) -> Result> { // 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, subquery_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 { + let subquery_ctx = SubqueryContext::default(); let mut builder = LoweredAggregateBuilder::new( e, logical_input_schema, physical_input_schema, execution_props, + &subquery_ctx, ) .with_human_display(human_display); @@ -2554,11 +2678,13 @@ pub fn create_aggregate_expr_and_maybe_filter( _ => (None, String::default(), e.clone()), }; + let subquery_ctx = SubqueryContext::default(); let mut builder = LoweredAggregateBuilder::new( &e, logical_input_schema, physical_input_schema, execution_props, + &subquery_ctx, ) .with_human_display(human_display); @@ -2953,6 +3079,7 @@ impl DefaultPhysicalPlanner { fn create_project_physical_exec_with_props( &self, execution_props: &ExecutionProps, + subquery_ctx: &SubqueryContext, input_exec: Arc, input: &Arc, 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, + subquery_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(), + &SubqueryContext::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(), + &SubqueryContext::default(), ); insta::assert_debug_snapshot!(cube, @r#" @@ -3744,6 +3881,7 @@ mod tests { logical_input_schema, physical_input_schema, session_state.execution_props(), + &SubqueryContext::default(), ); insta::assert_debug_snapshot!(rollup, @r#" @@ -3848,6 +3986,7 @@ mod tests { &col("a").not(), &dfschema, &make_session_state(), + &SubqueryContext::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], _session_state: &SessionState, + _subquery_ctx: &SubqueryContext, ) -> Result>> { internal_err!("BOOM") } @@ -4404,6 +4567,7 @@ mod tests { #[derive(PartialEq, Eq, Hash)] struct NoOpExtensionNode { schema: DFSchemaRef, + expressions: Vec, } impl Default for NoOpExtensionNode { @@ -4416,6 +4580,7 @@ mod tests { ) .unwrap(), ), + expressions: vec![], } } } @@ -4448,7 +4613,7 @@ mod tests { } fn expressions(&self) -> Vec { - 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, + exprs: Vec, _inputs: Vec, ) -> Result { - 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, - _children: Vec>, + children: Vec>, ) -> Result> { - 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], + session_state: &SessionState, + subquery_ctx: &SubqueryContext, + ) -> Result>> { + for expr in node.expressions() { + planner.create_physical_expr( + &expr, + node.schema(), + session_state, + subquery_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], _session_state: &SessionState, + _subquery_ctx: &SubqueryContext, ) -> Result>> { 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], _session_state: &SessionState, + _subquery_ctx: &SubqueryContext, ) -> Result>> { Ok(None) } @@ -5193,6 +5393,7 @@ digraph { _planner: &dyn PhysicalPlanner, scan: &TableScan, _session_state: &SessionState, + _subquery_ctx: &SubqueryContext, ) -> Result>> { if scan.source.is::() { 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 c53495421307b..a768022ea1bdb 100644 --- a/datafusion/core/src/test_util/parquet.rs +++ b/datafusion/core/src/test_util/parquet.rs @@ -28,7 +28,7 @@ use crate::datasource::listing::{ListingTableUrl, PartitionedFile}; 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::execution_props::{ExecutionProps, SubqueryContext}; use crate::logical_expr::simplify::SimplifyContext; use crate::optimizer::simplify_expressions::ExprSimplifier; use crate::physical_expr::create_physical_expr; @@ -172,8 +172,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(), + &SubqueryContext::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 a41803191ad05..2f6efe05f76c5 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::execution_props::SubqueryContext; 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, + &SubqueryContext::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 b837373632f07..0017a5d6a1e94 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::execution_props::SubqueryContext; 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], _session_state: &SessionState, + _subquery_ctx: &SubqueryContext, ) -> Result>> { Ok( if let Some(topk_node) = node.as_any().downcast_ref::() { diff --git a/datafusion/datasource/src/file_scan_config/mod.rs b/datafusion/datasource/src/file_scan_config/mod.rs index 91caabeee6a41..88ba8b456d9d5 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/datafusion/datasource/src/file_scan_config/mod.rs @@ -1655,7 +1655,7 @@ mod tests { fn test_split_groups_by_statistics() -> Result<()> { use chrono::TimeZone; use datafusion_common::DFSchema; - use datafusion_expr::execution_props::ExecutionProps; + use datafusion_expr::execution_props::{ExecutionProps, SubqueryContext}; use object_store::{ObjectMeta, path::Path}; struct File { @@ -1869,6 +1869,7 @@ mod tests { &expr, &DFSchema::try_from(Arc::clone(&table_schema))?, &ExecutionProps::default(), + &SubqueryContext::default(), ) }) .collect::>>()?, @@ -2235,7 +2236,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, SubqueryContext}, + }; let schema = Arc::new(Schema::new(vec![Field::new( "value", @@ -2249,7 +2253,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, + &SubqueryContext::default(), + ) + .unwrap() }) .into(); diff --git a/datafusion/datasource/src/projection.rs b/datafusion/datasource/src/projection.rs index 16207c086f7bc..c331d668bf5a8 100644 --- a/datafusion/datasource/src/projection.rs +++ b/datafusion/datasource/src/projection.rs @@ -297,7 +297,10 @@ mod test { use arrow::array::{AsArray, RecordBatch}; use arrow::datatypes::{DataType, Field, SchemaRef}; use datafusion_common::{DFSchema, ScalarValue, config::ConfigOptions, record_batch}; - use datafusion_expr::{Expr, ScalarUDF, col, execution_props::ExecutionProps}; + use datafusion_expr::{ + Expr, ScalarUDF, col, + execution_props::{ExecutionProps, SubqueryContext}, + }; use datafusion_functions::core::input_file_name::InputFileNameFunc; use datafusion_physical_expr::{ ScalarFunctionExpr, create_physical_exprs, projection::ProjectionExpr, @@ -324,8 +327,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(), + &SubqueryContext::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 649f74ed3997c..012e405d05a1f 100644 --- a/datafusion/expr/src/execution_props.rs +++ b/datafusion/expr/src/execution_props.rs @@ -64,12 +64,6 @@ pub struct ExecutionProps { pub config_options: Option>, /// Providers for scalar variables pub var_providers: Option>>, - /// Maps each logical `Subquery` to its index in `subquery_results`. - /// Populated by the physical planner before calling `create_physical_expr`. - pub subquery_indexes: HashMap, - /// 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 +84,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,6 +161,52 @@ impl ExecutionProps { } } +/// Per-plan context used by the physical planner when creating physical +/// expressions for [`Expr::ScalarSubquery`] nodes that read from a shared +/// [`ScalarSubqueryResults`] container. +/// +/// The physical planner builds this from the set of uncorrelated scalar +/// subqueries it has scheduled for execution. It is then passed explicitly +/// through `create_physical_expr` so that function can find the slot index for +/// each [`Subquery`]. +/// +/// An empty [`SubqueryContext`] (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 SubqueryContext { + indexes: HashMap, + results: ScalarSubqueryResults, +} + +impl SubqueryContext { + /// Create a [`SubqueryContext`] 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, + 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 { + 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); @@ -274,7 +312,7 @@ 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:?}") ); } diff --git a/datafusion/functions-nested/src/array_any_match.rs b/datafusion/functions-nested/src/array_any_match.rs index c8ba978881394..027b058642b28 100644 --- a/datafusion/functions-nested/src/array_any_match.rs +++ b/datafusion/functions-nested/src/array_any_match.rs @@ -273,7 +273,7 @@ mod tests { use datafusion_common::{DFSchema, Result}; use datafusion_expr::{ Expr, HigherOrderReturnFieldArgs, HigherOrderUDFImpl, ValueOrLambda, col, - execution_props::ExecutionProps, + execution_props::{ExecutionProps, SubqueryContext}, expr::{HigherOrderFunction, LambdaVariable}, lambda, lit, }; @@ -311,6 +311,7 @@ mod tests { )), &schema, &ExecutionProps::new(), + &SubqueryContext::default(), )? .evaluate(&RecordBatch::try_new( Arc::clone(schema.inner()), @@ -344,6 +345,7 @@ mod tests { )), &schema, &ExecutionProps::new(), + &SubqueryContext::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 0f208ce5d26b2..035456acb1c7d 100644 --- a/datafusion/functions-nested/src/lambda_utils.rs +++ b/datafusion/functions-nested/src/lambda_utils.rs @@ -137,7 +137,7 @@ pub(crate) mod test_utils { use datafusion_common::{DFSchema, Result}; use datafusion_expr::{ Expr, HigherOrderUDF, col, - execution_props::ExecutionProps, + execution_props::{ExecutionProps, SubqueryContext}, expr::{HigherOrderFunction, LambdaVariable}, lambda, }; @@ -174,6 +174,7 @@ pub(crate) mod test_utils { )), &schema, &ExecutionProps::new(), + &SubqueryContext::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 39c8541b51b2f..27f555fb622ff 100644 --- a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs +++ b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs @@ -39,6 +39,7 @@ use datafusion_common::{ metadata::FieldMetadata, tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRewriter}, }; +use datafusion_expr::execution_props::SubqueryContext; use datafusion_expr::expr::HigherOrderFunction; use datafusion_expr::{ BinaryExpr, Case, ColumnarValue, Expr, ExprSchemable, Like, Operator, Volatility, @@ -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, + &SubqueryContext::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 b29649e9ead49..87221c55565c7 100644 --- a/datafusion/optimizer/src/utils.rs +++ b/datafusion/optimizer/src/utils.rs @@ -27,6 +27,7 @@ use datafusion_common::cast::as_boolean_array; use datafusion_common::tree_node::{TransformedResult, TreeNode, TreeNodeRecursion}; use datafusion_common::{Column, DFSchema, Result, ScalarValue}; use datafusion_expr::execution_props::ExecutionProps; +use datafusion_expr::execution_props::SubqueryContext; use datafusion_expr::expr::{Exists, InSubquery, SetComparison}; use datafusion_expr::expr_rewriter::replace_col; use datafusion_expr::{ColumnarValue, Expr, logical_plan::LogicalPlan}; @@ -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, + &SubqueryContext::default(), + )? + .evaluate(&input_batch) } fn coerce(expr: Expr, schema: &DFSchema) -> Result { diff --git a/datafusion/physical-expr/src/aggregate.rs b/datafusion/physical-expr/src/aggregate.rs index e5d55aba4f51c..550085683894b 100644 --- a/datafusion/physical-expr/src/aggregate.rs +++ b/datafusion/physical-expr/src/aggregate.rs @@ -47,7 +47,7 @@ use datafusion_common::metadata::FieldMetadata; use datafusion_common::{ DFSchema, Result, ScalarValue, assert_or_internal_err, internal_err, not_impl_err, }; -use datafusion_expr::execution_props::ExecutionProps; +use datafusion_expr::execution_props::{ExecutionProps, SubqueryContext}; use datafusion_expr::expr::{ AggregateFunction, AggregateFunctionParams, NullTreatment, physical_name, }; @@ -423,6 +423,7 @@ pub struct LoweredAggregateBuilder<'a> { logical_input_schema: &'a DFSchema, physical_input_schema: &'a Schema, execution_props: &'a ExecutionProps, + subquery_ctx: &'a SubqueryContext, } impl<'a> LoweredAggregateBuilder<'a> { @@ -430,12 +431,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. `subquery_ctx` is used when creating + /// physical expressions that reference uncorrelated scalar subqueries. + /// Callers creating physical aggregates outside of physical planning should + /// pass `&SubqueryContext::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, + subquery_ctx: &'a SubqueryContext, ) -> Self { Self { expr, @@ -446,6 +452,7 @@ impl<'a> LoweredAggregateBuilder<'a> { logical_input_schema, physical_input_schema, execution_props, + subquery_ctx, } } @@ -484,6 +491,7 @@ impl<'a> LoweredAggregateBuilder<'a> { logical_input_schema, physical_input_schema, execution_props, + subquery_ctx, } = self; let (name, human_display, output_metadata, expr) = lower_aggregate_display( @@ -515,16 +523,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, + subquery_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, + subquery_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, + subquery_ctx, + )?; let ignore_nulls = null_treatment.unwrap_or(NullTreatment::RespectNulls) == NullTreatment::IgnoreNulls; @@ -1162,6 +1183,7 @@ mod tests { &logical_schema, &schema, &ExecutionProps::new(), + &SubqueryContext::default(), ) .build()?; @@ -1185,6 +1207,7 @@ mod tests { &logical_schema, &schema, &ExecutionProps::new(), + &SubqueryContext::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 1dca36b75f9f5..1e983ac9df8ff 100644 --- a/datafusion/physical-expr/src/analysis.rs +++ b/datafusion/physical-expr/src/analysis.rs @@ -349,7 +349,10 @@ mod tests { use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::{DFSchema, ScalarValue, assert_contains, stats::Precision}; use datafusion_expr::{ - Expr, col, execution_props::ExecutionProps, interval_arithmetic::Interval, lit, + Expr, col, + execution_props::{ExecutionProps, SubqueryContext}, + interval_arithmetic::Interval, + lit, }; use crate::{AnalysisContext, create_physical_expr, expressions::Column}; @@ -412,8 +415,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(), + &SubqueryContext::default(), + ) + .unwrap(); let analysis_result = analyze( &physical_expr, AnalysisContext::new(boundaries), @@ -453,8 +461,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(), + &SubqueryContext::default(), + ) + .unwrap(); let analysis_result = analyze( &physical_expr, AnalysisContext::new(boundaries), @@ -475,8 +488,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(), + &SubqueryContext::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 cfc9866fc8c3f..7635f5cc2f027 100644 --- a/datafusion/physical-expr/src/physical_expr.rs +++ b/datafusion/physical-expr/src/physical_expr.rs @@ -25,7 +25,7 @@ use arrow::datatypes::{DataType, Schema, SchemaRef}; 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::execution_props::{ExecutionProps, SubqueryContext}; use datafusion_expr::{Expr, Partitioning as LogicalPartitioning, SortExpr}; use datafusion_expr_common::casts::try_cast_literal_to_type; @@ -190,47 +190,68 @@ pub fn create_lex_ordering( exprs, &df_schema, execution_props, + &SubqueryContext::default(), )?)); } Ok(all_sort_orders) } /// Create a physical sort expression from a logical expression +/// +/// See [`create_physical_expr`] for details on the `subquery_ctx` argument. pub fn create_physical_sort_expr( e: &SortExpr, input_dfschema: &DFSchema, execution_props: &ExecutionProps, + subquery_ctx: &SubqueryContext, ) -> Result { - 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, subquery_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 `subquery_ctx` argument. pub fn create_physical_sort_exprs( exprs: &[SortExpr], input_dfschema: &DFSchema, execution_props: &ExecutionProps, + subquery_ctx: &SubqueryContext, ) -> Result> { exprs .iter() - .map(|e| create_physical_sort_expr(e, input_dfschema, execution_props)) + .map(|e| { + create_physical_sort_expr(e, input_dfschema, execution_props, subquery_ctx) + }) .collect() } /// Create physical partitioning from logical partitioning. +/// +/// See [`create_physical_expr`] for details on the `subquery_ctx` argument. pub fn create_physical_partitioning( partitioning: &LogicalPartitioning, input_dfschema: &DFSchema, execution_props: &ExecutionProps, + subquery_ctx: &SubqueryContext, ) -> Result { 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, + subquery_ctx, + ) + }) .collect::>>()?; Ok(Partitioning::Hash(exprs, *partition_count)) } @@ -239,6 +260,7 @@ pub fn create_physical_partitioning( range.ordering(), input_dfschema, execution_props, + subquery_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 d0d0508a106a5..8456764dade69 100644 --- a/datafusion/physical-expr/src/planner.rs +++ b/datafusion/physical-expr/src/planner.rs @@ -32,7 +32,7 @@ use datafusion_common::{ DFSchema, Result, ScalarValue, TableReference, ToDFSchema, exec_err, internal_datafusion_err, not_impl_err, plan_datafusion_err, plan_err, }; -use datafusion_expr::execution_props::ExecutionProps; +use datafusion_expr::execution_props::{ExecutionProps, SubqueryContext}; use datafusion_expr::expr::{ Alias, Cast, HigherOrderFunction, InList, Lambda, LambdaVariable, Placeholder, ScalarFunction, @@ -62,7 +62,7 @@ use datafusion_expr::{ /// # use datafusion_common::DFSchema; /// # use datafusion_expr::{Expr, col, lit}; /// # use datafusion_physical_expr::create_physical_expr; -/// # use datafusion_expr::execution_props::ExecutionProps; +/// # use datafusion_expr::execution_props::{ExecutionProps, SubqueryContext}; /// // 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 +70,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 `SubqueryContext`: +/// let physical_expr = +/// create_physical_expr(&expr, &df_schema, &props, &SubqueryContext::default()) +/// .unwrap(); /// ``` /// /// # Example: Executing a PhysicalExpr to obtain [ColumnarValue] @@ -82,13 +85,15 @@ use datafusion_expr::{ /// # use datafusion_common::{assert_batches_eq, DFSchema}; /// # use datafusion_expr::{Expr, col, lit, ColumnarValue}; /// # use datafusion_physical_expr::create_physical_expr; -/// # use datafusion_expr::execution_props::ExecutionProps; +/// # use datafusion_expr::execution_props::{ExecutionProps, SubqueryContext}; /// # 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, &SubqueryContext::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 +116,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. +/// * `subquery_ctx` - The [`SubqueryContext`] 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 +/// `&SubqueryContext::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, + subquery_ctx: &SubqueryContext, ) -> Result> { let input_schema = input_dfschema.as_arrow(); @@ -131,7 +145,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, + subquery_ctx, + )?) } } Expr::Column(c) => { @@ -167,12 +186,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, + subquery_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, + subquery_ctx, + ) } Expr::IsFalse(expr) => { let binary_op = binary_expr( @@ -180,12 +209,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, + subquery_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, + subquery_ctx, + ) } Expr::IsUnknown(expr) => { let binary_op = binary_expr( @@ -193,7 +232,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, + subquery_ctx, + ) } Expr::IsNotUnknown(expr) => { let binary_op = binary_expr( @@ -201,12 +245,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, + subquery_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, + subquery_ctx, + )?; + let rhs = create_physical_expr( + right, + input_dfschema, + execution_props, + subquery_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 +288,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, + subquery_ctx, + )?; + let physical_pattern = create_physical_expr( + pattern, + input_dfschema, + execution_props, + subquery_ctx, + )?; like( *negated, *case_insensitive, @@ -251,10 +318,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, + subquery_ctx, + )?; + let physical_pattern = create_physical_expr( + pattern, + input_dfschema, + execution_props, + subquery_ctx, + )?; similar_to(*negated, *case_insensitive, physical_expr, physical_pattern) } Expr::Case(case) => { @@ -263,6 +338,7 @@ pub fn create_physical_expr( e.as_ref(), input_dfschema, execution_props, + subquery_ctx, )?) } else { None @@ -272,10 +348,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, + subquery_ctx, + )?; + let then_expr = create_physical_exprs( + then_expr, + input_dfschema, + execution_props, + subquery_ctx, + )?; let when_then_expr: Vec<(Arc, Arc)> = when_expr .iter() @@ -288,6 +372,7 @@ pub fn create_physical_expr( e.as_ref(), input_dfschema, execution_props, + subquery_ctx, )?) } else { None @@ -295,7 +380,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, subquery_ctx)?, input_schema, Arc::clone(field), None, @@ -314,31 +399,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, + subquery_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, + subquery_ctx, + )?), Expr::Negative(expr) => expressions::negative( - create_physical_expr(expr, input_dfschema, execution_props)?, + create_physical_expr(expr, input_dfschema, execution_props, subquery_ctx)?, input_schema, ), Expr::IsNull(expr) => expressions::is_null(create_physical_expr( expr, input_dfschema, execution_props, + subquery_ctx, )?), Expr::IsNotNull(expr) => expressions::is_not_null(create_physical_expr( expr, input_dfschema, execution_props, + subquery_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, + subquery_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 +456,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, + subquery_ctx, + )?; + let low_expr = + create_physical_expr(low, input_dfschema, execution_props, subquery_ctx)?; + let high_expr = create_physical_expr( + high, + input_dfschema, + execution_props, + subquery_ctx, + )?; // rewrite the between into the two binary operators let binary_expr = binary( @@ -394,17 +504,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, + subquery_ctx, + )?; - let list_exprs = - create_physical_exprs(list, input_dfschema, execution_props)?; + let list_exprs = create_physical_exprs( + list, + input_dfschema, + execution_props, + subquery_ctx, + )?; expressions::in_list(value_expr, list_exprs, negated, input_schema) } }, Expr::ScalarSubquery(sq) => { - match execution_props.subquery_indexes.get(sq) { - Some(&index) => { + match subquery_ctx.index_of(sq) { + Some(index) => { let schema = sq.subquery.schema(); if schema.fields().len() != 1 { return plan_err!( @@ -418,7 +536,7 @@ pub fn create_physical_expr( dt, nullable, index, - execution_props.subquery_results.clone(), + subquery_ctx.results().clone(), ))) } None => { @@ -495,9 +613,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, + subquery_ctx, + ) } - _ => create_physical_expr(arg, input_dfschema, execution_props), + _ => create_physical_expr( + arg, + input_dfschema, + execution_props, + subquery_ctx, + ), }) .collect::>()?; @@ -515,7 +643,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, subquery_ctx)?, ), Expr::LambdaVariable(LambdaVariable { name, @@ -572,17 +700,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 `subquery_ctx` argument. pub fn create_physical_exprs<'a, I>( exprs: I, input_dfschema: &DFSchema, execution_props: &ExecutionProps, + subquery_ctx: &SubqueryContext, ) -> Result>> where I: IntoIterator, { exprs .into_iter() - .map(|expr| create_physical_expr(expr, input_dfschema, execution_props)) + .map(|expr| { + create_physical_expr(expr, input_dfschema, execution_props, subquery_ctx) + }) .collect() } @@ -591,7 +724,13 @@ pub fn logical2physical(expr: &Expr, schema: &Schema) -> Arc { // 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, + &SubqueryContext::default(), + ) + .unwrap() } #[cfg(test)] @@ -608,7 +747,12 @@ mod tests { fn lower_cast_expr(expr: &Expr, schema: &Schema) -> Result> { let df_schema = DFSchema::try_from(schema.clone())?; - create_physical_expr(expr, &df_schema, &ExecutionProps::new()) + create_physical_expr( + expr, + &df_schema, + &ExecutionProps::new(), + &SubqueryContext::default(), + ) } fn as_planner_cast(physical: &Arc) -> &expressions::CastExpr { @@ -623,7 +767,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(), + &SubqueryContext::default(), + )?; let batch = RecordBatch::try_new( Arc::new(schema), @@ -728,8 +877,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(), + &SubqueryContext::default(), + )?; Ok(()) } 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 f7d2745549c99..2c1397c28c6d3 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -591,3 +591,88 @@ metadata and schema fingerprint match. - Pass the current schema fingerprint to `CachedFileMetadata::is_valid_for`. See [PR #23201](https://github.com/apache/datafusion/pull/23201) for details. + +### Scalar-subquery state moved from `ExecutionProps` to an explicit `SubqueryContext` parameter + +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` values from logical `Expr` values. + +That state is now carried by a dedicated +`datafusion_expr::execution_props::SubqueryContext` passed explicitly to those +functions and the planner traits. 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 `subquery_ctx: &SubqueryContext` +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 `subquery_ctx: &SubqueryContext` +- `ExtensionPlanner::plan_extension` and `plan_table_scan` receive + `subquery_ctx: &SubqueryContext` 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 `&SubqueryContext::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 `SubqueryContext` instead. + +**Migration guide:** + +When creating a physical expression outside of physical planning, pass an empty +context: + +```rust,ignore +use datafusion_expr::execution_props::SubqueryContext; +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, &SubqueryContext::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], + session_state: &SessionState, + subquery_ctx: &SubqueryContext, // new parameter +) -> Result>> { + for expr in node.expressions() { + // Forward the context so scalar subqueries in this node's + // expressions resolve against the plan's subquery state + planner.create_physical_expr(&expr, node.schema(), session_state, subquery_ctx)?; + } + // ... +} +``` + +See [PR #23649](https://github.com/apache/datafusion/pull/23649) for details.