From 805a3705472941fee3236749104f72e0b18d9fa1 Mon Sep 17 00:00:00 2001 From: Albert Skalt Date: Wed, 11 Feb 2026 11:19:45 +0300 Subject: [PATCH 1/6] feat: make execution transfomation application rule generic Currently, it depends on the particular execution transformation that resolves placeholders. However, it would be more convenient to be able to specify the rule which is required to be applied to the plan. --- ...pr_resolver.rs => exec_transform_apply.rs} | 31 ++++--- .../core/tests/physical_optimizer/mod.rs | 2 +- .../tests/physical_optimizer/test_utils.rs | 2 +- datafusion/datasource/src/values.rs | 6 +- ...pr_resolver.rs => exec_transform_apply.rs} | 82 +++++++++---------- datafusion/physical-optimizer/src/lib.rs | 2 +- .../physical-optimizer/src/optimizer.rs | 9 +- .../physical-plan/benches/plan_transformer.rs | 25 ++---- .../physical-plan/src/plan_transformer/mod.rs | 74 ++++++++--------- .../plan_transformer/resolve_placeholders.rs | 6 +- datafusion/proto/src/physical_plan/mod.rs | 8 +- .../sqllogictest/test_files/explain.slt | 18 ++-- 12 files changed, 126 insertions(+), 139 deletions(-) rename datafusion/core/tests/physical_optimizer/{physical_expr_resolver.rs => exec_transform_apply.rs} (90%) rename datafusion/physical-optimizer/src/{physical_expr_resolver.rs => exec_transform_apply.rs} (52%) diff --git a/datafusion/core/tests/physical_optimizer/physical_expr_resolver.rs b/datafusion/core/tests/physical_optimizer/exec_transform_apply.rs similarity index 90% rename from datafusion/core/tests/physical_optimizer/physical_expr_resolver.rs rename to datafusion/core/tests/physical_optimizer/exec_transform_apply.rs index 43013cd62c65f..22d8b7c7d3881 100644 --- a/datafusion/core/tests/physical_optimizer/physical_expr_resolver.rs +++ b/datafusion/core/tests/physical_optimizer/exec_transform_apply.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! Integration tests for [`PhysicalExprResolver`] optimizer rule. +//! Integration tests for [`ExecutionTransformationApplier`] optimizer rule. use std::{collections::HashMap, sync::Arc}; @@ -29,11 +29,14 @@ use datafusion_physical_expr::{ expressions::{BinaryExpr, col, lit, placeholder}, }; use datafusion_physical_optimizer::{ - PhysicalOptimizerRule, physical_expr_resolver::PhysicalExprResolver, + PhysicalOptimizerRule, exec_transform_apply::ExecutionTransformationApplier, }; use datafusion_physical_plan::{ - ExecutionPlan, filter::FilterExec, get_plan_string, - plan_transformer::TransformPlanExec, repartition::RepartitionExec, + ExecutionPlan, + filter::FilterExec, + get_plan_string, + plan_transformer::{ResolvePlaceholdersRule, TransformPlanExec}, + repartition::RepartitionExec, }; use crate::physical_optimizer::test_utils::{ @@ -85,6 +88,12 @@ fn repartition_exec( )?)) } +fn placeholder_resolver() -> ExecutionTransformationApplier { + ExecutionTransformationApplier::new_post_optimization(Arc::new( + ResolvePlaceholdersRule::new(), + )) +} + #[test] fn test_noop_if_no_placeholders_found() -> Result<()> { let schema = create_schema(); @@ -105,8 +114,7 @@ fn test_noop_if_no_placeholders_found() -> Result<()> { assert_eq!(initial, expected_initial); - let after_optimize = PhysicalExprResolver::new_post_optimization() - .optimize(plan, &ConfigOptions::new())?; + let after_optimize = placeholder_resolver().optimize(plan, &ConfigOptions::new())?; let optimized_plan_string = get_plan_string(&after_optimize); assert_eq!(initial, optimized_plan_string); @@ -134,8 +142,7 @@ fn test_wrap_plan_with_transformer() -> Result<()> { assert_eq!(initial, expected_initial); - let after_optimize = PhysicalExprResolver::new_post_optimization() - .optimize(plan, &ConfigOptions::new())?; + let after_optimize = placeholder_resolver().optimize(plan, &ConfigOptions::new())?; let expected_optimized = [ "TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1]", @@ -201,7 +208,7 @@ fn test_remove_useless_transformers() -> Result<()> { assert_eq!(initial, expected_initial); let after_optimize = - PhysicalExprResolver::new().optimize(plan, &ConfigOptions::new())?; + ExecutionTransformationApplier::new().optimize(plan, &ConfigOptions::new())?; let expected_optimized = [ "GlobalLimitExec: skip=0, fetch=5", @@ -244,7 +251,7 @@ fn test_combine_transformers() -> Result<()> { assert_eq!(initial, expected_initial); let after_pre_optimization = - PhysicalExprResolver::new().optimize(plan, &ConfigOptions::new())?; + ExecutionTransformationApplier::new().optimize(plan, &ConfigOptions::new())?; let expected_optimized = [ "GlobalLimitExec: skip=0, fetch=5", @@ -257,8 +264,8 @@ fn test_combine_transformers() -> Result<()> { let optimized_plan_string = get_plan_string(&after_pre_optimization); assert_eq!(optimized_plan_string, expected_optimized); - let after_post_optimization = PhysicalExprResolver::new_post_optimization() - .optimize(after_pre_optimization, &ConfigOptions::new())?; + let after_post_optimization = + placeholder_resolver().optimize(after_pre_optimization, &ConfigOptions::new())?; let expected_optimized = [ "TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1]", diff --git a/datafusion/core/tests/physical_optimizer/mod.rs b/datafusion/core/tests/physical_optimizer/mod.rs index 7d22a8c25f209..011fe7af0b615 100644 --- a/datafusion/core/tests/physical_optimizer/mod.rs +++ b/datafusion/core/tests/physical_optimizer/mod.rs @@ -31,7 +31,7 @@ mod limit_pushdown; mod limited_distinct_aggregation; mod partition_statistics; #[expect(clippy::needless_pass_by_value)] -mod physical_expr_resolver; +mod exec_transform_apply; mod projection_pushdown; mod pushdown_sort; mod replace_with_order_preserving_variants; diff --git a/datafusion/core/tests/physical_optimizer/test_utils.rs b/datafusion/core/tests/physical_optimizer/test_utils.rs index 582506457560c..91f46b070879c 100644 --- a/datafusion/core/tests/physical_optimizer/test_utils.rs +++ b/datafusion/core/tests/physical_optimizer/test_utils.rs @@ -400,7 +400,7 @@ pub fn resolve_placeholders_exec( input: Arc, ) -> Arc { Arc::new( - TransformPlanExec::try_new(input, vec![Box::new(ResolvePlaceholdersRule::new())]) + TransformPlanExec::try_new(input, vec![Arc::new(ResolvePlaceholdersRule::new())]) .unwrap(), ) } diff --git a/datafusion/datasource/src/values.rs b/datafusion/datasource/src/values.rs index bea1f871d76a2..8e2b7b4b80750 100644 --- a/datafusion/datasource/src/values.rs +++ b/datafusion/datasource/src/values.rs @@ -437,7 +437,7 @@ mod tests { // Should be ValuesSource because of placeholder. assert!(values_exec.data_source().as_any().is::()); - let rules = vec![Box::new(ResolvePlaceholdersRule::new()) as Box<_>]; + let rules = vec![Arc::new(ResolvePlaceholdersRule::new()) as _]; let exec = Arc::new(TransformPlanExec::try_new(values_exec, rules)?); let task_ctx = Arc::new(TaskContext::default().with_param_values( ParamValues::List(vec![ScalarValue::Int32(Some(10)).into()]), @@ -470,7 +470,7 @@ mod tests { ]]; let values_exec = ValuesSource::try_new_exec(Arc::clone(&schema), data)?; - let rules = vec![Box::new(ResolvePlaceholdersRule::new()) as Box<_>]; + let rules = vec![Arc::new(ResolvePlaceholdersRule::new()) as _]; let exec = Arc::new(TransformPlanExec::try_new(values_exec, rules)?) as Arc<_>; let task_ctx = Arc::new(TaskContext::default().with_param_values( @@ -528,7 +528,7 @@ mod tests { vec![vec![lit(10), placeholder("$foo", DataType::Int32)]]; let values_exec = ValuesSource::try_new_exec(Arc::clone(&schema), data)?; - let rules = vec![Box::new(ResolvePlaceholdersRule::new()) as Box<_>]; + let rules = vec![Arc::new(ResolvePlaceholdersRule::new()) as _]; let exec = Arc::new(TransformPlanExec::try_new(values_exec, rules)?) as Arc<_>; let task_ctx = Arc::new(TaskContext::default()); diff --git a/datafusion/physical-optimizer/src/physical_expr_resolver.rs b/datafusion/physical-optimizer/src/exec_transform_apply.rs similarity index 52% rename from datafusion/physical-optimizer/src/physical_expr_resolver.rs rename to datafusion/physical-optimizer/src/exec_transform_apply.rs index a18a6d4d6bbc9..5312334ac1f82 100644 --- a/datafusion/physical-optimizer/src/physical_expr_resolver.rs +++ b/datafusion/physical-optimizer/src/exec_transform_apply.rs @@ -15,11 +15,10 @@ // specific language governing permissions and limitations // under the License. -//! [`PhysicalExprResolver`] ensures that the physical plan is prepared for placeholder resolution -//! by wrapping it in a [`TransformPlanExec`] with a [`ResolvePlaceholdersRule`] if the plan -//! contains any unresolved placeholders. The actual resolution happens during execution. +//! [`ExecutionTransformationApplier`] ensures that the required execution transformations +//! are applied to the physical plan. -use std::sync::Arc; +use std::{borrow::Cow, sync::Arc}; use datafusion_common::{ Result, @@ -28,62 +27,63 @@ use datafusion_common::{ }; use datafusion_physical_plan::{ ExecutionPlan, - plan_transformer::{ResolvePlaceholdersRule, TransformPlanExec}, + plan_transformer::{ExecutionTransformationRule, TransformPlanExec}, }; use crate::PhysicalOptimizerRule; -/// The phase in which the [`PhysicalExprResolver`] rule is applied. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum PhysicalExprResolverPhase { +/// The phase in which the [`ExecutionTransformationApplier`] rule is applied. +#[derive(Debug)] +pub enum ExecutionTransformationApplierPhase { /// Optimization that happens before most other optimizations. /// This optimization removes all [`TransformPlanExec`] execution plans from the plan /// tree. Pre, /// Optimization that happens after most other optimizations. - /// This optimization checks for the presence of placeholders in the optimized plan, and if - /// they are present, wraps the plan in a [`TransformPlanExec`] with a [`ResolvePlaceholdersRule`]. - Post, + /// This optimization checks if `rule` requires to transform the plan and wraps the plan with + /// [`TransformPlanExec`] if it so, or adds rule to the existing transformation node. + Post { + rule: Arc, + }, } -/// Physical optimizer rule that prepares the plan for placeholder resolution during execution. +/// Physical optimizer rule that wraps the plan with a certain execution-stage transformation. #[derive(Debug)] -pub struct PhysicalExprResolver { - phase: PhysicalExprResolverPhase, +pub struct ExecutionTransformationApplier { + phase: ExecutionTransformationApplierPhase, + name: Cow<'static, str>, } -impl PhysicalExprResolver { - /// Creates a new [`PhysicalExprResolver`] optimizer rule that runs in the pre-optimization - /// phase. In this phase, the rule removes any existing [`TransformPlanExec`] from the - /// plan tree. +impl ExecutionTransformationApplier { + /// Creates a new [`ExecutionTransformationApplier`] optimizer rule that runs in the + /// pre-optimization phase. pub fn new() -> Self { Self { - phase: PhysicalExprResolverPhase::Pre, + phase: ExecutionTransformationApplierPhase::Pre, + name: Cow::Borrowed("ExecutionTransformationApplier"), } } - /// Creates a new [`PhysicalExprResolver`] optimizer rule that runs in the post-optimization - /// phase. In this phase, the rule wraps the physical plan in a [`TransformPlanExec`] with a - /// [`ResolvePlaceholdersRule`] if the plan contains any unresolved placeholders. - pub fn new_post_optimization() -> Self { + /// Creates a new [`ExecutionTransformationApplier`] optimizer rule that runs in the + /// post-optimization phase. + pub fn new_post_optimization(rule: Arc) -> Self { + let name = format!("ExecutionTransformationApplier({})", rule.name()); Self { - phase: PhysicalExprResolverPhase::Post, + phase: ExecutionTransformationApplierPhase::Post { rule }, + name: name.into(), } } } -impl Default for PhysicalExprResolver { +impl Default for ExecutionTransformationApplier { fn default() -> Self { Self::new() } } -impl PhysicalOptimizerRule for PhysicalExprResolver { +impl PhysicalOptimizerRule for ExecutionTransformationApplier { fn name(&self) -> &str { - match self.phase { - PhysicalExprResolverPhase::Pre => "PhysicalExprResolver", - PhysicalExprResolverPhase::Post => "PhysicalExprResolver(Post)", - } + &self.name } fn optimize( @@ -91,8 +91,8 @@ impl PhysicalOptimizerRule for PhysicalExprResolver { plan: Arc, _config: &ConfigOptions, ) -> Result> { - match self.phase { - PhysicalExprResolverPhase::Pre => plan + match &self.phase { + ExecutionTransformationApplierPhase::Pre => plan .transform_up(|plan| { if let Some(plan) = plan.as_any().downcast_ref::() { @@ -102,26 +102,22 @@ impl PhysicalOptimizerRule for PhysicalExprResolver { } }) .map(|t| t.data), - PhysicalExprResolverPhase::Post => { + ExecutionTransformationApplierPhase::Post { rule } => { if let Some(transformer) = plan.as_any().downcast_ref::() { - let resolves_placeholders = - transformer.has_rule::(); - - if resolves_placeholders { + let has_rule = transformer.has_dyn_rule(rule); + if has_rule { + // Rule is already applied. Ok(plan) } else { transformer - .add_rule(Box::new(ResolvePlaceholdersRule::new())) + .add_rule(Arc::clone(rule)) .map(|r| Arc::new(r) as Arc<_>) } } else { - let transformer = TransformPlanExec::try_new( - plan, - vec![Box::new(ResolvePlaceholdersRule::new())], - )?; - + let transformer = + TransformPlanExec::try_new(plan, vec![Arc::clone(rule)])?; if transformer.plans_to_transform() > 0 { Ok(Arc::new(transformer)) } else { diff --git a/datafusion/physical-optimizer/src/lib.rs b/datafusion/physical-optimizer/src/lib.rs index 195b77c9b4b70..a5b22befb9e9e 100644 --- a/datafusion/physical-optimizer/src/lib.rs +++ b/datafusion/physical-optimizer/src/lib.rs @@ -30,6 +30,7 @@ pub mod combine_partial_final_agg; pub mod enforce_distribution; pub mod enforce_sorting; pub mod ensure_coop; +pub mod exec_transform_apply; pub mod filter_pushdown; pub mod join_selection; pub mod limit_pushdown; @@ -37,7 +38,6 @@ pub mod limit_pushdown_past_window; pub mod limited_distinct_aggregation; pub mod optimizer; pub mod output_requirements; -pub mod physical_expr_resolver; pub mod projection_pushdown; pub use datafusion_pruning as pruning; pub mod pushdown_sort; diff --git a/datafusion/physical-optimizer/src/optimizer.rs b/datafusion/physical-optimizer/src/optimizer.rs index f8a10e3a2727e..665bfa62d9c2d 100644 --- a/datafusion/physical-optimizer/src/optimizer.rs +++ b/datafusion/physical-optimizer/src/optimizer.rs @@ -25,12 +25,12 @@ use crate::combine_partial_final_agg::CombinePartialFinalAggregate; use crate::enforce_distribution::EnforceDistribution; use crate::enforce_sorting::EnforceSorting; use crate::ensure_coop::EnsureCooperative; +use crate::exec_transform_apply::ExecutionTransformationApplier; use crate::filter_pushdown::FilterPushdown; use crate::join_selection::JoinSelection; use crate::limit_pushdown::LimitPushdown; use crate::limited_distinct_aggregation::LimitedDistinctAggregation; use crate::output_requirements::OutputRequirements; -use crate::physical_expr_resolver::PhysicalExprResolver; use crate::projection_pushdown::ProjectionPushdown; use crate::sanity_checker::SanityCheckPlan; use crate::topk_aggregation::TopKAggregation; @@ -41,6 +41,7 @@ use crate::pushdown_sort::PushdownSort; use datafusion_common::Result; use datafusion_common::config::ConfigOptions; use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::plan_transformer::ResolvePlaceholdersRule; /// `PhysicalOptimizerRule` transforms one ['ExecutionPlan'] into another which /// computes the same results, but in a potentially more efficient way. @@ -88,7 +89,7 @@ impl PhysicalOptimizer { // this information is not lost across different rules during optimization. Arc::new(OutputRequirements::new_add_mode()), // This rule removes all existing `TransformPlanExec` nodes from the plan tree. - Arc::new(PhysicalExprResolver::new()), + Arc::new(ExecutionTransformationApplier::new()), Arc::new(AggregateStatistics::new()), // Statistics-based join selection will change the Auto mode to a real join implementation, // like collect left, or hash join, or future sort merge join, which will influence the @@ -151,7 +152,9 @@ impl PhysicalOptimizer { // This rule prepares the physical plan for placeholder resolution by wrapping it in a // `TransformPlanExec` with a `ResolvePlaceholdersRule` if it contains any unresolved // placeholders. - Arc::new(PhysicalExprResolver::new_post_optimization()), + Arc::new(ExecutionTransformationApplier::new_post_optimization( + Arc::new(ResolvePlaceholdersRule::new()), + )), // This FilterPushdown handles dynamic filters that may have references to the source ExecutionPlan. // Therefore it should be run at the end of the optimization process since any changes to the plan may break the dynamic filter's references. // See `FilterPushdownPhase` for more details. diff --git a/datafusion/physical-plan/benches/plan_transformer.rs b/datafusion/physical-plan/benches/plan_transformer.rs index 42816bd92d5cc..03b54b872176b 100644 --- a/datafusion/physical-plan/benches/plan_transformer.rs +++ b/datafusion/physical-plan/benches/plan_transformer.rs @@ -46,11 +46,7 @@ impl ExecutionTransformationRule for ResetAllRule { self } - fn clone_box(&self) -> Box { - Box::new(self.clone()) - } - - fn matches(&mut self, _node: &Arc) -> Result { + fn matches(&self, _node: &Arc) -> Result { Ok(true) } @@ -77,11 +73,7 @@ impl ExecutionTransformationRule for ResetByNameRule { self } - fn clone_box(&self) -> Box { - Box::new(self.clone()) - } - - fn matches(&mut self, node: &Arc) -> Result { + fn matches(&self, node: &Arc) -> Result { Ok(node.name() == self.node_name) } @@ -166,7 +158,7 @@ fn benchmark_with_transformer_exec( group: &mut BenchmarkGroup<'_, WallTime>, batch_label: &str, plan: &Arc, - rules: &[Box], + rules: &[Arc], batch_size: BatchSize, ) { let ctx = Arc::new(TaskContext::default()); @@ -182,7 +174,7 @@ fn benchmark_with_transformer_exec( || { ( Arc::clone(plan), - rules.iter().map(|r| r.clone_box()).collect(), + rules.to_vec(), ) }, |(plan, rules)| { @@ -201,8 +193,7 @@ fn benchmark_with_transformer_exec( ), |b| { let plan = Arc::clone(plan); - let rules = rules.iter().map(|r| r.clone_box()).collect(); - let transformer = Arc::new(TransformPlanExec::try_new(plan, rules).unwrap()); + let transformer = Arc::new(TransformPlanExec::try_new(plan, rules.to_vec()).unwrap()); b.iter_batched( || Arc::clone(&transformer), @@ -257,14 +248,14 @@ fn criterion_benchmark(c: &mut Criterion) { for count in rules_count { let reset_all_rules = (0..count) - .map(|_| Box::new(ResetAllRule {}) as Box<_>) + .map(|_| Arc::new(ResetAllRule {}) as _) .collect::>(); let reset_one_rules = (0..count) .map(|_| { - Box::new(ResetByNameRule { + Arc::new(ResetByNameRule { node_name: "EmptyExec".to_string(), - }) as Box<_> + }) as _ }) .collect::>(); diff --git a/datafusion/physical-plan/src/plan_transformer/mod.rs b/datafusion/physical-plan/src/plan_transformer/mod.rs index 9834e3c5d0d84..3237255c4da0f 100644 --- a/datafusion/physical-plan/src/plan_transformer/mod.rs +++ b/datafusion/physical-plan/src/plan_transformer/mod.rs @@ -60,11 +60,8 @@ pub trait ExecutionTransformationRule: Send + Sync + Debug { /// Returns the rule as [`Any`] so that it can be downcast to a specific implementation. fn as_any(&self) -> &dyn Any; - /// Clones this rule. - fn clone_box(&self) -> Box; - /// Checks if the given [`ExecutionPlan`] node matches the criteria for this rule. - fn matches(&mut self, _node: &Arc) -> Result { + fn matches(&self, _node: &Arc) -> Result { Ok(false) } @@ -91,12 +88,12 @@ struct TransformationPlan { /// Helper for building transformation plans for an [`ExecutionPlan`] tree during a single pass. struct TransformationPlanner { cursor: usize, - rules: Vec>, + rules: Vec>, plans: Vec, } impl TransformationPlanner { - fn new(rules: Vec>) -> Self { + fn new(rules: Vec>) -> Self { Self { cursor: 0, rules, @@ -140,14 +137,14 @@ impl<'n> TreeNodeVisitor<'n> for TransformationPlanner { /// rules. struct TransformationApplier<'rules, 'plans, 'ctx> { cursor: usize, - rules: &'rules [Box], + rules: &'rules [Arc], plans: &'plans [TransformationPlan], ctx: &'ctx TaskContext, } impl<'rules, 'plans, 'ctx> TransformationApplier<'rules, 'plans, 'ctx> { fn new( - rules: &'rules [Box], + rules: &'rules [Arc], plans: &'plans [TransformationPlan], ctx: &'ctx TaskContext, ) -> Self { @@ -205,7 +202,7 @@ pub struct TransformPlanExec { /// The input execution plan. input: Arc, /// The transformation rules to apply. - rules: Vec>, + rules: Vec>, /// The pre-calculated transformation plans. plans: Vec, /// Execution metrics. @@ -219,7 +216,7 @@ impl TransformPlanExec { /// initial traversal of the input plan. pub fn try_new( input: Arc, - rules: Vec>, + rules: Vec>, ) -> Result { let mut planner = TransformationPlanner::new(rules); input.visit(&mut planner)?; @@ -252,7 +249,7 @@ impl TransformPlanExec { } /// Returns the transformation rules. - pub fn rules(&self) -> &[Box] { + pub fn rules(&self) -> &[Arc] { &self.rules } @@ -261,10 +258,15 @@ impl TransformPlanExec { self.rules.iter().any(|r| r.as_any().is::()) } + /// Checks if the transformation rules contains a specific rule. + pub fn has_dyn_rule(&self, rule: &Arc) -> bool { + self.rules.iter().any(|r| Arc::ptr_eq(r, rule)) + } + /// Adds a new transformation rule and recalculates transformation plans. pub fn add_rule( &self, - new_rule: Box, + new_rule: Arc, ) -> Result { self.add_rules(vec![new_rule]) } @@ -272,18 +274,14 @@ impl TransformPlanExec { /// Adds new transformation rules and recalculates transformation plans. pub fn add_rules( &self, - new_rules: Vec>, + new_rules: Vec>, ) -> Result { let mut planner = TransformationPlanner::new(new_rules); self.input.visit(&mut planner)?; let new_rules = planner.rules; let new_plans = planner.plans; - let mut current_rules = self - .rules - .iter() - .map(|rule| rule.clone_box()) - .collect::>(); + let mut current_rules = self.rules.clone(); let offset = current_rules.len(); current_rules.extend(new_rules); @@ -394,10 +392,9 @@ impl ExecutionPlan for TransformPlanExec { self: Arc, children: Vec>, ) -> Result> { - let rules = self.rules.iter().map(|r| r.clone_box()).collect(); Ok(Arc::new(TransformPlanExec::try_new( Arc::clone(&children[0]), - rules, + self.rules.clone(), )?)) } @@ -459,11 +456,7 @@ mod tests { self } - fn clone_box(&self) -> Box { - Box::new(self.clone()) - } - - fn matches(&mut self, node: &Arc) -> Result { + fn matches(&self, node: &Arc) -> Result { Ok(node.name() == self.match_name || self.match_name == "all") } @@ -481,12 +474,12 @@ mod tests { let schema = Arc::new(Schema::empty()); let input = Arc::new(EmptyExec::new(schema)); - let resolve_rule = Box::new(ResolvePlaceholdersRule::new()); + let resolve_rule = Arc::new(ResolvePlaceholdersRule::new()); let exec = TransformPlanExec::try_new(input, vec![resolve_rule])?; assert!(exec.has_rule::()); assert!(!exec.has_rule::()); - let exec = exec.add_rule(Box::new(MockRule { + let exec = exec.add_rule(Arc::new(MockRule { name: "mock".to_string(), match_name: "any".to_string(), }))?; @@ -508,22 +501,23 @@ mod tests { // Node 1: CoalescePartitionsExec // Node 2: EmptyExec - let rule_a = Box::new(MockRule { + let rule_a: Arc = Arc::new(MockRule { name: "ruleA".to_string(), match_name: "CoalescePartitionsExec".to_string(), }); - let exec = TransformPlanExec::try_new(Arc::clone(&input), vec![rule_a.clone()])?; + let exec = + TransformPlanExec::try_new(Arc::clone(&input), vec![Arc::clone(&rule_a)])?; assert_eq!(exec.rules.len(), 1); assert_eq!(exec.plans.len(), 1); assert_eq!(exec.plans[0].node_index, 1); assert_eq!(exec.plans[0].rule_indices, vec![0]); - let rule_b = Box::new(MockRule { + let rule_b: Arc = Arc::new(MockRule { name: "ruleB".to_string(), match_name: "GlobalLimitExec".to_string(), }); - let exec = exec.add_rules(vec![rule_b.clone()])?; + let exec = exec.add_rules(vec![Arc::clone(&rule_b)])?; assert_eq!(exec.rules.len(), 2); assert_eq!(exec.plans.len(), 2); @@ -532,11 +526,11 @@ mod tests { assert_eq!(exec.plans[1].node_index, 1); assert_eq!(exec.plans[1].rule_indices, vec![0]); - let rule_c = Box::new(MockRule { + let rule_c: Arc = Arc::new(MockRule { name: "ruleC".to_string(), match_name: "EmptyExec".to_string(), }); - let exec = exec.add_rules(vec![rule_c.clone()])?; + let exec = exec.add_rules(vec![Arc::clone(&rule_c)])?; assert_eq!(exec.rules.len(), 3); assert_eq!(exec.plans.len(), 3); @@ -547,12 +541,12 @@ mod tests { assert_eq!(exec.plans[2].node_index, 2); assert_eq!(exec.plans[2].rule_indices, vec![2]); - let rule_d = Box::new(MockRule { + let rule_d: Arc = Arc::new(MockRule { name: "ruleD".to_string(), match_name: "all".to_string(), }); - let exec = exec.add_rules(vec![rule_d.clone()])?; + let exec = exec.add_rules(vec![Arc::clone(&rule_d)])?; let check_full_plan = |exec: TransformPlanExec| { assert_eq!(exec.rules.len(), 4); assert_eq!(exec.plans.len(), 3); @@ -569,10 +563,10 @@ mod tests { let exec = TransformPlanExec::try_new( Arc::clone(&input), vec![ - rule_a.clone(), - rule_b.clone(), - rule_c.clone(), - rule_d.clone(), + Arc::clone(&rule_a), + Arc::clone(&rule_b), + Arc::clone(&rule_c), + Arc::clone(&rule_d), ], )?; @@ -610,7 +604,7 @@ mod tests { let projection = ProjectionExec::try_new(vec![projection_expr], row)?; let transformer = TransformPlanExec::try_new( Arc::new(projection), - vec![Box::new(ResolvePlaceholdersRule::new())], + vec![Arc::new(ResolvePlaceholdersRule::new())], )?; let param_values = ParamValues::List(vec![ScalarValue::Int32(Some(20)).into()]); diff --git a/datafusion/physical-plan/src/plan_transformer/resolve_placeholders.rs b/datafusion/physical-plan/src/plan_transformer/resolve_placeholders.rs index 115e9cc70ff4e..cefb57420dda3 100644 --- a/datafusion/physical-plan/src/plan_transformer/resolve_placeholders.rs +++ b/datafusion/physical-plan/src/plan_transformer/resolve_placeholders.rs @@ -56,11 +56,7 @@ impl ExecutionTransformationRule for ResolvePlaceholdersRule { self } - fn clone_box(&self) -> Box { - Box::new(Self {}) - } - - fn matches(&mut self, node: &Arc) -> Result { + fn matches(&self, node: &Arc) -> Result { let Some(exprs) = node.physical_expressions() else { return Ok(false); }; diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 4d0b3645d5025..500145058be61 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -2231,13 +2231,13 @@ impl protobuf::PhysicalPlanNode { let input: Arc = into_physical_plan(&transform_plan.input, ctx, codec, proto_converter)?; - let mut rules: Vec> = + let mut rules: Vec> = Vec::with_capacity(transform_plan.rules.len()); for rule in transform_plan.rules.iter() { match &rule.rule_type { Some(RuleType::ResolvePlaceholders(_)) => { - rules.push(Box::new(ResolvePlaceholdersRule::new())) + rules.push(Arc::new(ResolvePlaceholdersRule::new())) } Some(RuleType::Extension(ext)) => { rules.push(codec.try_decode_transformation_rule(ext)?) @@ -3799,7 +3799,7 @@ pub trait PhysicalExtensionCodec: Debug + Send + Sync { fn try_decode_transformation_rule( &self, _buf: &[u8], - ) -> Result> { + ) -> Result> { not_impl_err!("PhysicalExtensionCodec is not provided") } @@ -4248,7 +4248,7 @@ impl PhysicalExtensionCodec for ComposedPhysicalExtensionCodec { fn try_decode_transformation_rule( &self, buf: &[u8], - ) -> Result> { + ) -> Result> { self.decode_protobuf(buf, |codec, data| { codec.try_decode_transformation_rule(data) }) diff --git a/datafusion/sqllogictest/test_files/explain.slt b/datafusion/sqllogictest/test_files/explain.slt index 0f911f77ea2e0..e23d8db2cf8d3 100644 --- a/datafusion/sqllogictest/test_files/explain.slt +++ b/datafusion/sqllogictest/test_files/explain.slt @@ -227,7 +227,7 @@ initial_physical_plan_with_schema DataSourceExec: file_groups={1 group: [[WORKSP physical_plan after OutputRequirements 01)OutputRequirementExec: order_by=[], dist_by=Unspecified 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], file_type=csv, has_header=true -physical_plan after PhysicalExprResolver SAME TEXT AS ABOVE +physical_plan after ExecutionTransformationApplier SAME TEXT AS ABOVE physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE @@ -244,7 +244,7 @@ physical_plan after LimitPushdown SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after PushdownSort SAME TEXT AS ABOVE physical_plan after EnsureCooperative SAME TEXT AS ABOVE -physical_plan after PhysicalExprResolver(Post) SAME TEXT AS ABOVE +physical_plan after ExecutionTransformationApplier(ResolvePlaceholders) SAME TEXT AS ABOVE physical_plan after FilterPushdown(Post) SAME TEXT AS ABOVE physical_plan after SanityCheckPlan SAME TEXT AS ABOVE physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], file_type=csv, has_header=true @@ -307,7 +307,7 @@ physical_plan after OutputRequirements 01)OutputRequirementExec: order_by=[], dist_by=Unspecified, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]: ScanBytes=Exact(32)),(Col[1]: ScanBytes=Inexact(24)),(Col[2]: ScanBytes=Exact(32)),(Col[3]: ScanBytes=Exact(32)),(Col[4]: ScanBytes=Exact(32)),(Col[5]: ScanBytes=Exact(64)),(Col[6]: ScanBytes=Exact(32)),(Col[7]: ScanBytes=Exact(64)),(Col[8]: ScanBytes=Inexact(88)),(Col[9]: ScanBytes=Inexact(49)),(Col[10]: ScanBytes=Exact(64))]] 02)--GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]: ScanBytes=Exact(32)),(Col[1]: ScanBytes=Inexact(24)),(Col[2]: ScanBytes=Exact(32)),(Col[3]: ScanBytes=Exact(32)),(Col[4]: ScanBytes=Exact(32)),(Col[5]: ScanBytes=Exact(64)),(Col[6]: ScanBytes=Exact(32)),(Col[7]: ScanBytes=Exact(64)),(Col[8]: ScanBytes=Inexact(88)),(Col[9]: ScanBytes=Inexact(49)),(Col[10]: ScanBytes=Exact(64))]] 03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]: ScanBytes=Exact(32)),(Col[1]: ScanBytes=Inexact(24)),(Col[2]: ScanBytes=Exact(32)),(Col[3]: ScanBytes=Exact(32)),(Col[4]: ScanBytes=Exact(32)),(Col[5]: ScanBytes=Exact(64)),(Col[6]: ScanBytes=Exact(32)),(Col[7]: ScanBytes=Exact(64)),(Col[8]: ScanBytes=Inexact(88)),(Col[9]: ScanBytes=Inexact(49)),(Col[10]: ScanBytes=Exact(64))]] -physical_plan after PhysicalExprResolver SAME TEXT AS ABOVE +physical_plan after ExecutionTransformationApplier SAME TEXT AS ABOVE physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE @@ -326,7 +326,7 @@ physical_plan after LimitPushdown DataSourceExec: file_groups={1 group: [[WORKSP physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after PushdownSort SAME TEXT AS ABOVE physical_plan after EnsureCooperative SAME TEXT AS ABOVE -physical_plan after PhysicalExprResolver(Post) SAME TEXT AS ABOVE +physical_plan after ExecutionTransformationApplier(ResolvePlaceholders) SAME TEXT AS ABOVE physical_plan after FilterPushdown(Post) SAME TEXT AS ABOVE physical_plan after SanityCheckPlan SAME TEXT AS ABOVE physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]: ScanBytes=Exact(32)),(Col[1]: ScanBytes=Inexact(24)),(Col[2]: ScanBytes=Exact(32)),(Col[3]: ScanBytes=Exact(32)),(Col[4]: ScanBytes=Exact(32)),(Col[5]: ScanBytes=Exact(64)),(Col[6]: ScanBytes=Exact(32)),(Col[7]: ScanBytes=Exact(64)),(Col[8]: ScanBytes=Inexact(88)),(Col[9]: ScanBytes=Inexact(49)),(Col[10]: ScanBytes=Exact(64))]] @@ -353,7 +353,7 @@ physical_plan after OutputRequirements 01)OutputRequirementExec: order_by=[], dist_by=Unspecified 02)--GlobalLimitExec: skip=0, fetch=10 03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet -physical_plan after PhysicalExprResolver SAME TEXT AS ABOVE +physical_plan after ExecutionTransformationApplier SAME TEXT AS ABOVE physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE @@ -372,7 +372,7 @@ physical_plan after LimitPushdown DataSourceExec: file_groups={1 group: [[WORKSP physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after PushdownSort SAME TEXT AS ABOVE physical_plan after EnsureCooperative SAME TEXT AS ABOVE -physical_plan after PhysicalExprResolver(Post) SAME TEXT AS ABOVE +physical_plan after ExecutionTransformationApplier(ResolvePlaceholders) SAME TEXT AS ABOVE physical_plan after FilterPushdown(Post) SAME TEXT AS ABOVE physical_plan after SanityCheckPlan SAME TEXT AS ABOVE physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet @@ -431,7 +431,7 @@ query TT explain select a from t1 where exists (select count(*) from t2); ---- logical_plan -01)LeftSemi Join: +01)LeftSemi Join: 02)--TableScan: t1 projection=[a] 03)--SubqueryAlias: __correlated_sq_1 04)----EmptyRelation: rows=1 @@ -594,7 +594,7 @@ initial_physical_plan_with_schema DataSourceExec: file_groups={1 group: [[WORKSP physical_plan after OutputRequirements 01)OutputRequirementExec: order_by=[], dist_by=Unspecified 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], file_type=csv, has_header=true -physical_plan after PhysicalExprResolver SAME TEXT AS ABOVE +physical_plan after ExecutionTransformationApplier SAME TEXT AS ABOVE physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE @@ -611,7 +611,7 @@ physical_plan after LimitPushdown SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after PushdownSort SAME TEXT AS ABOVE physical_plan after EnsureCooperative SAME TEXT AS ABOVE -physical_plan after PhysicalExprResolver(Post) SAME TEXT AS ABOVE +physical_plan after ExecutionTransformationApplier(ResolvePlaceholders) SAME TEXT AS ABOVE physical_plan after FilterPushdown(Post) SAME TEXT AS ABOVE physical_plan after SanityCheckPlan SAME TEXT AS ABOVE physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], file_type=csv, has_header=true From f9defa2f23c57c7da6240b581bb41089b00fe084 Mon Sep 17 00:00:00 2001 From: Albert Skalt Date: Thu, 12 Feb 2026 09:39:36 +0300 Subject: [PATCH 2/6] refactor(physical plan): simplify placeholder resolving flow Move it out from the separate physical plan into a function that prepares plan for execution. --- datafusion/core/benches/reset_plan_states.rs | 39 +- .../exec_transform_apply.rs | 283 -------- .../core/tests/physical_optimizer/mod.rs | 2 - .../tests/physical_optimizer/test_utils.rs | 12 - datafusion/core/tests/sql/select.rs | 28 +- datafusion/datasource/src/values.rs | 58 +- datafusion/execution/src/task.rs | 19 +- .../physical-expr/src/expressions/mod.rs | 4 +- .../src/expressions/placeholder.rs | 33 +- .../src/exec_transform_apply.rs | 134 ---- datafusion/physical-optimizer/src/lib.rs | 1 - .../physical-optimizer/src/optimizer.rs | 10 - datafusion/physical-plan/Cargo.toml | 4 - .../physical-plan/benches/plan_transformer.rs | 304 --------- .../physical-plan/src/execution_plan.rs | 50 +- datafusion/physical-plan/src/lib.rs | 1 - .../physical-plan/src/plan_transformer/mod.rs | 623 ------------------ .../plan_transformer/resolve_placeholders.rs | 122 ---- datafusion/proto/proto/datafusion.proto | 17 +- datafusion/proto/src/generated/pbjson.rs | 304 --------- datafusion/proto/src/generated/prost.rs | 28 +- datafusion/proto/src/physical_plan/mod.rs | 120 ---- .../sqllogictest/test_files/explain.slt | 8 - .../sqllogictest/test_files/placeholders.slt | 239 +++---- 24 files changed, 266 insertions(+), 2177 deletions(-) delete mode 100644 datafusion/core/tests/physical_optimizer/exec_transform_apply.rs delete mode 100644 datafusion/physical-optimizer/src/exec_transform_apply.rs delete mode 100644 datafusion/physical-plan/benches/plan_transformer.rs delete mode 100644 datafusion/physical-plan/src/plan_transformer/mod.rs delete mode 100644 datafusion/physical-plan/src/plan_transformer/resolve_placeholders.rs diff --git a/datafusion/core/benches/reset_plan_states.rs b/datafusion/core/benches/reset_plan_states.rs index f2f81f755b96e..0988702e6f950 100644 --- a/datafusion/core/benches/reset_plan_states.rs +++ b/datafusion/core/benches/reset_plan_states.rs @@ -23,6 +23,7 @@ use datafusion::prelude::SessionContext; use datafusion_catalog::MemTable; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::displayable; +use datafusion_physical_plan::execution_plan::prepare_execution; use datafusion_physical_plan::execution_plan::reset_plan_states; use tokio::runtime::Runtime; @@ -194,5 +195,41 @@ fn bench_reset_plan_states(c: &mut Criterion) { c.bench_function("query3", bench_query!(query3)); } -criterion_group!(benches, bench_reset_plan_states); +fn run_prepare_execution(b: &mut criterion::Bencher, plan: &Arc) { + b.iter(|| std::hint::black_box(prepare_execution(Arc::clone(plan), None).unwrap())); +} + +/// Benchmark is intended to measure overhead of actions, required to perform +/// making an independent instance of the execution plan to re-execute it with placeholders, +/// avoiding re-planning stage. +fn bench_prepare_execution(c: &mut Criterion) { + let rt = Runtime::new().unwrap(); + let ctx = SessionContext::new(); + ctx.register_table( + "t", + Arc::new(MemTable::try_new(Arc::clone(&SCHEMA), vec![vec![], vec![]]).unwrap()), + ) + .unwrap(); + + ctx.register_table( + "v", + Arc::new(MemTable::try_new(Arc::clone(&SCHEMA), vec![vec![], vec![]]).unwrap()), + ) + .unwrap(); + + macro_rules! bench_query { + ($query_producer: expr) => {{ + let sql = $query_producer(); + let plan = physical_plan(&ctx, &rt, &sql); + log::debug!("plan:\n{}", displayable(plan.as_ref()).indent(true)); + move |b| run_prepare_execution(b, &plan) + }}; + } + + c.bench_function("query1", bench_query!(query1)); + c.bench_function("query2", bench_query!(query2)); + c.bench_function("query3", bench_query!(query3)); +} + +criterion_group!(benches, bench_reset_plan_states, bench_prepare_execution); criterion_main!(benches); diff --git a/datafusion/core/tests/physical_optimizer/exec_transform_apply.rs b/datafusion/core/tests/physical_optimizer/exec_transform_apply.rs deleted file mode 100644 index 22d8b7c7d3881..0000000000000 --- a/datafusion/core/tests/physical_optimizer/exec_transform_apply.rs +++ /dev/null @@ -1,283 +0,0 @@ -// 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. - -//! Integration tests for [`ExecutionTransformationApplier`] optimizer rule. - -use std::{collections::HashMap, sync::Arc}; - -use arrow_schema::{DataType, Field, Schema, SchemaRef}; -use datafusion::config::ConfigOptions; -use datafusion_common::{ParamValues, Result, ScalarValue}; -use datafusion_execution::TaskContext; -use datafusion_expr::Operator; -use datafusion_physical_expr::{ - Partitioning, - expressions::{BinaryExpr, col, lit, placeholder}, -}; -use datafusion_physical_optimizer::{ - PhysicalOptimizerRule, exec_transform_apply::ExecutionTransformationApplier, -}; -use datafusion_physical_plan::{ - ExecutionPlan, - filter::FilterExec, - get_plan_string, - plan_transformer::{ResolvePlaceholdersRule, TransformPlanExec}, - repartition::RepartitionExec, -}; - -use crate::physical_optimizer::test_utils::{ - coalesce_partitions_exec, global_limit_exec, resolve_placeholders_exec, stream_exec, -}; - -fn create_schema() -> SchemaRef { - Arc::new(Schema::new(vec![ - Field::new("c1", DataType::Int32, true), - Field::new("c2", DataType::Int32, true), - Field::new("c3", DataType::Int32, true), - ])) -} - -fn filter_exec( - schema: SchemaRef, - input: Arc, -) -> Result> { - Ok(Arc::new(FilterExec::try_new( - Arc::new(BinaryExpr::new( - col("c3", schema.as_ref()).unwrap(), - Operator::Gt, - lit(0), - )), - input, - )?)) -} - -fn filter_exec_with_placeholders( - schema: SchemaRef, - input: Arc, -) -> Result> { - Ok(Arc::new(FilterExec::try_new( - Arc::new(BinaryExpr::new( - col("c3", schema.as_ref()).unwrap(), - Operator::Gt, - placeholder("$foo", DataType::Int32), - )), - input, - )?)) -} - -fn repartition_exec( - streaming_table: Arc, -) -> Result> { - Ok(Arc::new(RepartitionExec::try_new( - streaming_table, - Partitioning::RoundRobinBatch(8), - )?)) -} - -fn placeholder_resolver() -> ExecutionTransformationApplier { - ExecutionTransformationApplier::new_post_optimization(Arc::new( - ResolvePlaceholdersRule::new(), - )) -} - -#[test] -fn test_noop_if_no_placeholders_found() -> Result<()> { - let schema = create_schema(); - let streaming_table = stream_exec(&schema); - let repartition = repartition_exec(streaming_table)?; - let filter = filter_exec(schema, repartition)?; - let coalesce_partitions = coalesce_partitions_exec(filter); - let plan = global_limit_exec(coalesce_partitions, 0, Some(5)); - - let initial = get_plan_string(&plan); - let expected_initial = [ - "GlobalLimitExec: skip=0, fetch=5", - " CoalescePartitionsExec", - " FilterExec: c3@2 > 0", - " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", - " StreamingTableExec: partition_sizes=1, projection=[c1, c2, c3], infinite_source=true", - ]; - - assert_eq!(initial, expected_initial); - - let after_optimize = placeholder_resolver().optimize(plan, &ConfigOptions::new())?; - - let optimized_plan_string = get_plan_string(&after_optimize); - assert_eq!(initial, optimized_plan_string); - - Ok(()) -} - -#[test] -fn test_wrap_plan_with_transformer() -> Result<()> { - let schema = create_schema(); - let streaming_table = stream_exec(&schema); - let repartition = repartition_exec(streaming_table)?; - let filter = filter_exec_with_placeholders(schema, repartition)?; - let coalesce_partitions = coalesce_partitions_exec(filter); - let plan = global_limit_exec(coalesce_partitions, 0, Some(5)); - - let initial = get_plan_string(&plan); - let expected_initial = [ - "GlobalLimitExec: skip=0, fetch=5", - " CoalescePartitionsExec", - " FilterExec: c3@2 > $foo", - " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", - " StreamingTableExec: partition_sizes=1, projection=[c1, c2, c3], infinite_source=true", - ]; - - assert_eq!(initial, expected_initial); - - let after_optimize = placeholder_resolver().optimize(plan, &ConfigOptions::new())?; - - let expected_optimized = [ - "TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1]", - " GlobalLimitExec: skip=0, fetch=5", - " CoalescePartitionsExec", - " FilterExec: c3@2 > $foo", - " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", - " StreamingTableExec: partition_sizes=1, projection=[c1, c2, c3], infinite_source=true", - ]; - - let optimized_plan_string = get_plan_string(&after_optimize); - assert_eq!(optimized_plan_string, expected_optimized); - - let transformer = after_optimize - .as_ref() - .as_any() - .downcast_ref::() - .expect("should downcast"); - - let param_values = ParamValues::Map(HashMap::from_iter([( - "foo".to_string(), - ScalarValue::Int32(Some(100)).into(), - )])); - - let ctx = Arc::new(TaskContext::default().with_param_values(param_values)); - let resolved_plan = transformer.transform(&ctx)?; - let resolved_plan_string = get_plan_string(&resolved_plan); - let expected_resolved = [ - "GlobalLimitExec: skip=0, fetch=5", - " CoalescePartitionsExec", - " FilterExec: c3@2 > 100", - " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", - " StreamingTableExec: partition_sizes=1, projection=[c1, c2, c3], infinite_source=true", - ]; - - assert_eq!(resolved_plan_string, expected_resolved); - - Ok(()) -} - -#[test] -fn test_remove_useless_transformers() -> Result<()> { - let schema = create_schema(); - let streaming_table = stream_exec(&schema); - let repartition = repartition_exec(streaming_table)?; - let filter = filter_exec(schema, repartition)?; - let transformer = resolve_placeholders_exec(filter); - let coalesce_partitions = coalesce_partitions_exec(transformer); - let global_limit = global_limit_exec(coalesce_partitions, 0, Some(5)); - let plan = resolve_placeholders_exec(global_limit); - - let initial = get_plan_string(&plan); - let expected_initial = [ - "TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=0]", - " GlobalLimitExec: skip=0, fetch=5", - " CoalescePartitionsExec", - " TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=0]", - " FilterExec: c3@2 > 0", - " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", - " StreamingTableExec: partition_sizes=1, projection=[c1, c2, c3], infinite_source=true", - ]; - - assert_eq!(initial, expected_initial); - - let after_optimize = - ExecutionTransformationApplier::new().optimize(plan, &ConfigOptions::new())?; - - let expected_optimized = [ - "GlobalLimitExec: skip=0, fetch=5", - " CoalescePartitionsExec", - " FilterExec: c3@2 > 0", - " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", - " StreamingTableExec: partition_sizes=1, projection=[c1, c2, c3], infinite_source=true", - ]; - - let optimized_plan_string = get_plan_string(&after_optimize); - assert_eq!(optimized_plan_string, expected_optimized); - - Ok(()) -} - -#[test] -fn test_combine_transformers() -> Result<()> { - let schema = create_schema(); - let streaming_table = stream_exec(&schema); - let repartition = repartition_exec(streaming_table)?; - let transformer = resolve_placeholders_exec(repartition); - let filter = filter_exec_with_placeholders(schema, transformer)?; - let transformer = resolve_placeholders_exec(filter); - let coalesce_partitions = coalesce_partitions_exec(transformer); - let global_limit = global_limit_exec(coalesce_partitions, 0, Some(5)); - let plan = resolve_placeholders_exec(global_limit); - - let initial = get_plan_string(&plan); - let expected_initial = [ - "TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=0]", - " GlobalLimitExec: skip=0, fetch=5", - " CoalescePartitionsExec", - " TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1]", - " FilterExec: c3@2 > $foo", - " TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=0]", - " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", - " StreamingTableExec: partition_sizes=1, projection=[c1, c2, c3], infinite_source=true", - ]; - - assert_eq!(initial, expected_initial); - - let after_pre_optimization = - ExecutionTransformationApplier::new().optimize(plan, &ConfigOptions::new())?; - - let expected_optimized = [ - "GlobalLimitExec: skip=0, fetch=5", - " CoalescePartitionsExec", - " FilterExec: c3@2 > $foo", - " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", - " StreamingTableExec: partition_sizes=1, projection=[c1, c2, c3], infinite_source=true", - ]; - - let optimized_plan_string = get_plan_string(&after_pre_optimization); - assert_eq!(optimized_plan_string, expected_optimized); - - let after_post_optimization = - placeholder_resolver().optimize(after_pre_optimization, &ConfigOptions::new())?; - - let expected_optimized = [ - "TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1]", - " GlobalLimitExec: skip=0, fetch=5", - " CoalescePartitionsExec", - " FilterExec: c3@2 > $foo", - " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", - " StreamingTableExec: partition_sizes=1, projection=[c1, c2, c3], infinite_source=true", - ]; - - let optimized_plan_string = get_plan_string(&after_post_optimization); - assert_eq!(optimized_plan_string, expected_optimized); - - Ok(()) -} diff --git a/datafusion/core/tests/physical_optimizer/mod.rs b/datafusion/core/tests/physical_optimizer/mod.rs index 011fe7af0b615..cf179cb727cf1 100644 --- a/datafusion/core/tests/physical_optimizer/mod.rs +++ b/datafusion/core/tests/physical_optimizer/mod.rs @@ -30,8 +30,6 @@ mod join_selection; mod limit_pushdown; mod limited_distinct_aggregation; mod partition_statistics; -#[expect(clippy::needless_pass_by_value)] -mod exec_transform_apply; mod projection_pushdown; mod pushdown_sort; mod replace_with_order_preserving_variants; diff --git a/datafusion/core/tests/physical_optimizer/test_utils.rs b/datafusion/core/tests/physical_optimizer/test_utils.rs index 91f46b070879c..feac8190ffde4 100644 --- a/datafusion/core/tests/physical_optimizer/test_utils.rs +++ b/datafusion/core/tests/physical_optimizer/test_utils.rs @@ -59,9 +59,6 @@ use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::utils::{JoinFilter, JoinOn}; use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode, SortMergeJoinExec}; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; -use datafusion_physical_plan::plan_transformer::{ - ResolvePlaceholdersRule, TransformPlanExec, -}; use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::sorts::sort::SortExec; @@ -396,15 +393,6 @@ pub fn projection_exec( Ok(Arc::new(ProjectionExec::try_new(proj_exprs, input)?)) } -pub fn resolve_placeholders_exec( - input: Arc, -) -> Arc { - Arc::new( - TransformPlanExec::try_new(input, vec![Arc::new(ResolvePlaceholdersRule::new())]) - .unwrap(), - ) -} - /// A test [`ExecutionPlan`] whose requirements can be configured. #[derive(Debug)] pub struct RequirementsTestExec { diff --git a/datafusion/core/tests/sql/select.rs b/datafusion/core/tests/sql/select.rs index 5ba8545766e32..2fe542bc70a2d 100644 --- a/datafusion/core/tests/sql/select.rs +++ b/datafusion/core/tests/sql/select.rs @@ -19,7 +19,7 @@ use std::collections::HashMap; use super::*; use datafusion_common::{ParamValues, ScalarValue, metadata::ScalarAndMetadata}; -use datafusion_execution::TaskContext; +use datafusion_physical_plan::execution_plan::prepare_execution; use insta::assert_snapshot; #[tokio::test] @@ -446,8 +446,8 @@ async fn test_resolve_window_function() -> Result<()> { .await?; let param_values = ParamValues::List(vec![ScalarValue::Int32(Some(100)).into()]); - let task_ctx = Arc::new(TaskContext::from(&ctx).with_param_values(param_values)); - let batches = collect(plan, task_ctx).await?; + let plan = prepare_execution(plan, Some(¶m_values))?; + let batches = collect(plan, ctx.task_ctx()).await?; assert_snapshot!(batches_to_sort_string(&batches), @r" +----+--------------------------------------------------------------------------------------------------------------------------+ @@ -486,8 +486,8 @@ async fn test_resolve_join() -> Result<()> { .await?; let param_values = ParamValues::List(vec![ScalarValue::Int32(Some(8)).into()]); - let task_ctx = Arc::new(TaskContext::from(&ctx).with_param_values(param_values)); - let batches = collect(plan, task_ctx).await?; + let plan = prepare_execution(plan, Some(¶m_values))?; + let batches = collect(plan, ctx.task_ctx()).await?; assert_snapshot!(batches_to_sort_string(&batches), @r" +------+-----+ @@ -512,15 +512,12 @@ async fn test_resolve_cast() -> Result<()> { let param_values = ParamValues::List(vec![ ScalarValue::Utf8(Some("not a number".to_string())).into(), ]); - - let task_ctx = Arc::new(TaskContext::from(&ctx).with_param_values(param_values)); - let result = collect(Arc::clone(&plan), task_ctx).await; - assert!(result.is_err()); + assert!(prepare_execution(Arc::clone(&plan), Some(¶m_values)).is_err()); let param_values = ParamValues::List(vec![ScalarValue::Utf8(Some("200".to_string())).into()]); - let task_ctx = Arc::new(TaskContext::from(&ctx).with_param_values(param_values)); - let batches = collect(plan, task_ctx).await?; + let plan = prepare_execution(plan, Some(¶m_values))?; + let batches = collect(plan, ctx.task_ctx()).await?; assert_snapshot!(batches_to_sort_string(&batches), @r" +-----+ @@ -545,9 +542,8 @@ async fn test_resolve_try_cast() -> Result<()> { let param_values = ParamValues::List(vec![ ScalarValue::Utf8(Some("not a number".to_string())).into(), ]); - - let task_ctx = Arc::new(TaskContext::from(&ctx).with_param_values(param_values)); - let batches = collect(Arc::clone(&plan), task_ctx).await?; + let plan1 = prepare_execution(Arc::clone(&plan), Some(¶m_values))?; + let batches = collect(plan1, ctx.task_ctx()).await?; assert_snapshot!(batches_to_sort_string(&batches), @r" +----+ @@ -559,8 +555,8 @@ async fn test_resolve_try_cast() -> Result<()> { let param_values = ParamValues::List(vec![ScalarValue::Utf8(Some("200".to_string())).into()]); - let task_ctx = Arc::new(TaskContext::from(&ctx).with_param_values(param_values)); - let batches = collect(plan, task_ctx).await?; + let plan2 = prepare_execution(plan, Some(¶m_values))?; + let batches = collect(plan2, ctx.task_ctx()).await?; assert_snapshot!(batches_to_sort_string(&batches), @r" +-----+ diff --git a/datafusion/datasource/src/values.rs b/datafusion/datasource/src/values.rs index 8e2b7b4b80750..75ce506eeb649 100644 --- a/datafusion/datasource/src/values.rs +++ b/datafusion/datasource/src/values.rs @@ -341,8 +341,7 @@ mod tests { use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{BinaryExpr, lit, placeholder}; use datafusion_physical_plan::{ - ExecutionPlan, collect, - plan_transformer::{ResolvePlaceholdersRule, TransformPlanExec}, + ExecutionPlan, collect, execution_plan::prepare_execution, }; #[test] @@ -437,11 +436,13 @@ mod tests { // Should be ValuesSource because of placeholder. assert!(values_exec.data_source().as_any().is::()); - let rules = vec![Arc::new(ResolvePlaceholdersRule::new()) as _]; - let exec = Arc::new(TransformPlanExec::try_new(values_exec, rules)?); - let task_ctx = Arc::new(TaskContext::default().with_param_values( - ParamValues::List(vec![ScalarValue::Int32(Some(10)).into()]), - )); + let exec = prepare_execution( + values_exec, + Some(&ParamValues::List(vec![ + ScalarValue::Int32(Some(10)).into(), + ])), + )?; + let task_ctx = Arc::new(TaskContext::default()); let batch = collect(exec, task_ctx).await?; let expected = [ @@ -469,18 +470,18 @@ mod tests { placeholder("$2", DataType::Int32), ]]; - let values_exec = ValuesSource::try_new_exec(Arc::clone(&schema), data)?; - let rules = vec![Arc::new(ResolvePlaceholdersRule::new()) as _]; - let exec = Arc::new(TransformPlanExec::try_new(values_exec, rules)?) as Arc<_>; - - let task_ctx = Arc::new(TaskContext::default().with_param_values( - ParamValues::List(vec![ + let values_exec = ValuesSource::try_new_exec(Arc::clone(&schema), data)? as _; + let exec = prepare_execution( + Arc::clone(&values_exec), + Some(&ParamValues::List(vec![ ScalarValue::Int32(Some(10)).into(), ScalarValue::Int32(Some(20)).into(), - ]), - )); + ])), + )?; - let batch = collect(Arc::clone(&exec), task_ctx).await?; + let task_ctx = Arc::new(TaskContext::default()); + + let batch = collect(Arc::clone(&exec), Arc::clone(&task_ctx)).await?; let expected = [ "+----+----+", "| a | b |", @@ -490,12 +491,13 @@ mod tests { ]; assert_batches_eq!(expected, &batch); - let task_ctx = Arc::new(TaskContext::default().with_param_values( - ParamValues::List(vec![ + let exec = prepare_execution( + values_exec, + Some(&ParamValues::List(vec![ ScalarValue::Int32(Some(30)).into(), ScalarValue::Int32(Some(40)).into(), - ]), - )); + ])), + )?; let batch = collect(exec, task_ctx).await?; let expected = [ @@ -527,21 +529,21 @@ mod tests { let data: Vec>> = vec![vec![lit(10), placeholder("$foo", DataType::Int32)]]; - let values_exec = ValuesSource::try_new_exec(Arc::clone(&schema), data)?; - let rules = vec![Arc::new(ResolvePlaceholdersRule::new()) as _]; - let exec = Arc::new(TransformPlanExec::try_new(values_exec, rules)?) as Arc<_>; + let values_exec = ValuesSource::try_new_exec(Arc::clone(&schema), data)? as _; let task_ctx = Arc::new(TaskContext::default()); - let result = collect(Arc::clone(&exec), task_ctx).await; + let result = collect(Arc::clone(&values_exec), task_ctx).await; assert!(result.is_err()); - let task_ctx = Arc::new(TaskContext::default().with_param_values( - ParamValues::Map(HashMap::from_iter([( + let exec = prepare_execution( + values_exec, + Some(&ParamValues::Map(HashMap::from_iter([( "foo".to_string(), ScalarValue::Int32(Some(20)).into(), - )])), - )); + )]))), + )?; + let task_ctx = Arc::new(TaskContext::default()); let batch = collect(Arc::clone(&exec), task_ctx).await?; let expected = [ "+----+----+", diff --git a/datafusion/execution/src/task.rs b/datafusion/execution/src/task.rs index 2c191fbc03dd4..38f31cf4629eb 100644 --- a/datafusion/execution/src/task.rs +++ b/datafusion/execution/src/task.rs @@ -19,9 +19,7 @@ use crate::{ config::SessionConfig, memory_pool::MemoryPool, registry::FunctionRegistry, runtime_env::RuntimeEnv, }; -use datafusion_common::{ - ParamValues, Result, internal_datafusion_err, plan_datafusion_err, -}; +use datafusion_common::{Result, internal_datafusion_err, plan_datafusion_err}; use datafusion_expr::planner::ExprPlanner; use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF}; use std::collections::HashSet; @@ -50,8 +48,6 @@ pub struct TaskContext { window_functions: HashMap>, /// Runtime environment associated with this task context runtime: Arc, - /// External query parameters - param_values: Option, } impl Default for TaskContext { @@ -67,7 +63,6 @@ impl Default for TaskContext { aggregate_functions: HashMap::new(), window_functions: HashMap::new(), runtime, - param_values: None, } } } @@ -95,7 +90,6 @@ impl TaskContext { aggregate_functions, window_functions, runtime, - param_values: None, } } @@ -124,11 +118,6 @@ impl TaskContext { Arc::clone(&self.runtime) } - /// Return param values associated with this [`TaskContext`] - pub fn param_values(&self) -> &Option { - &self.param_values - } - pub fn scalar_functions(&self) -> &HashMap> { &self.scalar_functions } @@ -152,12 +141,6 @@ impl TaskContext { self.runtime = runtime; self } - - /// Update the [`ParamValues`] - pub fn with_param_values(mut self, param_values: ParamValues) -> Self { - self.param_values = Some(param_values); - self - } } impl FunctionRegistry for TaskContext { diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs index 34314f1c80d5f..8312d6df1f48c 100644 --- a/datafusion/physical-expr/src/expressions/mod.rs +++ b/datafusion/physical-expr/src/expressions/mod.rs @@ -55,6 +55,8 @@ pub use literal::{Literal, lit}; pub use negative::{NegativeExpr, negative}; pub use no_op::NoOp; pub use not::{NotExpr, not}; -pub use placeholder::{PlaceholderExpr, has_placeholders, placeholder}; +pub use placeholder::{ + PlaceholderExpr, has_placeholders, placeholder, resolve_expr_placeholders, +}; pub use try_cast::{TryCastExpr, try_cast}; pub use unknown_column::UnKnownColumn; diff --git a/datafusion/physical-expr/src/expressions/placeholder.rs b/datafusion/physical-expr/src/expressions/placeholder.rs index 9149517a0a57a..74cb81550ffc9 100644 --- a/datafusion/physical-expr/src/expressions/placeholder.rs +++ b/datafusion/physical-expr/src/expressions/placeholder.rs @@ -28,12 +28,43 @@ use arrow::{ datatypes::{DataType, Field, FieldRef, Schema}, }; use datafusion_common::{ - DataFusionError, Result, exec_datafusion_err, tree_node::TreeNode, + DataFusionError, ParamValues, Result, exec_datafusion_err, exec_err, + tree_node::{Transformed, TreeNode}, }; use datafusion_expr::ColumnarValue; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use std::hash::Hash; +use crate::{expressions::Literal, simplifier::const_evaluator::simplify_const_expr}; + +/// Resolves [`PlaceholderExpr`] in the physical expression using the provided [`ParamValues`]. +pub fn resolve_expr_placeholders( + expr: Arc, + param_values: Option<&ParamValues>, +) -> Result> { + let expr = expr.transform_up(|node| { + let Some(placeholder) = node.as_any().downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + let Some(ref field) = placeholder.field else { + return Ok(Transformed::no(node)); + }; + let Some(param_values) = param_values else { + return exec_err!("value for placeholder {} is not found", placeholder.id); + }; + let scalar = param_values.get_placeholders_with_values(&placeholder.id)?; + let value = scalar.value.cast_to(field.data_type())?; + let literal = Literal::new_with_metadata(value, scalar.metadata); + Ok(Transformed::yes(Arc::new(literal))) + })?; + + if expr.transformed { + simplify_const_expr(expr.data).map(|t| t.data) + } else { + Ok(expr.data) + } +} + /// Physical expression representing a placeholder parameter (e.g., $1, $2, or named parameters) in /// the physical plan. /// diff --git a/datafusion/physical-optimizer/src/exec_transform_apply.rs b/datafusion/physical-optimizer/src/exec_transform_apply.rs deleted file mode 100644 index 5312334ac1f82..0000000000000 --- a/datafusion/physical-optimizer/src/exec_transform_apply.rs +++ /dev/null @@ -1,134 +0,0 @@ -// 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. - -//! [`ExecutionTransformationApplier`] ensures that the required execution transformations -//! are applied to the physical plan. - -use std::{borrow::Cow, sync::Arc}; - -use datafusion_common::{ - Result, - config::ConfigOptions, - tree_node::{Transformed, TreeNode}, -}; -use datafusion_physical_plan::{ - ExecutionPlan, - plan_transformer::{ExecutionTransformationRule, TransformPlanExec}, -}; - -use crate::PhysicalOptimizerRule; - -/// The phase in which the [`ExecutionTransformationApplier`] rule is applied. -#[derive(Debug)] -pub enum ExecutionTransformationApplierPhase { - /// Optimization that happens before most other optimizations. - /// This optimization removes all [`TransformPlanExec`] execution plans from the plan - /// tree. - Pre, - /// Optimization that happens after most other optimizations. - /// This optimization checks if `rule` requires to transform the plan and wraps the plan with - /// [`TransformPlanExec`] if it so, or adds rule to the existing transformation node. - Post { - rule: Arc, - }, -} - -/// Physical optimizer rule that wraps the plan with a certain execution-stage transformation. -#[derive(Debug)] -pub struct ExecutionTransformationApplier { - phase: ExecutionTransformationApplierPhase, - name: Cow<'static, str>, -} - -impl ExecutionTransformationApplier { - /// Creates a new [`ExecutionTransformationApplier`] optimizer rule that runs in the - /// pre-optimization phase. - pub fn new() -> Self { - Self { - phase: ExecutionTransformationApplierPhase::Pre, - name: Cow::Borrowed("ExecutionTransformationApplier"), - } - } - - /// Creates a new [`ExecutionTransformationApplier`] optimizer rule that runs in the - /// post-optimization phase. - pub fn new_post_optimization(rule: Arc) -> Self { - let name = format!("ExecutionTransformationApplier({})", rule.name()); - Self { - phase: ExecutionTransformationApplierPhase::Post { rule }, - name: name.into(), - } - } -} - -impl Default for ExecutionTransformationApplier { - fn default() -> Self { - Self::new() - } -} - -impl PhysicalOptimizerRule for ExecutionTransformationApplier { - fn name(&self) -> &str { - &self.name - } - - fn optimize( - &self, - plan: Arc, - _config: &ConfigOptions, - ) -> Result> { - match &self.phase { - ExecutionTransformationApplierPhase::Pre => plan - .transform_up(|plan| { - if let Some(plan) = plan.as_any().downcast_ref::() - { - Ok(Transformed::yes(Arc::clone(plan.input()))) - } else { - Ok(Transformed::no(plan)) - } - }) - .map(|t| t.data), - ExecutionTransformationApplierPhase::Post { rule } => { - if let Some(transformer) = - plan.as_any().downcast_ref::() - { - let has_rule = transformer.has_dyn_rule(rule); - if has_rule { - // Rule is already applied. - Ok(plan) - } else { - transformer - .add_rule(Arc::clone(rule)) - .map(|r| Arc::new(r) as Arc<_>) - } - } else { - let transformer = - TransformPlanExec::try_new(plan, vec![Arc::clone(rule)])?; - if transformer.plans_to_transform() > 0 { - Ok(Arc::new(transformer)) - } else { - Ok(Arc::clone(transformer.input())) - } - } - } - } - } - - fn schema_check(&self) -> bool { - true - } -} diff --git a/datafusion/physical-optimizer/src/lib.rs b/datafusion/physical-optimizer/src/lib.rs index a5b22befb9e9e..3a0d79ae2d234 100644 --- a/datafusion/physical-optimizer/src/lib.rs +++ b/datafusion/physical-optimizer/src/lib.rs @@ -30,7 +30,6 @@ pub mod combine_partial_final_agg; pub mod enforce_distribution; pub mod enforce_sorting; pub mod ensure_coop; -pub mod exec_transform_apply; pub mod filter_pushdown; pub mod join_selection; pub mod limit_pushdown; diff --git a/datafusion/physical-optimizer/src/optimizer.rs b/datafusion/physical-optimizer/src/optimizer.rs index 665bfa62d9c2d..ff71c9ec64385 100644 --- a/datafusion/physical-optimizer/src/optimizer.rs +++ b/datafusion/physical-optimizer/src/optimizer.rs @@ -25,7 +25,6 @@ use crate::combine_partial_final_agg::CombinePartialFinalAggregate; use crate::enforce_distribution::EnforceDistribution; use crate::enforce_sorting::EnforceSorting; use crate::ensure_coop::EnsureCooperative; -use crate::exec_transform_apply::ExecutionTransformationApplier; use crate::filter_pushdown::FilterPushdown; use crate::join_selection::JoinSelection; use crate::limit_pushdown::LimitPushdown; @@ -41,7 +40,6 @@ use crate::pushdown_sort::PushdownSort; use datafusion_common::Result; use datafusion_common::config::ConfigOptions; use datafusion_physical_plan::ExecutionPlan; -use datafusion_physical_plan::plan_transformer::ResolvePlaceholdersRule; /// `PhysicalOptimizerRule` transforms one ['ExecutionPlan'] into another which /// computes the same results, but in a potentially more efficient way. @@ -88,8 +86,6 @@ impl PhysicalOptimizer { // If there is a output requirement of the query, make sure that // this information is not lost across different rules during optimization. Arc::new(OutputRequirements::new_add_mode()), - // This rule removes all existing `TransformPlanExec` nodes from the plan tree. - Arc::new(ExecutionTransformationApplier::new()), Arc::new(AggregateStatistics::new()), // Statistics-based join selection will change the Auto mode to a real join implementation, // like collect left, or hash join, or future sort merge join, which will influence the @@ -149,12 +145,6 @@ impl PhysicalOptimizer { // PushdownSort: Detect sorts that can be pushed down to data sources. Arc::new(PushdownSort::new()), Arc::new(EnsureCooperative::new()), - // This rule prepares the physical plan for placeholder resolution by wrapping it in a - // `TransformPlanExec` with a `ResolvePlaceholdersRule` if it contains any unresolved - // placeholders. - Arc::new(ExecutionTransformationApplier::new_post_optimization( - Arc::new(ResolvePlaceholdersRule::new()), - )), // This FilterPushdown handles dynamic filters that may have references to the source ExecutionPlan. // Therefore it should be run at the end of the optimization process since any changes to the plan may break the dynamic filter's references. // See `FilterPushdownPhase` for more details. diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 27bd8610b46e3..13f91fd7d4ea2 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -102,7 +102,3 @@ name = "sort_preserving_merge" harness = false name = "aggregate_vectorized" required-features = ["test_utils"] - -[[bench]] -harness = false -name = "plan_transformer" diff --git a/datafusion/physical-plan/benches/plan_transformer.rs b/datafusion/physical-plan/benches/plan_transformer.rs deleted file mode 100644 index 03b54b872176b..0000000000000 --- a/datafusion/physical-plan/benches/plan_transformer.rs +++ /dev/null @@ -1,304 +0,0 @@ -// 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::any::Any; -use std::sync::Arc; - -use arrow_schema::Schema; -use criterion::measurement::WallTime; -use datafusion_common::Result; -use datafusion_common::tree_node::{ - Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter, -}; -use datafusion_execution::TaskContext; -use datafusion_physical_plan::ExecutionPlan; -use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; -use datafusion_physical_plan::empty::EmptyExec; -use datafusion_physical_plan::plan_transformer::{ - ExecutionTransformationRule, TransformPlanExec, -}; - -use criterion::{BatchSize, BenchmarkGroup, Criterion, criterion_group, criterion_main}; - -#[derive(Debug, Clone)] -struct ResetAllRule {} - -impl ExecutionTransformationRule for ResetAllRule { - fn name(&self) -> &str { - "ResetAllRule" - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn matches(&self, _node: &Arc) -> Result { - Ok(true) - } - - fn rewrite( - &self, - node: Arc, - _ctx: &TaskContext, - ) -> Result>> { - node.reset_state().map(Transformed::yes) - } -} - -#[derive(Debug, Clone)] -struct ResetByNameRule { - node_name: String, -} - -impl ExecutionTransformationRule for ResetByNameRule { - fn name(&self) -> &str { - "ResetByNameRule" - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn matches(&self, node: &Arc) -> Result { - Ok(node.name() == self.node_name) - } - - fn rewrite( - &self, - node: Arc, - _ctx: &TaskContext, - ) -> Result>> { - node.reset_state().map(Transformed::yes) - } -} - -struct ResetAllRewriter; - -impl TreeNodeRewriter for ResetAllRewriter { - type Node = Arc; - - fn f_up(&mut self, node: Self::Node) -> Result> { - // This part is needed for handling nested rewriters. - if node.as_any().is::() { - return Ok(Transformed::new(node, false, TreeNodeRecursion::Jump)); - } - node.reset_state().map(Transformed::yes) - } -} - -struct ResetByNameRewriter { - node_name: String, -} - -impl TreeNodeRewriter for ResetByNameRewriter { - type Node = Arc; - - fn f_up(&mut self, node: Self::Node) -> Result> { - // This part is needed for handling nested rewriters. - if node.as_any().is::() { - return Ok(Transformed::new(node, false, TreeNodeRecursion::Jump)); - } - - if node.name() == self.node_name { - node.reset_state().map(Transformed::yes) - } else { - Ok(Transformed::no(node)) - } - } -} - -struct MockRewriterExec { - input: Arc, -} - -impl MockRewriterExec { - fn execute( - &self, - rewriter: &mut impl TreeNodeRewriter>, - ) -> Result> { - let input = Arc::clone(&self.input); - input.rewrite(rewriter).map(|t| t.data) - } -} - -fn create_deep_tree(depth: usize) -> Arc { - let schema = Arc::new(Schema::empty()); - let mut node: Arc = Arc::new(EmptyExec::new(schema)); - for _ in 0..depth { - node = Arc::new(CoalescePartitionsExec::new(node)); - } - node -} - -fn nodes_amount(plan: &Arc) -> usize { - let mut amount = 0; - plan.apply(|_| { - amount += 1; - Ok(TreeNodeRecursion::Continue) - }) - .unwrap(); - amount -} - -fn benchmark_with_transformer_exec( - group: &mut BenchmarkGroup<'_, WallTime>, - batch_label: &str, - plan: &Arc, - rules: &[Arc], - batch_size: BatchSize, -) { - let ctx = Arc::new(TaskContext::default()); - let nodes_amount = nodes_amount(plan); - - group.bench_function( - format!( - "transform_plan_exec_two_phases_{batch_label}_{nodes_amount}_nodes_{}_rule(s)", - rules.len() - ), - |b| { - b.iter_batched( - || { - ( - Arc::clone(plan), - rules.to_vec(), - ) - }, - |(plan, rules)| { - let transformer = TransformPlanExec::try_new(plan, rules).unwrap(); - transformer.transform(&ctx).unwrap(); - }, - batch_size, - ) - }, - ); - - group.bench_function( - format!( - "transform_plan_exec_second_phase_{batch_label}_{nodes_amount}_nodes_{}_rule(s)", - rules.len() - ), - |b| { - let plan = Arc::clone(plan); - let transformer = Arc::new(TransformPlanExec::try_new(plan, rules.to_vec()).unwrap()); - - b.iter_batched( - || Arc::clone(&transformer), - |transformer| { - transformer.transform(&ctx).unwrap(); - }, - batch_size, - ) - }, - ); -} - -fn benchmark_with_tree_node_rewriter( - group: &mut BenchmarkGroup<'_, WallTime>, - batch_label: &str, - plan: &Arc, - mut rewriter: impl TreeNodeRewriter>, - rewrites_amount: usize, - batch_size: BatchSize, -) { - let nodes_amount = nodes_amount(plan); - let mock_rewriter_exec = Arc::new(MockRewriterExec { - input: Arc::clone(plan), - }); - - group.bench_function( - format!( - "tree_node_rewriter_{batch_label}_{nodes_amount}_nodes_{rewrites_amount}_iteration(s)" - ), - |b| { - b.iter_batched( - || Arc::clone(&mock_rewriter_exec), - |plan| { - for _ in 0..rewrites_amount { - plan.execute(&mut rewriter).unwrap(); - } - }, - batch_size, - ) - }, - ); -} - -fn criterion_benchmark(c: &mut Criterion) { - let depths = [5, 30]; - let rules_count = [1, 2]; - let batch_size = BatchSize::SmallInput; - let mut group = c.benchmark_group("plan_transformation"); - - for depth in depths { - let plan = create_deep_tree(depth); - - for count in rules_count { - let reset_all_rules = (0..count) - .map(|_| Arc::new(ResetAllRule {}) as _) - .collect::>(); - - let reset_one_rules = (0..count) - .map(|_| { - Arc::new(ResetByNameRule { - node_name: "EmptyExec".to_string(), - }) as _ - }) - .collect::>(); - - benchmark_with_transformer_exec( - &mut group, - "reset_all", - &plan, - &reset_all_rules, - batch_size, - ); - - benchmark_with_tree_node_rewriter( - &mut group, - "reset_all", - &plan, - ResetAllRewriter, - count, - batch_size, - ); - - benchmark_with_transformer_exec( - &mut group, - "reset_one", - &plan, - &reset_one_rules, - batch_size, - ); - - benchmark_with_tree_node_rewriter( - &mut group, - "reset_one", - &plan, - ResetByNameRewriter { - node_name: "EmptyExec".to_string(), - }, - count, - batch_size, - ); - } - } - - group.finish(); -} - -criterion_group!(benches, criterion_benchmark); -criterion_main!(benches); diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index abfa1df828224..1c16d1e4b1864 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -31,6 +31,7 @@ pub use datafusion_common::utils::project_schema; pub use datafusion_common::{ColumnStatistics, Statistics, internal_err}; pub use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream}; pub use datafusion_expr::{Accumulator, ColumnarValue}; +use datafusion_physical_expr::expressions::resolve_expr_placeholders; pub use datafusion_physical_expr::window::WindowExpr; pub use datafusion_physical_expr::{ Distribution, Partitioning, PhysicalExpr, expressions, @@ -50,7 +51,7 @@ use arrow::array::{Array, RecordBatch}; use arrow::datatypes::SchemaRef; use datafusion_common::config::ConfigOptions; use datafusion_common::{ - Constraints, DataFusionError, Result, assert_eq_or_internal_err, + Constraints, DataFusionError, ParamValues, Result, assert_eq_or_internal_err, assert_or_internal_err, exec_err, }; use datafusion_common_runtime::JoinSet; @@ -1556,6 +1557,53 @@ pub fn reset_plan_states(plan: Arc) -> Result, + param_values: Option<&ParamValues>, +) -> Result> { + plan.transform_up(|plan| { + let plan = if let Some(iter) = plan.physical_expressions() { + let mut has_placeholders = false; + let exprs = iter + .map(|expr| { + let resolved_expr = + resolve_expr_placeholders(Arc::clone(&expr), param_values)?; + has_placeholders |= !Arc::ptr_eq(&expr, &resolved_expr); + Ok(resolved_expr) + }) + .collect::>>()?; + if !has_placeholders { + Arc::clone(&plan).reset_state()? + } else { + // `with_physical_expressions` resets plan state. + let Some(plan) = + plan.with_physical_expressions(ReplacePhysicalExpr { exprs })? + else { + return exec_err!( + "plan {} does not support expression substitution", + plan.name() + ); + }; + plan + } + } else { + plan.reset_state()? + }; + Ok(Transformed::yes(plan)) + }) + .map(|tnr| tnr.data) +} + /// Utility function yielding a string representation of the given [`ExecutionPlan`]. pub fn get_plan_string(plan: &Arc) -> Vec { let formatted = displayable(plan.as_ref()).indent(true).to_string(); diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index da28951e6d897..6467d7a2e389d 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -81,7 +81,6 @@ pub mod limit; pub mod memory; pub mod metrics; pub mod placeholder_row; -pub mod plan_transformer; pub mod projection; pub mod recursive_query; pub mod repartition; diff --git a/datafusion/physical-plan/src/plan_transformer/mod.rs b/datafusion/physical-plan/src/plan_transformer/mod.rs deleted file mode 100644 index 3237255c4da0f..0000000000000 --- a/datafusion/physical-plan/src/plan_transformer/mod.rs +++ /dev/null @@ -1,623 +0,0 @@ -// 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. - -//! This module provides a mechanism for two-pass [`ExecutionPlan`] tree transformations. -//! -//! In the first pass, the tree is visited (top-down) to identify nodes that require transformation -//! based on certain criteria. -//! -//! In the second pass, the stored plans are applied to the tree (top-down) using a -//! [`TreeNodeRewriter`]. -//! -//! This approach is beneficial because it allows making multiple transformations during a single -//! pass and caches node indices, ensuring that only nodes matching the criteria are transformed, -//! which can improve performance. - -mod resolve_placeholders; - -pub use resolve_placeholders::ResolvePlaceholdersRule; - -use std::any::Any; -use std::fmt::{self, Debug}; -use std::sync::Arc; - -use datafusion_common::{ - Result, - tree_node::{ - Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter, TreeNodeVisitor, - }, -}; -use datafusion_execution::{SendableRecordBatchStream, TaskContext}; -use itertools::Itertools; - -use crate::metrics::{ExecutionPlanMetricsSet, MetricBuilder, MetricsSet}; -use crate::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, Statistics}; - -/// A rule for transforming an [`ExecutionPlan`] node. -/// -/// This trait is used in the two-pass tree transformation process. Both passes use a top-down -/// traversal. In the first pass, the [`ExecutionTransformationRule::matches`] method is used to -/// identify nodes that require transformation. In the second pass, the -/// [`ExecutionTransformationRule::rewrite`] method is used to apply the transformation. -pub trait ExecutionTransformationRule: Send + Sync + Debug { - /// Returns the rule name. - fn name(&self) -> &str; - - /// Returns the rule as [`Any`] so that it can be downcast to a specific implementation. - fn as_any(&self) -> &dyn Any; - - /// Checks if the given [`ExecutionPlan`] node matches the criteria for this rule. - fn matches(&self, _node: &Arc) -> Result { - Ok(false) - } - - /// Transforms the given [`ExecutionPlan`] node. - fn rewrite( - &self, - node: Arc, - _ctx: &TaskContext, - ) -> Result>> { - Ok(Transformed::no(node)) - } -} - -/// Stores the transformation operations to be applied to an [`ExecutionPlan`] node at a specific -/// index. -#[derive(Debug, Clone)] -struct TransformationPlan { - /// The index of the node in the tree traversal. - pub node_index: usize, - /// The list of transformation rule indices to apply. - pub rule_indices: Vec, -} - -/// Helper for building transformation plans for an [`ExecutionPlan`] tree during a single pass. -struct TransformationPlanner { - cursor: usize, - rules: Vec>, - plans: Vec, -} - -impl TransformationPlanner { - fn new(rules: Vec>) -> Self { - Self { - cursor: 0, - rules, - plans: Vec::new(), - } - } -} - -impl<'n> TreeNodeVisitor<'n> for TransformationPlanner { - type Node = Arc; - - fn f_down(&mut self, node: &'n Self::Node) -> Result { - if node.as_any().is::() { - return Ok(TreeNodeRecursion::Jump); - } - - let index = self.cursor; - self.cursor += 1; - - let mut rule_indices = Vec::new(); - for (rule_index, rule) in self.rules.iter_mut().enumerate() { - if rule.matches(node)? { - rule_indices.push(rule_index); - } - } - - if !rule_indices.is_empty() { - self.plans.push(TransformationPlan { - node_index: index, - rule_indices, - }); - } - - Ok(TreeNodeRecursion::Continue) - } -} - -/// Helper for applying transformation plans to an [`ExecutionPlan`] tree during a single pass. -/// -/// This applier uses a [`TaskContext`] to provide necessary information for the transformation -/// rules. -struct TransformationApplier<'rules, 'plans, 'ctx> { - cursor: usize, - rules: &'rules [Arc], - plans: &'plans [TransformationPlan], - ctx: &'ctx TaskContext, -} - -impl<'rules, 'plans, 'ctx> TransformationApplier<'rules, 'plans, 'ctx> { - fn new( - rules: &'rules [Arc], - plans: &'plans [TransformationPlan], - ctx: &'ctx TaskContext, - ) -> Self { - Self { - cursor: 0, - rules, - plans, - ctx, - } - } -} - -impl<'rules, 'plans, 'ctx> TreeNodeRewriter - for TransformationApplier<'rules, 'plans, 'ctx> -{ - type Node = Arc; - - fn f_down(&mut self, mut node: Self::Node) -> Result> { - let Some(plan) = self.plans.first() else { - return Ok(Transformed::new(node, false, TreeNodeRecursion::Stop)); - }; - - if node.as_any().is::() { - return Ok(Transformed::new(node, false, TreeNodeRecursion::Jump)); - } - - let index = self.cursor; - self.cursor += 1; - - if index != plan.node_index { - return Ok(Transformed::no(node)); - } - - self.plans = &self.plans[1..]; - - let mut transformed = false; - for rule_index in plan.rule_indices.iter() { - let rule = &self.rules[*rule_index]; - let transform = rule.rewrite(node, self.ctx)?; - node = transform.data; - transformed |= transform.transformed; - } - - Ok(Transformed::new( - node, - transformed, - TreeNodeRecursion::Continue, - )) - } -} - -/// An [`ExecutionPlan`] that applies transformation rules during execution. -#[derive(Debug)] -pub struct TransformPlanExec { - /// The input execution plan. - input: Arc, - /// The transformation rules to apply. - rules: Vec>, - /// The pre-calculated transformation plans. - plans: Vec, - /// Execution metrics. - metrics: ExecutionPlanMetricsSet, -} - -impl TransformPlanExec { - /// Create a new [TransformPlanExec]. - /// - /// This method returns an error if any of the transformation rules return an error during the - /// initial traversal of the input plan. - pub fn try_new( - input: Arc, - rules: Vec>, - ) -> Result { - let mut planner = TransformationPlanner::new(rules); - input.visit(&mut planner)?; - - Ok(Self { - input, - rules: planner.rules, - plans: planner.plans, - metrics: ExecutionPlanMetricsSet::new(), - }) - } - - /// Returns the number of nodes that match at least one transformation rule. - pub fn plans_to_transform(&self) -> usize { - self.plans.len() - } - - pub fn transform( - &self, - context: &Arc, - ) -> Result> { - let mut applier = TransformationApplier::new(&self.rules, &self.plans, context); - let input = Arc::clone(&self.input); - input.rewrite(&mut applier).map(|t| t.data) - } - - /// Returns the input plan. - pub fn input(&self) -> &Arc { - &self.input - } - - /// Returns the transformation rules. - pub fn rules(&self) -> &[Arc] { - &self.rules - } - - /// Checks if the transformation rules contains a rule of a specific type. - pub fn has_rule(&self) -> bool { - self.rules.iter().any(|r| r.as_any().is::()) - } - - /// Checks if the transformation rules contains a specific rule. - pub fn has_dyn_rule(&self, rule: &Arc) -> bool { - self.rules.iter().any(|r| Arc::ptr_eq(r, rule)) - } - - /// Adds a new transformation rule and recalculates transformation plans. - pub fn add_rule( - &self, - new_rule: Arc, - ) -> Result { - self.add_rules(vec![new_rule]) - } - - /// Adds new transformation rules and recalculates transformation plans. - pub fn add_rules( - &self, - new_rules: Vec>, - ) -> Result { - let mut planner = TransformationPlanner::new(new_rules); - self.input.visit(&mut planner)?; - let new_rules = planner.rules; - let new_plans = planner.plans; - - let mut current_rules = self.rules.clone(); - - let offset = current_rules.len(); - current_rules.extend(new_rules); - - let mut merged_plans = Vec::with_capacity(self.plans.len() + new_plans.len()); - let mut old_plans_iter = self.plans.iter().peekable(); - let mut new_plans_iter = new_plans.into_iter().peekable(); - - while old_plans_iter.peek().is_some() || new_plans_iter.peek().is_some() { - match (old_plans_iter.peek(), new_plans_iter.peek()) { - (Some(&old), Some(new)) if old.node_index == new.node_index => { - let mut merged_plan = old.clone(); - for &rule_index in &new.rule_indices { - merged_plan.rule_indices.push(offset + rule_index); - } - merged_plans.push(merged_plan); - old_plans_iter.next(); - new_plans_iter.next(); - } - (Some(&old), Some(new)) if old.node_index < new.node_index => { - merged_plans.push(old.clone()); - old_plans_iter.next(); - } - (Some(_), Some(_)) => { - // old_node_index > new.node_index - let mut new_plan = new_plans_iter.next().unwrap(); - for rule_index in new_plan.rule_indices.iter_mut() { - *rule_index += offset; - } - merged_plans.push(new_plan); - } - (Some(&old), None) => { - merged_plans.push(old.clone()); - old_plans_iter.next(); - } - (None, Some(_)) => { - let mut new_plan = new_plans_iter.next().unwrap(); - for rule_index in new_plan.rule_indices.iter_mut() { - *rule_index += offset; - } - merged_plans.push(new_plan); - } - (None, None) => unreachable!(), - } - } - - Ok(Self { - input: Arc::clone(&self.input), - rules: current_rules, - plans: merged_plans, - metrics: ExecutionPlanMetricsSet::new(), - }) - } -} - -impl DisplayAs for TransformPlanExec { - fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { - match t { - DisplayFormatType::Default | DisplayFormatType::Verbose => { - let mut rule_to_nodes_count = vec![0; self.rules.len()]; - for plan in self.plans.iter() { - for rule_index in plan.rule_indices.iter() { - rule_to_nodes_count[*rule_index] += 1; - } - } - - let rules = rule_to_nodes_count - .into_iter() - .enumerate() - .map(|(rule_index, nodes_count)| { - let rule_name = self.rules[rule_index].name(); - format!("{rule_name}: plans_to_modify={nodes_count}") - }) - .join(", "); - - write!(f, "TransformPlanExec: rules=[{rules}]") - } - DisplayFormatType::TreeRender => Ok(()), - } - } -} - -impl ExecutionPlan for TransformPlanExec { - fn static_name() -> &'static str - where - Self: Sized, - { - "TransformPlanExec" - } - - fn name(&self) -> &str { - Self::static_name() - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn properties(&self) -> &PlanProperties { - self.input.properties() - } - - fn children(&self) -> Vec<&Arc> { - vec![&self.input] - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> Result> { - Ok(Arc::new(TransformPlanExec::try_new( - Arc::clone(&children[0]), - self.rules.clone(), - )?)) - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> Result { - let metric = MetricBuilder::new(&self.metrics).elapsed_compute(partition); - let _timer = metric.timer(); - - let transformed = self.transform(&context)?; - transformed.execute(partition, context) - } - - fn metrics(&self) -> Option { - Some(self.metrics.clone_inner()) - } - - fn statistics(&self) -> Result { - self.partition_statistics(None) - } - - fn partition_statistics(&self, partition: Option) -> Result { - self.input.partition_statistics(partition) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::coalesce_partitions::CoalescePartitionsExec; - use crate::empty::EmptyExec; - use crate::limit::GlobalLimitExec; - use crate::placeholder_row::PlaceholderRowExec; - use crate::plan_transformer::resolve_placeholders::ResolvePlaceholdersRule; - use crate::projection::ProjectionExec; - use crate::{get_plan_string, test}; - use arrow_schema::{DataType, Schema}; - use datafusion_common::{ParamValues, ScalarValue, tree_node::Transformed}; - use datafusion_expr::Operator; - use datafusion_physical_expr::expressions::{binary, lit, placeholder}; - use datafusion_physical_expr::projection::ProjectionExpr; - use insta::assert_snapshot; - use std::sync::Arc; - - #[derive(Debug, Clone)] - struct MockRule { - name: String, - match_name: String, - } - - impl ExecutionTransformationRule for MockRule { - fn name(&self) -> &str { - &self.name - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn matches(&self, node: &Arc) -> Result { - Ok(node.name() == self.match_name || self.match_name == "all") - } - - fn rewrite( - &self, - node: Arc, - _ctx: &TaskContext, - ) -> Result>> { - Ok(Transformed::no(node)) - } - } - - #[test] - fn test_has_rule() -> Result<()> { - let schema = Arc::new(Schema::empty()); - let input = Arc::new(EmptyExec::new(schema)); - - let resolve_rule = Arc::new(ResolvePlaceholdersRule::new()); - let exec = TransformPlanExec::try_new(input, vec![resolve_rule])?; - assert!(exec.has_rule::()); - assert!(!exec.has_rule::()); - - let exec = exec.add_rule(Arc::new(MockRule { - name: "mock".to_string(), - match_name: "any".to_string(), - }))?; - - assert!(exec.has_rule::()); - assert!(exec.has_rule::()); - - Ok(()) - } - - #[test] - fn test_add_rules_merge() -> Result<()> { - let schema = Arc::new(Schema::empty()); - let empty = Arc::new(EmptyExec::new(schema)); - let coalesce = Arc::new(CoalescePartitionsExec::new(empty)); - let input = Arc::new(GlobalLimitExec::new(coalesce, 0, None)) as Arc<_>; - - // Node 0: GlobalLimitExec - // Node 1: CoalescePartitionsExec - // Node 2: EmptyExec - - let rule_a: Arc = Arc::new(MockRule { - name: "ruleA".to_string(), - match_name: "CoalescePartitionsExec".to_string(), - }); - let exec = - TransformPlanExec::try_new(Arc::clone(&input), vec![Arc::clone(&rule_a)])?; - - assert_eq!(exec.rules.len(), 1); - assert_eq!(exec.plans.len(), 1); - assert_eq!(exec.plans[0].node_index, 1); - assert_eq!(exec.plans[0].rule_indices, vec![0]); - - let rule_b: Arc = Arc::new(MockRule { - name: "ruleB".to_string(), - match_name: "GlobalLimitExec".to_string(), - }); - let exec = exec.add_rules(vec![Arc::clone(&rule_b)])?; - - assert_eq!(exec.rules.len(), 2); - assert_eq!(exec.plans.len(), 2); - assert_eq!(exec.plans[0].node_index, 0); - assert_eq!(exec.plans[0].rule_indices, vec![1]); - assert_eq!(exec.plans[1].node_index, 1); - assert_eq!(exec.plans[1].rule_indices, vec![0]); - - let rule_c: Arc = Arc::new(MockRule { - name: "ruleC".to_string(), - match_name: "EmptyExec".to_string(), - }); - let exec = exec.add_rules(vec![Arc::clone(&rule_c)])?; - - assert_eq!(exec.rules.len(), 3); - assert_eq!(exec.plans.len(), 3); - assert_eq!(exec.plans[0].node_index, 0); - assert_eq!(exec.plans[0].rule_indices, vec![1]); - assert_eq!(exec.plans[1].node_index, 1); - assert_eq!(exec.plans[1].rule_indices, vec![0]); - assert_eq!(exec.plans[2].node_index, 2); - assert_eq!(exec.plans[2].rule_indices, vec![2]); - - let rule_d: Arc = Arc::new(MockRule { - name: "ruleD".to_string(), - match_name: "all".to_string(), - }); - - let exec = exec.add_rules(vec![Arc::clone(&rule_d)])?; - let check_full_plan = |exec: TransformPlanExec| { - assert_eq!(exec.rules.len(), 4); - assert_eq!(exec.plans.len(), 3); - assert_eq!(exec.plans[0].node_index, 0); - assert_eq!(exec.plans[0].rule_indices, vec![1, 3]); - assert_eq!(exec.plans[1].node_index, 1); - assert_eq!(exec.plans[1].rule_indices, vec![0, 3]); - assert_eq!(exec.plans[2].node_index, 2); - assert_eq!(exec.plans[2].rule_indices, vec![2, 3]); - }; - - check_full_plan(exec); - - let exec = TransformPlanExec::try_new( - Arc::clone(&input), - vec![ - Arc::clone(&rule_a), - Arc::clone(&rule_b), - Arc::clone(&rule_c), - Arc::clone(&rule_d), - ], - )?; - - check_full_plan(exec); - - let exec = TransformPlanExec::try_new(Arc::clone(&input), vec![rule_a])? - .add_rules(vec![rule_b, rule_c, rule_d])?; - - check_full_plan(exec); - - Ok(()) - } - - #[tokio::test] - async fn test_consts_evaluation() -> Result<()> { - let schema = test::aggr_test_schema(); - let expr = binary( - lit(10), - Operator::Plus, - binary( - lit(5), - Operator::Multiply, - placeholder("$1", DataType::Int32), - &schema, - )?, - &schema, - )?; - - let projection_expr = ProjectionExpr { - expr, - alias: "sum".to_string(), - }; - - let row = Arc::new(PlaceholderRowExec::new(schema)); - let projection = ProjectionExec::try_new(vec![projection_expr], row)?; - let transformer = TransformPlanExec::try_new( - Arc::new(projection), - vec![Arc::new(ResolvePlaceholdersRule::new())], - )?; - - let param_values = ParamValues::List(vec![ScalarValue::Int32(Some(20)).into()]); - let task_ctx = Arc::new(TaskContext::default().with_param_values(param_values)); - - let plan = transformer.transform(&task_ctx)?; - let plan_string = get_plan_string(&plan).join("\n"); - - assert_snapshot!(plan_string, @r" - ProjectionExec: expr=[110 as sum] - PlaceholderRowExec - "); - - Ok(()) - } -} diff --git a/datafusion/physical-plan/src/plan_transformer/resolve_placeholders.rs b/datafusion/physical-plan/src/plan_transformer/resolve_placeholders.rs deleted file mode 100644 index cefb57420dda3..0000000000000 --- a/datafusion/physical-plan/src/plan_transformer/resolve_placeholders.rs +++ /dev/null @@ -1,122 +0,0 @@ -// 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. - -//! Rule to resolve placeholders in the physical plan with actual values. - -use std::sync::Arc; - -use datafusion_common::{ - ParamValues, Result, exec_err, - tree_node::{Transformed, TreeNode}, -}; -use datafusion_execution::TaskContext; -use datafusion_physical_expr::{ - PhysicalExpr, - expressions::{Literal, PlaceholderExpr, has_placeholders}, - simplifier::const_evaluator::simplify_const_expr, -}; - -use crate::{ExecutionPlan, plan_transformer::ExecutionTransformationRule}; - -/// A transformation rule that replaces [`PlaceholderExpr`] with actual values. -/// -/// This rule is applied to the physical plan when actual parameter values are provided in the -/// [`TaskContext`]. It traverses the plan and replaces any placeholders found in physical -/// expressions with their corresponding literal values. -#[derive(Debug, Clone, Default)] -pub struct ResolvePlaceholdersRule {} - -impl ResolvePlaceholdersRule { - /// Create a new [`ResolvePlaceholdersRule`]. - pub fn new() -> Self { - Self {} - } -} - -impl ExecutionTransformationRule for ResolvePlaceholdersRule { - fn name(&self) -> &str { - "ResolvePlaceholders" - } - - fn as_any(&self) -> &dyn std::any::Any { - self - } - - fn matches(&self, node: &Arc) -> Result { - let Some(exprs) = node.physical_expressions() else { - return Ok(false); - }; - - for expr in exprs { - if has_placeholders(&expr) { - return Ok(true); - } - } - - Ok(false) - } - - fn rewrite( - &self, - node: Arc, - ctx: &TaskContext, - ) -> Result>> { - let Some(param_values) = ctx.param_values() else { - return Ok(Transformed::no(node)); - }; - - let Some(exprs) = node.physical_expressions() else { - return exec_err!("no physical expressions found"); - }; - - let new_exprs = exprs - .map(|expr| resolve_expr_placeholders(expr, param_values)) - .collect::>>()?; - - match node.with_physical_expressions(new_exprs.into())? { - Some(transformed_plan) => Ok(Transformed::yes(transformed_plan)), - None => exec_err!("failed to rewrite execution plan"), - } - } -} - -/// Resolves [`PlaceholderExpr`] in the physical expression using the provided [`ParamValues`]. -pub fn resolve_expr_placeholders( - expr: Arc, - param_values: &ParamValues, -) -> Result> { - let expr = expr.transform_up(|node| { - let Some(placeholder) = node.as_any().downcast_ref::() else { - return Ok(Transformed::no(node)); - }; - - if let Some(ref field) = placeholder.field { - let scalar = param_values.get_placeholders_with_values(&placeholder.id)?; - let value = scalar.value.cast_to(field.data_type())?; - let literal = Literal::new_with_metadata(value, scalar.metadata); - Ok(Transformed::yes(Arc::new(literal))) - } else { - Ok(Transformed::no(node)) - } - })?; - - if expr.transformed { - simplify_const_expr(expr.data).map(|t| t.data) - } else { - Ok(expr.data) - } -} diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto index 8fe68a264c3cc..9a8c25831051f 100644 --- a/datafusion/proto/proto/datafusion.proto +++ b/datafusion/proto/proto/datafusion.proto @@ -751,8 +751,7 @@ message PhysicalPlanNode { MemoryScanExecNode memory_scan = 35; AsyncFuncExecNode async_func = 36; BufferExecNode buffer = 37; - TransformPlanExecNode transform_plan = 38; - ValuesExecNode values_scan = 39; + ValuesExecNode values_scan = 38; } } @@ -1458,17 +1457,3 @@ message BufferExecNode { PhysicalPlanNode input = 1; uint64 capacity = 2; } - -message TransformPlanExecNode { - PhysicalPlanNode input = 1; - repeated TransformationRule rules = 2; -} - -message TransformationRule { - oneof rule_type { - bytes extension = 1; - ResolvePlaceholdersRule resolve_placeholders = 2; - } -} - -message ResolvePlaceholdersRule {} diff --git a/datafusion/proto/src/generated/pbjson.rs b/datafusion/proto/src/generated/pbjson.rs index f579478f976e3..3644140ed0538 100644 --- a/datafusion/proto/src/generated/pbjson.rs +++ b/datafusion/proto/src/generated/pbjson.rs @@ -17905,9 +17905,6 @@ impl serde::Serialize for PhysicalPlanNode { physical_plan_node::PhysicalPlanType::Buffer(v) => { struct_ser.serialize_field("buffer", v)?; } - physical_plan_node::PhysicalPlanType::TransformPlan(v) => { - struct_ser.serialize_field("transformPlan", v)?; - } physical_plan_node::PhysicalPlanType::ValuesScan(v) => { struct_ser.serialize_field("valuesScan", v)?; } @@ -17979,8 +17976,6 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { "async_func", "asyncFunc", "buffer", - "transform_plan", - "transformPlan", "values_scan", "valuesScan", ]; @@ -18023,7 +18018,6 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { MemoryScan, AsyncFunc, Buffer, - TransformPlan, ValuesScan, } impl<'de> serde::Deserialize<'de> for GeneratedField { @@ -18082,7 +18076,6 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { "memoryScan" | "memory_scan" => Ok(GeneratedField::MemoryScan), "asyncFunc" | "async_func" => Ok(GeneratedField::AsyncFunc), "buffer" => Ok(GeneratedField::Buffer), - "transformPlan" | "transform_plan" => Ok(GeneratedField::TransformPlan), "valuesScan" | "values_scan" => Ok(GeneratedField::ValuesScan), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } @@ -18356,13 +18349,6 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { return Err(serde::de::Error::duplicate_field("buffer")); } physical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_plan_node::PhysicalPlanType::Buffer) -; - } - GeneratedField::TransformPlan => { - if physical_plan_type__.is_some() { - return Err(serde::de::Error::duplicate_field("transformPlan")); - } - physical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_plan_node::PhysicalPlanType::TransformPlan) ; } GeneratedField::ValuesScan => { @@ -20942,77 +20928,6 @@ impl<'de> serde::Deserialize<'de> for RepartitionNode { deserializer.deserialize_struct("datafusion.RepartitionNode", FIELDS, GeneratedVisitor) } } -impl serde::Serialize for ResolvePlaceholdersRule { - #[allow(deprecated)] - fn serialize(&self, serializer: S) -> std::result::Result - where - S: serde::Serializer, - { - use serde::ser::SerializeStruct; - let len = 0; - let struct_ser = serializer.serialize_struct("datafusion.ResolvePlaceholdersRule", len)?; - struct_ser.end() - } -} -impl<'de> serde::Deserialize<'de> for ResolvePlaceholdersRule { - #[allow(deprecated)] - fn deserialize(deserializer: D) -> std::result::Result - where - D: serde::Deserializer<'de>, - { - const FIELDS: &[&str] = &[ - ]; - - #[allow(clippy::enum_variant_names)] - enum GeneratedField { - } - impl<'de> serde::Deserialize<'de> for GeneratedField { - fn deserialize(deserializer: D) -> std::result::Result - where - D: serde::Deserializer<'de>, - { - struct GeneratedVisitor; - - impl serde::de::Visitor<'_> for GeneratedVisitor { - type Value = GeneratedField; - - fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(formatter, "expected one of: {:?}", &FIELDS) - } - - #[allow(unused_variables)] - fn visit_str(self, value: &str) -> std::result::Result - where - E: serde::de::Error, - { - Err(serde::de::Error::unknown_field(value, FIELDS)) - } - } - deserializer.deserialize_identifier(GeneratedVisitor) - } - } - struct GeneratedVisitor; - impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { - type Value = ResolvePlaceholdersRule; - - fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str("struct datafusion.ResolvePlaceholdersRule") - } - - fn visit_map(self, mut map_: V) -> std::result::Result - where - V: serde::de::MapAccess<'de>, - { - while map_.next_key::()?.is_some() { - let _ = map_.next_value::()?; - } - Ok(ResolvePlaceholdersRule { - }) - } - } - deserializer.deserialize_struct("datafusion.ResolvePlaceholdersRule", FIELDS, GeneratedVisitor) - } -} impl serde::Serialize for RollupNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -23141,225 +23056,6 @@ impl<'de> serde::Deserialize<'de> for TableReference { deserializer.deserialize_struct("datafusion.TableReference", FIELDS, GeneratedVisitor) } } -impl serde::Serialize for TransformPlanExecNode { - #[allow(deprecated)] - fn serialize(&self, serializer: S) -> std::result::Result - where - S: serde::Serializer, - { - use serde::ser::SerializeStruct; - let mut len = 0; - if self.input.is_some() { - len += 1; - } - if !self.rules.is_empty() { - len += 1; - } - let mut struct_ser = serializer.serialize_struct("datafusion.TransformPlanExecNode", len)?; - if let Some(v) = self.input.as_ref() { - struct_ser.serialize_field("input", v)?; - } - if !self.rules.is_empty() { - struct_ser.serialize_field("rules", &self.rules)?; - } - struct_ser.end() - } -} -impl<'de> serde::Deserialize<'de> for TransformPlanExecNode { - #[allow(deprecated)] - fn deserialize(deserializer: D) -> std::result::Result - where - D: serde::Deserializer<'de>, - { - const FIELDS: &[&str] = &[ - "input", - "rules", - ]; - - #[allow(clippy::enum_variant_names)] - enum GeneratedField { - Input, - Rules, - } - impl<'de> serde::Deserialize<'de> for GeneratedField { - fn deserialize(deserializer: D) -> std::result::Result - where - D: serde::Deserializer<'de>, - { - struct GeneratedVisitor; - - impl serde::de::Visitor<'_> for GeneratedVisitor { - type Value = GeneratedField; - - fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(formatter, "expected one of: {:?}", &FIELDS) - } - - #[allow(unused_variables)] - fn visit_str(self, value: &str) -> std::result::Result - where - E: serde::de::Error, - { - match value { - "input" => Ok(GeneratedField::Input), - "rules" => Ok(GeneratedField::Rules), - _ => Err(serde::de::Error::unknown_field(value, FIELDS)), - } - } - } - deserializer.deserialize_identifier(GeneratedVisitor) - } - } - struct GeneratedVisitor; - impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { - type Value = TransformPlanExecNode; - - fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str("struct datafusion.TransformPlanExecNode") - } - - fn visit_map(self, mut map_: V) -> std::result::Result - where - V: serde::de::MapAccess<'de>, - { - let mut input__ = None; - let mut rules__ = None; - while let Some(k) = map_.next_key()? { - match k { - GeneratedField::Input => { - if input__.is_some() { - return Err(serde::de::Error::duplicate_field("input")); - } - input__ = map_.next_value()?; - } - GeneratedField::Rules => { - if rules__.is_some() { - return Err(serde::de::Error::duplicate_field("rules")); - } - rules__ = Some(map_.next_value()?); - } - } - } - Ok(TransformPlanExecNode { - input: input__, - rules: rules__.unwrap_or_default(), - }) - } - } - deserializer.deserialize_struct("datafusion.TransformPlanExecNode", FIELDS, GeneratedVisitor) - } -} -impl serde::Serialize for TransformationRule { - #[allow(deprecated)] - fn serialize(&self, serializer: S) -> std::result::Result - where - S: serde::Serializer, - { - use serde::ser::SerializeStruct; - let mut len = 0; - if self.rule_type.is_some() { - len += 1; - } - let mut struct_ser = serializer.serialize_struct("datafusion.TransformationRule", len)?; - if let Some(v) = self.rule_type.as_ref() { - match v { - transformation_rule::RuleType::Extension(v) => { - #[allow(clippy::needless_borrow)] - #[allow(clippy::needless_borrows_for_generic_args)] - struct_ser.serialize_field("extension", pbjson::private::base64::encode(&v).as_str())?; - } - transformation_rule::RuleType::ResolvePlaceholders(v) => { - struct_ser.serialize_field("resolvePlaceholders", v)?; - } - } - } - struct_ser.end() - } -} -impl<'de> serde::Deserialize<'de> for TransformationRule { - #[allow(deprecated)] - fn deserialize(deserializer: D) -> std::result::Result - where - D: serde::Deserializer<'de>, - { - const FIELDS: &[&str] = &[ - "extension", - "resolve_placeholders", - "resolvePlaceholders", - ]; - - #[allow(clippy::enum_variant_names)] - enum GeneratedField { - Extension, - ResolvePlaceholders, - } - impl<'de> serde::Deserialize<'de> for GeneratedField { - fn deserialize(deserializer: D) -> std::result::Result - where - D: serde::Deserializer<'de>, - { - struct GeneratedVisitor; - - impl serde::de::Visitor<'_> for GeneratedVisitor { - type Value = GeneratedField; - - fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(formatter, "expected one of: {:?}", &FIELDS) - } - - #[allow(unused_variables)] - fn visit_str(self, value: &str) -> std::result::Result - where - E: serde::de::Error, - { - match value { - "extension" => Ok(GeneratedField::Extension), - "resolvePlaceholders" | "resolve_placeholders" => Ok(GeneratedField::ResolvePlaceholders), - _ => Err(serde::de::Error::unknown_field(value, FIELDS)), - } - } - } - deserializer.deserialize_identifier(GeneratedVisitor) - } - } - struct GeneratedVisitor; - impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { - type Value = TransformationRule; - - fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str("struct datafusion.TransformationRule") - } - - fn visit_map(self, mut map_: V) -> std::result::Result - where - V: serde::de::MapAccess<'de>, - { - let mut rule_type__ = None; - while let Some(k) = map_.next_key()? { - match k { - GeneratedField::Extension => { - if rule_type__.is_some() { - return Err(serde::de::Error::duplicate_field("extension")); - } - rule_type__ = map_.next_value::<::std::option::Option<::pbjson::private::BytesDeserialize<_>>>()?.map(|x| transformation_rule::RuleType::Extension(x.0)); - } - GeneratedField::ResolvePlaceholders => { - if rule_type__.is_some() { - return Err(serde::de::Error::duplicate_field("resolvePlaceholders")); - } - rule_type__ = map_.next_value::<::std::option::Option<_>>()?.map(transformation_rule::RuleType::ResolvePlaceholders) -; - } - } - } - Ok(TransformationRule { - rule_type: rule_type__, - }) - } - } - deserializer.deserialize_struct("datafusion.TransformationRule", FIELDS, GeneratedVisitor) - } -} impl serde::Serialize for TryCastNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result diff --git a/datafusion/proto/src/generated/prost.rs b/datafusion/proto/src/generated/prost.rs index 755d04c68111f..dd7bbcc2c9163 100644 --- a/datafusion/proto/src/generated/prost.rs +++ b/datafusion/proto/src/generated/prost.rs @@ -1079,7 +1079,7 @@ pub mod table_reference { pub struct PhysicalPlanNode { #[prost( oneof = "physical_plan_node::PhysicalPlanType", - tags = "1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39" + tags = "1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38" )] pub physical_plan_type: ::core::option::Option, } @@ -1162,8 +1162,6 @@ pub mod physical_plan_node { #[prost(message, tag = "37")] Buffer(::prost::alloc::boxed::Box), #[prost(message, tag = "38")] - TransformPlan(::prost::alloc::boxed::Box), - #[prost(message, tag = "39")] ValuesScan(super::ValuesExecNode), } } @@ -2184,30 +2182,6 @@ pub struct BufferExecNode { #[prost(uint64, tag = "2")] pub capacity: u64, } -#[derive(Clone, PartialEq, ::prost::Message)] -pub struct TransformPlanExecNode { - #[prost(message, optional, boxed, tag = "1")] - pub input: ::core::option::Option<::prost::alloc::boxed::Box>, - #[prost(message, repeated, tag = "2")] - pub rules: ::prost::alloc::vec::Vec, -} -#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] -pub struct TransformationRule { - #[prost(oneof = "transformation_rule::RuleType", tags = "1, 2")] - pub rule_type: ::core::option::Option, -} -/// Nested message and enum types in `TransformationRule`. -pub mod transformation_rule { - #[derive(Clone, PartialEq, Eq, Hash, ::prost::Oneof)] - pub enum RuleType { - #[prost(bytes, tag = "1")] - Extension(::prost::alloc::vec::Vec), - #[prost(message, tag = "2")] - ResolvePlaceholders(super::ResolvePlaceholdersRule), - } -} -#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] -pub struct ResolvePlaceholdersRule {} #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] #[repr(i32)] pub enum WindowFrameUnits { diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 500145058be61..8fa82eed966fa 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -76,9 +76,6 @@ use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::memory::LazyMemoryExec; use datafusion_physical_plan::metrics::MetricType; use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; -use datafusion_physical_plan::plan_transformer::{ - ExecutionTransformationRule, ResolvePlaceholdersRule, TransformPlanExec, -}; use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::sorts::sort::SortExec; @@ -106,7 +103,6 @@ use crate::physical_plan::to_proto::{ use crate::protobuf::physical_aggregate_expr_node::AggregateFunction; use crate::protobuf::physical_expr_node::ExprType; use crate::protobuf::physical_plan_node::PhysicalPlanType; -use crate::protobuf::transformation_rule::RuleType; use crate::protobuf::{ self, ListUnnest as ProtoListUnnest, SortExprNode, SortMergeJoinExecNode, proto_error, window_agg_exec_node, @@ -317,13 +313,6 @@ impl protobuf::PhysicalPlanNode { PhysicalPlanType::Buffer(buffer) => { self.try_into_buffer_physical_plan(buffer, ctx, codec, proto_converter) } - PhysicalPlanType::TransformPlan(transform_plan) => self - .try_into_transform_plan_exec( - transform_plan, - ctx, - codec, - proto_converter, - ), PhysicalPlanType::ValuesScan(scan) => { self.try_into_values_scan_exec(scan, ctx, codec, proto_converter) } @@ -575,14 +564,6 @@ impl protobuf::PhysicalPlanNode { ); } - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_transform_plan_exec( - exec, - codec, - proto_converter, - ); - } - let mut buf: Vec = vec![]; match codec.try_encode(Arc::clone(&plan_clone), &mut buf) { Ok(_) => { @@ -2221,37 +2202,6 @@ impl protobuf::PhysicalPlanNode { Ok(Arc::new(BufferExec::new(input, buffer.capacity as usize))) } - fn try_into_transform_plan_exec( - &self, - transform_plan: &protobuf::TransformPlanExecNode, - ctx: &TaskContext, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&transform_plan.input, ctx, codec, proto_converter)?; - - let mut rules: Vec> = - Vec::with_capacity(transform_plan.rules.len()); - - for rule in transform_plan.rules.iter() { - match &rule.rule_type { - Some(RuleType::ResolvePlaceholders(_)) => { - rules.push(Arc::new(ResolvePlaceholdersRule::new())) - } - Some(RuleType::Extension(ext)) => { - rules.push(codec.try_decode_transformation_rule(ext)?) - } - None => { - return internal_err!("Missing rule_type in TransformationRule"); - } - } - } - - let transformer = TransformPlanExec::try_new(input, rules)?; - Ok(Arc::new(transformer)) - } - fn try_into_values_scan_exec( &self, scan: &protobuf::ValuesExecNode, @@ -3667,42 +3617,6 @@ impl protobuf::PhysicalPlanNode { }) } - fn try_from_transform_plan_exec( - exec: &TransformPlanExec, - extension_codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - Arc::clone(exec.input()), - extension_codec, - proto_converter, - )?; - - let mut rules = Vec::with_capacity(exec.rules().len()); - for rule in exec.rules() { - let rule_type = if rule.as_any().is::() { - Some(RuleType::ResolvePlaceholders( - protobuf::ResolvePlaceholdersRule {}, - )) - } else { - let mut buf = vec![]; - extension_codec - .try_encode_transformation_rule(rule.as_ref(), &mut buf)?; - Some(RuleType::Extension(buf)) - }; - rules.push(protobuf::TransformationRule { rule_type }); - } - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::TransformPlan(Box::new( - protobuf::TransformPlanExecNode { - input: Some(Box::new(input)), - rules, - }, - ))), - }) - } - fn try_from_values_source( source: &ValuesSource, extension_codec: &dyn PhysicalExtensionCodec, @@ -3796,21 +3710,6 @@ pub trait PhysicalExtensionCodec: Debug + Send + Sync { Ok(()) } - fn try_decode_transformation_rule( - &self, - _buf: &[u8], - ) -> Result> { - not_impl_err!("PhysicalExtensionCodec is not provided") - } - - fn try_encode_transformation_rule( - &self, - _node: &dyn ExecutionTransformationRule, - _buf: &mut Vec, - ) -> Result<()> { - not_impl_err!("PhysicalExtensionCodec is not provided") - } - fn try_decode_udwf(&self, name: &str, _buf: &[u8]) -> Result> { not_impl_err!("PhysicalExtensionCodec is not provided for window function {name}") } @@ -4244,25 +4143,6 @@ impl PhysicalExtensionCodec for ComposedPhysicalExtensionCodec { fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec) -> Result<()> { self.encode_protobuf(buf, |codec, data| codec.try_encode_udaf(node, data)) } - - fn try_decode_transformation_rule( - &self, - buf: &[u8], - ) -> Result> { - self.decode_protobuf(buf, |codec, data| { - codec.try_decode_transformation_rule(data) - }) - } - - fn try_encode_transformation_rule( - &self, - node: &dyn ExecutionTransformationRule, - buf: &mut Vec, - ) -> Result<()> { - self.encode_protobuf(buf, |codec, data| { - codec.try_encode_transformation_rule(node, data) - }) - } } fn into_physical_plan( diff --git a/datafusion/sqllogictest/test_files/explain.slt b/datafusion/sqllogictest/test_files/explain.slt index e23d8db2cf8d3..7938c5af66b38 100644 --- a/datafusion/sqllogictest/test_files/explain.slt +++ b/datafusion/sqllogictest/test_files/explain.slt @@ -227,7 +227,6 @@ initial_physical_plan_with_schema DataSourceExec: file_groups={1 group: [[WORKSP physical_plan after OutputRequirements 01)OutputRequirementExec: order_by=[], dist_by=Unspecified 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], file_type=csv, has_header=true -physical_plan after ExecutionTransformationApplier SAME TEXT AS ABOVE physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE @@ -244,7 +243,6 @@ physical_plan after LimitPushdown SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after PushdownSort SAME TEXT AS ABOVE physical_plan after EnsureCooperative SAME TEXT AS ABOVE -physical_plan after ExecutionTransformationApplier(ResolvePlaceholders) SAME TEXT AS ABOVE physical_plan after FilterPushdown(Post) SAME TEXT AS ABOVE physical_plan after SanityCheckPlan SAME TEXT AS ABOVE physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], file_type=csv, has_header=true @@ -307,7 +305,6 @@ physical_plan after OutputRequirements 01)OutputRequirementExec: order_by=[], dist_by=Unspecified, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]: ScanBytes=Exact(32)),(Col[1]: ScanBytes=Inexact(24)),(Col[2]: ScanBytes=Exact(32)),(Col[3]: ScanBytes=Exact(32)),(Col[4]: ScanBytes=Exact(32)),(Col[5]: ScanBytes=Exact(64)),(Col[6]: ScanBytes=Exact(32)),(Col[7]: ScanBytes=Exact(64)),(Col[8]: ScanBytes=Inexact(88)),(Col[9]: ScanBytes=Inexact(49)),(Col[10]: ScanBytes=Exact(64))]] 02)--GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]: ScanBytes=Exact(32)),(Col[1]: ScanBytes=Inexact(24)),(Col[2]: ScanBytes=Exact(32)),(Col[3]: ScanBytes=Exact(32)),(Col[4]: ScanBytes=Exact(32)),(Col[5]: ScanBytes=Exact(64)),(Col[6]: ScanBytes=Exact(32)),(Col[7]: ScanBytes=Exact(64)),(Col[8]: ScanBytes=Inexact(88)),(Col[9]: ScanBytes=Inexact(49)),(Col[10]: ScanBytes=Exact(64))]] 03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]: ScanBytes=Exact(32)),(Col[1]: ScanBytes=Inexact(24)),(Col[2]: ScanBytes=Exact(32)),(Col[3]: ScanBytes=Exact(32)),(Col[4]: ScanBytes=Exact(32)),(Col[5]: ScanBytes=Exact(64)),(Col[6]: ScanBytes=Exact(32)),(Col[7]: ScanBytes=Exact(64)),(Col[8]: ScanBytes=Inexact(88)),(Col[9]: ScanBytes=Inexact(49)),(Col[10]: ScanBytes=Exact(64))]] -physical_plan after ExecutionTransformationApplier SAME TEXT AS ABOVE physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE @@ -326,7 +323,6 @@ physical_plan after LimitPushdown DataSourceExec: file_groups={1 group: [[WORKSP physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after PushdownSort SAME TEXT AS ABOVE physical_plan after EnsureCooperative SAME TEXT AS ABOVE -physical_plan after ExecutionTransformationApplier(ResolvePlaceholders) SAME TEXT AS ABOVE physical_plan after FilterPushdown(Post) SAME TEXT AS ABOVE physical_plan after SanityCheckPlan SAME TEXT AS ABOVE physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]: ScanBytes=Exact(32)),(Col[1]: ScanBytes=Inexact(24)),(Col[2]: ScanBytes=Exact(32)),(Col[3]: ScanBytes=Exact(32)),(Col[4]: ScanBytes=Exact(32)),(Col[5]: ScanBytes=Exact(64)),(Col[6]: ScanBytes=Exact(32)),(Col[7]: ScanBytes=Exact(64)),(Col[8]: ScanBytes=Inexact(88)),(Col[9]: ScanBytes=Inexact(49)),(Col[10]: ScanBytes=Exact(64))]] @@ -353,7 +349,6 @@ physical_plan after OutputRequirements 01)OutputRequirementExec: order_by=[], dist_by=Unspecified 02)--GlobalLimitExec: skip=0, fetch=10 03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet -physical_plan after ExecutionTransformationApplier SAME TEXT AS ABOVE physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE @@ -372,7 +367,6 @@ physical_plan after LimitPushdown DataSourceExec: file_groups={1 group: [[WORKSP physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after PushdownSort SAME TEXT AS ABOVE physical_plan after EnsureCooperative SAME TEXT AS ABOVE -physical_plan after ExecutionTransformationApplier(ResolvePlaceholders) SAME TEXT AS ABOVE physical_plan after FilterPushdown(Post) SAME TEXT AS ABOVE physical_plan after SanityCheckPlan SAME TEXT AS ABOVE physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet @@ -594,7 +588,6 @@ initial_physical_plan_with_schema DataSourceExec: file_groups={1 group: [[WORKSP physical_plan after OutputRequirements 01)OutputRequirementExec: order_by=[], dist_by=Unspecified 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], file_type=csv, has_header=true -physical_plan after ExecutionTransformationApplier SAME TEXT AS ABOVE physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE @@ -611,7 +604,6 @@ physical_plan after LimitPushdown SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after PushdownSort SAME TEXT AS ABOVE physical_plan after EnsureCooperative SAME TEXT AS ABOVE -physical_plan after ExecutionTransformationApplier(ResolvePlaceholders) SAME TEXT AS ABOVE physical_plan after FilterPushdown(Post) SAME TEXT AS ABOVE physical_plan after SanityCheckPlan SAME TEXT AS ABOVE physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], file_type=csv, has_header=true diff --git a/datafusion/sqllogictest/test_files/placeholders.slt b/datafusion/sqllogictest/test_files/placeholders.slt index 3dfa7eb572dbc..2ff1837206e2f 100644 --- a/datafusion/sqllogictest/test_files/placeholders.slt +++ b/datafusion/sqllogictest/test_files/placeholders.slt @@ -17,15 +17,6 @@ ########## ## Test physical plans with placeholders. -## -## These tests verify that the physical planner correctly produces a -## `TransformPlanExec` node with the expected number of plans with -## placeholders. -## -## A test is considered successful when the physical plan contains -## `TransformPlanExec`, indicating that the plan can be traversed to find -## placeholders and successfully replace them with actual parameter values -## during execution. ########## statement ok @@ -46,10 +37,9 @@ logical_plan 02)--Projection: column1 AS id, column2 AS name 03)----Values: ($1, Utf8View("Samanta") AS Utf8("Samanta")), (Int32(5) AS Int64(5), $2) physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--DataSinkExec: sink=MemoryTable (partitions=1) -03)----ProjectionExec: expr=[column1@0 as id, column2@1 as name] -04)------DataSourceExec: placeholders=2 +01)DataSinkExec: sink=MemoryTable (partitions=1) +02)--ProjectionExec: expr=[column1@0 as id, column2@1 as name] +03)----DataSourceExec: placeholders=2 # Filter with multiple placeholders query TT @@ -60,9 +50,8 @@ logical_plan 02)--Filter: t1.name = $1 OR t1.id = $2 03)----TableScan: t1 projection=[id, name] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--FilterExec: name@1 = $1 OR id@0 = $2, projection=[id@0] -03)----DataSourceExec: partitions=1, partition_sizes=[1] +01)FilterExec: name@1 = $1 OR id@0 = $2, projection=[id@0] +02)--DataSourceExec: partitions=1, partition_sizes=[1] # Projection with placeholder query TT @@ -72,9 +61,8 @@ logical_plan 01)Projection: t1.id + $1 02)--TableScan: t1 projection=[id] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[id@0 + $1 as t1.id + $1] -03)----DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[id@0 + $1 as t1.id + $1] +02)--DataSourceExec: partitions=1, partition_sizes=[1] # Projection and filter with placeholders query TT @@ -85,11 +73,10 @@ logical_plan 02)--Filter: t1.name = $2 03)----TableScan: t1 projection=[id, name] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=2] -02)--ProjectionExec: expr=[id@0 + $1 as t1.id + $1] -03)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -04)------FilterExec: name@1 = $2 -05)--------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[id@0 + $1 as t1.id + $1] +02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +03)----FilterExec: name@1 = $2 +04)------DataSourceExec: partitions=1, partition_sizes=[1] statement ok DROP TABLE t1 @@ -112,13 +99,12 @@ logical_plan 01)Aggregate: groupBy=[[]], aggr=[[array_agg(agg_order.c1 + $1) ORDER BY [agg_order.c2 DESC NULLS FIRST, agg_order.c3 ASC NULLS LAST]]] 02)--TableScan: agg_order projection=[c1, c2, c3] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=2] -02)--AggregateExec: mode=Final, gby=[], aggr=[array_agg(agg_order.c1 + $1) ORDER BY [agg_order.c2 DESC NULLS FIRST, agg_order.c3 ASC NULLS LAST]] -03)----CoalescePartitionsExec -04)------AggregateExec: mode=Partial, gby=[], aggr=[array_agg(agg_order.c1 + $1) ORDER BY [agg_order.c2 DESC NULLS FIRST, agg_order.c3 ASC NULLS LAST]] -05)--------SortExec: expr=[c2@1 DESC, c3@2 ASC NULLS LAST], preserve_partitioning=[true] -06)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -07)------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/aggregate_agg_multi_order.csv]]}, projection=[c1, c2, c3], file_type=csv, has_header=true +01)AggregateExec: mode=Final, gby=[], aggr=[array_agg(agg_order.c1 + $1) ORDER BY [agg_order.c2 DESC NULLS FIRST, agg_order.c3 ASC NULLS LAST]] +02)--CoalescePartitionsExec +03)----AggregateExec: mode=Partial, gby=[], aggr=[array_agg(agg_order.c1 + $1) ORDER BY [agg_order.c2 DESC NULLS FIRST, agg_order.c3 ASC NULLS LAST]] +04)------SortExec: expr=[c2@1 DESC, c3@2 ASC NULLS LAST], preserve_partitioning=[true] +05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/aggregate_agg_multi_order.csv]]}, projection=[c1, c2, c3], file_type=csv, has_header=true statement ok DROP TABLE agg_order @@ -140,9 +126,7 @@ logical_plan 01)Projection: alltypes_plain.smallint_col 02)--Filter: alltypes_plain.int_col = $1 03)----TableScan: alltypes_plain projection=[smallint_col, int_col], partial_filters=[alltypes_plain.int_col = $1] -physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[smallint_col], file_type=parquet, predicate=int_col@4 = $1, pruning_predicate=int_col_null_count@2 != row_count@3 AND int_col_min@0 <= $1 AND $1 <= int_col_max@1, required_guarantees=[] +physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[smallint_col], file_type=parquet, predicate=int_col@4 = $1, pruning_predicate=int_col_null_count@2 != row_count@3 AND int_col_min@0 <= $1 AND $1 <= int_col_max@1, required_guarantees=[] # Projection with placeholder on parquet table query TT @@ -151,9 +135,7 @@ EXPLAIN SELECT smallint_col + $1 FROM alltypes_plain; logical_plan 01)Projection: alltypes_plain.smallint_col + $1 02)--TableScan: alltypes_plain projection=[smallint_col] -physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[smallint_col@3 + $1 as alltypes_plain.smallint_col + $1], file_type=parquet +physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[smallint_col@3 + $1 as alltypes_plain.smallint_col + $1], file_type=parquet # Projection and filter with placeholders on parquet table query TT @@ -164,10 +146,9 @@ logical_plan 02)--Filter: alltypes_plain.int_col = $2 03)----TableScan: alltypes_plain projection=[smallint_col, int_col], partial_filters=[alltypes_plain.int_col = $2] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=2] -02)--ProjectionExec: expr=[smallint_col@0 + $1 as alltypes_plain.smallint_col + $1] -03)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -04)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[smallint_col, int_col], file_type=parquet, predicate=int_col@4 = $2, pruning_predicate=int_col_null_count@2 != row_count@3 AND int_col_min@0 <= $2 AND $2 <= int_col_max@1, required_guarantees=[] +01)ProjectionExec: expr=[smallint_col@0 + $1 as alltypes_plain.smallint_col + $1] +02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[smallint_col, int_col], file_type=parquet, predicate=int_col@4 = $2, pruning_predicate=int_col_null_count@2 != row_count@3 AND int_col_min@0 <= $2 AND $2 <= int_col_max@1, required_guarantees=[] statement ok DROP TABLE alltypes_plain; @@ -192,12 +173,11 @@ logical_plan 03)----TableScan: t1 projection=[id, name] 04)----TableScan: t2 projection=[id, age] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[name@1 as name, age@0 as age] -03)----HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, t1.id + $1@2)], projection=[age@1, name@3] -04)------DataSourceExec: partitions=1, partition_sizes=[1] -05)------ProjectionExec: expr=[id@0 as id, name@1 as name, id@0 + $1 as t1.id + $1] -06)--------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[name@1 as name, age@0 as age] +02)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, t1.id + $1@2)], projection=[age@1, name@3] +03)----DataSourceExec: partitions=1, partition_sizes=[1] +04)----ProjectionExec: expr=[id@0 as id, name@1 as name, id@0 + $1 as t1.id + $1] +05)------DataSourceExec: partitions=1, partition_sizes=[1] # Join with placeholder in filter query TT @@ -210,12 +190,11 @@ logical_plan 04)----Filter: t2.age > $1 05)------TableScan: t2 projection=[id, age] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[name@1 as name, age@0 as age] -03)----HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0)], projection=[age@1, name@3] -04)------FilterExec: age@1 > $1 -05)--------DataSourceExec: partitions=1, partition_sizes=[1] -06)------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[name@1 as name, age@0 as age] +02)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0)], projection=[age@1, name@3] +03)----FilterExec: age@1 > $1 +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)----DataSourceExec: partitions=1, partition_sizes=[1] # Join with placeholder in projection query TT @@ -227,13 +206,12 @@ logical_plan 03)----TableScan: t1 projection=[id, name] 04)----TableScan: t2 projection=[id, age] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[name@1 as name, age@3 + $1 as t2.age + $1] -03)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -04)------ProjectionExec: expr=[id@2 as id, name@3 as name, id@0 as id, age@1 as age] -05)--------HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0)] -06)----------DataSourceExec: partitions=1, partition_sizes=[1] -07)----------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[name@1 as name, age@3 + $1 as t2.age + $1] +02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +03)----ProjectionExec: expr=[id@2 as id, name@3 as name, id@0 as id, age@1 as age] +04)------HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0)] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] +06)--------DataSourceExec: partitions=1, partition_sizes=[1] # Join with placeholder in ON statement query TT @@ -245,11 +223,10 @@ logical_plan 03)----TableScan: t1 projection=[id, name] 04)----TableScan: t2 projection=[id, age] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[name@1 as name, age@0 as age] -03)----NestedLoopJoinExec: join_type=Inner, filter=id@0 + id@1 = $1, projection=[age@1, name@3] -04)------DataSourceExec: partitions=1, partition_sizes=[1] -05)------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[name@1 as name, age@0 as age] +02)--NestedLoopJoinExec: join_type=Inner, filter=id@0 + id@1 = $1, projection=[age@1, name@3] +03)----DataSourceExec: partitions=1, partition_sizes=[1] +04)----DataSourceExec: partitions=1, partition_sizes=[1] statement ok DROP TABLE t1; @@ -270,11 +247,10 @@ logical_plan 02)--WindowAggr: windowExpr=[[sum(CAST(t1.id AS Int64)) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] 03)----TableScan: t1 projection=[id, name] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[id@0 as id, sum(t1.id) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 + $1 as sum(t1.id) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW + $1] -03)----BoundedWindowAggExec: wdw=[sum(t1.id) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(t1.id) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -04)------SortExec: expr=[name@1 ASC NULLS LAST, id@0 ASC NULLS LAST], preserve_partitioning=[false] -05)--------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[id@0 as id, sum(t1.id) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 + $1 as sum(t1.id) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW + $1] +02)--BoundedWindowAggExec: wdw=[sum(t1.id) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(t1.id) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortExec: expr=[name@1 ASC NULLS LAST, id@0 ASC NULLS LAST], preserve_partitioning=[false] +04)------DataSourceExec: partitions=1, partition_sizes=[1] # Window function with placeholder # Here we only resolve BoundedWindowAggExec. @@ -286,11 +262,10 @@ logical_plan 02)--WindowAggr: windowExpr=[[sum(CAST(t1.id + $1 AS Int64)) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] 03)----TableScan: t1 projection=[id, name] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[id@0 as id, sum(t1.id + $1) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as sum(t1.id + $1) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] -03)----BoundedWindowAggExec: wdw=[sum(t1.id + $1) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(t1.id + $1) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -04)------SortExec: expr=[name@1 ASC NULLS LAST, id@0 ASC NULLS LAST], preserve_partitioning=[false] -05)--------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[id@0 as id, sum(t1.id + $1) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as sum(t1.id + $1) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] +02)--BoundedWindowAggExec: wdw=[sum(t1.id + $1) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(t1.id + $1) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortExec: expr=[name@1 ASC NULLS LAST, id@0 ASC NULLS LAST], preserve_partitioning=[false] +04)------DataSourceExec: partitions=1, partition_sizes=[1] # UNION with placeholder query TT @@ -303,12 +278,11 @@ logical_plan 04)--Filter: t1.id = $2 05)----TableScan: t1 projection=[id] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=2] -02)--UnionExec -03)----FilterExec: id@0 = $1 -04)------DataSourceExec: partitions=1, partition_sizes=[1] -05)----FilterExec: id@0 = $2 -06)------DataSourceExec: partitions=1, partition_sizes=[1] +01)UnionExec +02)--FilterExec: id@0 = $1 +03)----DataSourceExec: partitions=1, partition_sizes=[1] +04)--FilterExec: id@0 = $2 +05)----DataSourceExec: partitions=1, partition_sizes=[1] statement ok DROP TABLE t1; @@ -333,12 +307,11 @@ logical_plan 06)--------Filter: t1.id = $1 07)----------TableScan: t1 projection=[id, name] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[id@1 as id, name@0 as name] -03)----NestedLoopJoinExec: join_type=Right -04)------FilterExec: id@0 = $1, projection=[name@1] -05)--------DataSourceExec: partitions=1, partition_sizes=[1] -06)------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[id@1 as id, name@0 as name] +02)--NestedLoopJoinExec: join_type=Right +03)----FilterExec: id@0 = $1, projection=[name@1] +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)----DataSourceExec: partitions=1, partition_sizes=[1] # CTE with placeholder query TT @@ -349,9 +322,8 @@ logical_plan 02)--Filter: t1.id = $1 03)----TableScan: t1 projection=[id, name] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--FilterExec: id@0 = $1 -03)----DataSourceExec: partitions=1, partition_sizes=[1] +01)FilterExec: id@0 = $1 +02)--DataSourceExec: partitions=1, partition_sizes=[1] statement ok DROP TABLE t1; @@ -372,12 +344,11 @@ logical_plan 02)--Aggregate: groupBy=[[t1.id]], aggr=[[count(Int64(1))]] 03)----TableScan: t1 projection=[id] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[id@0 + $1 as t1.id + $1, count(Int64(1))@1 as count(*)] -03)----AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[count(Int64(1))] -04)------RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=1 -05)--------AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[count(Int64(1))] -06)----------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[id@0 + $1 as t1.id + $1, count(Int64(1))@1 as count(*)] +02)--AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[count(Int64(1))] +03)----RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=1 +04)------AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[count(Int64(1))] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] # Group by with placeholder in HAVING query TT @@ -389,13 +360,12 @@ logical_plan 03)----Aggregate: groupBy=[[t1.id]], aggr=[[count(Int64(1))]] 04)------TableScan: t1 projection=[id] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[id@0 as id, count(Int64(1))@1 as count(*)] -03)----FilterExec: count(Int64(1))@1 > $1 -04)------AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[count(Int64(1))] -05)--------RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=1 -06)----------AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[count(Int64(1))] -07)------------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[id@0 as id, count(Int64(1))@1 as count(*)] +02)--FilterExec: count(Int64(1))@1 > $1 +03)----AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[count(Int64(1))] +04)------RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=1 +05)--------AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[count(Int64(1))] +06)----------DataSourceExec: partitions=1, partition_sizes=[1] # Order by with placeholder query TT @@ -405,9 +375,8 @@ logical_plan 01)Sort: t1.id + CAST($1 AS Int32) ASC NULLS LAST 02)--TableScan: t1 projection=[id] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--SortExec: expr=[id@0 + $1 ASC NULLS LAST], preserve_partitioning=[false] -03)----DataSourceExec: partitions=1, partition_sizes=[1] +01)SortExec: expr=[id@0 + $1 ASC NULLS LAST], preserve_partitioning=[false] +02)--DataSourceExec: partitions=1, partition_sizes=[1] # Group by and Order by with placeholders query TT @@ -418,13 +387,12 @@ logical_plan 02)--Aggregate: groupBy=[[t1.name]], aggr=[[sum(CAST(t1.id AS Int64))]] 03)----TableScan: t1 projection=[id, name] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=2] -02)--SortPreservingMergeExec: [sum(t1.id)@1 + $1 ASC NULLS LAST] -03)----SortExec: expr=[sum(t1.id)@1 + $1 ASC NULLS LAST], preserve_partitioning=[true] -04)------AggregateExec: mode=FinalPartitioned, gby=[name@0 as name], aggr=[sum(t1.id)] -05)--------RepartitionExec: partitioning=Hash([name@0], 4), input_partitions=1 -06)----------AggregateExec: mode=Partial, gby=[name@1 as name], aggr=[sum(t1.id)] -07)------------DataSourceExec: partitions=1, partition_sizes=[1] +01)SortPreservingMergeExec: [sum(t1.id)@1 + $1 ASC NULLS LAST] +02)--SortExec: expr=[sum(t1.id)@1 + $1 ASC NULLS LAST], preserve_partitioning=[true] +03)----AggregateExec: mode=FinalPartitioned, gby=[name@0 as name], aggr=[sum(t1.id)] +04)------RepartitionExec: partitioning=Hash([name@0], 4), input_partitions=1 +05)--------AggregateExec: mode=Partial, gby=[name@1 as name], aggr=[sum(t1.id)] +06)----------DataSourceExec: partitions=1, partition_sizes=[1] statement ok DROP TABLE t1; @@ -441,9 +409,8 @@ logical_plan 01)Projection: CAST($1 AS Int32) 02)--EmptyRelation: rows=1 physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[$1] -03)----PlaceholderRowExec +01)ProjectionExec: expr=[$1] +02)--PlaceholderRowExec # CAST with placeholder query TT @@ -453,9 +420,8 @@ logical_plan 01)Projection: CAST($1 AS Int32) 02)--EmptyRelation: rows=1 physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[$1] -03)----PlaceholderRowExec +01)ProjectionExec: expr=[$1] +02)--PlaceholderRowExec # TRY_CAST with placeholder query TT @@ -465,9 +431,8 @@ logical_plan 01)Projection: TRY_CAST($1 AS Int32) 02)--EmptyRelation: rows=1 physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[TRY_CAST($1 AS Int32) as $1] -03)----PlaceholderRowExec +01)ProjectionExec: expr=[TRY_CAST($1 AS Int32) as $1] +02)--PlaceholderRowExec ########## ## IN and BETWEEN with placeholders @@ -484,9 +449,8 @@ logical_plan 01)Filter: t1.id = $1 OR t1.id = $2 OR t1.id = $3 02)--TableScan: t1 projection=[id] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--FilterExec: id@0 = $1 OR id@0 = $2 OR id@0 = $3 -03)----DataSourceExec: partitions=1, partition_sizes=[1] +01)FilterExec: id@0 = $1 OR id@0 = $2 OR id@0 = $3 +02)--DataSourceExec: partitions=1, partition_sizes=[1] # BETWEEN with placeholders query TT @@ -496,9 +460,8 @@ logical_plan 01)Filter: t1.id >= $1 AND t1.id <= $2 02)--TableScan: t1 projection=[id] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--FilterExec: id@0 >= $1 AND id@0 <= $2 -03)----DataSourceExec: partitions=1, partition_sizes=[1] +01)FilterExec: id@0 >= $1 AND id@0 <= $2 +02)--DataSourceExec: partitions=1, partition_sizes=[1] ########## ## String and Arithmetic operations with placeholders @@ -512,9 +475,8 @@ logical_plan 01)Projection: CAST($1 AS Utf8View) || CAST($2 AS Utf8View) 02)--EmptyRelation: rows=1 physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[$1 || $2] -03)----PlaceholderRowExec +01)ProjectionExec: expr=[$1 || $2] +02)--PlaceholderRowExec # Arithmetic with placeholders query TT @@ -524,9 +486,8 @@ logical_plan 01)Projection: $1 + CAST($2 AS Int64) * CAST($3 AS Int64) 02)--EmptyRelation: rows=1 physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[$1 + CAST($2 AS Int64) * CAST($3 AS Int64) as $1 + $2 * $3] -03)----PlaceholderRowExec +01)ProjectionExec: expr=[$1 + CAST($2 AS Int64) * CAST($3 AS Int64) as $1 + $2 * $3] +02)--PlaceholderRowExec ########## ## LIKE and Regex with placeholders @@ -540,9 +501,8 @@ logical_plan 01)Filter: t1.name LIKE $1 02)--TableScan: t1 projection=[id, name] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--FilterExec: name@1 LIKE $1 -03)----DataSourceExec: partitions=1, partition_sizes=[1] +01)FilterExec: name@1 LIKE $1 +02)--DataSourceExec: partitions=1, partition_sizes=[1] # Regex with placeholder query TT @@ -552,9 +512,8 @@ logical_plan 01)Filter: t1.name ~ $1 02)--TableScan: t1 projection=[id, name] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--FilterExec: name@1 ~ $1 -03)----DataSourceExec: partitions=1, partition_sizes=[1] +01)FilterExec: name@1 ~ $1 +02)--DataSourceExec: partitions=1, partition_sizes=[1] statement ok DROP TABLE t1; From 9f7eae0518d766e5cee8e19b49f4c02145353b43 Mon Sep 17 00:00:00 2001 From: Albert Skalt <133099191+askalt@users.noreply.github.com> Date: Thu, 5 Feb 2026 00:30:22 +0300 Subject: [PATCH 3/6] Wrap immutable plan parts into Arc (make creating `ExecutionPlan`s less costly) (#19893) - Closes https://github.com/apache/datafusion/issues/19852 Improve performance of query planning and plan state re-set by making node clone cheap. - Store projection as `Option>` instead of `Option>` in `FilterExec`, `HashJoinExec`, `NestedLoopJoinExec`. - Store exprs as `Arc<[ProjectionExpr]>` instead of Vec in `ProjectionExprs`. - Store arced aggregation, filter, group by expressions within `AggregateExec`. --- datafusion/common/src/stats.rs | 13 +- datafusion/common/src/utils/mod.rs | 4 +- datafusion/core/src/physical_planner.rs | 2 +- .../physical_optimizer/join_selection.rs | 2 +- datafusion/physical-expr/src/projection.rs | 95 ++++-- .../src/enforce_distribution.rs | 44 +-- .../src/enforce_sorting/sort_pushdown.rs | 2 +- .../physical-optimizer/src/join_selection.rs | 50 +--- .../src/projection_pushdown.rs | 2 +- .../physical-plan/src/aggregates/mod.rs | 43 +-- .../src/aggregates/no_grouping.rs | 6 +- .../physical-plan/src/aggregates/row_hash.rs | 12 +- .../src/aggregates/topk_stream.rs | 4 +- datafusion/physical-plan/src/common.rs | 2 +- datafusion/physical-plan/src/filter.rs | 60 ++-- .../physical-plan/src/joins/hash_join/exec.rs | 276 ++++++++++++------ .../physical-plan/src/joins/hash_join/mod.rs | 2 +- datafusion/physical-plan/src/joins/mod.rs | 6 +- .../src/joins/nested_loop_join.rs | 154 +++++++--- datafusion/physical-plan/src/joins/utils.rs | 4 +- datafusion/physical-plan/src/projection.rs | 18 +- datafusion/proto/src/physical_plan/mod.rs | 2 +- 22 files changed, 512 insertions(+), 291 deletions(-) diff --git a/datafusion/common/src/stats.rs b/datafusion/common/src/stats.rs index ba13ef392d912..cecf1d03418d7 100644 --- a/datafusion/common/src/stats.rs +++ b/datafusion/common/src/stats.rs @@ -391,8 +391,13 @@ impl Statistics { /// For example, if we had statistics for columns `{"a", "b", "c"}`, /// projecting to `vec![2, 1]` would return statistics for columns `{"c", /// "b"}`. - pub fn project(mut self, projection: Option<&Vec>) -> Self { - let Some(projection) = projection else { + pub fn project(self, projection: Option<&impl AsRef<[usize]>>) -> Self { + let projection = projection.map(AsRef::as_ref); + self.project_impl(projection) + } + + fn project_impl(mut self, projection: Option<&[usize]>) -> Self { + let Some(projection) = projection.map(AsRef::as_ref) else { return self; }; @@ -410,7 +415,7 @@ impl Statistics { .map(Slot::Present) .collect(); - for idx in projection { + for idx in projection.iter() { let next_idx = self.column_statistics.len(); let slot = std::mem::replace( columns.get_mut(*idx).expect("projection out of bounds"), @@ -1066,7 +1071,7 @@ mod tests { #[test] fn test_project_none() { - let projection = None; + let projection: Option> = None; let stats = make_stats(vec![10, 20, 30]).project(projection.as_ref()); assert_eq!(stats, make_stats(vec![10, 20, 30])); } diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index 03310a7bde193..016b188e3d6b8 100644 --- a/datafusion/common/src/utils/mod.rs +++ b/datafusion/common/src/utils/mod.rs @@ -70,10 +70,10 @@ use std::thread::available_parallelism; /// ``` pub fn project_schema( schema: &SchemaRef, - projection: Option<&Vec>, + projection: Option<&impl AsRef<[usize]>>, ) -> Result { let schema = match projection { - Some(columns) => Arc::new(schema.project(columns)?), + Some(columns) => Arc::new(schema.project(columns.as_ref())?), None => Arc::clone(schema), }; Ok(schema) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 54c5867c1804a..6d1bed069ff52 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -1007,7 +1007,7 @@ impl DefaultPhysicalPlanner { // project the output columns excluding the async functions // The async functions are always appended to the end of the schema. .apply_projection(Some( - (0..input.schema().fields().len()).collect(), + (0..input.schema().fields().len()).collect::>(), ))? .with_batch_size(session_state.config().batch_size()) .build()? diff --git a/datafusion/core/tests/physical_optimizer/join_selection.rs b/datafusion/core/tests/physical_optimizer/join_selection.rs index 9234a95591baa..b640159ca8463 100644 --- a/datafusion/core/tests/physical_optimizer/join_selection.rs +++ b/datafusion/core/tests/physical_optimizer/join_selection.rs @@ -762,7 +762,7 @@ async fn test_hash_join_swap_on_joins_with_projections( "ProjectionExec won't be added above if HashJoinExec contains embedded projection", ); - assert_eq!(swapped_join.projection, Some(vec![0_usize])); + assert_eq!(swapped_join.projection.as_deref().unwrap(), &[0_usize]); assert_eq!(swapped.schema().fields.len(), 1); assert_eq!(swapped.schema().fields[0].name(), "small_col"); Ok(()) diff --git a/datafusion/physical-expr/src/projection.rs b/datafusion/physical-expr/src/projection.rs index a84cb618333d3..c6adccc9b49de 100644 --- a/datafusion/physical-expr/src/projection.rs +++ b/datafusion/physical-expr/src/projection.rs @@ -29,7 +29,8 @@ use arrow::datatypes::{Field, Schema, SchemaRef}; use datafusion_common::stats::{ColumnStatistics, Precision}; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_common::{ - Result, ScalarValue, assert_or_internal_err, internal_datafusion_err, plan_err, + Result, ScalarValue, Statistics, assert_or_internal_err, internal_datafusion_err, + plan_err, }; use datafusion_physical_expr_common::metrics::ExecutionPlanMetricsSet; @@ -125,7 +126,8 @@ impl From for (Arc, String) { /// indices. #[derive(Debug, Clone, PartialEq, Eq)] pub struct ProjectionExprs { - exprs: Vec, + /// [`Arc`] used for a cheap clone, which improves physical plan optimization performance. + exprs: Arc<[ProjectionExpr]>, } impl std::fmt::Display for ProjectionExprs { @@ -137,14 +139,16 @@ impl std::fmt::Display for ProjectionExprs { impl From> for ProjectionExprs { fn from(value: Vec) -> Self { - Self { exprs: value } + Self { + exprs: value.into(), + } } } impl From<&[ProjectionExpr]> for ProjectionExprs { fn from(value: &[ProjectionExpr]) -> Self { Self { - exprs: value.to_vec(), + exprs: value.iter().cloned().collect(), } } } @@ -152,7 +156,7 @@ impl From<&[ProjectionExpr]> for ProjectionExprs { impl FromIterator for ProjectionExprs { fn from_iter>(exprs: T) -> Self { Self { - exprs: exprs.into_iter().collect::>(), + exprs: exprs.into_iter().collect(), } } } @@ -164,12 +168,17 @@ impl AsRef<[ProjectionExpr]> for ProjectionExprs { } impl ProjectionExprs { - pub fn new(exprs: I) -> Self - where - I: IntoIterator, - { + /// Make a new [`ProjectionExprs`] from expressions iterator. + pub fn new(exprs: impl IntoIterator) -> Self { + Self { + exprs: exprs.into_iter().collect(), + } + } + + /// Make a new [`ProjectionExprs`] from expressions. + pub fn from_expressions(exprs: impl Into>) -> Self { Self { - exprs: exprs.into_iter().collect::>(), + exprs: exprs.into(), } } @@ -285,13 +294,14 @@ impl ProjectionExprs { { let exprs = self .exprs - .into_iter() + .iter() + .cloned() .map(|mut proj| { proj.expr = f(proj.expr)?; Ok(proj) }) - .collect::>>()?; - Ok(Self::new(exprs)) + .collect::>>()?; + Ok(Self::from_expressions(exprs)) } /// Apply another projection on top of this projection, returning the combined projection. @@ -361,7 +371,7 @@ impl ProjectionExprs { /// applied on top of this projection. pub fn try_merge(&self, other: &ProjectionExprs) -> Result { let mut new_exprs = Vec::with_capacity(other.exprs.len()); - for proj_expr in &other.exprs { + for proj_expr in other.exprs.iter() { let new_expr = update_expr(&proj_expr.expr, &self.exprs, true)? .ok_or_else(|| { internal_datafusion_err!( @@ -614,12 +624,12 @@ impl ProjectionExprs { /// ``` pub fn project_statistics( &self, - mut stats: datafusion_common::Statistics, + mut stats: Statistics, output_schema: &Schema, - ) -> Result { + ) -> Result { let mut column_statistics = vec![]; - for proj_expr in &self.exprs { + for proj_expr in self.exprs.iter() { let expr = &proj_expr.expr; let col_stats = if let Some(col) = expr.as_any().downcast_ref::() { std::mem::take(&mut stats.column_statistics[col.index()]) @@ -766,13 +776,52 @@ impl Projector { } } -impl IntoIterator for ProjectionExprs { - type Item = ProjectionExpr; - type IntoIter = std::vec::IntoIter; +/// Describes an immutable reference counted projection. +/// +/// This structure represents projecting a set of columns by index. +/// [`Arc`] is used to make it cheap to clone. +pub type ProjectionRef = Arc<[usize]>; - fn into_iter(self) -> Self::IntoIter { - self.exprs.into_iter() - } +/// Combine two projections. +/// +/// If `p1` is [`None`] then there are no changes. +/// Otherwise, if passed `p2` is not [`None`] then it is remapped +/// according to the `p1`. Otherwise, there are no changes. +/// +/// # Example +/// +/// If stored projection is [0, 2] and we call `apply_projection([0, 2, 3])`, +/// then the resulting projection will be [0, 3]. +/// +/// # Error +/// +/// Returns an internal error if `p1` contains index that is greater than `p2` len. +/// +pub fn combine_projections( + p1: Option<&ProjectionRef>, + p2: Option<&ProjectionRef>, +) -> Result> { + let Some(p1) = p1 else { + return Ok(None); + }; + let Some(p2) = p2 else { + return Ok(Some(Arc::clone(p1))); + }; + + Ok(Some( + p1.iter() + .map(|i| { + let idx = *i; + assert_or_internal_err!( + idx < p2.len(), + "unable to apply projection: index {} is greater than new projection len {}", + idx, + p2.len(), + ); + Ok(p2[*i]) + }) + .collect::>>()?, + )) } /// The function operates in two modes: diff --git a/datafusion/physical-optimizer/src/enforce_distribution.rs b/datafusion/physical-optimizer/src/enforce_distribution.rs index acb1c588097ee..790669b5c9dbf 100644 --- a/datafusion/physical-optimizer/src/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/enforce_distribution.rs @@ -49,7 +49,7 @@ use datafusion_physical_plan::aggregates::{ use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::execution_plan::EmissionType; use datafusion_physical_plan::joins::{ - CrossJoinExec, HashJoinExec, PartitionMode, SortMergeJoinExec, + CrossJoinExec, HashJoinExec, HashJoinExecBuilder, PartitionMode, SortMergeJoinExec, }; use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion_physical_plan::repartition::RepartitionExec; @@ -305,18 +305,19 @@ pub fn adjust_input_keys_ordering( Vec<(PhysicalExprRef, PhysicalExprRef)>, Vec, )| { - HashJoinExec::try_new( + HashJoinExecBuilder::new( Arc::clone(left), Arc::clone(right), new_conditions.0, - filter.clone(), - join_type, - // TODO: although projection is not used in the join here, because projection pushdown is after enforce_distribution. Maybe we need to handle it later. Same as filter. - projection.clone(), - PartitionMode::Partitioned, - *null_equality, - *null_aware, + *join_type, ) + .with_filter(filter.clone()) + // TODO: although projection is not used in the join here, because projection pushdown is after enforce_distribution. Maybe we need to handle it later. Same as filter. + .with_projection_ref(projection.clone()) + .with_partition_mode(PartitionMode::Partitioned) + .with_null_equality(*null_equality) + .with_null_aware(*null_aware) + .build() .map(|e| Arc::new(e) as _) }; return reorder_partitioned_join_keys( @@ -638,17 +639,20 @@ pub fn reorder_join_keys_to_inputs( right_keys, } = join_keys; let new_join_on = new_join_conditions(&left_keys, &right_keys); - return Ok(Arc::new(HashJoinExec::try_new( - Arc::clone(left), - Arc::clone(right), - new_join_on, - filter.clone(), - join_type, - projection.clone(), - PartitionMode::Partitioned, - *null_equality, - *null_aware, - )?)); + return Ok(Arc::new( + HashJoinExecBuilder::new( + Arc::clone(left), + Arc::clone(right), + new_join_on, + *join_type, + ) + .with_filter(filter.clone()) + .with_projection_ref(projection.clone()) + .with_partition_mode(PartitionMode::Partitioned) + .with_null_equality(*null_equality) + .with_null_aware(*null_aware) + .build()?, + )); } } } else if let Some(SortMergeJoinExec { diff --git a/datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs index 698fdea8e766e..2dc61ba2453fb 100644 --- a/datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs +++ b/datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs @@ -723,7 +723,7 @@ fn handle_hash_join( .collect(); let column_indices = build_join_column_index(plan); - let projected_indices: Vec<_> = if let Some(projection) = &plan.projection { + let projected_indices: Vec<_> = if let Some(projection) = plan.projection.as_ref() { projection.iter().map(|&i| &column_indices[i]).collect() } else { column_indices.iter().collect() diff --git a/datafusion/physical-optimizer/src/join_selection.rs b/datafusion/physical-optimizer/src/join_selection.rs index 7412d0ba97812..02ef378d704a0 100644 --- a/datafusion/physical-optimizer/src/join_selection.rs +++ b/datafusion/physical-optimizer/src/join_selection.rs @@ -34,7 +34,7 @@ use datafusion_physical_expr::expressions::Column; use datafusion_physical_plan::execution_plan::EmissionType; use datafusion_physical_plan::joins::utils::ColumnIndex; use datafusion_physical_plan::joins::{ - CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, + CrossJoinExec, HashJoinExec, HashJoinExecBuilder, NestedLoopJoinExec, PartitionMode, StreamJoinPartitionMode, SymmetricHashJoinExec, }; use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; @@ -191,30 +191,18 @@ pub(crate) fn try_collect_left( { Ok(Some(hash_join.swap_inputs(PartitionMode::CollectLeft)?)) } else { - Ok(Some(Arc::new(HashJoinExec::try_new( - Arc::clone(left), - Arc::clone(right), - hash_join.on().to_vec(), - hash_join.filter().cloned(), - hash_join.join_type(), - hash_join.projection.clone(), - PartitionMode::CollectLeft, - hash_join.null_equality(), - hash_join.null_aware, - )?))) + Ok(Some(Arc::new( + HashJoinExecBuilder::from(hash_join) + .with_partition_mode(PartitionMode::CollectLeft) + .build()?, + ))) } } - (true, false) => Ok(Some(Arc::new(HashJoinExec::try_new( - Arc::clone(left), - Arc::clone(right), - hash_join.on().to_vec(), - hash_join.filter().cloned(), - hash_join.join_type(), - hash_join.projection.clone(), - PartitionMode::CollectLeft, - hash_join.null_equality(), - hash_join.null_aware, - )?))), + (true, false) => Ok(Some(Arc::new( + HashJoinExecBuilder::from(hash_join) + .with_partition_mode(PartitionMode::CollectLeft) + .build()?, + ))), (false, true) => { // Don't swap null-aware anti joins as they have specific side requirements if hash_join.join_type().supports_swap() && !hash_join.null_aware { @@ -254,17 +242,11 @@ pub(crate) fn partitioned_hash_join( PartitionMode::Partitioned }; - Ok(Arc::new(HashJoinExec::try_new( - Arc::clone(left), - Arc::clone(right), - hash_join.on().to_vec(), - hash_join.filter().cloned(), - hash_join.join_type(), - hash_join.projection.clone(), - partition_mode, - hash_join.null_equality(), - hash_join.null_aware, - )?)) + Ok(Arc::new( + HashJoinExecBuilder::from(hash_join) + .with_partition_mode(partition_mode) + .build()?, + )) } } diff --git a/datafusion/physical-optimizer/src/projection_pushdown.rs b/datafusion/physical-optimizer/src/projection_pushdown.rs index 99922ba075cc0..44d0926a8b250 100644 --- a/datafusion/physical-optimizer/src/projection_pushdown.rs +++ b/datafusion/physical-optimizer/src/projection_pushdown.rs @@ -135,7 +135,7 @@ fn try_push_down_join_filter( ); let new_lhs_length = lhs_rewrite.data.0.schema().fields.len(); - let projections = match projections { + let projections = match projections.as_ref() { None => match join.join_type() { JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full => { // Build projections that ignore the newly projected columns. diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 464e5ce6c7c1a..8915513a6eb60 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -626,11 +626,14 @@ pub struct AggregateExec { /// Aggregation mode (full, partial) mode: AggregateMode, /// Group by expressions - group_by: PhysicalGroupBy, + /// [`Arc`] used for a cheap clone, which improves physical plan optimization performance. + group_by: Arc, /// Aggregate expressions - aggr_expr: Vec>, + /// The same reason to [`Arc`] it as for [`Self::group_by`]. + aggr_expr: Arc<[Arc]>, /// FILTER (WHERE clause) expression for each aggregate expression - filter_expr: Vec>>, + /// The same reason to [`Arc`] it as for [`Self::group_by`]. + filter_expr: Arc<[Option>]>, /// Configuration for limit-based optimizations limit_options: Option, /// Input plan, could be a partial aggregate or the input to the aggregate @@ -664,18 +667,18 @@ impl AggregateExec { /// Rewrites aggregate exec with new aggregate expressions. pub fn with_new_aggr_exprs( &self, - aggr_expr: Vec>, + aggr_expr: impl Into]>>, ) -> Self { Self { - aggr_expr, + aggr_expr: aggr_expr.into(), // clone the rest of the fields required_input_ordering: self.required_input_ordering.clone(), metrics: ExecutionPlanMetricsSet::new(), input_order_mode: self.input_order_mode.clone(), cache: self.cache.clone(), mode: self.mode, - group_by: self.group_by.clone(), - filter_expr: self.filter_expr.clone(), + group_by: Arc::clone(&self.group_by), + filter_expr: Arc::clone(&self.filter_expr), limit_options: self.limit_options, input: Arc::clone(&self.input), schema: Arc::clone(&self.schema), @@ -694,9 +697,9 @@ impl AggregateExec { input_order_mode: self.input_order_mode.clone(), cache: self.cache.clone(), mode: self.mode, - group_by: self.group_by.clone(), - aggr_expr: self.aggr_expr.clone(), - filter_expr: self.filter_expr.clone(), + group_by: Arc::clone(&self.group_by), + aggr_expr: Arc::clone(&self.aggr_expr), + filter_expr: Arc::clone(&self.filter_expr), input: Arc::clone(&self.input), schema: Arc::clone(&self.schema), input_schema: Arc::clone(&self.input_schema), @@ -711,12 +714,13 @@ impl AggregateExec { /// Create a new hash aggregate execution plan pub fn try_new( mode: AggregateMode, - group_by: PhysicalGroupBy, + group_by: impl Into>, aggr_expr: Vec>, filter_expr: Vec>>, input: Arc, input_schema: SchemaRef, ) -> Result { + let group_by = group_by.into(); let schema = create_schema(&input.schema(), &group_by, &aggr_expr, mode)?; let schema = Arc::new(schema); @@ -741,13 +745,16 @@ impl AggregateExec { /// the schema in such cases. fn try_new_with_schema( mode: AggregateMode, - group_by: PhysicalGroupBy, + group_by: impl Into>, mut aggr_expr: Vec>, - filter_expr: Vec>>, + filter_expr: impl Into>]>>, input: Arc, input_schema: SchemaRef, schema: SchemaRef, ) -> Result { + let group_by = group_by.into(); + let filter_expr = filter_expr.into(); + // Make sure arguments are consistent in size assert_eq_or_internal_err!( aggr_expr.len(), @@ -814,13 +821,13 @@ impl AggregateExec { &group_expr_mapping, &mode, &input_order_mode, - aggr_expr.as_slice(), + aggr_expr.as_ref(), )?; let mut exec = AggregateExec { mode, group_by, - aggr_expr, + aggr_expr: aggr_expr.into(), filter_expr, input, schema, @@ -1370,9 +1377,9 @@ impl ExecutionPlan for AggregateExec { ) -> Result> { let mut me = AggregateExec::try_new_with_schema( self.mode, - self.group_by.clone(), - self.aggr_expr.clone(), - self.filter_expr.clone(), + Arc::clone(&self.group_by), + self.aggr_expr.to_vec(), + Arc::clone(&self.filter_expr), Arc::clone(&children[0]), Arc::clone(&self.input_schema), Arc::clone(&self.schema), diff --git a/datafusion/physical-plan/src/aggregates/no_grouping.rs b/datafusion/physical-plan/src/aggregates/no_grouping.rs index eb9b6766ab8ee..fe8942097ac93 100644 --- a/datafusion/physical-plan/src/aggregates/no_grouping.rs +++ b/datafusion/physical-plan/src/aggregates/no_grouping.rs @@ -62,7 +62,7 @@ struct AggregateStreamInner { mode: AggregateMode, input: SendableRecordBatchStream, aggregate_expressions: Vec>>, - filter_expressions: Vec>>, + filter_expressions: Arc<[Option>]>, // ==== Runtime States/Buffers ==== accumulators: Vec, @@ -277,7 +277,7 @@ impl AggregateStream { partition: usize, ) -> Result { let agg_schema = Arc::clone(&agg.schema); - let agg_filter_expr = agg.filter_expr.clone(); + let agg_filter_expr = Arc::clone(&agg.filter_expr); let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); let input = agg.input.execute(partition, Arc::clone(context))?; @@ -285,7 +285,7 @@ impl AggregateStream { let aggregate_expressions = aggregate_expressions(&agg.aggr_expr, &agg.mode, 0)?; let filter_expressions = match agg.mode.input_mode() { AggregateInputMode::Raw => agg_filter_expr, - AggregateInputMode::Partial => vec![None; agg.aggr_expr.len()], + AggregateInputMode::Partial => vec![None; agg.aggr_expr.len()].into(), }; let accumulators = create_accumulators(&agg.aggr_expr)?; diff --git a/datafusion/physical-plan/src/aggregates/row_hash.rs b/datafusion/physical-plan/src/aggregates/row_hash.rs index b2cf396b1500d..de857370ce285 100644 --- a/datafusion/physical-plan/src/aggregates/row_hash.rs +++ b/datafusion/physical-plan/src/aggregates/row_hash.rs @@ -377,10 +377,10 @@ pub(crate) struct GroupedHashAggregateStream { /// /// For example, for an aggregate like `SUM(x) FILTER (WHERE x >= 100)`, /// the filter expression is `x > 100`. - filter_expressions: Vec>>, + filter_expressions: Arc<[Option>]>, /// GROUP BY expressions - group_by: PhysicalGroupBy, + group_by: Arc, /// max rows in output RecordBatches batch_size: usize, @@ -465,8 +465,8 @@ impl GroupedHashAggregateStream { ) -> Result { debug!("Creating GroupedHashAggregateStream"); let agg_schema = Arc::clone(&agg.schema); - let agg_group_by = agg.group_by.clone(); - let agg_filter_expr = agg.filter_expr.clone(); + let agg_group_by = Arc::clone(&agg.group_by); + let agg_filter_expr = Arc::clone(&agg.filter_expr); let batch_size = context.session_config().batch_size(); let input = agg.input.execute(partition, Arc::clone(context))?; @@ -475,7 +475,7 @@ impl GroupedHashAggregateStream { let timer = baseline_metrics.elapsed_compute().timer(); - let aggregate_exprs = agg.aggr_expr.clone(); + let aggregate_exprs = Arc::clone(&agg.aggr_expr); // arguments for each aggregate, one vec of expressions per // aggregate @@ -493,7 +493,7 @@ impl GroupedHashAggregateStream { let filter_expressions = match agg.mode.input_mode() { AggregateInputMode::Raw => agg_filter_expr, - AggregateInputMode::Partial => vec![None; agg.aggr_expr.len()], + AggregateInputMode::Partial => vec![None; agg.aggr_expr.len()].into(), }; // Instantiate the accumulators diff --git a/datafusion/physical-plan/src/aggregates/topk_stream.rs b/datafusion/physical-plan/src/aggregates/topk_stream.rs index 72c5d0c86745d..4aa566ccfcd0a 100644 --- a/datafusion/physical-plan/src/aggregates/topk_stream.rs +++ b/datafusion/physical-plan/src/aggregates/topk_stream.rs @@ -50,7 +50,7 @@ pub struct GroupedTopKAggregateStream { baseline_metrics: BaselineMetrics, group_by_metrics: GroupByMetrics, aggregate_arguments: Vec>>, - group_by: PhysicalGroupBy, + group_by: Arc, priority_map: PriorityMap, } @@ -62,7 +62,7 @@ impl GroupedTopKAggregateStream { limit: usize, ) -> Result { let agg_schema = Arc::clone(&aggr.schema); - let group_by = aggr.group_by.clone(); + let group_by = Arc::clone(&aggr.group_by); let input = aggr.input.execute(partition, Arc::clone(context))?; let baseline_metrics = BaselineMetrics::new(&aggr.metrics, partition); let group_by_metrics = GroupByMetrics::new(&aggr.metrics, partition); diff --git a/datafusion/physical-plan/src/common.rs b/datafusion/physical-plan/src/common.rs index 32dc60b56ad48..590f6f09e8b9e 100644 --- a/datafusion/physical-plan/src/common.rs +++ b/datafusion/physical-plan/src/common.rs @@ -181,7 +181,7 @@ pub fn compute_record_batch_statistics( /// Checks if the given projection is valid for the given schema. pub fn can_project( schema: &arrow::datatypes::SchemaRef, - projection: Option<&Vec>, + projection: Option<&[usize]>, ) -> Result<()> { match projection { Some(columns) => { diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index 79e3c2960c2ad..dc5ca465ccc23 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -20,6 +20,7 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll, ready}; +use datafusion_physical_expr::projection::{ProjectionRef, combine_projections}; use itertools::Itertools; use super::{ @@ -86,7 +87,7 @@ pub struct FilterExec { /// Properties equivalence properties, partitioning, etc. cache: PlanProperties, /// The projection indices of the columns in the output schema of join - projection: Option>, + projection: Option, /// Target batch size for output batches batch_size: usize, /// Number of rows to fetch @@ -97,7 +98,7 @@ pub struct FilterExec { pub struct FilterExecBuilder { predicate: Arc, input: Arc, - projection: Option>, + projection: Option, default_selectivity: u8, batch_size: usize, fetch: Option, @@ -137,18 +138,19 @@ impl FilterExecBuilder { /// /// If no projection is currently set, the new projection is used directly. /// If `None` is passed, the projection is cleared. - pub fn apply_projection(mut self, projection: Option>) -> Result { + pub fn apply_projection(self, projection: Option>) -> Result { + let projection = projection.map(Into::into); + self.apply_projection_by_ref(projection.as_ref()) + } + + /// The same as [`Self::apply_projection`] but takes projection shared reference. + pub fn apply_projection_by_ref( + mut self, + projection: Option<&ProjectionRef>, + ) -> Result { // Check if the projection is valid against current output schema - can_project(&self.input.schema(), projection.as_ref())?; - self.projection = match projection { - Some(new_proj) => match &self.projection { - Some(existing_proj) => { - Some(new_proj.iter().map(|i| existing_proj[*i]).collect()) - } - None => Some(new_proj), - }, - None => None, - }; + can_project(&self.input.schema(), projection.map(AsRef::as_ref))?; + self.projection = combine_projections(projection, self.projection.as_ref())?; Ok(self) } @@ -190,16 +192,14 @@ impl FilterExecBuilder { } // Validate projection if provided - if let Some(ref proj) = self.projection { - can_project(&self.input.schema(), Some(proj))?; - } + can_project(&self.input.schema(), self.projection.as_deref())?; // Compute properties once with all parameters let cache = FilterExec::compute_properties( &self.input, &self.predicate, self.default_selectivity, - self.projection.as_ref(), + self.projection.as_deref(), )?; Ok(FilterExec { @@ -303,8 +303,8 @@ impl FilterExec { } /// Projection - pub fn projection(&self) -> Option<&Vec> { - self.projection.as_ref() + pub fn projection(&self) -> &Option { + &self.projection } /// Calculates `Statistics` for `FilterExec`, by applying selectivity (either default, or estimated) to input statistics. @@ -381,7 +381,7 @@ impl FilterExec { input: &Arc, predicate: &Arc, default_selectivity: u8, - projection: Option<&Vec>, + projection: Option<&[usize]>, ) -> Result { // Combine the equal predicates with the input equivalence properties // to construct the equivalence properties: @@ -420,7 +420,7 @@ impl FilterExec { if let Some(projection) = projection { let schema = eq_properties.schema(); let projection_mapping = ProjectionMapping::from_indices(projection, schema)?; - let out_schema = project_schema(schema, Some(projection))?; + let out_schema = project_schema(schema, Some(&projection))?; output_partitioning = output_partitioning.project(&projection_mapping, &eq_properties); eq_properties = eq_properties.project(&projection_mapping, out_schema); @@ -665,7 +665,7 @@ impl ExecutionPlan for FilterExec { let new_predicate = conjunction(unhandled_filters); let updated_node = if new_predicate.eq(&lit(true)) { // FilterExec is no longer needed, but we may need to leave a projection in place - match self.projection() { + match self.projection().as_ref() { Some(projection_indices) => { let filter_child_schema = filter_input.schema(); let proj_exprs = projection_indices @@ -701,7 +701,7 @@ impl ExecutionPlan for FilterExec { &filter_input, &new_predicate, self.default_selectivity, - self.projection.as_ref(), + self.projection.as_deref(), )?, projection: self.projection.clone(), batch_size: self.batch_size, @@ -839,7 +839,7 @@ struct FilterExecStream { /// Runtime metrics recording metrics: FilterExecMetrics, /// The projection indices of the columns in the input schema - projection: Option>, + projection: Option, /// Batch coalescer to combine small batches batch_coalescer: LimitedBatchCoalescer, } @@ -930,8 +930,8 @@ impl Stream for FilterExecStream { .evaluate(&batch) .and_then(|v| v.into_array(batch.num_rows())) .and_then(|array| { - Ok(match self.projection { - Some(ref projection) => { + Ok(match self.projection.as_ref() { + Some(projection) => { let projected_batch = batch.project(projection)?; (array, projected_batch) }, @@ -1752,7 +1752,7 @@ mod tests { .build()?; // Verify projection is set correctly - assert_eq!(filter.projection(), Some(&vec![0, 2])); + assert_eq!(filter.projection(), &Some([0, 2].into())); // Verify schema contains only projected columns let output_schema = filter.schema(); @@ -1782,7 +1782,7 @@ mod tests { let filter = FilterExecBuilder::new(predicate, input).build()?; // Verify no projection is set - assert_eq!(filter.projection(), None); + assert!(filter.projection().is_none()); // Verify schema contains all columns let output_schema = filter.schema(); @@ -1996,7 +1996,7 @@ mod tests { .build()?; // Verify composed projection is [0, 3] - assert_eq!(filter.projection(), Some(&vec![0, 3])); + assert_eq!(filter.projection(), &Some([0, 3].into())); // Verify schema contains only columns a and d let output_schema = filter.schema(); @@ -2030,7 +2030,7 @@ mod tests { .build()?; // Projection should be cleared - assert_eq!(filter.projection(), None); + assert_eq!(filter.projection(), &None); // Schema should have all columns let output_schema = filter.schema(); diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 98f7d04383402..44e67d190a1c3 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -83,6 +83,7 @@ use datafusion_physical_expr::equivalence::{ ProjectionMapping, join_equivalence_properties, }; use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit}; +use datafusion_physical_expr::projection::{ProjectionRef, combine_projections}; use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef}; use ahash::RandomState; @@ -248,6 +249,172 @@ impl JoinLeftData { } } +/// Helps to build [`HashJoinExec`]. +pub struct HashJoinExecBuilder { + left: Arc, + right: Arc, + on: Vec<(PhysicalExprRef, PhysicalExprRef)>, + join_type: JoinType, + filter: Option, + projection: Option, + partition_mode: PartitionMode, + null_equality: NullEquality, + null_aware: bool, +} + +impl HashJoinExecBuilder { + /// Make a new [`HashJoinExecBuilder`]. + pub fn new( + left: Arc, + right: Arc, + on: Vec<(PhysicalExprRef, PhysicalExprRef)>, + join_type: JoinType, + ) -> Self { + Self { + left, + right, + on, + filter: None, + projection: None, + partition_mode: PartitionMode::Auto, + join_type, + null_equality: NullEquality::NullEqualsNothing, + null_aware: false, + } + } + + /// Set projection from the vector. + pub fn with_projection(self, projection: Option>) -> Self { + self.with_projection_ref(projection.map(Into::into)) + } + + /// Set projection from the shared reference. + pub fn with_projection_ref(mut self, projection: Option) -> Self { + self.projection = projection; + self + } + + /// Set optional filter. + pub fn with_filter(mut self, filter: Option) -> Self { + self.filter = filter; + self + } + + /// Set partition mode. + pub fn with_partition_mode(mut self, mode: PartitionMode) -> Self { + self.partition_mode = mode; + self + } + + /// Set null equality property. + pub fn with_null_equality(mut self, null_equality: NullEquality) -> Self { + self.null_equality = null_equality; + self + } + + /// Set null aware property. + pub fn with_null_aware(mut self, null_aware: bool) -> Self { + self.null_aware = null_aware; + self + } + + /// Build resulting execution plan. + pub fn build(self) -> Result { + let Self { + left, + right, + on, + join_type, + filter, + projection, + partition_mode, + null_equality, + null_aware, + } = self; + + let left_schema = left.schema(); + let right_schema = right.schema(); + if on.is_empty() { + return plan_err!("On constraints in HashJoinExec should be non-empty"); + } + + check_join_is_valid(&left_schema, &right_schema, &on)?; + + // Validate null_aware flag + if null_aware { + if !matches!(join_type, JoinType::LeftAnti) { + return plan_err!( + "null_aware can only be true for LeftAnti joins, got {join_type}" + ); + } + if on.len() != 1 { + return plan_err!( + "null_aware anti join only supports single column join key, got {} columns", + on.len() + ); + } + } + + let (join_schema, column_indices) = + build_join_schema(&left_schema, &right_schema, &join_type); + + let random_state = HASH_JOIN_SEED; + + let join_schema = Arc::new(join_schema); + + // check if the projection is valid + can_project(&join_schema, projection.as_deref())?; + + let cache = HashJoinExec::compute_properties( + &left, + &right, + &join_schema, + join_type, + &on, + partition_mode, + projection.as_deref(), + )?; + + // Initialize both dynamic filter and bounds accumulator to None + // They will be set later if dynamic filtering is enabled + + Ok(HashJoinExec { + left, + right, + on, + filter, + join_type, + join_schema, + left_fut: Default::default(), + random_state, + mode: partition_mode, + metrics: ExecutionPlanMetricsSet::new(), + projection, + column_indices, + null_equality, + null_aware, + cache, + dynamic_filter: None, + }) + } +} + +impl From<&HashJoinExec> for HashJoinExecBuilder { + fn from(exec: &HashJoinExec) -> Self { + Self { + left: Arc::clone(exec.left()), + right: Arc::clone(exec.right()), + on: exec.on.clone(), + join_type: exec.join_type, + filter: exec.filter.clone(), + projection: exec.projection.clone(), + partition_mode: exec.mode, + null_equality: exec.null_equality, + null_aware: exec.null_aware, + } + } +} + #[expect(rustdoc::private_intra_doc_links)] /// Join execution plan: Evaluates equijoin predicates in parallel on multiple /// partitions using a hash table and an optional filter list to apply post @@ -468,7 +635,7 @@ pub struct HashJoinExec { /// Execution metrics metrics: ExecutionPlanMetricsSet, /// The projection indices of the columns in the output schema of join - pub projection: Option>, + pub projection: Option, /// Information of index and left / right placement of columns column_indices: Vec, /// The equality null-handling behavior of the join algorithm. @@ -537,70 +704,13 @@ impl HashJoinExec { null_equality: NullEquality, null_aware: bool, ) -> Result { - let left_schema = left.schema(); - let right_schema = right.schema(); - if on.is_empty() { - return plan_err!("On constraints in HashJoinExec should be non-empty"); - } - - check_join_is_valid(&left_schema, &right_schema, &on)?; - - // Validate null_aware flag - if null_aware { - if !matches!(join_type, JoinType::LeftAnti) { - return plan_err!( - "null_aware can only be true for LeftAnti joins, got {join_type}" - ); - } - if on.len() != 1 { - return plan_err!( - "null_aware anti join only supports single column join key, got {} columns", - on.len() - ); - } - } - - let (join_schema, column_indices) = - build_join_schema(&left_schema, &right_schema, join_type); - - let random_state = HASH_JOIN_SEED; - - let join_schema = Arc::new(join_schema); - - // check if the projection is valid - can_project(&join_schema, projection.as_ref())?; - - let cache = Self::compute_properties( - &left, - &right, - &join_schema, - *join_type, - &on, - partition_mode, - projection.as_ref(), - )?; - - // Initialize both dynamic filter and bounds accumulator to None - // They will be set later if dynamic filtering is enabled - - Ok(HashJoinExec { - left, - right, - on, - filter, - join_type: *join_type, - join_schema, - left_fut: Default::default(), - random_state, - mode: partition_mode, - metrics: ExecutionPlanMetricsSet::new(), - projection, - column_indices, - null_equality, - null_aware, - cache, - dynamic_filter: None, - }) + HashJoinExecBuilder::new(left, right, on, *join_type) + .with_filter(filter) + .with_projection(projection) + .with_partition_mode(partition_mode) + .with_null_equality(null_equality) + .with_null_aware(null_aware) + .build() } fn create_dynamic_filter(on: &JoinOn) -> Arc { @@ -689,26 +799,14 @@ impl HashJoinExec { /// Return new instance of [HashJoinExec] with the given projection. pub fn with_projection(&self, projection: Option>) -> Result { + let projection = projection.map(Into::into); // check if the projection is valid - can_project(&self.schema(), projection.as_ref())?; - let projection = match projection { - Some(projection) => match &self.projection { - Some(p) => Some(projection.iter().map(|i| p[*i]).collect()), - None => Some(projection), - }, - None => None, - }; - Self::try_new( - Arc::clone(&self.left), - Arc::clone(&self.right), - self.on.clone(), - self.filter.clone(), - &self.join_type, - projection, - self.mode, - self.null_equality, - self.null_aware, - ) + can_project(&self.schema(), projection.as_deref())?; + let projection = + combine_projections(projection.as_ref(), self.projection.as_ref())?; + HashJoinExecBuilder::from(self) + .with_projection_ref(projection) + .build() } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. @@ -719,7 +817,7 @@ impl HashJoinExec { join_type: JoinType, on: JoinOnRef, mode: PartitionMode, - projection: Option<&Vec>, + projection: Option<&[usize]>, ) -> Result { // Calculate equivalence properties: let mut eq_properties = join_equivalence_properties( @@ -771,7 +869,7 @@ impl HashJoinExec { if let Some(projection) = projection { // construct a map from the input expressions to the output expression of the Projection let projection_mapping = ProjectionMapping::from_indices(projection, schema)?; - let out_schema = project_schema(schema, Some(projection))?; + let out_schema = project_schema(schema, Some(&projection))?; output_partitioning = output_partitioning.project(&projection_mapping, &eq_properties); eq_properties = eq_properties.project(&projection_mapping, out_schema); @@ -826,7 +924,7 @@ impl HashJoinExec { swap_join_projection( left.schema().fields().len(), right.schema().fields().len(), - self.projection.as_ref(), + self.projection.as_deref(), self.join_type(), ), partition_mode, @@ -1022,7 +1120,7 @@ impl ExecutionPlan for HashJoinExec { self.join_type, &self.on, self.mode, - self.projection.as_ref(), + self.projection.as_deref(), )?, // Keep the dynamic filter, bounds accumulator will be reset dynamic_filter: self.dynamic_filter.clone(), @@ -1183,7 +1281,7 @@ impl ExecutionPlan for HashJoinExec { let right_stream = self.right.execute(partition, context)?; // update column indices to reflect the projection - let column_indices_after_projection = match &self.projection { + let column_indices_after_projection = match self.projection.as_ref() { Some(projection) => projection .iter() .map(|i| self.column_indices[*i].clone()) diff --git a/datafusion/physical-plan/src/joins/hash_join/mod.rs b/datafusion/physical-plan/src/joins/hash_join/mod.rs index 8592e1d968535..b915802ea4015 100644 --- a/datafusion/physical-plan/src/joins/hash_join/mod.rs +++ b/datafusion/physical-plan/src/joins/hash_join/mod.rs @@ -17,7 +17,7 @@ //! [`HashJoinExec`] Partitioned Hash Join Operator -pub use exec::HashJoinExec; +pub use exec::{HashJoinExec, HashJoinExecBuilder}; pub use partitioned_hash_eval::{HashExpr, HashTableLookupExpr, SeededRandomState}; mod exec; diff --git a/datafusion/physical-plan/src/joins/mod.rs b/datafusion/physical-plan/src/joins/mod.rs index 848d0472fe885..2cdfa1e6ac020 100644 --- a/datafusion/physical-plan/src/joins/mod.rs +++ b/datafusion/physical-plan/src/joins/mod.rs @@ -20,8 +20,10 @@ use arrow::array::BooleanBufferBuilder; pub use cross_join::CrossJoinExec; use datafusion_physical_expr::PhysicalExprRef; -pub use hash_join::{HashExpr, HashJoinExec, HashTableLookupExpr, SeededRandomState}; -pub use nested_loop_join::NestedLoopJoinExec; +pub use hash_join::{ + HashExpr, HashJoinExec, HashJoinExecBuilder, HashTableLookupExpr, SeededRandomState, +}; +pub use nested_loop_join::{NestedLoopJoinExec, NestedLoopJoinExecBuilder}; use parking_lot::Mutex; // Note: SortMergeJoin is not used in plans yet pub use piecewise_merge_join::PiecewiseMergeJoinExec; diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index d4c0d1ac4a870..796f5567e1eac 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -73,7 +73,10 @@ use datafusion_physical_expr::equivalence::{ ProjectionMapping, join_equivalence_properties, }; -use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::{ + PhysicalExpr, + projection::{ProjectionRef, combine_projections}, +}; use futures::{Stream, StreamExt, TryStreamExt}; use log::debug; use parking_lot::Mutex; @@ -195,7 +198,7 @@ pub struct NestedLoopJoinExec { /// Information of index and left / right placement of columns column_indices: Vec, /// Projection to apply to the output of the join - projection: Option>, + projection: Option, /// Execution metrics metrics: ExecutionPlanMetricsSet, @@ -203,34 +206,76 @@ pub struct NestedLoopJoinExec { cache: PlanProperties, } -impl NestedLoopJoinExec { - /// Try to create a new [`NestedLoopJoinExec`] - pub fn try_new( +/// Helps to build [`NestedLoopJoinExec`]. +pub struct NestedLoopJoinExecBuilder { + left: Arc, + right: Arc, + join_type: JoinType, + filter: Option, + projection: Option, +} + +impl NestedLoopJoinExecBuilder { + /// Make a new [`NestedLoopJoinExecBuilder`]. + pub fn new( left: Arc, right: Arc, - filter: Option, - join_type: &JoinType, - projection: Option>, - ) -> Result { + join_type: JoinType, + ) -> Self { + Self { + left, + right, + join_type, + filter: None, + projection: None, + } + } + + /// Set projection from the vector. + pub fn with_projection(self, projection: Option>) -> Self { + self.with_projection_ref(projection.map(Into::into)) + } + + /// Set projection from the shared reference. + pub fn with_projection_ref(mut self, projection: Option) -> Self { + self.projection = projection; + self + } + + /// Set optional filter. + pub fn with_filter(mut self, filter: Option) -> Self { + self.filter = filter; + self + } + + /// Build resulting execution plan. + pub fn build(self) -> Result { + let Self { + left, + right, + join_type, + filter, + projection, + } = self; + let left_schema = left.schema(); let right_schema = right.schema(); check_join_is_valid(&left_schema, &right_schema, &[])?; let (join_schema, column_indices) = - build_join_schema(&left_schema, &right_schema, join_type); + build_join_schema(&left_schema, &right_schema, &join_type); let join_schema = Arc::new(join_schema); - let cache = Self::compute_properties( + let cache = NestedLoopJoinExec::compute_properties( &left, &right, &join_schema, - *join_type, - projection.as_ref(), + join_type, + projection.as_deref(), )?; - Ok(NestedLoopJoinExec { left, right, filter, - join_type: *join_type, + join_type, join_schema, build_side_data: Default::default(), column_indices, @@ -239,6 +284,34 @@ impl NestedLoopJoinExec { cache, }) } +} + +impl From<&NestedLoopJoinExec> for NestedLoopJoinExecBuilder { + fn from(exec: &NestedLoopJoinExec) -> Self { + Self { + left: Arc::clone(exec.left()), + right: Arc::clone(exec.right()), + join_type: exec.join_type, + filter: exec.filter.clone(), + projection: exec.projection.clone(), + } + } +} + +impl NestedLoopJoinExec { + /// Try to create a new [`NestedLoopJoinExec`] + pub fn try_new( + left: Arc, + right: Arc, + filter: Option, + join_type: &JoinType, + projection: Option>, + ) -> Result { + NestedLoopJoinExecBuilder::new(left, right, *join_type) + .with_projection(projection) + .with_filter(filter) + .build() + } /// left side pub fn left(&self) -> &Arc { @@ -260,8 +333,8 @@ impl NestedLoopJoinExec { &self.join_type } - pub fn projection(&self) -> Option<&Vec> { - self.projection.as_ref() + pub fn projection(&self) -> &Option { + &self.projection } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. @@ -270,7 +343,7 @@ impl NestedLoopJoinExec { right: &Arc, schema: &SchemaRef, join_type: JoinType, - projection: Option<&Vec>, + projection: Option<&[usize]>, ) -> Result { // Calculate equivalence properties: let mut eq_properties = join_equivalence_properties( @@ -313,7 +386,7 @@ impl NestedLoopJoinExec { if let Some(projection) = projection { // construct a map from the input expressions to the output expression of the Projection let projection_mapping = ProjectionMapping::from_indices(projection, schema)?; - let out_schema = project_schema(schema, Some(projection))?; + let out_schema = project_schema(schema, Some(&projection))?; output_partitioning = output_partitioning.project(&projection_mapping, &eq_properties); eq_properties = eq_properties.project(&projection_mapping, out_schema); @@ -337,22 +410,14 @@ impl NestedLoopJoinExec { } pub fn with_projection(&self, projection: Option>) -> Result { + let projection = projection.map(Into::into); // check if the projection is valid - can_project(&self.schema(), projection.as_ref())?; - let projection = match projection { - Some(projection) => match &self.projection { - Some(p) => Some(projection.iter().map(|i| p[*i]).collect()), - None => Some(projection), - }, - None => None, - }; - Self::try_new( - Arc::clone(&self.left), - Arc::clone(&self.right), - self.filter.clone(), - &self.join_type, - projection, - ) + can_project(&self.schema(), projection.as_deref())?; + let projection = + combine_projections(projection.as_ref(), self.projection.as_ref())?; + NestedLoopJoinExecBuilder::from(self) + .with_projection_ref(projection) + .build() } /// Returns a new `ExecutionPlan` that runs NestedLoopsJoins with the left @@ -374,7 +439,7 @@ impl NestedLoopJoinExec { swap_join_projection( left.schema().fields().len(), right.schema().fields().len(), - self.projection.as_ref(), + self.projection.as_deref(), self.join_type(), ), )?; @@ -479,13 +544,16 @@ impl ExecutionPlan for NestedLoopJoinExec { self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(NestedLoopJoinExec::try_new( - Arc::clone(&children[0]), - Arc::clone(&children[1]), - self.filter.clone(), - &self.join_type, - self.projection.clone(), - )?)) + Ok(Arc::new( + NestedLoopJoinExecBuilder::new( + Arc::clone(&children[0]), + Arc::clone(&children[1]), + self.join_type, + ) + .with_filter(self.filter.clone()) + .with_projection_ref(self.projection.clone()) + .build()?, + )) } fn execute( @@ -524,7 +592,7 @@ impl ExecutionPlan for NestedLoopJoinExec { let probe_side_data = self.right.execute(partition, context)?; // update column indices to reflect the projection - let column_indices_after_projection = match &self.projection { + let column_indices_after_projection = match self.projection.as_ref() { Some(projection) => projection .iter() .map(|i| self.column_indices[*i].clone()) diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index a9243fe04e28d..e709703e07d45 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -1674,7 +1674,7 @@ fn swap_reverting_projection( pub fn swap_join_projection( left_schema_len: usize, right_schema_len: usize, - projection: Option<&Vec>, + projection: Option<&[usize]>, join_type: &JoinType, ) -> Option> { match join_type { @@ -1685,7 +1685,7 @@ pub fn swap_join_projection( | JoinType::RightAnti | JoinType::RightSemi | JoinType::LeftMark - | JoinType::RightMark => projection.cloned(), + | JoinType::RightMark => projection.map(|p| p.to_vec()), _ => projection.map(|p| { p.iter() .map(|i| { diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index a4de3f0101ade..905bd13309ed2 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -140,13 +140,19 @@ impl ProjectionExec { E: Into, { let input_schema = input.schema(); - // convert argument to Vec - let expr_vec = expr.into_iter().map(Into::into).collect::>(); - let projection = ProjectionExprs::new(expr_vec); + let expr_arc = expr.into_iter().map(Into::into).collect::>(); + let projection = ProjectionExprs::from_expressions(expr_arc); let projector = projection.make_projector(&input_schema)?; + Self::try_from_projector(projector, input) + } + fn try_from_projector( + projector: Projector, + input: Arc, + ) -> Result { // Construct a map from the input expressions to the output expression of the Projection - let projection_mapping = projection.projection_mapping(&input_schema)?; + let projection_mapping = + projector.projection().projection_mapping(&input.schema())?; let cache = Self::compute_properties( &input, &projection_mapping, @@ -307,8 +313,8 @@ impl ExecutionPlan for ProjectionExec { self: Arc, mut children: Vec>, ) -> Result> { - ProjectionExec::try_new( - self.projector.projection().clone(), + ProjectionExec::try_from_projector( + self.projector.clone(), children.swap_remove(0), ) .map(|p| Arc::new(p) as _) diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 8fa82eed966fa..55646f937cf1d 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -3214,7 +3214,7 @@ impl protobuf::PhysicalPlanNode { right: Some(Box::new(right)), join_type: join_type.into(), filter, - projection: exec.projection().map_or_else(Vec::new, |v| { + projection: exec.projection().as_ref().map_or_else(Vec::new, |v| { v.iter().map(|x| *x as u32).collect::>() }), }, From c7497ebad2189f8876d4b18e23ded19885adbfa0 Mon Sep 17 00:00:00 2001 From: Albert Skalt Date: Tue, 13 Jan 2026 17:10:23 +0300 Subject: [PATCH 4/6] add fast-path for `with_new_children` This patch aims to implement a fast-path for the ExecutionPlan::with_new_children function for some plans, moving closer to a physical plan re-use implementation and improving planning performance. If the passed children properties are the same as in self, we do not actually recompute self's properties (which could be costly if projection mapping is required). Instead, we just replace the children and re-use self's properties as-is. To be able to compare two different properties -- ExecutionPlan::properties(...) signature is modified and now returns `&Arc`. If `children` properties are the same in `with_new_children` -- we clone our properties arc and then a parent plan will consider our properties as unchanged, doing the same. - Return `&Arc` from `ExecutionPlan::properties(...)` instead of a reference. - Implement `with_new_children` fast-path if there is no children properties changes for all major plans. Note: currently, `reset_plan_states` does not allow to re-use plan in general: it is not supported for dynamic filters and recursive queries features, as in this case state reset should update pointers in the children plans. Closes https://github.com/apache/datafusion/issues/19796 --- .../custom_data_source/custom_datasource.rs | 6 +- .../memory_pool_execution_plan.rs | 4 +- .../proto/composed_extension_codec.rs | 4 +- .../examples/relation_planner/table_sample.rs | 6 +- datafusion/catalog/src/memory/table.rs | 6 +- datafusion/core/benches/reset_plan_states.rs | 2 + datafusion/core/src/physical_planner.rs | 14 +-- .../core/tests/custom_sources_cases/mod.rs | 9 +- .../provider_filter_pushdown.rs | 9 +- .../tests/custom_sources_cases/statistics.rs | 6 +- datafusion/core/tests/fuzz_cases/once_exec.rs | 6 +- .../enforce_distribution.rs | 6 +- .../physical_optimizer/join_selection.rs | 12 +-- .../physical_optimizer/pushdown_utils.rs | 2 +- .../tests/physical_optimizer/test_utils.rs | 8 +- .../tests/user_defined/insert_operation.rs | 20 ++-- .../tests/user_defined/user_defined_plan.rs | 10 +- datafusion/datasource/src/sink.rs | 6 +- datafusion/datasource/src/source.rs | 20 ++-- datafusion/ffi/src/execution_plan.rs | 20 ++-- datafusion/ffi/src/tests/async_provider.rs | 8 +- .../src/equivalence/properties/mod.rs | 7 +- .../physical-optimizer/src/ensure_coop.rs | 6 +- .../src/output_requirements.rs | 6 +- .../physical-plan/src/aggregates/mod.rs | 38 ++++++-- datafusion/physical-plan/src/analyze.rs | 6 +- datafusion/physical-plan/src/async_func.rs | 26 +++-- datafusion/physical-plan/src/buffer.rs | 27 ++++-- .../physical-plan/src/coalesce_batches.rs | 27 ++++-- .../physical-plan/src/coalesce_partitions.rs | 30 ++++-- datafusion/physical-plan/src/coop.rs | 26 +++-- datafusion/physical-plan/src/display.rs | 2 +- datafusion/physical-plan/src/empty.rs | 8 +- .../physical-plan/src/execution_plan.rs | 57 +++++++++-- datafusion/physical-plan/src/explain.rs | 6 +- datafusion/physical-plan/src/filter.rs | 31 ++++-- .../physical-plan/src/joins/cross_join.rs | 32 +++++-- .../physical-plan/src/joins/hash_join/exec.rs | 35 ++++--- .../src/joins/nested_loop_join.rs | 30 +++++- .../src/joins/piecewise_merge_join/exec.rs | 94 +++++++++++++------ .../src/joins/sort_merge_join/exec.rs | 24 ++++- .../src/joins/symmetric_hash_join.rs | 23 ++++- datafusion/physical-plan/src/limit.rs | 51 +++++++--- datafusion/physical-plan/src/memory.rs | 19 ++-- .../physical-plan/src/placeholder_row.rs | 8 +- datafusion/physical-plan/src/projection.rs | 24 ++++- .../physical-plan/src/recursive_query.rs | 6 +- .../physical-plan/src/repartition/mod.rs | 34 +++++-- .../physical-plan/src/sorts/partial_sort.rs | 28 ++++-- datafusion/physical-plan/src/sorts/sort.rs | 50 +++++----- .../src/sorts/sort_preserving_merge.rs | 34 +++++-- datafusion/physical-plan/src/streaming.rs | 8 +- datafusion/physical-plan/src/test.rs | 14 +-- datafusion/physical-plan/src/test/exec.rs | 38 ++++---- datafusion/physical-plan/src/union.rs | 41 ++++++-- datafusion/physical-plan/src/unnest.rs | 25 +++-- .../src/windows/bounded_window_agg_exec.rs | 22 ++++- .../src/windows/window_agg_exec.rs | 28 ++++-- datafusion/physical-plan/src/work_table.rs | 8 +- .../custom-table-providers.md | 6 +- docs/source/library-user-guide/upgrading.md | 55 +++++++++++ 61 files changed, 859 insertions(+), 365 deletions(-) diff --git a/datafusion-examples/examples/custom_data_source/custom_datasource.rs b/datafusion-examples/examples/custom_data_source/custom_datasource.rs index b276ae32cf247..7abb39e1a7130 100644 --- a/datafusion-examples/examples/custom_data_source/custom_datasource.rs +++ b/datafusion-examples/examples/custom_data_source/custom_datasource.rs @@ -192,7 +192,7 @@ impl TableProvider for CustomDataSource { struct CustomExec { db: CustomDataSource, projected_schema: SchemaRef, - cache: PlanProperties, + cache: Arc, } impl CustomExec { @@ -207,7 +207,7 @@ impl CustomExec { Self { db, projected_schema, - cache, + cache: Arc::new(cache), } } @@ -238,7 +238,7 @@ impl ExecutionPlan for CustomExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs b/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs index 48475acbb1542..6f377ea1ce3b9 100644 --- a/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs +++ b/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs @@ -199,7 +199,7 @@ impl ExternalBatchBufferer { struct BufferingExecutionPlan { schema: SchemaRef, input: Arc, - properties: PlanProperties, + properties: Arc, } impl BufferingExecutionPlan { @@ -233,7 +233,7 @@ impl ExecutionPlan for BufferingExecutionPlan { self.schema.clone() } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.properties } diff --git a/datafusion-examples/examples/proto/composed_extension_codec.rs b/datafusion-examples/examples/proto/composed_extension_codec.rs index f3910d461b6a8..b4f3d4f098996 100644 --- a/datafusion-examples/examples/proto/composed_extension_codec.rs +++ b/datafusion-examples/examples/proto/composed_extension_codec.rs @@ -106,7 +106,7 @@ impl ExecutionPlan for ParentExec { self } - fn properties(&self) -> &datafusion::physical_plan::PlanProperties { + fn properties(&self) -> &Arc { unreachable!() } @@ -182,7 +182,7 @@ impl ExecutionPlan for ChildExec { self } - fn properties(&self) -> &datafusion::physical_plan::PlanProperties { + fn properties(&self) -> &Arc { unreachable!() } diff --git a/datafusion-examples/examples/relation_planner/table_sample.rs b/datafusion-examples/examples/relation_planner/table_sample.rs index 657432ef31362..895f2fdd4ff3a 100644 --- a/datafusion-examples/examples/relation_planner/table_sample.rs +++ b/datafusion-examples/examples/relation_planner/table_sample.rs @@ -618,7 +618,7 @@ pub struct SampleExec { upper_bound: f64, seed: u64, metrics: ExecutionPlanMetricsSet, - cache: PlanProperties, + cache: Arc, } impl SampleExec { @@ -656,7 +656,7 @@ impl SampleExec { upper_bound, seed, metrics: ExecutionPlanMetricsSet::new(), - cache, + cache: Arc::new(cache), }) } @@ -686,7 +686,7 @@ impl ExecutionPlan for SampleExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/catalog/src/memory/table.rs b/datafusion/catalog/src/memory/table.rs index 7865eb016bee1..484b5f805e547 100644 --- a/datafusion/catalog/src/memory/table.rs +++ b/datafusion/catalog/src/memory/table.rs @@ -549,7 +549,7 @@ fn evaluate_filters_to_mask( struct DmlResultExec { rows_affected: u64, schema: SchemaRef, - properties: PlanProperties, + properties: Arc, } impl DmlResultExec { @@ -570,7 +570,7 @@ impl DmlResultExec { Self { rows_affected, schema, - properties, + properties: Arc::new(properties), } } } @@ -604,7 +604,7 @@ impl ExecutionPlan for DmlResultExec { Arc::clone(&self.schema) } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.properties } diff --git a/datafusion/core/benches/reset_plan_states.rs b/datafusion/core/benches/reset_plan_states.rs index 0988702e6f950..a4f9e716bd761 100644 --- a/datafusion/core/benches/reset_plan_states.rs +++ b/datafusion/core/benches/reset_plan_states.rs @@ -167,6 +167,8 @@ fn run_reset_states(b: &mut criterion::Bencher, plan: &Arc) { /// making an independent instance of the execution plan to re-execute it, avoiding /// re-planning stage. fn bench_reset_plan_states(c: &mut Criterion) { + env_logger::init(); + let rt = Runtime::new().unwrap(); let ctx = SessionContext::new(); ctx.register_table( diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 6d1bed069ff52..ce070b4b0203e 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -3722,13 +3722,15 @@ mod tests { #[derive(Debug)] struct NoOpExecutionPlan { - cache: PlanProperties, + cache: Arc, } impl NoOpExecutionPlan { fn new(schema: SchemaRef) -> Self { let cache = Self::compute_properties(schema); - Self { cache } + Self { + cache: Arc::new(cache), + } } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. @@ -3766,7 +3768,7 @@ mod tests { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -3920,7 +3922,7 @@ digraph { fn children(&self) -> Vec<&Arc> { self.0.iter().collect::>() } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { unimplemented!() } fn execute( @@ -3969,7 +3971,7 @@ digraph { fn children(&self) -> Vec<&Arc> { unimplemented!() } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { unimplemented!() } fn execute( @@ -4090,7 +4092,7 @@ digraph { fn children(&self) -> Vec<&Arc> { vec![] } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { unimplemented!() } fn execute( diff --git a/datafusion/core/tests/custom_sources_cases/mod.rs b/datafusion/core/tests/custom_sources_cases/mod.rs index 8453615c2886b..6dd09ebd8832f 100644 --- a/datafusion/core/tests/custom_sources_cases/mod.rs +++ b/datafusion/core/tests/custom_sources_cases/mod.rs @@ -79,7 +79,7 @@ struct CustomTableProvider; #[derive(Debug, Clone)] struct CustomExecutionPlan { projection: Option>, - cache: PlanProperties, + cache: Arc, } impl CustomExecutionPlan { @@ -88,7 +88,10 @@ impl CustomExecutionPlan { let schema = project_schema(&schema, projection.as_ref()).expect("projected schema"); let cache = Self::compute_properties(schema); - Self { projection, cache } + Self { + projection, + cache: Arc::new(cache), + } } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. @@ -157,7 +160,7 @@ impl ExecutionPlan for CustomExecutionPlan { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs b/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs index ca1eaa1f958ea..7a624d0636530 100644 --- a/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs +++ b/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs @@ -62,13 +62,16 @@ fn create_batch(value: i32, num_rows: usize) -> Result { #[derive(Debug)] struct CustomPlan { batches: Vec, - cache: PlanProperties, + cache: Arc, } impl CustomPlan { fn new(schema: SchemaRef, batches: Vec) -> Self { let cache = Self::compute_properties(schema); - Self { batches, cache } + Self { + batches, + cache: Arc::new(cache), + } } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. @@ -109,7 +112,7 @@ impl ExecutionPlan for CustomPlan { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/core/tests/custom_sources_cases/statistics.rs b/datafusion/core/tests/custom_sources_cases/statistics.rs index 820c2a470b376..6964d72eacffe 100644 --- a/datafusion/core/tests/custom_sources_cases/statistics.rs +++ b/datafusion/core/tests/custom_sources_cases/statistics.rs @@ -45,7 +45,7 @@ use async_trait::async_trait; struct StatisticsValidation { stats: Statistics, schema: Arc, - cache: PlanProperties, + cache: Arc, } impl StatisticsValidation { @@ -59,7 +59,7 @@ impl StatisticsValidation { Self { stats, schema, - cache, + cache: Arc::new(cache), } } @@ -158,7 +158,7 @@ impl ExecutionPlan for StatisticsValidation { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/core/tests/fuzz_cases/once_exec.rs b/datafusion/core/tests/fuzz_cases/once_exec.rs index 49e2caaa7417c..69edf9be1d825 100644 --- a/datafusion/core/tests/fuzz_cases/once_exec.rs +++ b/datafusion/core/tests/fuzz_cases/once_exec.rs @@ -32,7 +32,7 @@ use std::sync::{Arc, Mutex}; pub struct OnceExec { /// the results to send back stream: Mutex>, - cache: PlanProperties, + cache: Arc, } impl Debug for OnceExec { @@ -46,7 +46,7 @@ impl OnceExec { let cache = Self::compute_properties(stream.schema()); Self { stream: Mutex::new(Some(stream)), - cache, + cache: Arc::new(cache), } } @@ -83,7 +83,7 @@ impl ExecutionPlan for OnceExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 94ae82a9ad755..b3a3d29d070b1 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -120,7 +120,7 @@ macro_rules! assert_plan { struct SortRequiredExec { input: Arc, expr: LexOrdering, - cache: PlanProperties, + cache: Arc, } impl SortRequiredExec { @@ -132,7 +132,7 @@ impl SortRequiredExec { Self { input, expr: requirement, - cache, + cache: Arc::new(cache), } } @@ -174,7 +174,7 @@ impl ExecutionPlan for SortRequiredExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/core/tests/physical_optimizer/join_selection.rs b/datafusion/core/tests/physical_optimizer/join_selection.rs index b640159ca8463..ab35f399e2b11 100644 --- a/datafusion/core/tests/physical_optimizer/join_selection.rs +++ b/datafusion/core/tests/physical_optimizer/join_selection.rs @@ -979,7 +979,7 @@ impl RecordBatchStream for UnboundedStream { pub struct UnboundedExec { batch_produce: Option, batch: RecordBatch, - cache: PlanProperties, + cache: Arc, } impl UnboundedExec { @@ -995,7 +995,7 @@ impl UnboundedExec { Self { batch_produce, batch, - cache, + cache: Arc::new(cache), } } @@ -1052,7 +1052,7 @@ impl ExecutionPlan for UnboundedExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -1091,7 +1091,7 @@ pub enum SourceType { pub struct StatisticsExec { stats: Statistics, schema: Arc, - cache: PlanProperties, + cache: Arc, } impl StatisticsExec { @@ -1105,7 +1105,7 @@ impl StatisticsExec { Self { stats, schema: Arc::new(schema), - cache, + cache: Arc::new(cache), } } @@ -1153,7 +1153,7 @@ impl ExecutionPlan for StatisticsExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/core/tests/physical_optimizer/pushdown_utils.rs b/datafusion/core/tests/physical_optimizer/pushdown_utils.rs index 6db9d816c40aa..02a3f7d0f67dd 100644 --- a/datafusion/core/tests/physical_optimizer/pushdown_utils.rs +++ b/datafusion/core/tests/physical_optimizer/pushdown_utils.rs @@ -484,7 +484,7 @@ impl ExecutionPlan for TestNode { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { self.input.properties() } diff --git a/datafusion/core/tests/physical_optimizer/test_utils.rs b/datafusion/core/tests/physical_optimizer/test_utils.rs index feac8190ffde4..f8c91ba272a9f 100644 --- a/datafusion/core/tests/physical_optimizer/test_utils.rs +++ b/datafusion/core/tests/physical_optimizer/test_utils.rs @@ -454,7 +454,7 @@ impl ExecutionPlan for RequirementsTestExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { self.input.properties() } @@ -825,7 +825,7 @@ pub fn sort_expr_named(name: &str, index: usize) -> PhysicalSortExpr { pub struct TestScan { schema: SchemaRef, output_ordering: Vec, - plan_properties: PlanProperties, + plan_properties: Arc, // Store the requested ordering for display requested_ordering: Option, } @@ -859,7 +859,7 @@ impl TestScan { Self { schema, output_ordering, - plan_properties, + plan_properties: Arc::new(plan_properties), requested_ordering: None, } } @@ -915,7 +915,7 @@ impl ExecutionPlan for TestScan { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.plan_properties } diff --git a/datafusion/core/tests/user_defined/insert_operation.rs b/datafusion/core/tests/user_defined/insert_operation.rs index 7ad00dece1b24..4d2a31ca1f960 100644 --- a/datafusion/core/tests/user_defined/insert_operation.rs +++ b/datafusion/core/tests/user_defined/insert_operation.rs @@ -122,20 +122,22 @@ impl TableProvider for TestInsertTableProvider { #[derive(Debug)] struct TestInsertExec { op: InsertOp, - plan_properties: PlanProperties, + plan_properties: Arc, } impl TestInsertExec { fn new(op: InsertOp) -> Self { Self { op, - plan_properties: PlanProperties::new( - EquivalenceProperties::new(make_count_schema()), - Partitioning::UnknownPartitioning(1), - EmissionType::Incremental, - Boundedness::Bounded, - ) - .with_scheduling_type(SchedulingType::Cooperative), + plan_properties: Arc::new( + PlanProperties::new( + EquivalenceProperties::new(make_count_schema()), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + ) + .with_scheduling_type(SchedulingType::Cooperative), + ), } } } @@ -159,7 +161,7 @@ impl ExecutionPlan for TestInsertExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.plan_properties } diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs b/datafusion/core/tests/user_defined/user_defined_plan.rs index d53e076739608..c2533e73d2be9 100644 --- a/datafusion/core/tests/user_defined/user_defined_plan.rs +++ b/datafusion/core/tests/user_defined/user_defined_plan.rs @@ -653,13 +653,17 @@ struct TopKExec { input: Arc, /// The maximum number of values k: usize, - cache: PlanProperties, + cache: Arc, } impl TopKExec { fn new(input: Arc, k: usize) -> Self { let cache = Self::compute_properties(input.schema()); - Self { input, k, cache } + Self { + input, + k, + cache: Arc::new(cache), + } } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. @@ -704,7 +708,7 @@ impl ExecutionPlan for TopKExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index 5acc89722b200..f149109dff5cc 100644 --- a/datafusion/datasource/src/sink.rs +++ b/datafusion/datasource/src/sink.rs @@ -89,7 +89,7 @@ pub struct DataSinkExec { count_schema: SchemaRef, /// Optional required sort order for output data. sort_order: Option, - cache: PlanProperties, + cache: Arc, } impl Debug for DataSinkExec { @@ -117,7 +117,7 @@ impl DataSinkExec { sink, count_schema: make_count_schema(), sort_order, - cache, + cache: Arc::new(cache), } } @@ -174,7 +174,7 @@ impl ExecutionPlan for DataSinkExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index 75413c743b8e1..cffb4f41e61c9 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -267,7 +267,7 @@ pub struct DataSourceExec { /// The source of the data -- for example, `FileScanConfig` or `MemorySourceConfig` data_source: Arc, /// Cached plan properties such as sort order - cache: PlanProperties, + cache: Arc, } impl DisplayAs for DataSourceExec { @@ -291,7 +291,7 @@ impl ExecutionPlan for DataSourceExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -361,7 +361,7 @@ impl ExecutionPlan for DataSourceExec { fn with_fetch(&self, limit: Option) -> Option> { let data_source = self.data_source.with_fetch(limit)?; - let cache = self.cache.clone(); + let cache = Arc::clone(&self.cache); Some(Arc::new(Self { data_source, cache })) } @@ -405,7 +405,8 @@ impl ExecutionPlan for DataSourceExec { let mut new_node = self.clone(); new_node.data_source = data_source; // Re-compute properties since we have new filters which will impact equivalence info - new_node.cache = Self::compute_properties(&new_node.data_source); + new_node.cache = + Arc::new(Self::compute_properties(&new_node.data_source)); Ok(FilterPushdownPropagation { filters: res.filters, @@ -470,7 +471,10 @@ impl DataSourceExec { // Default constructor for `DataSourceExec`, setting the `cooperative` flag to `true`. pub fn new(data_source: Arc) -> Self { let cache = Self::compute_properties(&data_source); - Self { data_source, cache } + Self { + data_source, + cache: Arc::new(cache), + } } /// Return the source object @@ -479,20 +483,20 @@ impl DataSourceExec { } pub fn with_data_source(mut self, data_source: Arc) -> Self { - self.cache = Self::compute_properties(&data_source); + self.cache = Arc::new(Self::compute_properties(&data_source)); self.data_source = data_source; self } /// Assign constraints pub fn with_constraints(mut self, constraints: Constraints) -> Self { - self.cache = self.cache.with_constraints(constraints); + Arc::make_mut(&mut self.cache).set_constraints(constraints); self } /// Assign output partitioning pub fn with_partitioning(mut self, partitioning: Partitioning) -> Self { - self.cache = self.cache.with_partitioning(partitioning); + Arc::make_mut(&mut self.cache).partitioning = partitioning; self } diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index c879b022067c3..ec8f7538fee1a 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/datafusion/ffi/src/execution_plan.rs @@ -90,7 +90,7 @@ impl FFI_ExecutionPlan { unsafe extern "C" fn properties_fn_wrapper( plan: &FFI_ExecutionPlan, ) -> FFI_PlanProperties { - plan.inner().properties().into() + plan.inner().properties().as_ref().into() } unsafe extern "C" fn children_fn_wrapper( @@ -192,7 +192,7 @@ impl Drop for FFI_ExecutionPlan { pub struct ForeignExecutionPlan { name: String, plan: FFI_ExecutionPlan, - properties: PlanProperties, + properties: Arc, children: Vec>, } @@ -244,7 +244,7 @@ impl TryFrom<&FFI_ExecutionPlan> for Arc { let plan = ForeignExecutionPlan { name, plan: plan.clone(), - properties, + properties: Arc::new(properties), children, }; @@ -262,7 +262,7 @@ impl ExecutionPlan for ForeignExecutionPlan { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.properties } @@ -278,7 +278,7 @@ impl ExecutionPlan for ForeignExecutionPlan { plan: self.plan.clone(), name: self.name.clone(), children, - properties: self.properties.clone(), + properties: Arc::clone(&self.properties), })) } @@ -305,19 +305,19 @@ pub(crate) mod tests { #[derive(Debug)] pub struct EmptyExec { - props: PlanProperties, + props: Arc, children: Vec>, } impl EmptyExec { pub fn new(schema: arrow::datatypes::SchemaRef) -> Self { Self { - props: PlanProperties::new( + props: Arc::new(PlanProperties::new( datafusion::physical_expr::EquivalenceProperties::new(schema), Partitioning::UnknownPartitioning(3), EmissionType::Incremental, Boundedness::Bounded, - ), + )), children: Vec::default(), } } @@ -342,7 +342,7 @@ pub(crate) mod tests { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.props } @@ -355,7 +355,7 @@ pub(crate) mod tests { children: Vec>, ) -> Result> { Ok(Arc::new(EmptyExec { - props: self.props.clone(), + props: Arc::clone(&self.props), children, })) } diff --git a/datafusion/ffi/src/tests/async_provider.rs b/datafusion/ffi/src/tests/async_provider.rs index 6149736c58555..8370cf19e6589 100644 --- a/datafusion/ffi/src/tests/async_provider.rs +++ b/datafusion/ffi/src/tests/async_provider.rs @@ -162,7 +162,7 @@ impl Drop for AsyncTableProvider { #[derive(Debug)] struct AsyncTestExecutionPlan { - properties: datafusion_physical_plan::PlanProperties, + properties: Arc, batch_request: mpsc::Sender, batch_receiver: broadcast::Receiver>, } @@ -173,12 +173,12 @@ impl AsyncTestExecutionPlan { batch_receiver: broadcast::Receiver>, ) -> Self { Self { - properties: datafusion_physical_plan::PlanProperties::new( + properties: Arc::new(datafusion_physical_plan::PlanProperties::new( EquivalenceProperties::new(super::create_test_schema()), Partitioning::UnknownPartitioning(3), datafusion_physical_plan::execution_plan::EmissionType::Incremental, datafusion_physical_plan::execution_plan::Boundedness::Bounded, - ), + )), batch_request, batch_receiver, } @@ -194,7 +194,7 @@ impl ExecutionPlan for AsyncTestExecutionPlan { self } - fn properties(&self) -> &datafusion_physical_plan::PlanProperties { + fn properties(&self) -> &Arc { &self.properties } diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index 996bc4b08fcd2..a98341b10765a 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -207,8 +207,13 @@ impl EquivalenceProperties { } /// Adds constraints to the properties. - pub fn with_constraints(mut self, constraints: Constraints) -> Self { + pub fn set_constraints(&mut self, constraints: Constraints) { self.constraints = constraints; + } + + /// Adds constraints to the properties. + pub fn with_constraints(mut self, constraints: Constraints) -> Self { + self.set_constraints(constraints); self } diff --git a/datafusion/physical-optimizer/src/ensure_coop.rs b/datafusion/physical-optimizer/src/ensure_coop.rs index 5d00d00bce21d..ef8946f9a49d1 100644 --- a/datafusion/physical-optimizer/src/ensure_coop.rs +++ b/datafusion/physical-optimizer/src/ensure_coop.rs @@ -281,7 +281,7 @@ mod tests { input: Arc, scheduling_type: SchedulingType, evaluation_type: EvaluationType, - properties: PlanProperties, + properties: Arc, } impl DummyExec { @@ -305,7 +305,7 @@ mod tests { input, scheduling_type, evaluation_type, - properties, + properties: Arc::new(properties), } } } @@ -327,7 +327,7 @@ mod tests { fn as_any(&self) -> &dyn Any { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.properties } fn children(&self) -> Vec<&Arc> { diff --git a/datafusion/physical-optimizer/src/output_requirements.rs b/datafusion/physical-optimizer/src/output_requirements.rs index 0dc6a25fbc0b7..9c4169ec654f8 100644 --- a/datafusion/physical-optimizer/src/output_requirements.rs +++ b/datafusion/physical-optimizer/src/output_requirements.rs @@ -98,7 +98,7 @@ pub struct OutputRequirementExec { input: Arc, order_requirement: Option, dist_requirement: Distribution, - cache: PlanProperties, + cache: Arc, fetch: Option, } @@ -114,7 +114,7 @@ impl OutputRequirementExec { input, order_requirement: requirements, dist_requirement, - cache, + cache: Arc::new(cache), fetch, } } @@ -200,7 +200,7 @@ impl ExecutionPlan for OutputRequirementExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 8915513a6eb60..e4c5a86d6158f 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -25,7 +25,9 @@ use crate::aggregates::{ no_grouping::AggregateStream, row_hash::GroupedHashAggregateStream, topk_stream::GroupedTopKAggregateStream, }; -use crate::execution_plan::{CardinalityEffect, EmissionType, ReplacePhysicalExpr}; +use crate::execution_plan::{ + CardinalityEffect, EmissionType, ReplacePhysicalExpr, has_same_children_properties, +}; use crate::filter_pushdown::{ ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, PushedDownPredicate, @@ -33,7 +35,7 @@ use crate::filter_pushdown::{ use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use crate::{ DisplayFormatType, Distribution, ExecutionPlan, InputOrderMode, - SendableRecordBatchStream, Statistics, + SendableRecordBatchStream, Statistics, check_if_same_properties, }; use datafusion_common::config::ConfigOptions; use datafusion_physical_expr::utils::collect_columns; @@ -651,7 +653,7 @@ pub struct AggregateExec { required_input_ordering: Option, /// Describes how the input is ordered relative to the group by columns input_order_mode: InputOrderMode, - cache: PlanProperties, + cache: Arc, /// During initialization, if the plan supports dynamic filtering (see [`AggrDynFilter`]), /// it is set to `Some(..)` regardless of whether it can be pushed down to a child node. /// @@ -675,7 +677,7 @@ impl AggregateExec { required_input_ordering: self.required_input_ordering.clone(), metrics: ExecutionPlanMetricsSet::new(), input_order_mode: self.input_order_mode.clone(), - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), mode: self.mode, group_by: Arc::clone(&self.group_by), filter_expr: Arc::clone(&self.filter_expr), @@ -695,7 +697,7 @@ impl AggregateExec { required_input_ordering: self.required_input_ordering.clone(), metrics: ExecutionPlanMetricsSet::new(), input_order_mode: self.input_order_mode.clone(), - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), mode: self.mode, group_by: Arc::clone(&self.group_by), aggr_expr: Arc::clone(&self.aggr_expr), @@ -836,7 +838,7 @@ impl AggregateExec { required_input_ordering, limit_options: None, input_order_mode, - cache, + cache: Arc::new(cache), dynamic_filter: None, }; @@ -1194,6 +1196,17 @@ impl AggregateExec { _ => Precision::Absent, } } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for AggregateExec { @@ -1332,7 +1345,7 @@ impl ExecutionPlan for AggregateExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -1375,6 +1388,8 @@ impl ExecutionPlan for AggregateExec { self: Arc, children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); + let mut me = AggregateExec::try_new_with_schema( self.mode, Arc::clone(&self.group_by), @@ -2514,14 +2529,17 @@ mod tests { struct TestYieldingExec { /// True if this exec should yield back to runtime the first time it is polled pub yield_first: bool, - cache: PlanProperties, + cache: Arc, } impl TestYieldingExec { fn new(yield_first: bool) -> Self { let schema = some_data().0; let cache = Self::compute_properties(schema); - Self { yield_first, cache } + Self { + yield_first, + cache: Arc::new(cache), + } } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. @@ -2562,7 +2580,7 @@ mod tests { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/physical-plan/src/analyze.rs b/datafusion/physical-plan/src/analyze.rs index 1fb8f93a38782..eca31ea0e194f 100644 --- a/datafusion/physical-plan/src/analyze.rs +++ b/datafusion/physical-plan/src/analyze.rs @@ -51,7 +51,7 @@ pub struct AnalyzeExec { pub(crate) input: Arc, /// The output schema for RecordBatches of this exec node schema: SchemaRef, - cache: PlanProperties, + cache: Arc, } impl AnalyzeExec { @@ -70,7 +70,7 @@ impl AnalyzeExec { metric_types, input, schema, - cache, + cache: Arc::new(cache), } } @@ -131,7 +131,7 @@ impl ExecutionPlan for AnalyzeExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/physical-plan/src/async_func.rs b/datafusion/physical-plan/src/async_func.rs index a61fd95949d1a..93701a25b3bf4 100644 --- a/datafusion/physical-plan/src/async_func.rs +++ b/datafusion/physical-plan/src/async_func.rs @@ -16,10 +16,12 @@ // under the License. use crate::coalesce::LimitedBatchCoalescer; +use crate::execution_plan::has_same_children_properties; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use crate::stream::RecordBatchStreamAdapter; use crate::{ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + check_if_same_properties, }; use arrow::array::RecordBatch; use arrow_schema::{Fields, Schema, SchemaRef}; @@ -45,12 +47,12 @@ use std::task::{Context, Poll, ready}; /// /// The schema of the output of the AsyncFuncExec is: /// Input columns followed by one column for each async expression -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct AsyncFuncExec { /// The async expressions to evaluate async_exprs: Vec>, input: Arc, - cache: PlanProperties, + cache: Arc, metrics: ExecutionPlanMetricsSet, } @@ -84,7 +86,7 @@ impl AsyncFuncExec { Ok(Self { input, async_exprs, - cache, + cache: Arc::new(cache), metrics: ExecutionPlanMetricsSet::new(), }) } @@ -113,6 +115,17 @@ impl AsyncFuncExec { pub fn input(&self) -> &Arc { &self.input } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for AsyncFuncExec { @@ -149,7 +162,7 @@ impl ExecutionPlan for AsyncFuncExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -159,16 +172,17 @@ impl ExecutionPlan for AsyncFuncExec { fn with_new_children( self: Arc, - children: Vec>, + mut children: Vec>, ) -> Result> { assert_eq_or_internal_err!( children.len(), 1, "AsyncFuncExec wrong number of children" ); + check_if_same_properties!(self, children); Ok(Arc::new(AsyncFuncExec::try_new( self.async_exprs.clone(), - Arc::clone(&children[0]), + children.swap_remove(0), )?)) } diff --git a/datafusion/physical-plan/src/buffer.rs b/datafusion/physical-plan/src/buffer.rs index 3b80f9924e311..ae8d015988228 100644 --- a/datafusion/physical-plan/src/buffer.rs +++ b/datafusion/physical-plan/src/buffer.rs @@ -18,7 +18,9 @@ //! [`BufferExec`] decouples production and consumption on messages by buffering the input in the //! background up to a certain capacity. -use crate::execution_plan::{CardinalityEffect, SchedulingType}; +use crate::execution_plan::{ + CardinalityEffect, SchedulingType, has_same_children_properties, +}; use crate::filter_pushdown::{ ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, @@ -27,6 +29,7 @@ use crate::projection::ProjectionExec; use crate::stream::RecordBatchStreamAdapter; use crate::{ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SortOrderPushdownResult, + check_if_same_properties, }; use arrow::array::RecordBatch; use datafusion_common::config::ConfigOptions; @@ -92,7 +95,7 @@ use tokio::sync::{OwnedSemaphorePermit, Semaphore}; #[derive(Debug, Clone)] pub struct BufferExec { input: Arc, - properties: PlanProperties, + properties: Arc, capacity: usize, metrics: ExecutionPlanMetricsSet, } @@ -100,14 +103,12 @@ pub struct BufferExec { impl BufferExec { /// Builds a new [BufferExec] with the provided capacity in bytes. pub fn new(input: Arc, capacity: usize) -> Self { - let properties = input - .properties() - .clone() + let properties = PlanProperties::clone(input.properties()) .with_scheduling_type(SchedulingType::Cooperative); Self { input, - properties, + properties: Arc::new(properties), capacity, metrics: ExecutionPlanMetricsSet::new(), } @@ -122,6 +123,17 @@ impl BufferExec { pub fn capacity(&self) -> usize { self.capacity } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for BufferExec { @@ -146,7 +158,7 @@ impl ExecutionPlan for BufferExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.properties } @@ -166,6 +178,7 @@ impl ExecutionPlan for BufferExec { self: Arc, mut children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); if children.len() != 1 { return plan_err!("BufferExec can only have one child"); } diff --git a/datafusion/physical-plan/src/coalesce_batches.rs b/datafusion/physical-plan/src/coalesce_batches.rs index dfcd3cb0bcae7..27a521ec6bc5d 100644 --- a/datafusion/physical-plan/src/coalesce_batches.rs +++ b/datafusion/physical-plan/src/coalesce_batches.rs @@ -27,6 +27,7 @@ use super::{DisplayAs, ExecutionPlanProperties, PlanProperties, Statistics}; use crate::projection::ProjectionExec; use crate::{ DisplayFormatType, ExecutionPlan, RecordBatchStream, SendableRecordBatchStream, + check_if_same_properties, }; use arrow::datatypes::SchemaRef; @@ -36,7 +37,7 @@ use datafusion_execution::TaskContext; use datafusion_physical_expr::PhysicalExpr; use crate::coalesce::{LimitedBatchCoalescer, PushBatchStatus}; -use crate::execution_plan::CardinalityEffect; +use crate::execution_plan::{CardinalityEffect, has_same_children_properties}; use crate::filter_pushdown::{ ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, @@ -71,7 +72,7 @@ pub struct CoalesceBatchesExec { fetch: Option, /// Execution metrics metrics: ExecutionPlanMetricsSet, - cache: PlanProperties, + cache: Arc, } #[expect(deprecated)] @@ -84,7 +85,7 @@ impl CoalesceBatchesExec { target_batch_size, fetch: None, metrics: ExecutionPlanMetricsSet::new(), - cache, + cache: Arc::new(cache), } } @@ -115,6 +116,17 @@ impl CoalesceBatchesExec { input.boundedness(), ) } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } #[expect(deprecated)] @@ -159,7 +171,7 @@ impl ExecutionPlan for CoalesceBatchesExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -177,10 +189,11 @@ impl ExecutionPlan for CoalesceBatchesExec { fn with_new_children( self: Arc, - children: Vec>, + mut children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); Ok(Arc::new( - CoalesceBatchesExec::new(Arc::clone(&children[0]), self.target_batch_size) + CoalesceBatchesExec::new(children.swap_remove(0), self.target_batch_size) .with_fetch(self.fetch), )) } @@ -222,7 +235,7 @@ impl ExecutionPlan for CoalesceBatchesExec { target_batch_size: self.target_batch_size, fetch: limit, metrics: self.metrics.clone(), - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), })) } diff --git a/datafusion/physical-plan/src/coalesce_partitions.rs b/datafusion/physical-plan/src/coalesce_partitions.rs index 22dcc85d6ea3a..a7a2ce8a18b5c 100644 --- a/datafusion/physical-plan/src/coalesce_partitions.rs +++ b/datafusion/physical-plan/src/coalesce_partitions.rs @@ -27,11 +27,13 @@ use super::{ DisplayAs, ExecutionPlanProperties, PlanProperties, SendableRecordBatchStream, Statistics, }; -use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType}; +use crate::execution_plan::{ + CardinalityEffect, EvaluationType, SchedulingType, has_same_children_properties, +}; use crate::filter_pushdown::{FilterDescription, FilterPushdownPhase}; use crate::projection::{ProjectionExec, make_with_child}; use crate::sort_pushdown::SortOrderPushdownResult; -use crate::{DisplayFormatType, ExecutionPlan, Partitioning}; +use crate::{DisplayFormatType, ExecutionPlan, Partitioning, check_if_same_properties}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_common::config::ConfigOptions; @@ -47,7 +49,7 @@ pub struct CoalescePartitionsExec { input: Arc, /// Execution metrics metrics: ExecutionPlanMetricsSet, - cache: PlanProperties, + cache: Arc, /// Optional number of rows to fetch. Stops producing rows after this fetch pub(crate) fetch: Option, } @@ -59,7 +61,7 @@ impl CoalescePartitionsExec { CoalescePartitionsExec { input, metrics: ExecutionPlanMetricsSet::new(), - cache, + cache: Arc::new(cache), fetch: None, } } @@ -100,6 +102,17 @@ impl CoalescePartitionsExec { .with_evaluation_type(drive) .with_scheduling_type(scheduling) } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for CoalescePartitionsExec { @@ -135,7 +148,7 @@ impl ExecutionPlan for CoalescePartitionsExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -149,9 +162,10 @@ impl ExecutionPlan for CoalescePartitionsExec { fn with_new_children( self: Arc, - children: Vec>, + mut children: Vec>, ) -> Result> { - let mut plan = CoalescePartitionsExec::new(Arc::clone(&children[0])); + check_if_same_properties!(self, children); + let mut plan = CoalescePartitionsExec::new(children.swap_remove(0)); plan.fetch = self.fetch; Ok(Arc::new(plan)) } @@ -274,7 +288,7 @@ impl ExecutionPlan for CoalescePartitionsExec { input: Arc::clone(&self.input), fetch: limit, metrics: self.metrics.clone(), - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), })) } diff --git a/datafusion/physical-plan/src/coop.rs b/datafusion/physical-plan/src/coop.rs index ce54a451ac4d1..acc79ee009690 100644 --- a/datafusion/physical-plan/src/coop.rs +++ b/datafusion/physical-plan/src/coop.rs @@ -87,14 +87,14 @@ use crate::filter_pushdown::{ use crate::projection::ProjectionExec; use crate::{ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, SortOrderPushdownResult, + SendableRecordBatchStream, SortOrderPushdownResult, check_if_same_properties, }; use arrow::record_batch::RecordBatch; use arrow_schema::Schema; use datafusion_common::{Result, Statistics, assert_eq_or_internal_err}; use datafusion_execution::TaskContext; -use crate::execution_plan::SchedulingType; +use crate::execution_plan::{SchedulingType, has_same_children_properties}; use crate::stream::RecordBatchStreamAdapter; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use futures::{Stream, StreamExt}; @@ -217,16 +217,15 @@ where #[derive(Debug, Clone)] pub struct CooperativeExec { input: Arc, - properties: PlanProperties, + properties: Arc, } impl CooperativeExec { /// Creates a new `CooperativeExec` operator that wraps the given input execution plan. pub fn new(input: Arc) -> Self { - let properties = input - .properties() - .clone() - .with_scheduling_type(SchedulingType::Cooperative); + let properties = PlanProperties::clone(input.properties()) + .with_scheduling_type(SchedulingType::Cooperative) + .into(); Self { input, properties } } @@ -235,6 +234,16 @@ impl CooperativeExec { pub fn input(&self) -> &Arc { &self.input } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + ..Self::clone(self) + } + } } impl DisplayAs for CooperativeExec { @@ -260,7 +269,7 @@ impl ExecutionPlan for CooperativeExec { self.input.schema() } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.properties } @@ -281,6 +290,7 @@ impl ExecutionPlan for CooperativeExec { 1, "CooperativeExec requires exactly one child" ); + check_if_same_properties!(self, children); Ok(Arc::new(CooperativeExec::new(children.swap_remove(0)))) } diff --git a/datafusion/physical-plan/src/display.rs b/datafusion/physical-plan/src/display.rs index 52c37a106b39e..d525f44541d5f 100644 --- a/datafusion/physical-plan/src/display.rs +++ b/datafusion/physical-plan/src/display.rs @@ -1153,7 +1153,7 @@ mod tests { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { unimplemented!() } diff --git a/datafusion/physical-plan/src/empty.rs b/datafusion/physical-plan/src/empty.rs index fcfbcfa3e8277..d3e62f91c0570 100644 --- a/datafusion/physical-plan/src/empty.rs +++ b/datafusion/physical-plan/src/empty.rs @@ -43,7 +43,7 @@ pub struct EmptyExec { schema: SchemaRef, /// Number of partitions partitions: usize, - cache: PlanProperties, + cache: Arc, } impl EmptyExec { @@ -53,7 +53,7 @@ impl EmptyExec { EmptyExec { schema, partitions: 1, - cache, + cache: Arc::new(cache), } } @@ -62,7 +62,7 @@ impl EmptyExec { self.partitions = partitions; // Changing partitions may invalidate output partitioning, so update it: let output_partitioning = Self::output_partitioning_helper(self.partitions); - self.cache = self.cache.with_partitioning(output_partitioning); + Arc::make_mut(&mut self.cache).partitioning = output_partitioning; self } @@ -114,7 +114,7 @@ impl ExecutionPlan for EmptyExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 1c16d1e4b1864..21a558b80bce0 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -129,7 +129,7 @@ pub trait ExecutionPlan: Debug + DisplayAs + Send + Sync { /// /// This information is available via methods on [`ExecutionPlanProperties`] /// trait, which is implemented for all `ExecutionPlan`s. - fn properties(&self) -> &PlanProperties; + fn properties(&self) -> &Arc; /// Returns an error if this individual node does not conform to its invariants. /// These invariants are typically only checked in debug mode. @@ -1106,12 +1106,17 @@ impl PlanProperties { self } - /// Overwrite equivalence properties with its new value. - pub fn with_eq_properties(mut self, eq_properties: EquivalenceProperties) -> Self { + /// Set equivalence properties having mut reference. + pub fn set_eq_properties(&mut self, eq_properties: EquivalenceProperties) { // Changing equivalence properties also changes output ordering, so // make sure to overwrite it: self.output_ordering = eq_properties.output_ordering(); self.eq_properties = eq_properties; + } + + /// Overwrite equivalence properties with its new value. + pub fn with_eq_properties(mut self, eq_properties: EquivalenceProperties) -> Self { + self.set_eq_properties(eq_properties); self } @@ -1143,9 +1148,14 @@ impl PlanProperties { self } + /// Set constraints having mut reference. + pub fn set_constraints(&mut self, constraints: Constraints) { + self.eq_properties.set_constraints(constraints); + } + /// Overwrite constraints with its new value. pub fn with_constraints(mut self, constraints: Constraints) -> Self { - self.eq_properties = self.eq_properties.with_constraints(constraints); + self.set_constraints(constraints); self } @@ -1604,6 +1614,41 @@ pub fn prepare_execution( .map(|tnr| tnr.data) } +/// Check if the `plan` children has the same properties as passed `children`. +/// In this case plan can avoid self properties re-computation when its children +/// replace is requested. +/// The size of `children` must be equal to the size of `ExecutionPlan::children()`. +pub fn has_same_children_properties( + plan: &Arc, + children: &[Arc], +) -> Result { + let old_children = plan.children(); + assert_eq_or_internal_err!( + children.len(), + old_children.len(), + "Wrong number of children" + ); + for (lhs, rhs) in plan.children().iter().zip(children.iter()) { + if !Arc::ptr_eq(lhs.properties(), rhs.properties()) { + return Ok(false); + } + } + Ok(true) +} + +/// Helper macro to avoid properties re-computation if passed children properties +/// the same as plan already has. Could be used to implement fast-path for method +/// [`ExecutionPlan::with_new_children`]. +#[macro_export] +macro_rules! check_if_same_properties { + ($plan: expr, $children: expr) => { + if has_same_children_properties(&$plan, &$children)? { + let plan = $plan.with_new_children_and_same_properties($children); + return Ok(Arc::new(plan)); + } + }; +} + /// Utility function yielding a string representation of the given [`ExecutionPlan`]. pub fn get_plan_string(plan: &Arc) -> Vec { let formatted = displayable(plan.as_ref()).indent(true).to_string(); @@ -1666,7 +1711,7 @@ mod tests { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { unimplemented!() } @@ -1733,7 +1778,7 @@ mod tests { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { unimplemented!() } diff --git a/datafusion/physical-plan/src/explain.rs b/datafusion/physical-plan/src/explain.rs index aa3c0afefe8b5..bf21b0484689a 100644 --- a/datafusion/physical-plan/src/explain.rs +++ b/datafusion/physical-plan/src/explain.rs @@ -44,7 +44,7 @@ pub struct ExplainExec { stringified_plans: Vec, /// control which plans to print verbose: bool, - cache: PlanProperties, + cache: Arc, } impl ExplainExec { @@ -59,7 +59,7 @@ impl ExplainExec { schema, stringified_plans, verbose, - cache, + cache: Arc::new(cache), } } @@ -112,7 +112,7 @@ impl ExecutionPlan for ExplainExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index dc5ca465ccc23..5393753aa509e 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -27,9 +27,12 @@ use super::{ ColumnStatistics, DisplayAs, ExecutionPlanProperties, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, }; +use crate::check_if_same_properties; use crate::coalesce::{LimitedBatchCoalescer, PushBatchStatus}; use crate::common::can_project; -use crate::execution_plan::{CardinalityEffect, ReplacePhysicalExpr}; +use crate::execution_plan::{ + CardinalityEffect, ReplacePhysicalExpr, has_same_children_properties, +}; use crate::filter_pushdown::{ ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, PushedDown, PushedDownPredicate, @@ -85,7 +88,7 @@ pub struct FilterExec { /// Selectivity for statistics. 0 = no rows, 100 = all rows default_selectivity: u8, /// Properties equivalence properties, partitioning, etc. - cache: PlanProperties, + cache: Arc, /// The projection indices of the columns in the output schema of join projection: Option, /// Target batch size for output batches @@ -207,7 +210,7 @@ impl FilterExecBuilder { input: self.input, metrics: ExecutionPlanMetricsSet::new(), default_selectivity: self.default_selectivity, - cache, + cache: Arc::new(cache), projection: self.projection, batch_size: self.batch_size, fetch: self.fetch, @@ -280,7 +283,7 @@ impl FilterExec { input: Arc::clone(&self.input), metrics: self.metrics.clone(), default_selectivity: self.default_selectivity, - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), projection: self.projection.clone(), batch_size, fetch: self.fetch, @@ -433,6 +436,17 @@ impl FilterExec { input.boundedness(), )) } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for FilterExec { @@ -487,7 +501,7 @@ impl ExecutionPlan for FilterExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -504,6 +518,7 @@ impl ExecutionPlan for FilterExec { self: Arc, mut children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); let new_input = children.swap_remove(0); FilterExecBuilder::from(&*self) .with_input(new_input) @@ -697,12 +712,12 @@ impl ExecutionPlan for FilterExec { input: Arc::clone(&filter_input), metrics: self.metrics.clone(), default_selectivity: self.default_selectivity, - cache: Self::compute_properties( + cache: Arc::new(Self::compute_properties( &filter_input, &new_predicate, self.default_selectivity, self.projection.as_deref(), - )?, + )?), projection: self.projection.clone(), batch_size: self.batch_size, fetch: self.fetch, @@ -722,7 +737,7 @@ impl ExecutionPlan for FilterExec { input: Arc::clone(&self.input), metrics: self.metrics.clone(), default_selectivity: self.default_selectivity, - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), projection: self.projection.clone(), batch_size: self.batch_size, fetch, diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index 7ada14be66543..c25408e49db64 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -25,7 +25,9 @@ use super::utils::{ OnceAsync, OnceFut, StatefulStreamResult, adjust_right_output_partitioning, reorder_output_after_swap, }; -use crate::execution_plan::{EmissionType, boundedness_from_children}; +use crate::execution_plan::{ + EmissionType, boundedness_from_children, has_same_children_properties, +}; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use crate::projection::{ ProjectionExec, join_allows_pushdown, join_table_borders, new_join_children, @@ -34,7 +36,7 @@ use crate::projection::{ use crate::{ ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, Statistics, handle_state, + SendableRecordBatchStream, Statistics, check_if_same_properties, handle_state, }; use arrow::array::{RecordBatch, RecordBatchOptions}; @@ -94,7 +96,7 @@ pub struct CrossJoinExec { /// Execution plan metrics metrics: ExecutionPlanMetricsSet, /// Properties such as schema, equivalence properties, ordering, partitioning, etc. - cache: PlanProperties, + cache: Arc, } impl CrossJoinExec { @@ -125,7 +127,7 @@ impl CrossJoinExec { schema, left_fut: Default::default(), metrics: ExecutionPlanMetricsSet::default(), - cache, + cache: Arc::new(cache), } } @@ -192,6 +194,23 @@ impl CrossJoinExec { &self.right.schema(), ) } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + + Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + left_fut: Default::default(), + cache: Arc::clone(&self.cache), + schema: Arc::clone(&self.schema), + } + } } /// Asynchronously collect the result of the left child @@ -256,7 +275,7 @@ impl ExecutionPlan for CrossJoinExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -272,6 +291,7 @@ impl ExecutionPlan for CrossJoinExec { self: Arc, children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); Ok(Arc::new(CrossJoinExec::new( Arc::clone(&children[0]), Arc::clone(&children[1]), @@ -285,7 +305,7 @@ impl ExecutionPlan for CrossJoinExec { schema: Arc::clone(&self.schema), left_fut: Default::default(), // reset the build side! metrics: ExecutionPlanMetricsSet::default(), - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), }; Ok(Arc::new(new_exec)) } diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 44e67d190a1c3..2647468376120 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -24,6 +24,7 @@ use std::{any::Any, vec}; use crate::ExecutionPlanProperties; use crate::execution_plan::{ EmissionType, ReplacePhysicalExpr, boundedness_from_children, + has_same_children_properties, }; use crate::filter_pushdown::{ ChildPushdownResult, FilterDescription, FilterPushdownPhase, @@ -393,7 +394,7 @@ impl HashJoinExecBuilder { column_indices, null_equality, null_aware, - cache, + cache: Arc::new(cache), dynamic_filter: None, }) } @@ -643,7 +644,7 @@ pub struct HashJoinExec { /// Flag to indicate if this is a null-aware anti join pub null_aware: bool, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, /// Dynamic filter for pushing down to the probe side /// Set when dynamic filter pushdown is detected in handle_child_pushdown_result. /// HashJoinExec also needs to keep a shared bounds accumulator for coordinating updates. @@ -1037,7 +1038,7 @@ impl ExecutionPlan for HashJoinExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -1098,6 +1099,20 @@ impl ExecutionPlan for HashJoinExec { self: Arc, children: Vec>, ) -> Result> { + let cache = if has_same_children_properties(&self, &children)? { + Arc::clone(&self.cache) + } else { + Arc::new(Self::compute_properties( + &children[0], + &children[1], + &self.join_schema, + self.join_type, + &self.on, + self.mode, + self.projection.as_deref(), + )?) + }; + Ok(Arc::new(HashJoinExec { left: Arc::clone(&children[0]), right: Arc::clone(&children[1]), @@ -1113,15 +1128,7 @@ impl ExecutionPlan for HashJoinExec { column_indices: self.column_indices.clone(), null_equality: self.null_equality, null_aware: self.null_aware, - cache: Self::compute_properties( - &children[0], - &children[1], - &self.join_schema, - self.join_type, - &self.on, - self.mode, - self.projection.as_deref(), - )?, + cache, // Keep the dynamic filter, bounds accumulator will be reset dynamic_filter: self.dynamic_filter.clone(), })) @@ -1144,7 +1151,7 @@ impl ExecutionPlan for HashJoinExec { column_indices: self.column_indices.clone(), null_equality: self.null_equality, null_aware: self.null_aware, - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), // Reset dynamic filter and bounds accumulator to initial state dynamic_filter: None, })) @@ -1473,7 +1480,7 @@ impl ExecutionPlan for HashJoinExec { column_indices: self.column_indices.clone(), null_equality: self.null_equality, null_aware: self.null_aware, - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), dynamic_filter: Some(HashJoinExecDynamicFilter { filter: dynamic_filter, build_accumulator: OnceLock::new(), diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index 796f5567e1eac..92ead98d439e4 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -31,6 +31,7 @@ use super::utils::{ use crate::common::can_project; use crate::execution_plan::{ EmissionType, ReplacePhysicalExpr, boundedness_from_children, + has_same_children_properties, }; use crate::joins::SharedBitmapBuilder; use crate::joins::utils::{ @@ -48,6 +49,7 @@ use crate::projection::{ use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, PlanProperties, RecordBatchStream, SendableRecordBatchStream, + check_if_same_properties, }; use arrow::array::{ @@ -203,7 +205,7 @@ pub struct NestedLoopJoinExec { /// Execution metrics metrics: ExecutionPlanMetricsSet, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } /// Helps to build [`NestedLoopJoinExec`]. @@ -281,7 +283,7 @@ impl NestedLoopJoinExecBuilder { column_indices, projection, metrics: Default::default(), - cache, + cache: Arc::new(cache), }) } } @@ -467,6 +469,27 @@ impl NestedLoopJoinExec { Ok(plan) } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + + Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + build_side_data: Default::default(), + cache: Arc::clone(&self.cache), + filter: self.filter.clone(), + join_type: self.join_type, + join_schema: Arc::clone(&self.join_schema), + column_indices: self.column_indices.clone(), + projection: self.projection.clone(), + } + } } impl DisplayAs for NestedLoopJoinExec { @@ -521,7 +544,7 @@ impl ExecutionPlan for NestedLoopJoinExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -544,6 +567,7 @@ impl ExecutionPlan for NestedLoopJoinExec { self: Arc, children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); Ok(Arc::new( NestedLoopJoinExecBuilder::new( Arc::clone(&children[0]), diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs index d7ece845e943c..014d4cb13e2b7 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs @@ -41,7 +41,9 @@ use std::fmt::Formatter; use std::sync::Arc; use std::sync::atomic::AtomicUsize; -use crate::execution_plan::{EmissionType, boundedness_from_children}; +use crate::execution_plan::{ + EmissionType, boundedness_from_children, has_same_children_properties, +}; use crate::joins::piecewise_merge_join::classic_join::{ ClassicPWMJStream, PiecewiseMergeJoinStreamState, @@ -51,7 +53,9 @@ use crate::joins::piecewise_merge_join::utils::{ }; use crate::joins::utils::asymmetric_join_output_partitioning; use crate::metrics::MetricsSet; -use crate::{DisplayAs, DisplayFormatType, ExecutionPlanProperties}; +use crate::{ + DisplayAs, DisplayFormatType, ExecutionPlanProperties, check_if_same_properties, +}; use crate::{ ExecutionPlan, PlanProperties, joins::{ @@ -86,7 +90,7 @@ use crate::{ /// Both sides are sorted so that we can iterate from index 0 to the end on each side. This ordering ensures /// that when we find the first matching pair of rows, we can emit the current stream row joined with all remaining /// probe rows from the match position onward, without rescanning earlier probe rows. -/// +/// /// For `<` and `<=` operators, both inputs are sorted in **descending** order, while for `>` and `>=` operators /// they are sorted in **ascending** order. This choice ensures that the pointer on the buffered side can advance /// monotonically as we stream new batches from the stream side. @@ -129,34 +133,34 @@ use crate::{ /// /// Processing Row 1: /// -/// Sorted Buffered Side Sorted Streamed Side -/// ┌──────────────────┐ ┌──────────────────┐ -/// 1 │ 100 │ 1 │ 100 │ -/// ├──────────────────┤ ├──────────────────┤ -/// 2 │ 200 │ ─┐ 2 │ 200 │ -/// ├──────────────────┤ │ For row 1 on streamed side with ├──────────────────┤ -/// 3 │ 200 │ │ value 100, we emit rows 2 - 5. 3 │ 500 │ +/// Sorted Buffered Side Sorted Streamed Side +/// ┌──────────────────┐ ┌──────────────────┐ +/// 1 │ 100 │ 1 │ 100 │ +/// ├──────────────────┤ ├──────────────────┤ +/// 2 │ 200 │ ─┐ 2 │ 200 │ +/// ├──────────────────┤ │ For row 1 on streamed side with ├──────────────────┤ +/// 3 │ 200 │ │ value 100, we emit rows 2 - 5. 3 │ 500 │ /// ├──────────────────┤ │ as matches when the operator is └──────────────────┘ /// 4 │ 300 │ │ `Operator::Lt` (<) Emitting all /// ├──────────────────┤ │ rows after the first match (row /// 5 │ 400 │ ─┘ 2 buffered side; 100 < 200) -/// └──────────────────┘ +/// └──────────────────┘ /// /// Processing Row 2: /// By sorting the streamed side we know /// -/// Sorted Buffered Side Sorted Streamed Side -/// ┌──────────────────┐ ┌──────────────────┐ -/// 1 │ 100 │ 1 │ 100 │ -/// ├──────────────────┤ ├──────────────────┤ -/// 2 │ 200 │ <- Start here when probing for the 2 │ 200 │ -/// ├──────────────────┤ streamed side row 2. ├──────────────────┤ -/// 3 │ 200 │ 3 │ 500 │ +/// Sorted Buffered Side Sorted Streamed Side +/// ┌──────────────────┐ ┌──────────────────┐ +/// 1 │ 100 │ 1 │ 100 │ +/// ├──────────────────┤ ├──────────────────┤ +/// 2 │ 200 │ <- Start here when probing for the 2 │ 200 │ +/// ├──────────────────┤ streamed side row 2. ├──────────────────┤ +/// 3 │ 200 │ 3 │ 500 │ /// ├──────────────────┤ └──────────────────┘ -/// 4 │ 300 │ -/// ├──────────────────┤ +/// 4 │ 300 │ +/// ├──────────────────┤ /// 5 │ 400 │ -/// └──────────────────┘ +/// └──────────────────┘ /// ``` /// /// ## Existence Joins (Semi, Anti, Mark) @@ -202,10 +206,10 @@ use crate::{ /// 1 │ 100 │ 1 │ 500 │ /// ├──────────────────┤ ├──────────────────┤ /// 2 │ 200 │ 2 │ 200 │ -/// ├──────────────────┤ ├──────────────────┤ +/// ├──────────────────┤ ├──────────────────┤ /// 3 │ 200 │ 3 │ 300 │ /// ├──────────────────┤ └──────────────────┘ -/// 4 │ 300 │ ─┐ +/// 4 │ 300 │ ─┐ /// ├──────────────────┤ | We emit matches for row 4 - 5 /// 5 │ 400 │ ─┘ on the buffered side. /// └──────────────────┘ @@ -236,11 +240,11 @@ use crate::{ /// /// # Mark Join: /// Sorts the probe side, then computes the min/max range of the probe keys and scans the buffered side only -/// within that range. +/// within that range. /// Complexity: `O(|S| + scan(R[range]))`. /// /// ## Nested Loop Join -/// Compares every row from `S` with every row from `R`. +/// Compares every row from `S` with every row from `R`. /// Complexity: `O(|S| * |R|)`. /// /// ## Nested Loop Join @@ -273,13 +277,12 @@ pub struct PiecewiseMergeJoinExec { left_child_plan_required_order: LexOrdering, /// The right sort order, descending for `<`, `<=` operations + ascending for `>`, `>=` operations /// Unsorted for mark joins - #[expect(dead_code)] right_batch_required_orders: LexOrdering, /// This determines the sort order of all join columns used in sorting the stream and buffered execution plans. sort_options: SortOptions, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, /// Number of partitions to process num_partitions: usize, } @@ -373,7 +376,7 @@ impl PiecewiseMergeJoinExec { left_child_plan_required_order, right_batch_required_orders, sort_options, - cache, + cache: Arc::new(cache), num_partitions, }) } @@ -466,6 +469,31 @@ impl PiecewiseMergeJoinExec { pub fn swap_inputs(&self) -> Result> { todo!() } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + let buffered = children.swap_remove(0); + let streamed = children.swap_remove(0); + Self { + buffered, + streamed, + on: self.on.clone(), + operator: self.operator, + join_type: self.join_type, + schema: Arc::clone(&self.schema), + left_child_plan_required_order: self.left_child_plan_required_order.clone(), + right_batch_required_orders: self.right_batch_required_orders.clone(), + sort_options: self.sort_options, + cache: Arc::clone(&self.cache), + num_partitions: self.num_partitions, + + // Re-set state. + metrics: ExecutionPlanMetricsSet::new(), + buffered_fut: Default::default(), + } + } } impl ExecutionPlan for PiecewiseMergeJoinExec { @@ -477,7 +505,7 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -511,6 +539,7 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { self: Arc, children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); match &children[..] { [left, right] => Ok(Arc::new(PiecewiseMergeJoinExec::try_new( Arc::clone(left), @@ -527,6 +556,13 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { } } + fn reset_state(self: Arc) -> Result> { + Ok(Arc::new(self.with_new_children_and_same_properties(vec![ + Arc::clone(&self.buffered), + Arc::clone(&self.streamed), + ]))) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index b3e87daff360a..83457806ca801 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -25,6 +25,7 @@ use std::sync::Arc; use crate::execution_plan::{ EmissionType, ReplacePhysicalExpr, boundedness_from_children, + has_same_children_properties, }; use crate::expressions::PhysicalSortExpr; use crate::joins::sort_merge_join::metrics::SortMergeJoinMetrics; @@ -41,7 +42,7 @@ use crate::projection::{ }; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - PlanProperties, SendableRecordBatchStream, Statistics, + PlanProperties, SendableRecordBatchStream, Statistics, check_if_same_properties, }; use arrow::compute::SortOptions; @@ -130,7 +131,7 @@ pub struct SortMergeJoinExec { /// Defines the null equality for the join. pub null_equality: NullEquality, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } impl SortMergeJoinExec { @@ -201,7 +202,7 @@ impl SortMergeJoinExec { right_sort_exprs, sort_options, null_equality, - cache, + cache: Arc::new(cache), }) } @@ -343,6 +344,20 @@ impl SortMergeJoinExec { reorder_output_after_swap(Arc::new(new_join), &left.schema(), &right.schema()) } } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for SortMergeJoinExec { @@ -408,7 +423,7 @@ impl ExecutionPlan for SortMergeJoinExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -443,6 +458,7 @@ impl ExecutionPlan for SortMergeJoinExec { self: Arc, children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); match &children[..] { [left, right] => Ok(Arc::new(SortMergeJoinExec::try_new( Arc::clone(left), diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 3adef3a2dd93f..41fed11f4dc0a 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -32,9 +32,11 @@ use std::sync::Arc; use std::task::{Context, Poll}; use std::vec; +use crate::check_if_same_properties; use crate::common::SharedMemoryReservation; use crate::execution_plan::{ ReplacePhysicalExpr, boundedness_from_children, emission_type_from_children, + has_same_children_properties, }; use crate::joins::stream_join_utils::{ PruningJoinHashMap, SortedFilterExpr, StreamJoinMetrics, @@ -200,7 +202,7 @@ pub struct SymmetricHashJoinExec { /// Partition Mode mode: StreamJoinPartitionMode, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } impl SymmetricHashJoinExec { @@ -256,7 +258,7 @@ impl SymmetricHashJoinExec { left_sort_exprs, right_sort_exprs, mode, - cache, + cache: Arc::new(cache), }) } @@ -363,6 +365,20 @@ impl SymmetricHashJoinExec { } Ok(false) } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for SymmetricHashJoinExec { @@ -414,7 +430,7 @@ impl ExecutionPlan for SymmetricHashJoinExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -456,6 +472,7 @@ impl ExecutionPlan for SymmetricHashJoinExec { self: Arc, children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); Ok(Arc::new(SymmetricHashJoinExec::try_new( Arc::clone(&children[0]), Arc::clone(&children[1]), diff --git a/datafusion/physical-plan/src/limit.rs b/datafusion/physical-plan/src/limit.rs index fea7acb221304..e1bd606fa25ba 100644 --- a/datafusion/physical-plan/src/limit.rs +++ b/datafusion/physical-plan/src/limit.rs @@ -27,8 +27,13 @@ use super::{ DisplayAs, ExecutionPlanProperties, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, }; -use crate::execution_plan::{Boundedness, CardinalityEffect}; -use crate::{DisplayFormatType, Distribution, ExecutionPlan, Partitioning}; +use crate::execution_plan::{ + Boundedness, CardinalityEffect, has_same_children_properties, +}; +use crate::{ + DisplayFormatType, Distribution, ExecutionPlan, Partitioning, + check_if_same_properties, +}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; @@ -51,10 +56,10 @@ pub struct GlobalLimitExec { fetch: Option, /// Execution metrics metrics: ExecutionPlanMetricsSet, - cache: PlanProperties, /// Does the limit have to preserve the order of its input, and if so what is it? /// Some optimizations may reorder the input if no particular sort is required required_ordering: Option, + cache: Arc, } impl GlobalLimitExec { @@ -66,8 +71,8 @@ impl GlobalLimitExec { skip, fetch, metrics: ExecutionPlanMetricsSet::new(), - cache, required_ordering: None, + cache: Arc::new(cache), } } @@ -106,6 +111,17 @@ impl GlobalLimitExec { pub fn set_required_ordering(&mut self, required_ordering: Option) { self.required_ordering = required_ordering; } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for GlobalLimitExec { @@ -144,7 +160,7 @@ impl ExecutionPlan for GlobalLimitExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -166,10 +182,11 @@ impl ExecutionPlan for GlobalLimitExec { fn with_new_children( self: Arc, - children: Vec>, + mut children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); Ok(Arc::new(GlobalLimitExec::new( - Arc::clone(&children[0]), + children.swap_remove(0), self.skip, self.fetch, ))) @@ -229,7 +246,7 @@ impl ExecutionPlan for GlobalLimitExec { } /// LocalLimitExec applies a limit to a single partition -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct LocalLimitExec { /// Input execution plan input: Arc, @@ -237,10 +254,10 @@ pub struct LocalLimitExec { fetch: usize, /// Execution metrics metrics: ExecutionPlanMetricsSet, - cache: PlanProperties, /// If the child plan is a sort node, after the sort node is removed during /// physical optimization, we should add the required ordering to the limit node required_ordering: Option, + cache: Arc, } impl LocalLimitExec { @@ -251,8 +268,8 @@ impl LocalLimitExec { input, fetch, metrics: ExecutionPlanMetricsSet::new(), - cache, required_ordering: None, + cache: Arc::new(cache), } } @@ -286,6 +303,17 @@ impl LocalLimitExec { pub fn set_required_ordering(&mut self, required_ordering: Option) { self.required_ordering = required_ordering; } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for LocalLimitExec { @@ -315,7 +343,7 @@ impl ExecutionPlan for LocalLimitExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -335,6 +363,7 @@ impl ExecutionPlan for LocalLimitExec { self: Arc, children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); match children.len() { 1 => Ok(Arc::new(LocalLimitExec::new( Arc::clone(&children[0]), diff --git a/datafusion/physical-plan/src/memory.rs b/datafusion/physical-plan/src/memory.rs index 4a406ca648d57..5c684bec98f9b 100644 --- a/datafusion/physical-plan/src/memory.rs +++ b/datafusion/physical-plan/src/memory.rs @@ -161,7 +161,7 @@ pub struct LazyMemoryExec { /// Functions to generate batches for each partition batch_generators: Vec>>, /// Plan properties cache storing equivalence properties, partitioning, and execution mode - cache: PlanProperties, + cache: Arc, /// Execution metrics metrics: ExecutionPlanMetricsSet, } @@ -200,7 +200,8 @@ impl LazyMemoryExec { EmissionType::Incremental, boundedness, ) - .with_scheduling_type(SchedulingType::Cooperative); + .with_scheduling_type(SchedulingType::Cooperative) + .into(); Ok(Self { schema, @@ -215,9 +216,9 @@ impl LazyMemoryExec { match projection.as_ref() { Some(columns) => { let projected = Arc::new(self.schema.project(columns).unwrap()); - self.cache = self.cache.with_eq_properties(EquivalenceProperties::new( - Arc::clone(&projected), - )); + Arc::make_mut(&mut self.cache).set_eq_properties( + EquivalenceProperties::new(Arc::clone(&projected)), + ); self.schema = projected; self.projection = projection; self @@ -236,12 +237,12 @@ impl LazyMemoryExec { partition_count, generator_count ); - self.cache.partitioning = partitioning; + Arc::make_mut(&mut self.cache).partitioning = partitioning; Ok(()) } pub fn add_ordering(&mut self, ordering: impl IntoIterator) { - self.cache + Arc::make_mut(&mut self.cache) .eq_properties .add_orderings(std::iter::once(ordering)); } @@ -306,7 +307,7 @@ impl ExecutionPlan for LazyMemoryExec { Arc::clone(&self.schema) } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -365,7 +366,7 @@ impl ExecutionPlan for LazyMemoryExec { Ok(Arc::new(LazyMemoryExec { schema: Arc::clone(&self.schema), batch_generators: generators, - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), metrics: ExecutionPlanMetricsSet::new(), projection: self.projection.clone(), })) diff --git a/datafusion/physical-plan/src/placeholder_row.rs b/datafusion/physical-plan/src/placeholder_row.rs index 4d00b73cff39c..f95b10771c02f 100644 --- a/datafusion/physical-plan/src/placeholder_row.rs +++ b/datafusion/physical-plan/src/placeholder_row.rs @@ -43,7 +43,7 @@ pub struct PlaceholderRowExec { schema: SchemaRef, /// Number of partitions partitions: usize, - cache: PlanProperties, + cache: Arc, } impl PlaceholderRowExec { @@ -54,7 +54,7 @@ impl PlaceholderRowExec { PlaceholderRowExec { schema, partitions, - cache, + cache: Arc::new(cache), } } @@ -63,7 +63,7 @@ impl PlaceholderRowExec { self.partitions = partitions; // Update output partitioning when updating partitions: let output_partitioning = Self::output_partitioning_helper(self.partitions); - self.cache = self.cache.with_partitioning(output_partitioning); + Arc::make_mut(&mut self.cache).partitioning = output_partitioning; self } @@ -132,7 +132,7 @@ impl ExecutionPlan for PlaceholderRowExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index 905bd13309ed2..628a5722950bd 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -27,13 +27,15 @@ use super::{ SendableRecordBatchStream, SortOrderPushdownResult, Statistics, }; use crate::column_rewriter::PhysicalColumnRewriter; -use crate::execution_plan::{CardinalityEffect, ReplacePhysicalExpr}; +use crate::execution_plan::{ + CardinalityEffect, ReplacePhysicalExpr, has_same_children_properties, +}; use crate::filter_pushdown::{ ChildFilterDescription, ChildPushdownResult, FilterColumnChecker, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, PushedDownPredicate, }; use crate::joins::utils::{ColumnIndex, JoinFilter, JoinOn, JoinOnRef}; -use crate::{DisplayFormatType, ExecutionPlan, PhysicalExpr}; +use crate::{DisplayFormatType, ExecutionPlan, PhysicalExpr, check_if_same_properties}; use std::any::Any; use std::collections::HashMap; use std::pin::Pin; @@ -81,7 +83,7 @@ pub struct ProjectionExec { /// Execution metrics metrics: ExecutionPlanMetricsSet, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } impl ProjectionExec { @@ -162,7 +164,7 @@ impl ProjectionExec { projector, input, metrics: ExecutionPlanMetricsSet::new(), - cache, + cache: Arc::new(cache), }) } @@ -225,6 +227,17 @@ impl ProjectionExec { } Ok(alias_map) } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for ProjectionExec { @@ -278,7 +291,7 @@ impl ExecutionPlan for ProjectionExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -313,6 +326,7 @@ impl ExecutionPlan for ProjectionExec { self: Arc, mut children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); ProjectionExec::try_from_projector( self.projector.clone(), children.swap_remove(0), diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index 936a02581e89c..f8847cbacefb5 100644 --- a/datafusion/physical-plan/src/recursive_query.rs +++ b/datafusion/physical-plan/src/recursive_query.rs @@ -74,7 +74,7 @@ pub struct RecursiveQueryExec { /// Execution metrics metrics: ExecutionPlanMetricsSet, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } impl RecursiveQueryExec { @@ -97,7 +97,7 @@ impl RecursiveQueryExec { is_distinct, work_table, metrics: ExecutionPlanMetricsSet::new(), - cache, + cache: Arc::new(cache), }) } @@ -143,7 +143,7 @@ impl ExecutionPlan for RecursiveQueryExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 612c7bb27ddf4..2e02cd3210786 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -31,7 +31,9 @@ use super::{ DisplayAs, ExecutionPlanProperties, RecordBatchStream, SendableRecordBatchStream, }; use crate::coalesce::LimitedBatchCoalescer; -use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType}; +use crate::execution_plan::{ + CardinalityEffect, EvaluationType, SchedulingType, has_same_children_properties, +}; use crate::hash_utils::create_hashes; use crate::metrics::{BaselineMetrics, SpillMetrics}; use crate::projection::{ProjectionExec, all_columns, make_with_child, update_expr}; @@ -39,7 +41,10 @@ use crate::sorts::streaming_merge::StreamingMergeBuilder; use crate::spill::spill_manager::SpillManager; use crate::spill::spill_pool::{self, SpillPoolWriter}; use crate::stream::RecordBatchStreamAdapter; -use crate::{DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, Statistics}; +use crate::{ + DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, Statistics, + check_if_same_properties, +}; use arrow::array::{PrimitiveArray, RecordBatch, RecordBatchOptions}; use arrow::compute::take_arrays; @@ -763,7 +768,7 @@ pub struct RepartitionExec { /// `SortPreservingRepartitionExec`, false means `RepartitionExec`. preserve_order: bool, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } #[derive(Debug, Clone)] @@ -832,6 +837,18 @@ impl RepartitionExec { pub fn name(&self) -> &str { "RepartitionExec" } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + state: Default::default(), + ..Self::clone(self) + } + } } impl DisplayAs for RepartitionExec { @@ -891,7 +908,7 @@ impl ExecutionPlan for RepartitionExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -903,6 +920,7 @@ impl ExecutionPlan for RepartitionExec { self: Arc, mut children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); let mut repartition = RepartitionExec::try_new( children.swap_remove(0), self.partitioning().clone(), @@ -1204,7 +1222,7 @@ impl ExecutionPlan for RepartitionExec { _config: &ConfigOptions, ) -> Result>> { use Partitioning::*; - let mut new_properties = self.cache.clone(); + let mut new_properties = PlanProperties::clone(&self.cache); new_properties.partitioning = match new_properties.partitioning { RoundRobinBatch(_) => RoundRobinBatch(target_partitions), Hash(hash, _) => Hash(hash, target_partitions), @@ -1215,7 +1233,7 @@ impl ExecutionPlan for RepartitionExec { state: Arc::clone(&self.state), metrics: self.metrics.clone(), preserve_order: self.preserve_order, - cache: new_properties, + cache: new_properties.into(), }))) } } @@ -1235,7 +1253,7 @@ impl RepartitionExec { state: Default::default(), metrics: ExecutionPlanMetricsSet::new(), preserve_order, - cache, + cache: Arc::new(cache), }) } @@ -1296,7 +1314,7 @@ impl RepartitionExec { // to maintain order self.input.output_partitioning().partition_count() > 1; let eq_properties = Self::eq_properties_helper(&self.input, self.preserve_order); - self.cache = self.cache.with_eq_properties(eq_properties); + Arc::make_mut(&mut self.cache).set_eq_properties(eq_properties); self } diff --git a/datafusion/physical-plan/src/sorts/partial_sort.rs b/datafusion/physical-plan/src/sorts/partial_sort.rs index e12aefd3003e1..f094258c8a9f5 100644 --- a/datafusion/physical-plan/src/sorts/partial_sort.rs +++ b/datafusion/physical-plan/src/sorts/partial_sort.rs @@ -57,12 +57,13 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; -use crate::execution_plan::ReplacePhysicalExpr; +use crate::execution_plan::{ReplacePhysicalExpr, has_same_children_properties}; use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; use crate::sorts::sort::sort_batch; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, Partitioning, PlanProperties, SendableRecordBatchStream, Statistics, + check_if_same_properties, }; use arrow::compute::concat_batches; @@ -94,7 +95,7 @@ pub struct PartialSortExec { /// Fetch highest/lowest n results fetch: Option, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } impl PartialSortExec { @@ -115,7 +116,7 @@ impl PartialSortExec { metrics_set: ExecutionPlanMetricsSet::new(), preserve_partitioning, fetch: None, - cache, + cache: Arc::new(cache), } } @@ -133,12 +134,8 @@ impl PartialSortExec { /// input partitions producing a single, sorted partition. pub fn with_preserve_partitioning(mut self, preserve_partitioning: bool) -> Self { self.preserve_partitioning = preserve_partitioning; - self.cache = self - .cache - .with_partitioning(Self::output_partitioning_helper( - &self.input, - self.preserve_partitioning, - )); + Arc::make_mut(&mut self.cache).partitioning = + Self::output_partitioning_helper(&self.input, self.preserve_partitioning); self } @@ -208,6 +205,16 @@ impl PartialSortExec { input.boundedness(), )) } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + ..Self::clone(self) + } + } } impl DisplayAs for PartialSortExec { @@ -256,7 +263,7 @@ impl ExecutionPlan for PartialSortExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -284,6 +291,7 @@ impl ExecutionPlan for PartialSortExec { self: Arc, children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); let new_partial_sort = PartialSortExec::new( self.expr.clone(), Arc::clone(&children[0]), diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 565e482592958..6e54721b0e01e 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -29,6 +29,7 @@ use parking_lot::RwLock; use crate::common::spawn_buffered; use crate::execution_plan::{ Boundedness, CardinalityEffect, EmissionType, ReplacePhysicalExpr, + has_same_children_properties, }; use crate::expressions::PhysicalSortExpr; use crate::filter_pushdown::{ @@ -953,7 +954,7 @@ pub struct SortExec { /// Normalized common sort prefix between the input and the sort expressions (only used with fetch) common_sort_prefix: Vec, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, /// Filter matching the state of the sort for dynamic filter pushdown. /// If `fetch` is `Some`, this will also be set and a TopK operator may be used. /// If `fetch` is `None`, this will be `None`. @@ -975,7 +976,7 @@ impl SortExec { preserve_partitioning, fetch: None, common_sort_prefix: sort_prefix, - cache, + cache: Arc::new(cache), filter: None, } } @@ -994,12 +995,8 @@ impl SortExec { /// input partitions producing a single, sorted partition. pub fn with_preserve_partitioning(mut self, preserve_partitioning: bool) -> Self { self.preserve_partitioning = preserve_partitioning; - self.cache = self - .cache - .with_partitioning(Self::output_partitioning_helper( - &self.input, - self.preserve_partitioning, - )); + Arc::make_mut(&mut self.cache).partitioning = + Self::output_partitioning_helper(&self.input, self.preserve_partitioning); self } @@ -1023,7 +1020,7 @@ impl SortExec { preserve_partitioning: self.preserve_partitioning, common_sort_prefix: self.common_sort_prefix.clone(), fetch: self.fetch, - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), filter: self.filter.clone(), } } @@ -1036,12 +1033,12 @@ impl SortExec { /// operation since rows that are not going to be included /// can be dropped. pub fn with_fetch(&self, fetch: Option) -> Self { - let mut cache = self.cache.clone(); + let mut cache = PlanProperties::clone(&self.cache); // If the SortExec can emit incrementally (that means the sort requirements // and properties of the input match), the SortExec can generate its result // without scanning the entire input when a fetch value exists. let is_pipeline_friendly = matches!( - self.cache.emission_type, + cache.emission_type, EmissionType::Incremental | EmissionType::Both ); if fetch.is_some() && is_pipeline_friendly { @@ -1053,7 +1050,7 @@ impl SortExec { }); let mut new_sort = self.cloned(); new_sort.fetch = fetch; - new_sort.cache = cache; + new_sort.cache = cache.into(); new_sort.filter = filter; new_sort } @@ -1208,7 +1205,7 @@ impl ExecutionPlan for SortExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -1237,14 +1234,17 @@ impl ExecutionPlan for SortExec { let mut new_sort = self.cloned(); assert_eq!(children.len(), 1, "SortExec should have exactly one child"); new_sort.input = Arc::clone(&children[0]); - // Recompute the properties based on the new input since they may have changed - let (cache, sort_prefix) = Self::compute_properties( - &new_sort.input, - new_sort.expr.clone(), - new_sort.preserve_partitioning, - )?; - new_sort.cache = cache; - new_sort.common_sort_prefix = sort_prefix; + + if !has_same_children_properties(&self, &children)? { + // Recompute the properties based on the new input since they may have changed + let (cache, sort_prefix) = Self::compute_properties( + &new_sort.input, + new_sort.expr.clone(), + new_sort.preserve_partitioning, + )?; + new_sort.cache = Arc::new(cache); + new_sort.common_sort_prefix = sort_prefix; + } Ok(Arc::new(new_sort)) } @@ -1509,7 +1509,7 @@ mod tests { pub struct SortedUnboundedExec { schema: Schema, batch_size: u64, - cache: PlanProperties, + cache: Arc, } impl DisplayAs for SortedUnboundedExec { @@ -1549,7 +1549,7 @@ mod tests { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -2302,7 +2302,9 @@ mod tests { let source = SortedUnboundedExec { schema: schema.clone(), batch_size: 2, - cache: SortedUnboundedExec::compute_properties(Arc::new(schema.clone())), + cache: Arc::new(SortedUnboundedExec::compute_properties(Arc::new( + schema.clone(), + ))), }; let mut plan = SortExec::new( [PhysicalSortExpr::new_default(Arc::new(Column::new( diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 91ab8906dec4d..b2889cc55c0e2 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -28,6 +28,7 @@ use crate::sorts::streaming_merge::StreamingMergeBuilder; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, Partitioning, PlanProperties, SendableRecordBatchStream, Statistics, + check_if_same_properties, }; use datafusion_common::{Result, assert_eq_or_internal_err, internal_err}; @@ -36,7 +37,9 @@ use datafusion_execution::memory_pool::MemoryConsumer; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements}; -use crate::execution_plan::{EvaluationType, ReplacePhysicalExpr, SchedulingType}; +use crate::execution_plan::{ + EvaluationType, ReplacePhysicalExpr, SchedulingType, has_same_children_properties, +}; use log::{debug, trace}; /// Sort preserving merge execution plan @@ -94,7 +97,7 @@ pub struct SortPreservingMergeExec { /// Optional number of rows to fetch. Stops producing rows after this fetch fetch: Option, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, /// Use round-robin selection of tied winners of loser tree /// /// See [`Self::with_round_robin_repartition`] for more information. @@ -110,7 +113,7 @@ impl SortPreservingMergeExec { expr, metrics: ExecutionPlanMetricsSet::new(), fetch: None, - cache, + cache: Arc::new(cache), enable_round_robin_repartition: true, } } @@ -181,6 +184,16 @@ impl SortPreservingMergeExec { .with_evaluation_type(drive) .with_scheduling_type(scheduling) } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + ..Self::clone(self) + } + } } impl DisplayAs for SortPreservingMergeExec { @@ -226,7 +239,7 @@ impl ExecutionPlan for SortPreservingMergeExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -241,7 +254,7 @@ impl ExecutionPlan for SortPreservingMergeExec { expr: self.expr.clone(), metrics: self.metrics.clone(), fetch: limit, - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), enable_round_robin_repartition: true, })) } @@ -281,10 +294,11 @@ impl ExecutionPlan for SortPreservingMergeExec { fn with_new_children( self: Arc, - children: Vec>, + mut children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); Ok(Arc::new( - SortPreservingMergeExec::new(self.expr.clone(), Arc::clone(&children[0])) + SortPreservingMergeExec::new(self.expr.clone(), children.swap_remove(0)) .with_fetch(self.fetch), )) } @@ -1393,7 +1407,7 @@ mod tests { #[derive(Debug, Clone)] struct CongestedExec { schema: Schema, - cache: PlanProperties, + cache: Arc, congestion: Arc, } @@ -1429,7 +1443,7 @@ mod tests { fn as_any(&self) -> &dyn Any { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } fn children(&self) -> Vec<&Arc> { @@ -1522,7 +1536,7 @@ mod tests { }; let source = CongestedExec { schema: schema.clone(), - cache: properties, + cache: Arc::new(properties), congestion: Arc::new(Congestion::new(partition_count)), }; let spm = SortPreservingMergeExec::new( diff --git a/datafusion/physical-plan/src/streaming.rs b/datafusion/physical-plan/src/streaming.rs index c8b8d95718cb8..1535482374110 100644 --- a/datafusion/physical-plan/src/streaming.rs +++ b/datafusion/physical-plan/src/streaming.rs @@ -67,7 +67,7 @@ pub struct StreamingTableExec { projected_output_ordering: Vec, infinite: bool, limit: Option, - cache: PlanProperties, + cache: Arc, metrics: ExecutionPlanMetricsSet, } @@ -111,7 +111,7 @@ impl StreamingTableExec { projected_output_ordering, infinite, limit, - cache, + cache: Arc::new(cache), metrics: ExecutionPlanMetricsSet::new(), }) } @@ -236,7 +236,7 @@ impl ExecutionPlan for StreamingTableExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -335,7 +335,7 @@ impl ExecutionPlan for StreamingTableExec { projected_output_ordering: self.projected_output_ordering.clone(), infinite: self.infinite, limit, - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), metrics: self.metrics.clone(), })) } diff --git a/datafusion/physical-plan/src/test.rs b/datafusion/physical-plan/src/test.rs index c6d0940c35480..aa079e73dd086 100644 --- a/datafusion/physical-plan/src/test.rs +++ b/datafusion/physical-plan/src/test.rs @@ -75,7 +75,7 @@ pub struct TestMemoryExec { /// The maximum number of records to read from this plan. If `None`, /// all records after filtering are returned. fetch: Option, - cache: PlanProperties, + cache: Arc, } impl DisplayAs for TestMemoryExec { @@ -134,7 +134,7 @@ impl ExecutionPlan for TestMemoryExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -239,7 +239,7 @@ impl TestMemoryExec { Ok(Self { partitions: partitions.to_vec(), schema, - cache: PlanProperties::new( + cache: Arc::new(PlanProperties::new( EquivalenceProperties::new_with_orderings( Arc::clone(&projected_schema), Vec::::new(), @@ -247,7 +247,7 @@ impl TestMemoryExec { Partitioning::UnknownPartitioning(partitions.len()), EmissionType::Incremental, Boundedness::Bounded, - ), + )), projected_schema, projection, sort_information: vec![], @@ -265,7 +265,7 @@ impl TestMemoryExec { ) -> Result> { let mut source = Self::try_new(partitions, schema, projection)?; let cache = source.compute_properties(); - source.cache = cache; + source.cache = Arc::new(cache); Ok(Arc::new(source)) } @@ -273,7 +273,7 @@ impl TestMemoryExec { pub fn update_cache(source: &Arc) -> TestMemoryExec { let cache = source.compute_properties(); let mut source = (**source).clone(); - source.cache = cache; + source.cache = Arc::new(cache); source } @@ -342,7 +342,7 @@ impl TestMemoryExec { } self.sort_information = sort_information; - self.cache = self.compute_properties(); + self.cache = Arc::new(self.compute_properties()); Ok(self) } diff --git a/datafusion/physical-plan/src/test/exec.rs b/datafusion/physical-plan/src/test/exec.rs index 4507cccba05a9..a8b21f70f7760 100644 --- a/datafusion/physical-plan/src/test/exec.rs +++ b/datafusion/physical-plan/src/test/exec.rs @@ -125,7 +125,7 @@ pub struct MockExec { /// if true (the default), sends data using a separate task to ensure the /// batches are not available without this stream yielding first use_task: bool, - cache: PlanProperties, + cache: Arc, } impl MockExec { @@ -142,7 +142,7 @@ impl MockExec { data, schema, use_task: true, - cache, + cache: Arc::new(cache), } } @@ -192,7 +192,7 @@ impl ExecutionPlan for MockExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -299,7 +299,7 @@ pub struct BarrierExec { /// all streams wait on this barrier to produce barrier: Arc, - cache: PlanProperties, + cache: Arc, } impl BarrierExec { @@ -312,7 +312,7 @@ impl BarrierExec { data, schema, barrier, - cache, + cache: Arc::new(cache), } } @@ -364,7 +364,7 @@ impl ExecutionPlan for BarrierExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -429,7 +429,7 @@ impl ExecutionPlan for BarrierExec { /// A mock execution plan that errors on a call to execute #[derive(Debug)] pub struct ErrorExec { - cache: PlanProperties, + cache: Arc, } impl Default for ErrorExec { @@ -446,7 +446,9 @@ impl ErrorExec { true, )])); let cache = Self::compute_properties(schema); - Self { cache } + Self { + cache: Arc::new(cache), + } } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. @@ -487,7 +489,7 @@ impl ExecutionPlan for ErrorExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -517,7 +519,7 @@ impl ExecutionPlan for ErrorExec { pub struct StatisticsExec { stats: Statistics, schema: Arc, - cache: PlanProperties, + cache: Arc, } impl StatisticsExec { pub fn new(stats: Statistics, schema: Schema) -> Self { @@ -530,7 +532,7 @@ impl StatisticsExec { Self { stats, schema: Arc::new(schema), - cache, + cache: Arc::new(cache), } } @@ -577,7 +579,7 @@ impl ExecutionPlan for StatisticsExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -623,7 +625,7 @@ pub struct BlockingExec { /// Ref-counting helper to check if the plan and the produced stream are still in memory. refs: Arc<()>, - cache: PlanProperties, + cache: Arc, } impl BlockingExec { @@ -633,7 +635,7 @@ impl BlockingExec { Self { schema, refs: Default::default(), - cache, + cache: Arc::new(cache), } } @@ -684,7 +686,7 @@ impl ExecutionPlan for BlockingExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -766,7 +768,7 @@ pub struct PanicExec { /// Number of output partitions. Each partition will produce this /// many empty output record batches prior to panicking batches_until_panics: Vec, - cache: PlanProperties, + cache: Arc, } impl PanicExec { @@ -778,7 +780,7 @@ impl PanicExec { Self { schema, batches_until_panics, - cache, + cache: Arc::new(cache), } } @@ -830,7 +832,7 @@ impl ExecutionPlan for PanicExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index b6f943886e309..7a3fa8eb3a3fe 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -32,9 +32,10 @@ use super::{ SendableRecordBatchStream, Statistics, metrics::{ExecutionPlanMetricsSet, MetricsSet}, }; +use crate::check_if_same_properties; use crate::execution_plan::{ InvariantLevel, boundedness_from_children, check_default_invariants, - emission_type_from_children, + emission_type_from_children, has_same_children_properties, }; use crate::filter_pushdown::{FilterDescription, FilterPushdownPhase}; use crate::metrics::BaselineMetrics; @@ -100,7 +101,7 @@ pub struct UnionExec { /// Execution metrics metrics: ExecutionPlanMetricsSet, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } impl UnionExec { @@ -118,7 +119,7 @@ impl UnionExec { UnionExec { inputs, metrics: ExecutionPlanMetricsSet::new(), - cache, + cache: Arc::new(cache), } } @@ -147,7 +148,7 @@ impl UnionExec { Ok(Arc::new(UnionExec { inputs, metrics: ExecutionPlanMetricsSet::new(), - cache, + cache: Arc::new(cache), })) } } @@ -183,6 +184,17 @@ impl UnionExec { boundedness_from_children(inputs), )) } + + fn with_new_children_and_same_properties( + &self, + children: Vec>, + ) -> Self { + Self { + inputs: children, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for UnionExec { @@ -210,7 +222,7 @@ impl ExecutionPlan for UnionExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -259,6 +271,7 @@ impl ExecutionPlan for UnionExec { self: Arc, children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); UnionExec::try_new(children) } @@ -411,7 +424,7 @@ pub struct InterleaveExec { /// Execution metrics metrics: ExecutionPlanMetricsSet, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } impl InterleaveExec { @@ -425,7 +438,7 @@ impl InterleaveExec { Ok(InterleaveExec { inputs, metrics: ExecutionPlanMetricsSet::new(), - cache, + cache: Arc::new(cache), }) } @@ -447,6 +460,17 @@ impl InterleaveExec { boundedness_from_children(inputs), )) } + + fn with_new_children_and_same_properties( + &self, + children: Vec>, + ) -> Self { + Self { + inputs: children, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for InterleaveExec { @@ -474,7 +498,7 @@ impl ExecutionPlan for InterleaveExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -495,6 +519,7 @@ impl ExecutionPlan for InterleaveExec { can_interleave(children.iter()), "Can not create InterleaveExec: new children can not be interleaved" ); + check_if_same_properties!(self, children); Ok(Arc::new(InterleaveExec::try_new(children)?)) } diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index 5fef754e80780..d8a8937c46b5e 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -26,9 +26,10 @@ use super::metrics::{ RecordOutput, }; use super::{DisplayAs, ExecutionPlanProperties, PlanProperties}; +use crate::execution_plan::has_same_children_properties; use crate::{ DisplayFormatType, Distribution, ExecutionPlan, RecordBatchStream, - SendableRecordBatchStream, + SendableRecordBatchStream, check_if_same_properties, }; use arrow::array::{ @@ -74,7 +75,7 @@ pub struct UnnestExec { /// Execution metrics metrics: ExecutionPlanMetricsSet, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } impl UnnestExec { @@ -100,7 +101,7 @@ impl UnnestExec { struct_column_indices, options, metrics: Default::default(), - cache, + cache: Arc::new(cache), }) } @@ -193,6 +194,17 @@ impl UnnestExec { pub fn options(&self) -> &UnnestOptions { &self.options } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for UnnestExec { @@ -221,7 +233,7 @@ impl ExecutionPlan for UnnestExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -231,10 +243,11 @@ impl ExecutionPlan for UnnestExec { fn with_new_children( self: Arc, - children: Vec>, + mut children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); Ok(Arc::new(UnnestExec::new( - Arc::clone(&children[0]), + children.swap_remove(0), self.list_column_indices.clone(), self.struct_column_indices.clone(), Arc::clone(&self.schema), diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index bf4c6a4a256b4..7daf42b71c747 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -28,7 +28,7 @@ use std::sync::Arc; use std::task::{Context, Poll}; use super::utils::create_schema; -use crate::execution_plan::ReplacePhysicalExpr; +use crate::execution_plan::{ReplacePhysicalExpr, has_same_children_properties}; use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; use crate::windows::{ calc_requirements, get_ordered_partition_by_indices, get_partition_by_sort_exprs, @@ -37,7 +37,7 @@ use crate::windows::{ use crate::{ ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, InputOrderMode, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, Statistics, WindowExpr, + SendableRecordBatchStream, Statistics, WindowExpr, check_if_same_properties, }; use arrow::compute::take_record_batch; @@ -95,7 +95,7 @@ pub struct BoundedWindowAggExec { // See `get_ordered_partition_by_indices` for more details. ordered_partition_by_indices: Vec, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, /// If `can_rerepartition` is false, partition_keys is always empty. can_repartition: bool, } @@ -136,7 +136,7 @@ impl BoundedWindowAggExec { metrics: ExecutionPlanMetricsSet::new(), input_order_mode, ordered_partition_by_indices, - cache, + cache: Arc::new(cache), can_repartition, }) } @@ -250,6 +250,17 @@ impl BoundedWindowAggExec { total_byte_size: Precision::Absent, }) } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for BoundedWindowAggExec { @@ -306,7 +317,7 @@ impl ExecutionPlan for BoundedWindowAggExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -341,6 +352,7 @@ impl ExecutionPlan for BoundedWindowAggExec { self: Arc, children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); Ok(Arc::new(BoundedWindowAggExec::try_new( self.window_expr.clone(), Arc::clone(&children[0]), diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs index d1a87ac594ee9..865ed666e5cf5 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -23,7 +23,9 @@ use std::sync::Arc; use std::task::{Context, Poll}; use super::utils::create_schema; -use crate::execution_plan::{EmissionType, ReplacePhysicalExpr}; +use crate::execution_plan::{ + EmissionType, ReplacePhysicalExpr, has_same_children_properties, +}; use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; use crate::windows::{ calc_requirements, get_ordered_partition_by_indices, get_partition_by_sort_exprs, @@ -32,7 +34,7 @@ use crate::windows::{ use crate::{ ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, PhysicalExpr, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, Statistics, WindowExpr, + SendableRecordBatchStream, Statistics, WindowExpr, check_if_same_properties, }; use arrow::array::ArrayRef; @@ -65,7 +67,7 @@ pub struct WindowAggExec { // see `get_ordered_partition_by_indices` for more details. ordered_partition_by_indices: Vec, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, /// If `can_partition` is false, partition_keys is always empty. can_repartition: bool, } @@ -89,7 +91,7 @@ impl WindowAggExec { schema, metrics: ExecutionPlanMetricsSet::new(), ordered_partition_by_indices, - cache, + cache: Arc::new(cache), can_repartition, }) } @@ -158,6 +160,17 @@ impl WindowAggExec { .unwrap_or_else(Vec::new) } } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for WindowAggExec { @@ -206,7 +219,7 @@ impl ExecutionPlan for WindowAggExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -242,11 +255,12 @@ impl ExecutionPlan for WindowAggExec { fn with_new_children( self: Arc, - children: Vec>, + mut children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); Ok(Arc::new(WindowAggExec::try_new( self.window_expr.clone(), - Arc::clone(&children[0]), + children.swap_remove(0), true, )?)) } diff --git a/datafusion/physical-plan/src/work_table.rs b/datafusion/physical-plan/src/work_table.rs index 1313909adbba2..b01ea513a8ede 100644 --- a/datafusion/physical-plan/src/work_table.rs +++ b/datafusion/physical-plan/src/work_table.rs @@ -109,7 +109,7 @@ pub struct WorkTableExec { /// Execution metrics metrics: ExecutionPlanMetricsSet, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } impl WorkTableExec { @@ -129,7 +129,7 @@ impl WorkTableExec { projection, work_table: Arc::new(WorkTable::new(name)), metrics: ExecutionPlanMetricsSet::new(), - cache, + cache: Arc::new(cache), }) } @@ -181,7 +181,7 @@ impl ExecutionPlan for WorkTableExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -263,7 +263,7 @@ impl ExecutionPlan for WorkTableExec { projection: self.projection.clone(), metrics: ExecutionPlanMetricsSet::new(), work_table, - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), })) } } diff --git a/docs/source/library-user-guide/custom-table-providers.md b/docs/source/library-user-guide/custom-table-providers.md index 8e1dee9e843ac..50005a7527da0 100644 --- a/docs/source/library-user-guide/custom-table-providers.md +++ b/docs/source/library-user-guide/custom-table-providers.md @@ -108,7 +108,7 @@ impl ExecutionPlan for CustomExec { } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { unreachable!() } @@ -232,7 +232,7 @@ The `scan` method of the `TableProvider` returns a `Result &PlanProperties { +# fn properties(&self) -> &Arc { # unreachable!() # } # @@ -424,7 +424,7 @@ This will allow you to use the custom table provider in DataFusion. For example, # } # # -# fn properties(&self) -> &PlanProperties { +# fn properties(&self) -> &Arc { # unreachable!() # } # diff --git a/docs/source/library-user-guide/upgrading.md b/docs/source/library-user-guide/upgrading.md index 182f2f0ef9f92..f62d3d8f948ac 100644 --- a/docs/source/library-user-guide/upgrading.md +++ b/docs/source/library-user-guide/upgrading.md @@ -61,6 +61,61 @@ FileSinkConfig { } ``` +### `ExecutionPlan` properties method return type + +Now `ExecutionPlan::properties()` must return `&Arc` instead of a reference. This was done to enable the comparison of properties and to determine that they have not changed within the +`with_new_children` method. To migrate, in all `ExecutionPlan` implementations, you need to wrap stored `PlanProperties` in an `Arc`: + +```diff +- cache: PlanProperties, ++ cache: Arc, + +... + +- fn properties(&self) -> &PlanProperties { ++ fn properties(&self) -> &Arc { + &self.cache + } +``` + +Note: The optimization for `with_new_children` can be implemented for any `ExecutionPlan`. This can reduce planning time as well as the time for resetting plan states. +To support it, you can use the macro: `check_if_same_properties`. For it to work, you need to implement the function: `with_new_children_and_same_properties` with semantics +identical to `with_new_children`, but operating under the assumption that the properties of the children plans have not changed. + +An example of supporting this optimization for `ProjectionExec`: + +```diff + impl ProjectionExec { ++ fn with_new_children_and_same_properties( ++ &self, ++ mut children: Vec>, ++ ) -> Self { ++ Self { ++ input: children.swap_remove(0), ++ metrics: ExecutionPlanMetricsSet::new(), ++ ..Self::clone(self) ++ } ++ } + } + + impl ExecutionPlan for ProjectionExec { + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { ++ check_if_same_properties!(self, children); + ProjectionExec::try_new( + self.projector.projection().into_iter().cloned(), + children.swap_remove(0), + ) + .map(|p| Arc::new(p) as _) + } + } + +... + +``` + ### `SimplifyInfo` trait removed, `SimplifyContext` now uses builder-style API The `SimplifyInfo` trait has been removed and replaced with the concrete `SimplifyContext` struct. This simplifies the expression simplification API and removes the need for trait objects. From bcb311f3cdc259d2e4d0b033d9faeb50927575e1 Mon Sep 17 00:00:00 2001 From: Albert Skalt Date: Thu, 12 Feb 2026 10:13:47 +0300 Subject: [PATCH 5/6] rebase fixes --- datafusion/physical-plan/src/aggregates/mod.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index e4c5a86d6158f..1d85703e9166f 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -1667,9 +1667,9 @@ impl ExecutionPlan for AggregateExec { } Ok(Some(Arc::new(Self { - group_by, - aggr_expr, - filter_expr, + group_by: Arc::new(group_by), + aggr_expr: aggr_expr.into(), + filter_expr: filter_expr.into(), dynamic_filter: None, metrics: ExecutionPlanMetricsSet::new(), ..self.clone() From 3aa6724a4ed0b5dc7eec08f9029228e8db7697ba Mon Sep 17 00:00:00 2001 From: Albert Skalt Date: Sat, 14 Feb 2026 19:31:02 +0300 Subject: [PATCH 6/6] feat(execution plan): add reusable plan This patch adds an execution plan wrapper that allows to bind parameters and reuse plans with placeholders. --- datafusion/core/benches/reset_plan_states.rs | 247 ++++++++++++------ datafusion/core/tests/sql/select.rs | 63 +++-- datafusion/datasource/src/values.rs | 68 ++--- .../physical-plan/src/execution_plan.rs | 50 +--- .../physical-plan/src/joins/hash_join/exec.rs | 2 +- .../src/joins/nested_loop_join.rs | 2 +- datafusion/physical-plan/src/lib.rs | 1 + datafusion/physical-plan/src/reuse.rs | 198 ++++++++++++++ 8 files changed, 439 insertions(+), 192 deletions(-) create mode 100644 datafusion/physical-plan/src/reuse.rs diff --git a/datafusion/core/benches/reset_plan_states.rs b/datafusion/core/benches/reset_plan_states.rs index a4f9e716bd761..56ad0a359d0eb 100644 --- a/datafusion/core/benches/reset_plan_states.rs +++ b/datafusion/core/benches/reset_plan_states.rs @@ -15,16 +15,20 @@ // specific language governing permissions and limitations // under the License. +use std::cell::OnceCell; use std::sync::{Arc, LazyLock}; use arrow_schema::{DataType, Field, Fields, Schema, SchemaRef}; +use criterion::measurement::WallTime; use criterion::{Criterion, criterion_group, criterion_main}; use datafusion::prelude::SessionContext; use datafusion_catalog::MemTable; +use datafusion_common::metadata::ScalarAndMetadata; +use datafusion_common::{ParamValues, ScalarValue}; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::displayable; -use datafusion_physical_plan::execution_plan::prepare_execution; use datafusion_physical_plan::execution_plan::reset_plan_states; +use datafusion_physical_plan::reuse::ReusableExecutionPlan; use tokio::runtime::Runtime; const NUM_FIELDS: usize = 1000; @@ -38,6 +42,44 @@ static SCHEMA: LazyLock = LazyLock::new(|| { )) }); +/// Decides when to generate placeholders, helping to form a query +/// with a certain placeholders percent. +struct PlaceholderCounter { + placeholders_percent: usize, + c: usize, + placeholder_idx: usize, + num_placeholders: usize, +} + +impl PlaceholderCounter { + fn new(placeholders_percent: usize) -> Self { + Self { + placeholders_percent, + c: 0, + placeholder_idx: 0, + num_placeholders: 0, + } + } + + fn placeholder(&mut self) -> Option { + let is_placeholder = self.c < self.placeholders_percent; + self.c += 1; + if self.c >= 100 { + self.c = 0; + } + if is_placeholder { + self.num_placeholders += 1; + Some("$1".to_owned()) + } else { + None + } + } + + fn placeholder_or(&mut self, f: impl FnOnce() -> String) -> String { + self.placeholder().unwrap_or_else(f) + } +} + fn col_name(i: usize) -> String { format!("x_{i}") } @@ -48,7 +90,7 @@ fn aggr_name(i: usize) -> String { fn physical_plan( ctx: &SessionContext, - rt: &Runtime, + rt: &tokio::runtime::Handle, sql: &str, ) -> Arc { rt.block_on(async { @@ -61,15 +103,16 @@ fn physical_plan( }) } -fn predicate(col_name: impl Fn(usize) -> String, len: usize) -> String { +fn predicate(mut comparee: impl FnMut(usize) -> (String, String), len: usize) -> String { let mut predicate = String::new(); for i in 0..len { if i > 0 { predicate.push_str(" AND "); } - predicate.push_str(&col_name(i)); + let (lhs, rhs) = comparee(i); + predicate.push_str(&lhs); predicate.push_str(" = "); - predicate.push_str(&i.to_string()); + predicate.push_str(&rhs); } predicate } @@ -84,7 +127,8 @@ fn predicate(col_name: impl Fn(usize) -> String, len: usize) -> String { /// /// Where `p1` and `p2` some long predicates. /// -fn query1() -> String { +fn query0(placeholders_percent: usize) -> (String, usize) { + let mut plc = PlaceholderCounter::new(placeholders_percent); let mut query = String::new(); query.push_str("SELECT "); for i in 0..NUM_FIELDS { @@ -92,15 +136,37 @@ fn query1() -> String { query.push_str(", "); } query.push_str("AVG("); - query.push_str(&col_name(i)); + + if let Some(placeholder) = plc.placeholder() { + query.push_str(&format!("{}+{}", placeholder, col_name(i))); + } else { + query.push_str(&col_name(i)); + } + query.push_str(") AS "); query.push_str(&aggr_name(i)); } query.push_str(" FROM t WHERE "); - query.push_str(&predicate(col_name, PREDICATE_LEN)); + query.push_str(&predicate( + |i| { + ( + plc.placeholder_or(|| col_name(i)), + plc.placeholder_or(|| col_name(i + 1)), + ) + }, + PREDICATE_LEN, + )); query.push_str(" HAVING "); - query.push_str(&predicate(aggr_name, PREDICATE_LEN)); - query + query.push_str(&predicate( + |i| { + ( + plc.placeholder_or(|| aggr_name(i)), + plc.placeholder_or(|| aggr_name(i + 1)), + ) + }, + PREDICATE_LEN, + )); + (query, plc.num_placeholders) } /// Returns a typical plan for the query like: @@ -110,27 +176,35 @@ fn query1() -> String { /// WHERE p1 /// ``` /// -fn query2() -> String { +fn query1(placeholders_percent: usize) -> (String, usize) { + let mut plc = PlaceholderCounter::new(placeholders_percent); let mut query = String::new(); query.push_str("SELECT "); for i in (0..NUM_FIELDS).step_by(2) { if i > 0 { query.push_str(", "); } - if (i / 2) % 2 == 0 { - query.push_str(&format!("t.{}", col_name(i))); + let col = if (i / 2) % 2 == 0 { + format!("t.{}", col_name(i)) } else { - query.push_str(&format!("v.{}", col_name(i))); - } + format!("v.{}", col_name(i)) + }; + let add = plc.placeholder_or(|| "1".to_owned()); + let proj = format!("{col} + {add}"); + query.push_str(&proj); } query.push_str(" FROM t JOIN v ON t.x_0 = v.x_0 WHERE "); - fn qualified_name(i: usize) -> String { - format!("t.{}", col_name(i)) - } - - query.push_str(&predicate(qualified_name, PREDICATE_LEN)); - query + query.push_str(&predicate( + |i| { + ( + plc.placeholder_or(|| format!("t.{}", col_name(i))), + plc.placeholder_or(|| i.to_string()), + ) + }, + PREDICATE_LEN, + )); + (query, plc.num_placeholders) } /// Returns a typical plan for the query like: @@ -140,7 +214,8 @@ fn query2() -> String { /// WHERE p /// ``` /// -fn query3() -> String { +fn query2(placeholders_percent: usize) -> (String, usize) { + let mut plc = PlaceholderCounter::new(placeholders_percent); let mut query = String::new(); query.push_str("SELECT "); @@ -151,24 +226,23 @@ fn query3() -> String { } query.push_str(&col_name(i * 2)); query.push_str(" + "); - query.push_str(&col_name(i * 2 + 1)); + query.push_str(&plc.placeholder_or(|| col_name(i * 2 + 1))); } query.push_str(" FROM t WHERE "); - query.push_str(&predicate(col_name, PREDICATE_LEN)); - query + query.push_str(&predicate( + |i| { + ( + plc.placeholder_or(|| col_name(i)), + plc.placeholder_or(|| i.to_string()), + ) + }, + PREDICATE_LEN, + )); + (query, plc.num_placeholders) } -fn run_reset_states(b: &mut criterion::Bencher, plan: &Arc) { - b.iter(|| std::hint::black_box(reset_plan_states(Arc::clone(plan)).unwrap())); -} - -/// Benchmark is intended to measure overhead of actions, required to perform -/// making an independent instance of the execution plan to re-execute it, avoiding -/// re-planning stage. -fn bench_reset_plan_states(c: &mut Criterion) { - env_logger::init(); - +fn init() -> (SessionContext, Runtime) { let rt = Runtime::new().unwrap(); let ctx = SessionContext::new(); ctx.register_table( @@ -182,56 +256,73 @@ fn bench_reset_plan_states(c: &mut Criterion) { Arc::new(MemTable::try_new(Arc::clone(&SCHEMA), vec![vec![], vec![]]).unwrap()), ) .unwrap(); - - macro_rules! bench_query { - ($query_producer: expr) => {{ - let sql = $query_producer(); - let plan = physical_plan(&ctx, &rt, &sql); - log::debug!("plan:\n{}", displayable(plan.as_ref()).indent(true)); - move |b| run_reset_states(b, &plan) - }}; - } - - c.bench_function("query1", bench_query!(query1)); - c.bench_function("query2", bench_query!(query2)); - c.bench_function("query3", bench_query!(query3)); + (ctx, rt) } -fn run_prepare_execution(b: &mut criterion::Bencher, plan: &Arc) { - b.iter(|| std::hint::black_box(prepare_execution(Arc::clone(plan), None).unwrap())); +/// Benchmark is intended to measure overhead of actions, required to perform +/// making an independent instance of the execution plan to re-execute it, avoiding +/// re-planning stage. +fn bench_reset( + g: &mut criterion::BenchmarkGroup<'_, WallTime>, + query_fn: impl FnOnce() -> String, +) { + let (ctx, rt) = init(); + let query = query_fn(); + let rt = rt.handle(); + let plan: OnceCell> = OnceCell::new(); + g.bench_function("reset", |b| { + let plan = plan.get_or_init(|| { + log::info!("sql:\n{}\n\n", query); + let plan = physical_plan(&ctx, rt, &query); + log::info!("plan:\n{}", displayable(plan.as_ref()).indent(true)); + plan + }); + b.iter(|| std::hint::black_box(reset_plan_states(Arc::clone(&plan)).unwrap())) + }); } -/// Benchmark is intended to measure overhead of actions, required to perform -/// making an independent instance of the execution plan to re-execute it with placeholders, -/// avoiding re-planning stage. -fn bench_prepare_execution(c: &mut Criterion) { - let rt = Runtime::new().unwrap(); - let ctx = SessionContext::new(); - ctx.register_table( - "t", - Arc::new(MemTable::try_new(Arc::clone(&SCHEMA), vec![vec![], vec![]]).unwrap()), - ) - .unwrap(); +/// The same as [`bench_reset`] for placeholdered plans. +/// `placeholders_percent` is a percent of placeholders that must be used in generated queries. +fn bench_bind( + g: &mut criterion::BenchmarkGroup<'_, WallTime>, + placeholders_percent: usize, + query_fn: impl FnOnce(usize) -> (String, usize), +) { + let (ctx, rt) = init(); + let params = ParamValues::List(vec![ScalarAndMetadata::new( + ScalarValue::Int64(Some(42)), + None, + )]); + let (query, num_placeholders) = query_fn(placeholders_percent); + let rt = rt.handle(); + let plan: OnceCell = OnceCell::new(); + g.bench_function(format!("{num_placeholders}_placeholders"), move |b| { + let plan = plan.get_or_init(|| { + log::info!("sql:\n{}\n\n", query); + let plan = physical_plan(&ctx, rt, &query); + log::info!("plan:\n{}", displayable(plan.as_ref()).indent(true)); + plan.into() + }); + b.iter(|| std::hint::black_box(plan.bind(Some(¶ms)))) + }); +} - ctx.register_table( - "v", - Arc::new(MemTable::try_new(Arc::clone(&SCHEMA), vec![vec![], vec![]]).unwrap()), - ) - .unwrap(); +fn criterion_benchmark(c: &mut Criterion) { + env_logger::init(); - macro_rules! bench_query { - ($query_producer: expr) => {{ - let sql = $query_producer(); - let plan = physical_plan(&ctx, &rt, &sql); - log::debug!("plan:\n{}", displayable(plan.as_ref()).indent(true)); - move |b| run_prepare_execution(b, &plan) - }}; + for (query_idx, query_fn) in [query0, query1, query2].iter().enumerate() { + { + let mut g = c.benchmark_group(format!("reset_query{query_idx}")); + bench_reset(&mut g, || query_fn(0).0); + } + { + let mut g = c.benchmark_group(format!("bind_query{query_idx}")); + for placeholders_percent in [0, 1, 10, 50, 100] { + bench_bind(&mut g, placeholders_percent, query_fn); + } + } } - - c.bench_function("query1", bench_query!(query1)); - c.bench_function("query2", bench_query!(query2)); - c.bench_function("query3", bench_query!(query3)); } -criterion_group!(benches, bench_reset_plan_states, bench_prepare_execution); +criterion_group!(benches, criterion_benchmark); criterion_main!(benches); diff --git a/datafusion/core/tests/sql/select.rs b/datafusion/core/tests/sql/select.rs index 2fe542bc70a2d..3c2e256a09bc4 100644 --- a/datafusion/core/tests/sql/select.rs +++ b/datafusion/core/tests/sql/select.rs @@ -19,7 +19,7 @@ use std::collections::HashMap; use super::*; use datafusion_common::{ParamValues, ScalarValue, metadata::ScalarAndMetadata}; -use datafusion_physical_plan::execution_plan::prepare_execution; +use datafusion_physical_plan::reuse::ReusableExecutionPlan; use insta::assert_snapshot; #[tokio::test] @@ -439,15 +439,15 @@ async fn test_resolve_window_function() -> Result<()> { let batch = record_batch!(("id", Int32, [1, 2]), ("name", Utf8, ["Alex", "Bob"]))?; ctx.register_batch("t1", batch)?; - let plan = ctx - .sql("SELECT id, SUM(id + $1) OVER (PARTITION BY name ORDER BY id) FROM t1") - .await? - .create_physical_plan() - .await?; + let plan = ReusableExecutionPlan::new( + ctx.sql("SELECT id, SUM(id + $1) OVER (PARTITION BY name ORDER BY id) FROM t1") + .await? + .create_physical_plan() + .await?, + ); let param_values = ParamValues::List(vec![ScalarValue::Int32(Some(100)).into()]); - let plan = prepare_execution(plan, Some(¶m_values))?; - let batches = collect(plan, ctx.task_ctx()).await?; + let batches = collect(plan.bind(Some(¶m_values))?.plan(), ctx.task_ctx()).await?; assert_snapshot!(batches_to_sort_string(&batches), @r" +----+--------------------------------------------------------------------------------------------------------------------------+ @@ -479,15 +479,15 @@ async fn test_resolve_join() -> Result<()> { ctx.register_batch("t1", batch_1)?; ctx.register_batch("t2", batch_2)?; - let plan = ctx - .sql("SELECT t1.name, t2.age FROM t1 JOIN t2 ON t1.id + $1 = t2.id;") - .await? - .create_physical_plan() - .await?; + let plan = ReusableExecutionPlan::new( + ctx.sql("SELECT t1.name, t2.age FROM t1 JOIN t2 ON t1.id + $1 = t2.id;") + .await? + .create_physical_plan() + .await?, + ); let param_values = ParamValues::List(vec![ScalarValue::Int32(Some(8)).into()]); - let plan = prepare_execution(plan, Some(¶m_values))?; - let batches = collect(plan, ctx.task_ctx()).await?; + let batches = collect(plan.bind(Some(¶m_values))?.plan(), ctx.task_ctx()).await?; assert_snapshot!(batches_to_sort_string(&batches), @r" +------+-----+ @@ -503,21 +503,21 @@ async fn test_resolve_join() -> Result<()> { #[tokio::test] async fn test_resolve_cast() -> Result<()> { let ctx = SessionContext::new(); - let plan = ctx - .sql("SELECT CAST($1 as INT)") - .await? - .create_physical_plan() - .await?; + let plan = ReusableExecutionPlan::new( + ctx.sql("SELECT CAST($1 as INT)") + .await? + .create_physical_plan() + .await?, + ); let param_values = ParamValues::List(vec![ ScalarValue::Utf8(Some("not a number".to_string())).into(), ]); - assert!(prepare_execution(Arc::clone(&plan), Some(¶m_values)).is_err()); + assert!(plan.bind(Some(¶m_values)).is_err()); let param_values = ParamValues::List(vec![ScalarValue::Utf8(Some("200".to_string())).into()]); - let plan = prepare_execution(plan, Some(¶m_values))?; - let batches = collect(plan, ctx.task_ctx()).await?; + let batches = collect(plan.bind(Some(¶m_values))?.plan(), ctx.task_ctx()).await?; assert_snapshot!(batches_to_sort_string(&batches), @r" +-----+ @@ -533,17 +533,17 @@ async fn test_resolve_cast() -> Result<()> { #[tokio::test] async fn test_resolve_try_cast() -> Result<()> { let ctx = SessionContext::new(); - let plan = ctx - .sql("SELECT TRY_CAST($1 as INT)") - .await? - .create_physical_plan() - .await?; + let plan = ReusableExecutionPlan::new( + ctx.sql("SELECT TRY_CAST($1 as INT)") + .await? + .create_physical_plan() + .await?, + ); let param_values = ParamValues::List(vec![ ScalarValue::Utf8(Some("not a number".to_string())).into(), ]); - let plan1 = prepare_execution(Arc::clone(&plan), Some(¶m_values))?; - let batches = collect(plan1, ctx.task_ctx()).await?; + let batches = collect(plan.bind(Some(¶m_values))?.plan(), ctx.task_ctx()).await?; assert_snapshot!(batches_to_sort_string(&batches), @r" +----+ @@ -555,8 +555,7 @@ async fn test_resolve_try_cast() -> Result<()> { let param_values = ParamValues::List(vec![ScalarValue::Utf8(Some("200".to_string())).into()]); - let plan2 = prepare_execution(plan, Some(¶m_values))?; - let batches = collect(plan2, ctx.task_ctx()).await?; + let batches = collect(plan.bind(Some(¶m_values))?.plan(), ctx.task_ctx()).await?; assert_snapshot!(batches_to_sort_string(&batches), @r" +-----+ diff --git a/datafusion/datasource/src/values.rs b/datafusion/datasource/src/values.rs index 75ce506eeb649..348e26fef177a 100644 --- a/datafusion/datasource/src/values.rs +++ b/datafusion/datasource/src/values.rs @@ -341,7 +341,7 @@ mod tests { use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{BinaryExpr, lit, placeholder}; use datafusion_physical_plan::{ - ExecutionPlan, collect, execution_plan::prepare_execution, + ExecutionPlan, collect, reuse::ReusableExecutionPlan, }; #[test] @@ -436,15 +436,17 @@ mod tests { // Should be ValuesSource because of placeholder. assert!(values_exec.data_source().as_any().is::()); - let exec = prepare_execution( - values_exec, - Some(&ParamValues::List(vec![ - ScalarValue::Int32(Some(10)).into(), - ])), - )?; + let exec = ReusableExecutionPlan::new(values_exec as _); let task_ctx = Arc::new(TaskContext::default()); - let batch = collect(exec, task_ctx).await?; + let batch = collect( + exec.bind(Some(&ParamValues::List(vec![ + ScalarValue::Int32(Some(10)).into(), + ])))? + .plan(), + task_ctx, + ) + .await?; let expected = [ "+-----+----+", "| a | b |", @@ -471,17 +473,19 @@ mod tests { ]]; let values_exec = ValuesSource::try_new_exec(Arc::clone(&schema), data)? as _; - let exec = prepare_execution( - Arc::clone(&values_exec), - Some(&ParamValues::List(vec![ - ScalarValue::Int32(Some(10)).into(), - ScalarValue::Int32(Some(20)).into(), - ])), - )?; + let exec = ReusableExecutionPlan::new(values_exec); let task_ctx = Arc::new(TaskContext::default()); - let batch = collect(Arc::clone(&exec), Arc::clone(&task_ctx)).await?; + let batch = collect( + exec.bind(Some(&ParamValues::List(vec![ + ScalarValue::Int32(Some(10)).into(), + ScalarValue::Int32(Some(20)).into(), + ])))? + .plan(), + Arc::clone(&task_ctx), + ) + .await?; let expected = [ "+----+----+", "| a | b |", @@ -491,15 +495,15 @@ mod tests { ]; assert_batches_eq!(expected, &batch); - let exec = prepare_execution( - values_exec, - Some(&ParamValues::List(vec![ + let batch = collect( + exec.bind(Some(&ParamValues::List(vec![ ScalarValue::Int32(Some(30)).into(), ScalarValue::Int32(Some(40)).into(), - ])), - )?; - - let batch = collect(exec, task_ctx).await?; + ])))? + .plan(), + task_ctx, + ) + .await?; let expected = [ "+----+----+", "| a | b |", @@ -535,16 +539,18 @@ mod tests { let result = collect(Arc::clone(&values_exec), task_ctx).await; assert!(result.is_err()); - let exec = prepare_execution( - values_exec, - Some(&ParamValues::Map(HashMap::from_iter([( - "foo".to_string(), - ScalarValue::Int32(Some(20)).into(), - )]))), - )?; + let exec = ReusableExecutionPlan::new(values_exec); let task_ctx = Arc::new(TaskContext::default()); - let batch = collect(Arc::clone(&exec), task_ctx).await?; + let batch = collect( + exec.bind(Some(&ParamValues::Map(HashMap::from_iter([( + "foo".to_string(), + ScalarValue::Int32(Some(20)).into(), + )]))))? + .plan(), + task_ctx, + ) + .await?; let expected = [ "+----+----+", "| a | b |", diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 21a558b80bce0..6ad1afa051e8d 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -31,7 +31,6 @@ pub use datafusion_common::utils::project_schema; pub use datafusion_common::{ColumnStatistics, Statistics, internal_err}; pub use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream}; pub use datafusion_expr::{Accumulator, ColumnarValue}; -use datafusion_physical_expr::expressions::resolve_expr_placeholders; pub use datafusion_physical_expr::window::WindowExpr; pub use datafusion_physical_expr::{ Distribution, Partitioning, PhysicalExpr, expressions, @@ -51,7 +50,7 @@ use arrow::array::{Array, RecordBatch}; use arrow::datatypes::SchemaRef; use datafusion_common::config::ConfigOptions; use datafusion_common::{ - Constraints, DataFusionError, ParamValues, Result, assert_eq_or_internal_err, + Constraints, DataFusionError, Result, assert_eq_or_internal_err, assert_or_internal_err, exec_err, }; use datafusion_common_runtime::JoinSet; @@ -1567,53 +1566,6 @@ pub fn reset_plan_states(plan: Arc) -> Result, - param_values: Option<&ParamValues>, -) -> Result> { - plan.transform_up(|plan| { - let plan = if let Some(iter) = plan.physical_expressions() { - let mut has_placeholders = false; - let exprs = iter - .map(|expr| { - let resolved_expr = - resolve_expr_placeholders(Arc::clone(&expr), param_values)?; - has_placeholders |= !Arc::ptr_eq(&expr, &resolved_expr); - Ok(resolved_expr) - }) - .collect::>>()?; - if !has_placeholders { - Arc::clone(&plan).reset_state()? - } else { - // `with_physical_expressions` resets plan state. - let Some(plan) = - plan.with_physical_expressions(ReplacePhysicalExpr { exprs })? - else { - return exec_err!( - "plan {} does not support expression substitution", - plan.name() - ); - }; - plan - } - } else { - plan.reset_state()? - }; - Ok(Transformed::yes(plan)) - }) - .map(|tnr| tnr.data) -} - /// Check if the `plan` children has the same properties as passed `children`. /// In this case plan can avoid self properties re-computation when its children /// replace is requested. diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 2647468376120..c3811083b94e4 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -1537,7 +1537,7 @@ impl ExecutionPlan for HashJoinExec { column_indices: self.column_indices.clone(), null_equality: self.null_equality, null_aware: self.null_aware, - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), // Reset dynamic filter and bounds accumulator to initial state dynamic_filter: None, }))) diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index 92ead98d439e4..4ecb58d9ce077 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -753,7 +753,7 @@ impl ExecutionPlan for NestedLoopJoinExec { column_indices: self.column_indices.clone(), projection: self.projection.clone(), metrics: Default::default(), - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), }))) } } diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 6467d7a2e389d..e65ec96ad84eb 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -84,6 +84,7 @@ pub mod placeholder_row; pub mod projection; pub mod recursive_query; pub mod repartition; +pub mod reuse; pub mod sort_pushdown; pub mod sorts; pub mod spill; diff --git a/datafusion/physical-plan/src/reuse.rs b/datafusion/physical-plan/src/reuse.rs new file mode 100644 index 0000000000000..5f3536ccc11f5 --- /dev/null +++ b/datafusion/physical-plan/src/reuse.rs @@ -0,0 +1,198 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use datafusion_common::{ + ParamValues, Result, exec_err, + tree_node::{Transformed, TreeNode, TreeNodeRecursion}, +}; +use datafusion_physical_expr::{ + PhysicalExpr, + expressions::{has_placeholders, resolve_expr_placeholders}, +}; + +use crate::{ExecutionPlan, execution_plan::ReplacePhysicalExpr}; + +/// Wraps an [`ExecutionPlan`] that may contain placeholders, making +/// it re-executable, avoiding re-planning stage. +/// +/// # Limitations +/// +/// Plan does not support re-execution if it (OR): +/// +/// * uses dynamic filters, +/// * represents a recursive query. +/// +/// This invariant is not enforced by [`ReusableExecutionPlan`] itself, so it +/// must be checked by user. +/// +pub struct ReusableExecutionPlan { + binder: Binder, + bound_plan: Option>, +} + +impl ReusableExecutionPlan { + /// Make a new [`ReusableExecutionPlan`] bound to the passed `plan`. + pub fn new(plan: Arc) -> Self { + let binder = Binder::new(plan); + Self { + binder, + bound_plan: None, + } + } + + /// Return a ready to execution instance of [`ReusableExecutionPlan`], + /// where placeholders are bound to the passed `params`. + pub fn bind(&self, params: Option<&ParamValues>) -> Result { + let bound_plan = self.binder.bind(params).map(Some)?; + Ok(Self { + binder: self.binder.clone(), + bound_plan, + }) + } + + /// Return an inner plan to execute. + /// + /// If this plan is a result of [`Self::bind`] call, then bound plan is returned. + /// Otherwise, an initial plan is returned. + pub fn plan(&self) -> Arc { + self.bound_plan + .clone() + .unwrap_or_else(|| Arc::clone(&self.binder.plan)) + } +} + +impl From> for ReusableExecutionPlan { + fn from(plan: Arc) -> Self { + Self::new(plan) + } +} + +#[derive(Debug, Clone)] +struct NodeWithPlaceholders { + /// The index of the node in the tree traversal. + idx: usize, + /// Positions of the placeholders among plan physical expressions. + placeholder_idx: Vec, +} + +impl NodeWithPlaceholders { + /// Returns [`Some`] if passed `node` contains placeholders and must + /// be resolved on binding stage. + fn new(node: &Arc, idx: usize) -> Option { + let placeholder_idx = if let Some(iter) = node.physical_expressions() { + iter.enumerate() + .filter_map(|(i, expr)| { + if has_placeholders(&expr) { + Some(i) + } else { + None + } + }) + .collect() + } else { + vec![] + }; + + if placeholder_idx.is_empty() { + None + } else { + Some(Self { + idx, + placeholder_idx, + }) + } + } + + fn resolve( + &self, + node: &Arc, + params: Option<&ParamValues>, + ) -> Result> { + let Some(expr) = node.physical_expressions() else { + return exec_err!("node {} does not support expressions export", node.name()); + }; + let mut exprs: Vec> = expr.collect(); + for idx in self.placeholder_idx.iter() { + exprs[*idx] = resolve_expr_placeholders(Arc::clone(&exprs[*idx]), params)?; + } + let Some(resolved_node) = + node.with_physical_expressions(ReplacePhysicalExpr { exprs })? + else { + return exec_err!( + "node {} does not support expressions replace", + node.name() + ); + }; + Ok(resolved_node) + } +} + +/// Helper to bound placeholders and reset plan nodes state. +#[derive(Clone)] +struct Binder { + /// Created during [`Binder`] construction. + /// This way we avoid runtime rebuild for expressions without placeholders. + nodes_to_resolve: Arc<[NodeWithPlaceholders]>, + plan: Arc, +} + +impl Binder { + fn new(plan: Arc) -> Self { + let mut nodes_to_resolve = vec![]; + let mut cursor = 0; + + plan.apply(|node| { + let idx = cursor; + cursor += 1; + if let Some(node) = NodeWithPlaceholders::new(node, idx) { + nodes_to_resolve.push(node); + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + + Self { + plan, + nodes_to_resolve: nodes_to_resolve.into(), + } + } + + fn bind(&self, params: Option<&ParamValues>) -> Result> { + let mut cursor = 0; + let mut resolve_node_idx = 0; + Arc::clone(&self.plan) + .transform_down(|node| { + let idx = cursor; + cursor += 1; + if resolve_node_idx < self.nodes_to_resolve.len() + && self.nodes_to_resolve[resolve_node_idx].idx == idx + { + // Note: `resolve` replaces plan expressions, which also resets a plan state. + let resolved_node = + self.nodes_to_resolve[resolve_node_idx].resolve(&node, params)?; + resolve_node_idx += 1; + Ok(Transformed::yes(resolved_node)) + } else { + // Reset state. + Ok(Transformed::yes(node.reset_state()?)) + } + }) + .map(|tnr| tnr.data) + } +}