diff --git a/datafusion-examples/examples/custom_data_source/custom_datasource.rs b/datafusion-examples/examples/custom_data_source/custom_datasource.rs index a5a38edf0b6f..6f176f8b4660 100644 --- a/datafusion-examples/examples/custom_data_source/custom_datasource.rs +++ b/datafusion-examples/examples/custom_data_source/custom_datasource.rs @@ -36,8 +36,8 @@ use datafusion::physical_expr::EquivalenceProperties; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::memory::MemoryStream; use datafusion::physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, - SendableRecordBatchStream, project_schema, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, + PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, project_schema, }; use datafusion::prelude::*; @@ -268,13 +268,24 @@ impl ExecutionPlan for CustomExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, 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 6decb84b55be..89eff74e0d73 100644 --- a/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs +++ b/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs @@ -39,7 +39,8 @@ use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::logical_expr::LogicalPlanBuilder; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ReplaceChildrenOptions, }; use datafusion::prelude::*; use futures::stream::{StreamExt, TryStreamExt}; @@ -237,9 +238,10 @@ impl ExecutionPlan for BufferingExecutionPlan { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _options: ReplaceChildrenOptions, ) -> Result> { if children.len() == 1 { Ok(Arc::new(BufferingExecutionPlan::new( @@ -251,6 +253,16 @@ impl ExecutionPlan for BufferingExecutionPlan { } } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, diff --git a/datafusion-examples/examples/proto/composed_extension_codec.rs b/datafusion-examples/examples/proto/composed_extension_codec.rs index d5197fe61bea..51c1bc7c5518 100644 --- a/datafusion-examples/examples/proto/composed_extension_codec.rs +++ b/datafusion-examples/examples/proto/composed_extension_codec.rs @@ -39,6 +39,7 @@ use datafusion::common::Result; use datafusion::common::internal_err; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::execution::TaskContext; +use datafusion::physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}; use datafusion::physical_plan::{DisplayAs, ExecutionPlan}; use datafusion::prelude::SessionContext; use datafusion_proto::physical_plan::{ @@ -111,13 +112,24 @@ impl ExecutionPlan for ParentExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unreachable!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -198,13 +210,24 @@ impl ExecutionPlan for ChildExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unreachable!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, diff --git a/datafusion-examples/examples/relation_planner/table_sample.rs b/datafusion-examples/examples/relation_planner/table_sample.rs index 388175ee3a17..7a8f533ac3a9 100644 --- a/datafusion-examples/examples/relation_planner/table_sample.rs +++ b/datafusion-examples/examples/relation_planner/table_sample.rs @@ -100,7 +100,6 @@ use futures::{ use rand::{Rng, SeedableRng, rngs::StdRng}; use tonic::async_trait; -use datafusion::optimizer::simplify_expressions::simplify_literal::parse_literal; use datafusion::{ catalog::Session, execution::{ @@ -116,6 +115,10 @@ use datafusion::{ physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner}, prelude::*, }; +use datafusion::{ + optimizer::simplify_expressions::simplify_literal::parse_literal, + physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}, +}; use datafusion_common::{ DFSchemaRef, DataFusionError, Result, Statistics, internal_err, not_impl_err, plan_datafusion_err, plan_err, tree_node::TreeNodeRecursion, @@ -698,9 +701,10 @@ impl ExecutionPlan for SampleExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(Arc::new(Self::try_new( children.swap_remove(0), @@ -710,6 +714,16 @@ impl ExecutionPlan for SampleExec { )?)) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, diff --git a/datafusion/catalog/src/memory/table.rs b/datafusion/catalog/src/memory/table.rs index ef5669a3a13f..4cf96cb364be 100644 --- a/datafusion/catalog/src/memory/table.rs +++ b/datafusion/catalog/src/memory/table.rs @@ -45,8 +45,8 @@ use datafusion_physical_expr::{ use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PhysicalExpr, - PlanProperties, collect_partitioned, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, + PhysicalExpr, PlanProperties, ReplaceChildrenOptions, collect_partitioned, }; use datafusion_session::Session; @@ -572,13 +572,24 @@ impl ExecutionPlan for DmlResultExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 1202f08a567e..3c1e7b50780a 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -3323,6 +3323,7 @@ mod tests { use datafusion_functions_aggregate::expr_fn::sum; use datafusion_physical_expr::EquivalenceProperties; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; + use datafusion_physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}; use datafusion_session::QueryPlanner; #[derive(Debug)] @@ -4932,9 +4933,10 @@ mod tests { vec![] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { if children.is_empty() { Ok(self) @@ -4943,6 +4945,16 @@ mod tests { } } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -5102,12 +5114,22 @@ digraph { fn name(&self) -> &str { "always ok" } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(Arc::new(Self(children))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } fn schema(&self) -> SchemaRef { Arc::new(Schema::empty()) } @@ -5157,12 +5179,22 @@ digraph { fn schema(&self) -> SchemaRef { Arc::new(Schema::empty()) } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } fn children(&self) -> Vec<&Arc> { unimplemented!() } @@ -5216,10 +5248,16 @@ digraph { // ok plan let ok_node: Arc = Arc::new(OkExtensionNode(vec![])); let child = Arc::clone(&ok_node); - let ok_plan = Arc::clone(&ok_node).with_new_children(vec![ - Arc::clone(&child).with_new_children(vec![Arc::clone(&child)])?, - Arc::clone(&child), - ])?; + let ok_plan = Arc::clone(&ok_node).replace_children( + vec![ + Arc::clone(&child).replace_children( + vec![Arc::clone(&child)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?, + Arc::clone(&child), + ], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; // Test: check should pass with same schema let equal_schema = ok_plan.schema(); @@ -5251,10 +5289,16 @@ digraph { // Test: should fail when descendent extension node fails let failing_node: Arc = Arc::new(InvariantFailsExtensionNode); - let invalid_plan = ok_node.with_new_children(vec![ - Arc::clone(&child).with_new_children(vec![Arc::clone(&failing_node)])?, - Arc::clone(&child), - ])?; + let invalid_plan = ok_node.replace_children( + vec![ + Arc::clone(&child).replace_children( + vec![Arc::clone(&failing_node)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?, + Arc::clone(&child), + ], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; let result = OptimizationInvariantChecker::new(&rule) .check(&invalid_plan, &ok_plan.schema()); if cfg!(debug_assertions) { @@ -5287,12 +5331,22 @@ digraph { fn schema(&self) -> SchemaRef { Arc::new(Schema::empty()) } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } fn children(&self) -> Vec<&Arc> { vec![] } @@ -5339,10 +5393,16 @@ digraph { let failing_node: Arc = Arc::new(ExecutableInvariantFails); let ok_node: Arc = Arc::new(OkExtensionNode(vec![])); let child = Arc::clone(&ok_node); - let plan = ok_node.with_new_children(vec![ - Arc::clone(&child).with_new_children(vec![Arc::clone(&failing_node)])?, - Arc::clone(&child), - ])?; + let plan = ok_node.replace_children( + vec![ + Arc::clone(&child).replace_children( + vec![Arc::clone(&failing_node)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?, + Arc::clone(&child), + ], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; let expected_err = InvariantChecker(InvariantLevel::Executable) .check(&plan) .unwrap_err(); diff --git a/datafusion/core/tests/custom_sources_cases/mod.rs b/datafusion/core/tests/custom_sources_cases/mod.rs index 43f4be05b24e..7abbcd6e9578 100644 --- a/datafusion/core/tests/custom_sources_cases/mod.rs +++ b/datafusion/core/tests/custom_sources_cases/mod.rs @@ -40,10 +40,12 @@ use datafusion_common::project_schema; use datafusion_common::stats::Precision; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_physical_expr::EquivalenceProperties; -use datafusion_physical_plan::PlanProperties; use datafusion_physical_plan::StatisticsArgs; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; +use datafusion_physical_plan::{ + ChildrenPropertiesMode, PlanProperties, ReplaceChildrenOptions, +}; use async_trait::async_trait; use futures::stream::Stream; @@ -165,13 +167,24 @@ impl ExecutionPlan for CustomExecutionPlan { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, 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 7437bbc5437c..a8f7f09ad016 100644 --- a/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs +++ b/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs @@ -41,6 +41,7 @@ use datafusion_expr::expr::{BinaryExpr, Cast}; use datafusion_functions_aggregate::expr_fn::count; use datafusion_physical_expr::EquivalenceProperties; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion_physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}; use async_trait::async_trait; @@ -117,9 +118,10 @@ impl ExecutionPlan for CustomPlan { vec![] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { // CustomPlan has no children if children.is_empty() { @@ -129,6 +131,16 @@ impl ExecutionPlan for CustomPlan { } } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/core/tests/custom_sources_cases/statistics.rs b/datafusion/core/tests/custom_sources_cases/statistics.rs index c14ca685b240..6213be8e2d24 100644 --- a/datafusion/core/tests/custom_sources_cases/statistics.rs +++ b/datafusion/core/tests/custom_sources_cases/statistics.rs @@ -37,7 +37,9 @@ use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{project_schema, stats::Precision}; use datafusion_physical_expr::EquivalenceProperties; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; -use datafusion_physical_plan::{StatisticsArgs, StatisticsContext}; +use datafusion_physical_plan::{ + ChildrenPropertiesMode, ReplaceChildrenOptions, StatisticsArgs, StatisticsContext, +}; use async_trait::async_trait; @@ -160,13 +162,24 @@ impl ExecutionPlan for StatisticsValidation { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/core/tests/fuzz_cases/once_exec.rs b/datafusion/core/tests/fuzz_cases/once_exec.rs index 638cbe4c9d41..c1db9a110d86 100644 --- a/datafusion/core/tests/fuzz_cases/once_exec.rs +++ b/datafusion/core/tests/fuzz_cases/once_exec.rs @@ -16,13 +16,14 @@ // under the License. use arrow_schema::SchemaRef; -use datafusion_common::internal_datafusion_err; use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{Result, internal_datafusion_err}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ReplaceChildrenOptions, }; use std::fmt::{Debug, Formatter}; use std::sync::{Arc, Mutex}; @@ -87,19 +88,30 @@ impl ExecutionPlan for OnceExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, - ) -> datafusion_common::Result> { + _: ReplaceChildrenOptions, + ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + /// Returns a stream which yields data fn execute( &self, partition: usize, _context: Arc, - ) -> datafusion_common::Result { + ) -> Result { assert_eq!(partition, 0); let stream = self.stream.lock().unwrap().take(); @@ -111,8 +123,8 @@ impl ExecutionPlan for OnceExec { &self, _f: &mut dyn FnMut( &Arc, - ) -> datafusion_common::Result, - ) -> datafusion_common::Result { + ) -> Result, + ) -> Result { Ok(TreeNodeRecursion::Continue) } } diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index e0b152d1f0aa..2dbacf1d898a 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -74,7 +74,8 @@ use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::union::UnionExec; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlanProperties, PlanProperties, displayable, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlanProperties, + PlanProperties, ReplaceChildrenOptions, displayable, }; use insta::Settings; @@ -193,9 +194,10 @@ impl ExecutionPlan for SortRequiredExec { vec![Some(OrderingRequirements::from(self.expr.clone()))] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { assert_eq!(children.len(), 1); let child = children.pop().unwrap(); @@ -205,6 +207,16 @@ impl ExecutionPlan for SortRequiredExec { ))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn apply_expressions( &self, _f: &mut dyn FnMut(&Arc) -> Result, @@ -290,15 +302,26 @@ impl ExecutionPlan for SinglePartitionMaintainsOrderExec { vec![false] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { assert_eq!(children.len(), 1); let child = children.pop().unwrap(); Ok(Arc::new(Self::new(child))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn apply_expressions( &self, _f: &mut dyn FnMut(&Arc) -> Result, diff --git a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs index 83fabcdff8da..86b60519da37 100644 --- a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs +++ b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs @@ -45,8 +45,9 @@ use datafusion_physical_plan::limit::GlobalLimitExec; use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::union::UnionExec; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, Partitioning, - PlanProperties, SendableRecordBatchStream, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + ExecutionPlanProperties, Partitioning, PlanProperties, ReplaceChildrenOptions, + SendableRecordBatchStream, }; use datafusion_physical_optimizer::output_requirements::OutputRequirementExec; @@ -130,18 +131,32 @@ impl ExecutionPlan for MockMultiPartitionExec { fn children(&self) -> Vec<&Arc> { vec![] } + + fn replace_children( + self: Arc, + _: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + Ok(self) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn apply_expressions( &self, _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( - self: Arc, - _children: Vec>, - ) -> Result> { - Ok(self) - } + fn execute( &self, _partition: usize, @@ -1022,17 +1037,27 @@ impl ExecutionPlan for MockReqExec { fn maintains_input_order(&self) -> Vec { vec![true] } - fn with_new_children( + fn replace_children( self: Arc, - mut c: Vec>, + mut children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { - assert_eq!(c.len(), 1); + assert_eq!(children.len(), 1); Ok(Arc::new(MockReqExec::new( - c.pop().expect("1 child"), + children.pop().expect("1 child"), self.dist.clone(), self.ord.clone(), ))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } fn execute( &self, _p: usize, diff --git a/datafusion/core/tests/physical_optimizer/join_selection.rs b/datafusion/core/tests/physical_optimizer/join_selection.rs index c7e3799842c8..265279a8ca1e 100644 --- a/datafusion/core/tests/physical_optimizer/join_selection.rs +++ b/datafusion/core/tests/physical_optimizer/join_selection.rs @@ -39,13 +39,15 @@ use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_optimizer::join_selection::JoinSelection; -use datafusion_physical_plan::ExecutionPlanProperties; use datafusion_physical_plan::displayable; use datafusion_physical_plan::joins::utils::ColumnIndex; use datafusion_physical_plan::joins::utils::JoinFilter; use datafusion_physical_plan::joins::{HashJoinExec, NestedLoopJoinExec, PartitionMode}; use datafusion_physical_plan::projection::ProjectionExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; +use datafusion_physical_plan::{ + ChildrenPropertiesMode, ExecutionPlanProperties, ReplaceChildrenOptions, +}; use datafusion_physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, StatisticsArgs, StatisticsContext, @@ -1108,13 +1110,24 @@ impl ExecutionPlan for UnboundedExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -1212,13 +1225,24 @@ impl ExecutionPlan for StatisticsExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/core/tests/physical_optimizer/pushdown_utils.rs b/datafusion/core/tests/physical_optimizer/pushdown_utils.rs index 4f8b9ad42b6c..0c2286527dbc 100644 --- a/datafusion/core/tests/physical_optimizer/pushdown_utils.rs +++ b/datafusion/core/tests/physical_optimizer/pushdown_utils.rs @@ -30,6 +30,7 @@ use datafusion_physical_expr_common::physical_expr::fmt_sql; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_plan::filter::batch_filter; use datafusion_physical_plan::filter_pushdown::{FilterPushdownPhase, PushedDown}; +use datafusion_physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}; use datafusion_physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, displayable, filter::FilterExec, @@ -489,9 +490,10 @@ impl ExecutionPlan for TestNode { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { assert!(children.len() == 1); Ok(Arc::new(TestNode::new( @@ -501,6 +503,16 @@ impl ExecutionPlan for TestNode { ))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/core/tests/physical_optimizer/test_utils.rs b/datafusion/core/tests/physical_optimizer/test_utils.rs index 833077fe491b..0835497f3451 100644 --- a/datafusion/core/tests/physical_optimizer/test_utils.rs +++ b/datafusion/core/tests/physical_optimizer/test_utils.rs @@ -70,9 +70,9 @@ use datafusion_physical_plan::tree_node::PlanContext; use datafusion_physical_plan::union::UnionExec; use datafusion_physical_plan::windows::{BoundedWindowAggExec, create_window_expr}; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, InputDistributionRequirements, - InputOrderMode, Partitioning, PlanProperties, SortOrderPushdownResult, - StatisticsArgs, displayable, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + InputDistributionRequirements, InputOrderMode, Partitioning, PlanProperties, + ReplaceChildrenOptions, SortOrderPushdownResult, StatisticsArgs, displayable, }; /// Create a non sorted parquet exec @@ -526,9 +526,10 @@ impl ExecutionPlan for RequirementsTestExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { assert_eq!(children.len(), 1); Ok(RequirementsTestExec::new(Arc::clone(&children[0])) @@ -538,6 +539,16 @@ impl ExecutionPlan for RequirementsTestExec { .into_arc()) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -1025,9 +1036,10 @@ impl ExecutionPlan for TestScan { vec![] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { if children.is_empty() { Ok(self) @@ -1036,6 +1048,16 @@ impl ExecutionPlan for TestScan { } } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/core/tests/user_defined/insert_operation.rs b/datafusion/core/tests/user_defined/insert_operation.rs index c61fe018aa74..0eefcdb551a6 100644 --- a/datafusion/core/tests/user_defined/insert_operation.rs +++ b/datafusion/core/tests/user_defined/insert_operation.rs @@ -28,7 +28,9 @@ use datafusion_common::config::Dialect; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_expr::{Expr, TableType, dml::InsertOp}; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; -use datafusion_physical_plan::execution_plan::SchedulingType; +use datafusion_physical_plan::{ + ChildrenPropertiesMode, ReplaceChildrenOptions, execution_plan::SchedulingType, +}; use datafusion_physical_plan::{ DisplayAs, ExecutionPlan, PlanProperties, execution_plan::{Boundedness, EmissionType}, @@ -162,14 +164,25 @@ impl ExecutionPlan for TestInsertExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { assert!(children.is_empty()); Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs b/datafusion/core/tests/user_defined/user_defined_plan.rs index 2b042b613dbc..da7fdd88793e 100644 --- a/datafusion/core/tests/user_defined/user_defined_plan.rs +++ b/datafusion/core/tests/user_defined/user_defined_plan.rs @@ -99,6 +99,7 @@ use datafusion_expr::{FetchType, InvariantLevel, Projection, SortExpr}; use datafusion_optimizer::AnalyzerRule; use datafusion_optimizer::optimizer::ApplyOrder; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion_physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}; use async_trait::async_trait; use datafusion_common::cast::as_string_view_array; @@ -725,13 +726,24 @@ impl ExecutionPlan for TopKExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(Arc::new(TopKExec::new(children[0].clone(), self.k))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + /// Execute one partition and return an iterator over RecordBatch fn execute( &self, diff --git a/datafusion/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index b89cf5d356f7..4bf04133b784 100644 --- a/datafusion/datasource/src/sink.rs +++ b/datafusion/datasource/src/sink.rs @@ -32,9 +32,9 @@ use datafusion_physical_expr_common::sort_expr::{LexRequirement, OrderingRequire use datafusion_physical_plan::metrics::MetricsSet; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, - InputDistributionRequirements, Partitioning, PlanProperties, - SendableRecordBatchStream, execute_input_stream, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + ExecutionPlanProperties, InputDistributionRequirements, Partitioning, PlanProperties, + ReplaceChildrenOptions, SendableRecordBatchStream, execute_input_stream, }; use async_trait::async_trait; @@ -305,9 +305,10 @@ impl ExecutionPlan for DataSinkExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(Arc::new(Self::new( Arc::clone(&children[0]), @@ -316,6 +317,16 @@ impl ExecutionPlan for DataSinkExec { ))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn apply_expressions( &self, _f: &mut dyn FnMut(&Arc) -> Result, diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index 929fed02b3eb..741010c59519 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -33,7 +33,8 @@ use datafusion_physical_plan::metrics::{ use datafusion_physical_plan::projection::ProjectionExec; use datafusion_physical_plan::stream::BatchSplitStream; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ReplaceChildrenOptions, }; use itertools::Itertools; @@ -400,6 +401,24 @@ impl ExecutionPlan for DataSourceExec { Vec::new() } + fn replace_children( + self: Arc, + _: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + Ok(self) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn apply_expressions( &self, f: &mut dyn FnMut(&Arc) -> Result, @@ -408,13 +427,6 @@ impl ExecutionPlan for DataSourceExec { self.data_source.apply_expressions(f) } - fn with_new_children( - self: Arc, - _: Vec>, - ) -> Result> { - Ok(self) - } - /// Implementation of [`ExecutionPlan::repartitioned`] which relies upon the inner [`DataSource::repartitioned`]. /// /// If the data source does not support changing its partitioning, returns `Ok(None)` (the default). Refer diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index a0dd5e6cb619..d7ee5dace30c 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/datafusion/ffi/src/execution_plan.rs @@ -25,8 +25,8 @@ use datafusion_common::{DataFusionError, Result, Statistics}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr_common::metrics::MetricsSet; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, StatisticsArgs, - StatisticsContext, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ReplaceChildrenOptions, StatisticsArgs, StatisticsContext, }; use stabby::string::String as SString; use stabby::vec::Vec as SVec; @@ -185,7 +185,10 @@ unsafe extern "C" fn with_new_children_fn_wrapper( .collect(); let children = sresult_return!(children); - let new_plan = sresult_return!(inner_plan.with_new_children(children)); + let new_plan = sresult_return!(inner_plan.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute) + )); FFI_Result::Ok(FFI_ExecutionPlan::new(new_plan, runtime)) } @@ -302,7 +305,7 @@ fn pass_runtime_to_children( // If the parent is foreign and the child is local to this library, then when // we called `children()` above we will get something other than a // `ForeignExecutionPlan`. In this case wrap the plan in a `ForeignExecutionPlan` - // because when we call `with_new_children` below it will extract the + // because when we call `replace_children` below it will extract the // FFI plan that does contain the runtime. if plan_is_foreign && !child.is::() { updated_children = true; @@ -315,7 +318,12 @@ fn pass_runtime_to_children( }) .collect::>>()?; if updated_children { - Arc::clone(plan).with_new_children(children).map(Some) + Arc::clone(plan) + .replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + .map(Some) } else { Ok(None) } @@ -453,9 +461,10 @@ impl ExecutionPlan for ForeignExecutionPlan { self.children.iter().collect() } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { let children = children .into_iter() @@ -467,6 +476,16 @@ impl ExecutionPlan for ForeignExecutionPlan { (&new_plan).try_into() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, @@ -621,9 +640,10 @@ pub mod tests { self.children.iter().collect() } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(Arc::new(EmptyExec { props: Arc::clone(&self.props), @@ -635,6 +655,16 @@ pub mod tests { })) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -777,7 +807,10 @@ pub mod tests { assert_eq!(parent_foreign.children().len(), 0); assert_eq!(child_foreign.children().len(), 0); - let parent_foreign = parent_foreign.with_new_children(vec![child_foreign])?; + let parent_foreign = parent_foreign.replace_children( + vec![child_foreign], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; assert_eq!(parent_foreign.children().len(), 1); // Version 2: Adding child to the local plan @@ -787,7 +820,10 @@ pub mod tests { let child_foreign = >::try_from(&child_local)?; let parent_plan = Arc::new(EmptyExec::new(Arc::clone(&schema))); - let parent_plan = parent_plan.with_new_children(vec![child_foreign])?; + let parent_plan = parent_plan.replace_children( + vec![child_foreign], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; let mut parent_local = FFI_ExecutionPlan::new(parent_plan, None); parent_local.library_marker_id = crate::mock_foreign_marker_id; let parent_foreign = >::try_from(&parent_local)?; diff --git a/datafusion/ffi/src/tests/async_provider.rs b/datafusion/ffi/src/tests/async_provider.rs index b9f353e89b8f..83057d8c45db 100644 --- a/datafusion/ffi/src/tests/async_provider.rs +++ b/datafusion/ffi/src/tests/async_provider.rs @@ -37,7 +37,9 @@ use datafusion_common::{Result, exec_err}; use datafusion_execution::RecordBatchStream; use datafusion_expr::Expr; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; -use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::{ + ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions, +}; use datafusion_session::Session; use futures::Stream; use tokio::runtime::Handle; @@ -211,13 +213,24 @@ impl ExecutionPlan for AsyncTestExecutionPlan { Vec::default() } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/ffi/tests/ffi_execution_plan.rs b/datafusion/ffi/tests/ffi_execution_plan.rs index 946dca61e1b8..4067d7eb49b2 100644 --- a/datafusion/ffi/tests/ffi_execution_plan.rs +++ b/datafusion/ffi/tests/ffi_execution_plan.rs @@ -22,12 +22,15 @@ mod tests { use arrow_schema::DataType; use datafusion_common::DataFusionError; use datafusion_common::tree_node::TreeNodeRecursion; - use datafusion_ffi::execution_plan::FFI_ExecutionPlan; - use datafusion_ffi::execution_plan::ForeignExecutionPlan; - use datafusion_ffi::execution_plan::{ExecutionPlanPrivateData, tests::EmptyExec}; + use datafusion_ffi::execution_plan::{ + ExecutionPlanPrivateData, FFI_ExecutionPlan, ForeignExecutionPlan, + tests::EmptyExec, + }; use datafusion_ffi::tests::utils::get_module; - use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::execution_plan::InvariantLevel; + use datafusion_physical_plan::{ + ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions, + }; use std::sync::Arc; #[test] @@ -135,7 +138,10 @@ mod tests { let grandchild_plan = generate_local_plan(); - let child_plan = child_plan.with_new_children(vec![grandchild_plan])?; + let child_plan = child_plan.replace_children( + vec![grandchild_plan], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; unsafe { // Originally the runtime is not set. We go through the unsafe casting @@ -150,7 +156,10 @@ mod tests { assert!((*grandchild_private_data).runtime.is_none()); } - let parent_plan = generate_local_plan().with_new_children(vec![child_plan])?; + let parent_plan = generate_local_plan().replace_children( + vec![child_plan], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; // Adding the grandchild beneath this FFI plan should get the runtime passed down. let runtime = tokio::runtime::Builder::new_current_thread() diff --git a/datafusion/physical-optimizer/src/ensure_coop.rs b/datafusion/physical-optimizer/src/ensure_coop.rs index 10da9e4e7517..93862df3b423 100644 --- a/datafusion/physical-optimizer/src/ensure_coop.rs +++ b/datafusion/physical-optimizer/src/ensure_coop.rs @@ -130,7 +130,10 @@ impl PhysicalOptimizerRule for EnsureCooperative { #[cfg(test)] mod tests { use super::*; - use datafusion_physical_plan::{displayable, test::scan_partitioned}; + use datafusion_physical_plan::{ + ChildrenPropertiesMode, ReplaceChildrenOptions, displayable, + test::scan_partitioned, + }; use insta::assert_snapshot; #[tokio::test] @@ -328,9 +331,10 @@ mod tests { fn children(&self) -> Vec<&Arc> { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(Arc::new(DummyExec::new( &self.name, @@ -339,6 +343,15 @@ mod tests { self.evaluation_type, ))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } fn execute( &self, _: usize, diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index adbc3dde7a7d..07bc98b2db79 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -55,7 +55,9 @@ use datafusion_physical_plan::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, }; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; -use datafusion_physical_plan::execution_plan::EmissionType; +use datafusion_physical_plan::execution_plan::{ + EmissionType, replace_children_if_necessary, +}; use datafusion_physical_plan::joins::{ CrossJoinExec, HashJoinExec, PartitionMode, SortMergeJoinExec, }; @@ -69,7 +71,7 @@ use datafusion_physical_plan::windows::WindowAggExec; use datafusion_physical_plan::windows::{BoundedWindowAggExec, get_best_fitting_window}; use datafusion_physical_plan::{ ChildSatisfactionOptions, Distribution, ExecutionPlan, InputDistributionRequirements, - Partitioning, with_new_children_if_necessary, + Partitioning, }; use itertools::izip; @@ -760,8 +762,10 @@ fn preserving_order_enables_streaming( return Ok(false); } // Build parent with the ordered child - let with_ordered = - Arc::clone(parent).with_new_children(vec![Arc::clone(ordered_child)])?; + let with_ordered = replace_children_if_necessary( + Arc::clone(parent), + vec![Arc::clone(ordered_child)], + )?; if with_ordered.pipeline_behavior() == EmissionType::Final { // Parent is blocking even with ordering — no benefit return Ok(false); @@ -769,7 +773,8 @@ fn preserving_order_enables_streaming( // Build parent with an unordered child via CoalescePartitionsExec. let unordered_child: Arc = Arc::new(CoalescePartitionsExec::new(Arc::clone(ordered_child))); - let without_ordered = Arc::clone(parent).with_new_children(vec![unordered_child])?; + let without_ordered = + replace_children_if_necessary(Arc::clone(parent), vec![unordered_child])?; Ok(without_ordered.pipeline_behavior() == EmissionType::Final) } @@ -1519,16 +1524,16 @@ pub fn ensure_distribution( // Data Arc::new(InterleaveExec::try_new(children_plans)?) } else { - // Route through `with_new_children_if_necessary` so the common + // Route through `replace_children_if_necessary` so the common // case where no child was replaced above skips the expensive - // `with_new_children` rebuild. For nodes like `ProjectionExec`, - // `with_new_children` recomputes schema / equivalence properties / + // `replace_children` rebuild. For nodes like `ProjectionExec`, + // `replace_children` recomputes schema / equivalence properties / // output ordering via `try_new` even when the input Arcs are // identical, which dominates `ensure_distribution` time on deep // projection stacks over plans where no distribution change // applies (point queries with no join / aggregate / unmet // ordering). - with_new_children_if_necessary(plan, children_plans)? + replace_children_if_necessary(plan, children_plans)? }; Ok(Transformed::yes(DistributionContext::new( diff --git a/datafusion/physical-optimizer/src/filter_pushdown.rs b/datafusion/physical-optimizer/src/filter_pushdown.rs index 06aa632a9d3f..18fe15100051 100644 --- a/datafusion/physical-optimizer/src/filter_pushdown.rs +++ b/datafusion/physical-optimizer/src/filter_pushdown.rs @@ -39,11 +39,12 @@ use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_common::{Result, assert_eq_or_internal_err, config::ConfigOptions}; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr_common::physical_expr::is_volatile; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::execution_plan::replace_children_if_necessary; use datafusion_physical_plan::filter_pushdown::{ ChildFilterPushdownResult, ChildPushdownResult, FilterPushdownPhase, FilterPushdownPropagation, PushedDown, }; -use datafusion_physical_plan::{ExecutionPlan, with_new_children_if_necessary}; use itertools::{Itertools, izip}; @@ -573,7 +574,7 @@ fn push_down_filters( } // Re-create this node with new children - let updated_node = with_new_children_if_necessary(Arc::clone(node), new_children)?; + let updated_node = replace_children_if_necessary(Arc::clone(node), new_children)?; // TODO: by calling `handle_child_pushdown_result` we are assuming that the // `ExecutionPlan` implementation will not change the plan itself. diff --git a/datafusion/physical-optimizer/src/hash_join_buffering.rs b/datafusion/physical-optimizer/src/hash_join_buffering.rs index 7a198cac13fc..dbdfd34a9a01 100644 --- a/datafusion/physical-optimizer/src/hash_join_buffering.rs +++ b/datafusion/physical-optimizer/src/hash_join_buffering.rs @@ -21,6 +21,7 @@ use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::buffer::BufferExec; +use datafusion_physical_plan::execution_plan::replace_children_if_necessary; use datafusion_physical_plan::joins::HashJoinExec; use std::sync::Arc; @@ -74,19 +75,25 @@ impl PhysicalOptimizerRule for HashJoinBuffering { if node.left.is::() { return Ok(Transformed::no(plan)); } - plan.with_new_children(vec![ - Arc::new(BufferExec::new(Arc::clone(&node.left), capacity)), - Arc::clone(&node.right), - ])? + replace_children_if_necessary( + plan, + vec![ + Arc::new(BufferExec::new(Arc::clone(&node.left), capacity)), + Arc::clone(&node.right), + ], + )? } else { // Do not stack BufferExec nodes together. if node.right.is::() { return Ok(Transformed::no(plan)); } - plan.with_new_children(vec![ - Arc::clone(&node.left), - Arc::new(BufferExec::new(Arc::clone(&node.right), capacity)), - ])? + replace_children_if_necessary( + plan, + vec![ + Arc::clone(&node.left), + Arc::new(BufferExec::new(Arc::clone(&node.right), capacity)), + ], + )? }, )) }) diff --git a/datafusion/physical-optimizer/src/limit_pushdown.rs b/datafusion/physical-optimizer/src/limit_pushdown.rs index 01a288f7f163..f88a2be14e98 100644 --- a/datafusion/physical-optimizer/src/limit_pushdown.rs +++ b/datafusion/physical-optimizer/src/limit_pushdown.rs @@ -72,6 +72,7 @@ use datafusion_common::tree_node::{Transformed, TreeNodeRecursion}; use datafusion_common::utils::combine_limit; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::empty::EmptyExec; +use datafusion_physical_plan::execution_plan::replace_children_if_necessary; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; use datafusion_physical_plan::projection::ProjectionExec; @@ -403,7 +404,7 @@ pub(crate) fn pushdown_limits( .collect::>()?; if changed { - new_node.data.with_new_children(new_children) + replace_children_if_necessary(new_node.data, new_children) } else { Ok(new_node.data) } diff --git a/datafusion/physical-optimizer/src/output_requirements.rs b/datafusion/physical-optimizer/src/output_requirements.rs index fc8bf490b9f0..541981270169 100644 --- a/datafusion/physical-optimizer/src/output_requirements.rs +++ b/datafusion/physical-optimizer/src/output_requirements.rs @@ -34,7 +34,9 @@ use datafusion_common::{Result, Statistics, internal_err}; use datafusion_execution::TaskContext; use datafusion_physical_expr::Distribution; use datafusion_physical_expr_common::sort_expr::OrderingRequirements; -use datafusion_physical_plan::execution_plan::Boundedness; +use datafusion_physical_plan::execution_plan::{ + Boundedness, replace_children_if_necessary, +}; use datafusion_physical_plan::projection::{ ProjectionExec, make_with_child, update_expr, update_ordering_requirement, }; @@ -42,8 +44,9 @@ use datafusion_physical_plan::scalar_subquery::ScalarSubqueryExec; use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::{ - ChildStats, DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, - PlanProperties, SendableRecordBatchStream, StatisticsArgs, + ChildStats, ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + ExecutionPlanProperties, PlanProperties, ReplaceChildrenOptions, + SendableRecordBatchStream, StatisticsArgs, }; /// This rule either adds or removes [`OutputRequirements`]s to/from the physical @@ -233,9 +236,10 @@ impl ExecutionPlan for OutputRequirementExec { vec![self.order_requirement.clone()] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(Arc::new(Self::new( children.remove(0), // has a single child @@ -245,6 +249,16 @@ impl ExecutionPlan for OutputRequirementExec { ))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -471,7 +485,7 @@ fn require_top_ordering_helper( require_top_ordering_helper(Arc::clone(&children[idx]))?; if is_changed { children[idx] = new_child; - return Ok((plan.with_new_children(children)?, true)); + return Ok((replace_children_if_necessary(plan, children)?, true)); } } Ok((plan, false)) diff --git a/datafusion/physical-optimizer/src/topk_repartition.rs b/datafusion/physical-optimizer/src/topk_repartition.rs index 115bdc3cb535..d8fa1ac986f9 100644 --- a/datafusion/physical-optimizer/src/topk_repartition.rs +++ b/datafusion/physical-optimizer/src/topk_repartition.rs @@ -48,6 +48,7 @@ use crate::PhysicalOptimizerRule; use datafusion_common::Result; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; +use datafusion_physical_plan::execution_plan::replace_children_if_necessary; use std::sync::Arc; // CoalesceBatchesExec is deprecated on main (replaced by arrow-rs BatchCoalescer), // but older DataFusion versions may still insert it between SortExec and RepartitionExec. @@ -151,7 +152,7 @@ impl PhysicalOptimizerRule for TopKRepartition { // Rebuild the tree above the repartition let new_sort_input = if let Some(parent) = repart_parent { - parent.with_new_children(vec![new_repartition])? + replace_children_if_necessary(parent, vec![new_repartition])? } else { new_repartition }; diff --git a/datafusion/physical-optimizer/src/window_topn.rs b/datafusion/physical-optimizer/src/window_topn.rs index 29b8f4a46000..20bd8b0d38a1 100644 --- a/datafusion/physical-optimizer/src/window_topn.rs +++ b/datafusion/physical-optimizer/src/window_topn.rs @@ -60,6 +60,7 @@ use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal}; use datafusion_physical_expr::window::StandardWindowExpr; use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr}; use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::execution_plan::replace_children_if_necessary; use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::projection::ProjectionExec; use datafusion_physical_plan::repartition::RepartitionExec; @@ -192,13 +193,13 @@ impl WindowTopN { .ok()?; // Step 7: Rebuild window with PartitionedTopKExec as its child - let mut result = window_exec - .with_new_children(vec![Arc::new(partitioned_topk)]) - .ok()?; + let mut result = + replace_children_if_necessary(window_exec, vec![Arc::new(partitioned_topk)]) + .ok()?; // Step 8: Rebuild intermediate nodes (ProjectionExec/RepartitionExec) for node in intermediates.into_iter().rev() { - result = node.with_new_children(vec![result]).ok()?; + result = replace_children_if_necessary(node, vec![result]).ok()?; } Some(result) diff --git a/datafusion/physical-plan/benches/compute_statistics.rs b/datafusion/physical-plan/benches/compute_statistics.rs index 93c95ea4ba09..cddf4c2396f4 100644 --- a/datafusion/physical-plan/benches/compute_statistics.rs +++ b/datafusion/physical-plan/benches/compute_statistics.rs @@ -47,8 +47,8 @@ use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::CrossJoinExec; use datafusion_physical_plan::statistics::StatisticsArgs; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, Partitioning, SendableRecordBatchStream, - StatisticsContext, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Partitioning, + ReplaceChildrenOptions, SendableRecordBatchStream, StatisticsContext, }; /// Minimal leaf node for benchmarking @@ -98,18 +98,29 @@ impl ExecutionPlan for BenchLeaf { vec![] } - fn apply_expressions( - &self, - _f: &mut dyn FnMut(&Arc) -> Result, - ) -> Result { - Ok(TreeNodeRecursion::Continue) + fn replace_children( + self: Arc, + _: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + Ok(self) } fn with_new_children( self: Arc, - _children: Vec>, + children: Vec>, ) -> Result> { - Ok(self) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) } fn execute( diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 1c5b56ffadbc..a39c6f34862d 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -165,9 +165,10 @@ use crate::filter_pushdown::{ }; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions, validate_child_count}; use crate::{ DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, - InputOrderMode, SendableRecordBatchStream, Statistics, check_if_same_properties, + InputOrderMode, SendableRecordBatchStream, Statistics, }; use datafusion_common::config::ConfigOptions; use parking_lot::Mutex; @@ -2031,6 +2032,45 @@ impl ExecutionPlan for AggregateExec { vec![&self.input] } + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let mut me = AggregateExec::try_new_with_schema( + self.mode, + 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), + )?; + me.limit_options = self.limit_options; + me.dynamic_filter.clone_from(&self.dynamic_filter); + Ok(Arc::new(me)) + } + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn apply_expressions( &self, f: &mut dyn FnMut(&Arc) -> Result, @@ -2068,36 +2108,14 @@ impl ExecutionPlan for AggregateExec { .collect() } - fn with_new_children( - 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), - self.aggr_expr.to_vec(), - Arc::clone(&self.filter_expr), - Arc::clone(&children[0]), - Arc::clone(&self.input_schema), - Arc::clone(&self.schema), - )?; - me.limit_options = self.limit_options; - me.dynamic_filter.clone_from(&self.dynamic_filter); - - Ok(Arc::new(me)) - } - fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -3681,18 +3699,29 @@ mod tests { vec![] } - fn apply_expressions( - &self, - _f: &mut dyn FnMut(&Arc) -> Result, - ) -> Result { - Ok(TreeNodeRecursion::Continue) + fn replace_children( + self: Arc, + _: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + internal_err!("Children cannot be replaced in {self:?}") } fn with_new_children( self: Arc, - _: Vec>, + children: Vec>, ) -> Result> { - internal_err!("Children cannot be replaced in {self:?}") + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) } fn execute( @@ -5029,8 +5058,10 @@ mod tests { Arc::clone(&blocking_exec) as Arc, schema, )?); - let new_agg = - Arc::clone(&aggregate_exec).with_new_children(vec![blocking_exec])?; + let new_agg = Arc::clone(&aggregate_exec).replace_children( + vec![blocking_exec], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; assert_eq!(new_agg.schema(), aggregate_exec.schema()); Ok(()) } diff --git a/datafusion/physical-plan/src/analyze.rs b/datafusion/physical-plan/src/analyze.rs index 9a6951895338..d1519828c24b 100644 --- a/datafusion/physical-plan/src/analyze.rs +++ b/datafusion/physical-plan/src/analyze.rs @@ -27,7 +27,10 @@ use super::{ use crate::display::DisplayableExecutionPlan; use crate::execution_plan::EvaluationType; use crate::metrics::{MetricCategory, MetricType}; -use crate::{DisplayFormatType, ExecutionPlan, Partitioning}; +use crate::{ + ChildrenPropertiesMode, DisplayFormatType, ExecutionPlan, Partitioning, + ReplaceChildrenOptions, +}; use arrow::{array::StringBuilder, datatypes::SchemaRef, record_batch::RecordBatch}; use datafusion_common::format::ExplainFormat; @@ -228,9 +231,10 @@ impl ExecutionPlan for AnalyzeExec { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(Arc::new( AnalyzeExec::builder( @@ -246,6 +250,16 @@ impl ExecutionPlan for AnalyzeExec { )) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/async_func.rs b/datafusion/physical-plan/src/async_func.rs index 91531ec35c55..f3ef13d4fd3a 100644 --- a/datafusion/physical-plan/src/async_func.rs +++ b/datafusion/physical-plan/src/async_func.rs @@ -19,13 +19,14 @@ use crate::coalesce::LimitedBatchCoalescer; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use crate::stream::{EmptyRecordBatchStream, RecordBatchStreamAdapter}; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, - check_if_same_properties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + ExecutionPlanProperties, PlanProperties, ReplaceChildrenOptions, + validate_child_count, }; use arrow::array::RecordBatch; use arrow_schema::{FieldRef, Fields, Schema, SchemaRef}; +use datafusion_common::Result; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; -use datafusion_common::{Result, assert_eq_or_internal_err}; use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr::ScalarFunctionExpr; use datafusion_physical_expr::async_scalar_function::AsyncFuncExpr; @@ -170,31 +171,43 @@ impl ExecutionPlan for AsyncFuncExec { ) } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, ) -> 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(), - children.swap_remove(0), - )?)) + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => Ok(Arc::new(AsyncFuncExec::try_new( + self.async_exprs.clone(), + children.swap_remove(0), + )?)), + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -316,6 +329,7 @@ impl AsyncFuncExec { node: &datafusion_proto_models::protobuf::PhysicalPlanNode, ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, ) -> Result> { + use datafusion_common::assert_eq_or_internal_err; use datafusion_proto_models::protobuf; let async_func = crate::expect_plan_variant!( node, diff --git a/datafusion/physical-plan/src/buffer.rs b/datafusion/physical-plan/src/buffer.rs index a1c3c7ea0165..24cca6b0b17f 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, EvaluationType, SchedulingType}; +use crate::execution_plan::{ + CardinalityEffect, EvaluationType, SchedulingType, replace_children_if_necessary, +}; use crate::filter_pushdown::{ ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, @@ -27,13 +29,13 @@ use crate::projection::ProjectionExec; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::RecordBatchStreamAdapter; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SortOrderPushdownResult, - check_if_same_properties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ReplaceChildrenOptions, SortOrderPushdownResult, validate_child_count, }; use arrow::array::RecordBatch; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::TreeNodeRecursion; -use datafusion_common::{Result, Statistics, internal_err, plan_err}; +use datafusion_common::{Result, Statistics, internal_err}; use datafusion_common_runtime::SpawnedTask; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; @@ -167,26 +169,42 @@ impl ExecutionPlan for BufferExec { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { - check_if_same_properties!(self, children); - if children.len() != 1 { - return plan_err!("BufferExec can only have one child"); + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + Ok(Arc::new(Self::new(children.swap_remove(0), self.capacity))) + } } - Ok(Arc::new(Self::new(children.swap_remove(0), self.capacity))) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -271,9 +289,10 @@ impl ExecutionPlan for BufferExec { projection: &ProjectionExec, ) -> Result>> { match self.input.try_swapping_with_projection(projection)? { - Some(new_input) => Ok(Some( - Arc::new(self.clone()).with_new_children(vec![new_input])?, - )), + Some(new_input) => Ok(Some(replace_children_if_necessary( + Arc::new(self.clone()), + vec![new_input], + )?)), None => Ok(None), } } diff --git a/datafusion/physical-plan/src/coalesce_batches.rs b/datafusion/physical-plan/src/coalesce_batches.rs index 511e5d793b87..cb0f9b2ce4b3 100644 --- a/datafusion/physical-plan/src/coalesce_batches.rs +++ b/datafusion/physical-plan/src/coalesce_batches.rs @@ -27,8 +27,8 @@ use crate::projection::ProjectionExec; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::EmptyRecordBatchStream; use crate::{ - DisplayFormatType, ExecutionPlan, RecordBatchStream, SendableRecordBatchStream, - check_if_same_properties, + ChildrenPropertiesMode, DisplayFormatType, ExecutionPlan, RecordBatchStream, + ReplaceChildrenOptions, SendableRecordBatchStream, validate_child_count, }; use arrow::datatypes::SchemaRef; @@ -39,7 +39,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, replace_children_if_necessary}; use crate::filter_pushdown::{ ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, @@ -181,26 +181,43 @@ impl ExecutionPlan for CoalesceBatchesExec { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new( - CoalesceBatchesExec::new(children.swap_remove(0), self.target_batch_size) - .with_fetch(self.fetch), - )) + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => Ok(Arc::new( + CoalesceBatchesExec::new(children.swap_remove(0), self.target_batch_size) + .with_fetch(self.fetch), + )), + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -260,9 +277,10 @@ impl ExecutionPlan for CoalesceBatchesExec { projection: &ProjectionExec, ) -> Result>> { match self.input.try_swapping_with_projection(projection)? { - Some(new_input) => Ok(Some( - Arc::new(self.clone()).with_new_children(vec![new_input])?, - )), + Some(new_input) => Ok(Some(replace_children_if_necessary( + Arc::new(self.clone()), + vec![new_input], + )?)), None => Ok(None), } } diff --git a/datafusion/physical-plan/src/coalesce_partitions.rs b/datafusion/physical-plan/src/coalesce_partitions.rs index 5cd7a707c23b..6f58eb2f1e6b 100644 --- a/datafusion/physical-plan/src/coalesce_partitions.rs +++ b/datafusion/physical-plan/src/coalesce_partitions.rs @@ -26,12 +26,17 @@ use super::{ DisplayAs, ExecutionPlanProperties, PlanProperties, SendableRecordBatchStream, Statistics, }; -use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType}; +use crate::execution_plan::{ + CardinalityEffect, EvaluationType, SchedulingType, replace_children_if_necessary, +}; use crate::filter_pushdown::{FilterDescription, FilterPushdownPhase}; use crate::projection::{ProjectionExec, make_with_child}; use crate::sort_pushdown::SortOrderPushdownResult; use crate::statistics::{ChildStats, StatisticsArgs}; -use crate::{DisplayFormatType, ExecutionPlan, Partitioning, check_if_same_properties}; +use crate::{ + ChildrenPropertiesMode, DisplayFormatType, ExecutionPlan, Partitioning, + ReplaceChildrenOptions, validate_child_count, +}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_common::config::ConfigOptions; @@ -151,25 +156,44 @@ impl ExecutionPlan for CoalescePartitionsExec { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let mut plan = CoalescePartitionsExec::new(children.swap_remove(0)); + plan.fetch = self.fetch; + Ok(Arc::new(plan)) + } + } + } + + fn with_new_children( + self: Arc, + children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - let mut plan = CoalescePartitionsExec::new(children.swap_remove(0)); - plan.fetch = self.fetch; - Ok(Arc::new(plan)) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -304,8 +328,7 @@ impl ExecutionPlan for CoalescePartitionsExec { self.input .with_preserve_order(preserve_order) .and_then(|new_input| { - Arc::new(self.clone()) - .with_new_children(vec![new_input]) + replace_children_if_necessary(Arc::new(self.clone()), vec![new_input]) .ok() }) } diff --git a/datafusion/physical-plan/src/coop.rs b/datafusion/physical-plan/src/coop.rs index dc7d98891e11..9e27b26d6e7c 100644 --- a/datafusion/physical-plan/src/coop.rs +++ b/datafusion/physical-plan/src/coop.rs @@ -87,15 +87,16 @@ use crate::filter_pushdown::{ use crate::projection::ProjectionExec; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, SortOrderPushdownResult, check_if_same_properties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, + SortOrderPushdownResult, validate_child_count, }; use arrow::record_batch::RecordBatch; use arrow_schema::Schema; -use datafusion_common::{Result, Statistics, assert_eq_or_internal_err}; +use datafusion_common::{Result, Statistics}; use datafusion_execution::TaskContext; -use crate::execution_plan::SchedulingType; +use crate::execution_plan::{SchedulingType, replace_children_if_necessary}; use crate::stream::RecordBatchStreamAdapter; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use futures::{Stream, StreamExt}; @@ -275,27 +276,41 @@ impl ExecutionPlan for CooperativeExec { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { - assert_eq_or_internal_err!( - children.len(), - 1, - "CooperativeExec requires exactly one child" - ); - check_if_same_properties!(self, children); - Ok(Arc::new(CooperativeExec::new(children.swap_remove(0)))) + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + Ok(Arc::new(CooperativeExec::new(children.swap_remove(0)))) + } + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - ..Self::clone(&*self) - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -332,9 +347,10 @@ impl ExecutionPlan for CooperativeExec { projection: &ProjectionExec, ) -> Result>> { match self.input.try_swapping_with_projection(projection)? { - Some(new_input) => Ok(Some( - Arc::new(self.clone()).with_new_children(vec![new_input])?, - )), + Some(new_input) => Ok(Some(replace_children_if_necessary( + Arc::new(self.clone()), + vec![new_input], + )?)), None => Ok(None), } } @@ -365,11 +381,13 @@ impl ExecutionPlan for CooperativeExec { match child.try_pushdown_sort(order)? { SortOrderPushdownResult::Exact { inner } => { - let new_exec = Arc::new(self.clone()).with_new_children(vec![inner])?; + let new_exec = + replace_children_if_necessary(Arc::new(self.clone()), vec![inner])?; Ok(SortOrderPushdownResult::Exact { inner: new_exec }) } SortOrderPushdownResult::Inexact { inner } => { - let new_exec = Arc::new(self.clone()).with_new_children(vec![inner])?; + let new_exec = + replace_children_if_necessary(Arc::new(self.clone()), vec![inner])?; Ok(SortOrderPushdownResult::Inexact { inner: new_exec }) } SortOrderPushdownResult::Unsupported => { diff --git a/datafusion/physical-plan/src/display.rs b/datafusion/physical-plan/src/display.rs index 2370e3e6ce6e..d2bdcef2e97a 100644 --- a/datafusion/physical-plan/src/display.rs +++ b/datafusion/physical-plan/src/display.rs @@ -1511,7 +1511,10 @@ mod tests { use datafusion_physical_expr::PhysicalExpr; use crate::statistics::StatisticsArgs; - use crate::{DisplayAs, ExecutionPlan, PlanProperties}; + use crate::{ + ChildrenPropertiesMode, DisplayAs, ExecutionPlan, PlanProperties, + ReplaceChildrenOptions, + }; use super::DisplayableExecutionPlan; @@ -1545,9 +1548,10 @@ mod tests { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unimplemented!() } @@ -1559,6 +1563,16 @@ mod tests { Ok(TreeNodeRecursion::Continue) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _: usize, @@ -1636,10 +1650,11 @@ mod tests { use crate::empty::EmptyExec; use crate::filter::FilterExec; use crate::projection::ProjectionExec; + use crate::{ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions}; use datafusion_physical_expr::expressions::{binary, col, lit}; use datafusion_physical_expr::{Partitioning, PhysicalExpr}; - fn sample_plan() -> Arc { + fn sample_plan() -> Arc { let schema = Arc::new(Schema::new(vec![ Field::new("a", DataType::Int32, false), Field::new("b", DataType::Int32, false), @@ -1727,12 +1742,23 @@ mod tests { { Ok(datafusion_common::tree_node::TreeNodeRecursion::Continue) } - fn with_new_children( + + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } fn execute( &self, _: usize, diff --git a/datafusion/physical-plan/src/empty.rs b/datafusion/physical-plan/src/empty.rs index bd91ec742d48..dd08ff36a9d8 100644 --- a/datafusion/physical-plan/src/empty.rs +++ b/datafusion/physical-plan/src/empty.rs @@ -20,7 +20,10 @@ use std::sync::Arc; use crate::memory::MemoryStream; -use crate::{DisplayAs, PlanProperties, SendableRecordBatchStream, Statistics}; +use crate::{ + ChildrenPropertiesMode, DisplayAs, PlanProperties, ReplaceChildrenOptions, + SendableRecordBatchStream, Statistics, +}; use crate::{ DisplayFormatType, ExecutionPlan, Partitioning, execution_plan::{Boundedness, EmissionType}, @@ -127,13 +130,24 @@ impl ExecutionPlan for EmptyExec { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, @@ -247,8 +261,8 @@ impl EmptyExec { mod tests { use super::*; use crate::common; + use crate::execution_plan::replace_children_if_necessary; use crate::test; - use crate::with_new_children_if_necessary; #[tokio::test] async fn empty() -> Result<()> { @@ -271,7 +285,7 @@ mod tests { let schema = test::aggr_test_schema(); let empty = Arc::new(EmptyExec::new(Arc::clone(&schema))); - let empty2 = with_new_children_if_necessary( + let empty2 = replace_children_if_necessary( Arc::clone(&empty) as Arc, vec![], )?; @@ -279,7 +293,7 @@ mod tests { let too_many_kids = vec![empty2]; assert!( - with_new_children_if_necessary(empty, too_many_kids).is_err(), + replace_children_if_necessary(empty, too_many_kids).is_err(), "expected error when providing list of kids" ); Ok(()) diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index ff803746aa99..a4d081b3d9e7 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -267,6 +267,26 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// joins). fn children(&self) -> Vec<&Arc>; + /// Returns a clone of the existing plan with the children replaced, + /// skipping recomputation of plan properties when the options indicate + /// the new children's properties are unchanged. + /// + /// Callers should typically call [`replace_children_if_necessary`] and + /// not invoke this method directly. + fn replace_children( + self: Arc, + children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + #[expect(deprecated)] + match options.children_properties { + ChildrenPropertiesMode::Keep => { + self.with_new_children_and_same_properties(children) + } + ChildrenPropertiesMode::Recompute => self.with_new_children(children), + } + } + /// Apply a closure `f` to each root expression that this node owns and uses /// during execution, either by evaluating it or updating it dynamically. /// @@ -331,27 +351,98 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { f: &mut dyn FnMut(&Arc) -> Result, ) -> Result; - /// Returns a new `ExecutionPlan` where all existing children were replaced - /// by the `children`, in order + /// Deprecated. + /// + /// DataFusion will remove this method in the future in favor of + /// [`ExecutionPlan::replace_children`]. + /// + /// Note that this method is still required by the trait; implementations + /// should delegate to [`ExecutionPlan::replace_children`] with + /// [`ChildrenPropertiesMode::Recompute`]. + /// + /// # Example Implementation + /// ``` + /// # #![allow(deprecated)] + /// # use std::fmt; + /// # use std::sync::Arc; + /// # use datafusion_common::Result; + /// # use datafusion_common::tree_node::TreeNodeRecursion; + /// # use datafusion_execution::{SendableRecordBatchStream, TaskContext}; + /// # use datafusion_physical_expr::PhysicalExpr; + /// # use datafusion_physical_plan::{ + /// # ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + /// # PlanProperties, ReplaceChildrenOptions, + /// # }; + /// # #[derive(Debug)] + /// # struct MyExec { + /// # input: Arc, + /// # } + /// # impl DisplayAs for MyExec { + /// # fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + /// # write!(f, "MyExec") + /// # } + /// # } + /// impl ExecutionPlan for MyExec { + /// // ... + /// # fn name(&self) -> &'static str { + /// # "MyExec" + /// # } + /// # fn properties(&self) -> &Arc { + /// # self.input.properties() + /// # } + /// # fn children(&self) -> Vec<&Arc> { + /// # vec![&self.input] + /// # } + /// # fn apply_expressions( + /// # &self, + /// # _f: &mut dyn FnMut(&Arc) -> Result, + /// # ) -> Result { + /// # Ok(TreeNodeRecursion::Continue) + /// # } + /// # fn execute( + /// # &self, + /// # _partition: usize, + /// # _context: Arc, + /// # ) -> Result { + /// # unimplemented!() + /// # } + /// fn replace_children( + /// self: Arc, + /// mut children: Vec>, + /// _options: ReplaceChildrenOptions, + /// ) -> Result> { + /// Ok(Arc::new(MyExec { + /// input: children.swap_remove(0), + /// })) + /// } + /// + /// fn with_new_children( + /// self: Arc, + /// children: Vec>, + /// ) -> Result> { + /// // call into `replace_children` with `ReplaceChildrenOptions` + /// self.replace_children( + /// children, + /// ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + /// ) + /// } + /// } + /// ``` + #[deprecated( + since = "55.0.0", + note = "Use `ExecutionPlan::replace_children` with `ReplaceChildrenOptions`" + )] fn with_new_children( self: Arc, children: Vec>, ) -> Result>; - /// Fast-path used by [`with_new_children_if_necessary`] when the new - /// `children` are known to have the same [`PlanProperties`] as the current - /// children. Implementations should swap the children in without - /// recomputing this plan's `PlanProperties` (typically by cloning `self` - /// and replacing the child pointers). - /// - /// The default implementation falls back to - /// [`ExecutionPlan::with_new_children`] which is always correct but - /// forfeits the fast-path: implementations that own an expensive - /// `PlanProperties` (e.g. projection mapping, complex equivalence - /// classes) should override this method. - /// - /// Callers should route through [`with_new_children_if_necessary`] and - /// not invoke this method directly. + /// Deprecated. Implement [`ExecutionPlan::replace_children`] instead. + #[deprecated( + since = "55.0.0", + note = "Use `ExecutionPlan::replace_children` with `ReplaceChildrenOptions`" + )] + #[expect(deprecated)] fn with_new_children_and_same_properties( self: Arc, children: Vec>, @@ -362,11 +453,11 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// Reset any internal state within this [`ExecutionPlan`]. /// /// This method is called when an [`ExecutionPlan`] needs to be re-executed, - /// such as in recursive queries. Unlike [`ExecutionPlan::with_new_children`], this method + /// such as in recursive queries. Unlike [`ExecutionPlan::replace_children`], this method /// ensures that any stateful components (e.g., [`DynamicFilterPhysicalExpr`]) /// are reset to their initial state. /// - /// The default implementation simply calls [`ExecutionPlan::with_new_children`] with the existing children, + /// The default implementation simply calls [`ExecutionPlan::replace_children`] with the existing children, /// effectively creating a new instance of the [`ExecutionPlan`] with the same children but without /// necessarily resetting any internal state. Implementations that require resetting of some /// internal state should override this method to provide the necessary logic. @@ -375,13 +466,16 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// it will be called from within a walk of the execution plan tree so that it will be called on each child later /// or was already called on each child. /// - /// Note to implementers: unlike [`ExecutionPlan::with_new_children`] this method does not accept new children as an argument, + /// Note to implementers: unlike [`ExecutionPlan::replace_children`] this method does not accept new children as an argument, /// thus it is expected that any cached plan properties will remain valid after the reset. /// /// [`DynamicFilterPhysicalExpr`]: datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr fn reset_state(self: Arc) -> Result> { let children = self.children().into_iter().cloned().collect(); - self.with_new_children(children) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } /// If supported, attempt to increase the partitioning of this `ExecutionPlan` to @@ -390,7 +484,7 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// If the `ExecutionPlan` does not support changing its partitioning, /// returns `Ok(None)` (the default). /// - /// It is the `ExecutionPlan` can increase its partitioning, but not to the + /// If the `ExecutionPlan` can increase its partitioning, but not to /// `target_partitions`, it may return an ExecutionPlan with fewer /// partitions. This might happen, for example, if each new partition would /// be too small to be efficiently processed individually. @@ -511,7 +605,7 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// partition: usize, /// context: Arc, /// ) -> Result { - /// // use functions from futures crate convert the batch into a stream + /// // use functions from futures crate to convert the batch into a stream /// let fut = futures::future::ready(Ok(self.batch.clone())); /// let stream = futures::stream::once(fut); /// Ok(Box::pin(RecordBatchStreamAdapter::new( @@ -743,7 +837,7 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// up the plan that `DataSourceExec` can actually bind the filters. /// /// The default implementation bars all parent filters from being pushed down and adds no new filters. - /// This is the safest option, making filter pushdown opt-in on a per-node pasis. + /// This is the safest option, making filter pushdown opt-in on a per-node basis. /// /// There are two different phases in filter pushdown, which some operators may handle the same and some differently. /// Depending on the phase the operator may or may not be allowed to modify the plan. @@ -936,6 +1030,36 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { } } +/// Options for [`ExecutionPlan::replace_children`] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ReplaceChildrenOptions { + /// Describes how plan properties should be handled for the replacement + /// children. + pub children_properties: ChildrenPropertiesMode, +} + +impl ReplaceChildrenOptions { + /// Create new options for [`ExecutionPlan::replace_children`]. + pub const fn new(children_properties: ChildrenPropertiesMode) -> Self { + Self { + children_properties, + } + } +} + +/// Indicates whether the plan properties of the new children must be recomputed. +/// +/// Part of [`ReplaceChildrenOptions`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ChildrenPropertiesMode { + /// The plan properties of the new children are identical to the properties + /// of the existing children, so we can skip recomputation. + Keep, + /// The plan properties of the new children are different from the properties + /// of the existing children, so we must recompute the properties from scratch. + Recompute, +} + /// Allows a type to be treated as a reference to an /// [`Arc`]. /// @@ -1354,12 +1478,10 @@ pub(crate) fn emission_type_from_children<'a>( } } -/// Stores certain, often expensive to compute, plan properties used in query -/// optimization. +/// Stores plan properties used in query optimization. /// -/// These properties are stored a single structure to permit this information to -/// be computed once and then those cached results used multiple times without -/// recomputation (aka a cache) +/// Serves as a cache for these properties, which are often +/// expensive to compute. #[derive(Debug, Clone)] pub struct PlanProperties { /// See [ExecutionPlanProperties::equivalence_properties] @@ -1563,19 +1685,19 @@ pub fn need_data_exchange(plan: Arc) -> bool { /// /// 1. **Same child pointers** — if every `children[i]` is `Arc::ptr_eq` to the /// corresponding existing child, the original `plan` is returned -/// unchanged (no allocation, no [`ExecutionPlan::with_new_children`] +/// unchanged (no allocation, no [`ExecutionPlan::replace_children`] /// call). /// 2. **Same child properties** — if the children's `PlanProperties` Arcs /// match (via [`has_same_children_properties`]), the plan's own /// `PlanProperties` cache can be reused. This calls -/// [`ExecutionPlan::with_new_children_and_same_properties`], which -/// swaps the child pointers without recomputing `PlanProperties`. +/// [`ExecutionPlan::replace_children`] with [`ChildrenPropertiesMode::Keep`], +/// which swaps the child pointers without recomputing `PlanProperties`. /// 3. **Full recompute** — otherwise, delegate to -/// [`ExecutionPlan::with_new_children`], which recomputes -/// `PlanProperties` from scratch. +/// [`ExecutionPlan::replace_children`] with [`ChildrenPropertiesMode::Recompute`], +/// which recomputes `PlanProperties` from scratch. /// /// The size of `children` must be equal to the size of `ExecutionPlan::children()`. -pub fn with_new_children_if_necessary( +pub fn replace_children_if_necessary( plan: Arc, children: Vec>, ) -> Result> { @@ -1596,11 +1718,25 @@ pub fn with_new_children_if_necessary( } // Layer 2: same child properties → reuse `PlanProperties` cache. if has_same_children_properties(plan.as_ref(), &children)? { - return plan.with_new_children_and_same_properties(children); + return plan.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ); } } // Layer 3: full recompute. - plan.with_new_children(children) + plan.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) +} + +#[deprecated(since = "55.0.0", note = "Use `replace_children_if_necessary`")] +pub fn with_new_children_if_necessary( + plan: Arc, + children: Vec>, +) -> Result> { + replace_children_if_necessary(plan, children) } /// Return a [`DisplayableExecutionPlan`] wrapper around an @@ -1871,9 +2007,9 @@ pub fn has_same_children_properties( /// the same as plan already has. Could be used to implement fast-path for method /// [`ExecutionPlan::with_new_children`]. /// -/// New call sites should route through [`with_new_children_if_necessary`], +/// New call sites should route through [`replace_children_if_necessary`], /// which applies this check together with the child-pointer short-circuit -/// (see [`with_new_children_if_necessary`] for the layered policy). This +/// (see [`replace_children_if_necessary`] for the layered policy). This /// macro remains for direct-caller sites that have not been migrated yet. #[macro_export] macro_rules! check_if_same_properties { @@ -1888,6 +2024,22 @@ macro_rules! check_if_same_properties { }; } +/// Helper macro to validate that replacement children match a plan's existing +/// child count. +/// +/// This is useful for [`ExecutionPlan::replace_children`] implementations that +/// need to preserve the same child-count validation behavior. +#[macro_export] +macro_rules! validate_child_count { + ($plan: expr, $children: expr) => { + datafusion_common::assert_eq_or_internal_err!( + $children.len(), + $plan.children().len(), + "Wrong number of children" + ); + }; +} + /// 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(); @@ -1979,9 +2131,10 @@ mod tests { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unimplemented!() } @@ -1993,6 +2146,16 @@ mod tests { Ok(TreeNodeRecursion::Continue) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn dynamic_expressions_produced(&self) -> Vec> { self.dynamic_expressions.iter().map(Arc::clone).collect() } @@ -2086,13 +2249,24 @@ mod tests { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -2143,13 +2317,24 @@ mod tests { self.0.apply_expressions(f) } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn downcast_delegate(&self) -> Option<&dyn ExecutionPlan> { Some(self.0.as_ref()) } @@ -2208,12 +2393,23 @@ mod tests { ) -> Result { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } fn execute( &self, _partition: usize, @@ -2277,38 +2473,59 @@ mod tests { ) -> Result { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { - self.recompute_calls - .fetch_add(1, std::sync::atomic::Ordering::SeqCst); - // Full recompute: allocate a fresh `PlanProperties` Arc so this - // path is observable via `Arc::ptr_eq` on properties. - let new_input = children.swap_remove(0); - let cache = Arc::new(PlanProperties::new( - EquivalenceProperties::new(Arc::new(Schema::empty())), - Partitioning::UnknownPartitioning(1), - EmissionType::Final, - Boundedness::Bounded, - )); - Ok(Arc::new(Self { - input: new_input, - cache, - recompute_calls: Arc::clone(&self.recompute_calls), - fast_path_calls: Arc::clone(&self.fast_path_calls), - })) + match options.children_properties { + ChildrenPropertiesMode::Keep => { + self.fast_path_calls + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(Arc::new(Self { + input: children.swap_remove(0), + ..Self::clone(&*self) + })) + } + ChildrenPropertiesMode::Recompute => { + self.recompute_calls + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + // Full recompute: allocate a fresh `PlanProperties` Arc so this + // path is observable via `Arc::ptr_eq` on properties. + let new_input = children.swap_remove(0); + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + Ok(Arc::new(Self { + input: new_input, + cache, + recompute_calls: Arc::clone(&self.recompute_calls), + fast_path_calls: Arc::clone(&self.fast_path_calls), + })) + } + } + } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - self.fast_path_calls - .fetch_add(1, std::sync::atomic::Ordering::SeqCst); - Ok(Arc::new(Self { - input: children.swap_remove(0), - ..Self::clone(&*self) - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( &self, @@ -2406,9 +2623,9 @@ mod tests { } /// Cover the three short-circuit layers of - /// [`with_new_children_if_necessary`]. + /// [`replace_children_if_necessary`]. #[test] - fn test_with_new_children_if_necessary_layers() -> Result<()> { + fn test_replace_children_if_necessary_layers() -> Result<()> { use std::sync::atomic::Ordering; // Two leaves that share the same `PlanProperties` Arc but sit behind @@ -2438,7 +2655,7 @@ mod tests { let orig_props = Arc::clone(parent.properties()); // Layer 1: same child pointer → returns the original plan Arc verbatim. - let out = with_new_children_if_necessary( + let out = replace_children_if_necessary( Arc::clone(&parent_dyn), vec![Arc::clone(&leaf_a)], )?; @@ -2451,7 +2668,7 @@ mod tests { // Arc is reused (not reallocated). assert!(!Arc::ptr_eq(&leaf_a, &leaf_b)); assert!(Arc::ptr_eq(leaf_a.properties(), leaf_b.properties())); - let out = with_new_children_if_necessary( + let out = replace_children_if_necessary( Arc::clone(&parent_dyn), vec![Arc::clone(&leaf_b)], )?; @@ -2461,7 +2678,7 @@ mod tests { // Layer 3: child's `PlanProperties` Arc differs → full recompute. assert!(!Arc::ptr_eq(leaf_a.properties(), leaf_c.properties())); - let out = with_new_children_if_necessary( + let out = replace_children_if_necessary( Arc::clone(&parent_dyn), vec![Arc::clone(&leaf_c)], )?; @@ -2479,7 +2696,7 @@ mod tests { /// `with_new_children`, so downstream / external `ExecutionPlan` /// implementations keep the semantics-preserving path. #[test] - fn test_with_new_children_if_necessary_default_fallback() -> Result<()> { + fn test_replace_children_if_necessary_default_fallback() -> Result<()> { use std::sync::atomic::Ordering; let leaf_props = Arc::new(PlanProperties::new( @@ -2498,10 +2715,20 @@ mod tests { let parent = Arc::new(WithChildrenTestParentDefault::new(Arc::clone(&leaf_a))); let parent_dyn: Arc = Arc::clone(&parent) as _; - // Distinct child Arc but same `PlanProperties` Arc — the helper - // enters the "same properties" branch and calls the trait method, - // whose default forwards to `with_new_children`. - let out = with_new_children_if_necessary( + // Using the same child means we return the original plan Arc verbatim, so even when + // the `replace_children` `ChildrenPropertiesMode::Keep` path is not defined, + // we do not recompute. + let out = replace_children_if_necessary( + Arc::clone(&parent_dyn), + vec![Arc::clone(&leaf_a)], + )?; + assert!(Arc::ptr_eq(&out, &parent_dyn)); + assert_eq!(parent.recompute_calls.load(Ordering::SeqCst), 0); + + // Using a distinct child but the same `PlanProperties` Arc means the helper + // attempts to enter the Keep branch. If it does not exist, we fall back + // to recomputation. + let out = replace_children_if_necessary( Arc::clone(&parent_dyn), vec![Arc::clone(&leaf_b)], )?; diff --git a/datafusion/physical-plan/src/explain.rs b/datafusion/physical-plan/src/explain.rs index 72ceae81f372..3b31ee748b73 100644 --- a/datafusion/physical-plan/src/explain.rs +++ b/datafusion/physical-plan/src/explain.rs @@ -22,7 +22,10 @@ use std::sync::Arc; use super::{DisplayAs, PlanProperties, SendableRecordBatchStream}; use crate::execution_plan::{Boundedness, EmissionType}; use crate::stream::RecordBatchStreamAdapter; -use crate::{DisplayFormatType, ExecutionPlan, Partitioning}; +use crate::{ + ChildrenPropertiesMode, DisplayFormatType, ExecutionPlan, Partitioning, + ReplaceChildrenOptions, +}; use arrow::{array::StringBuilder, datatypes::SchemaRef, record_batch::RecordBatch}; use datafusion_common::display::StringifiedPlan; @@ -124,13 +127,24 @@ impl ExecutionPlan for ExplainExec { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index 65abfa259d33..5df5482fb75d 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -28,10 +28,9 @@ 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; +use crate::execution_plan::{CardinalityEffect, replace_children_if_necessary}; use crate::filter_pushdown::{ ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, PushedDown, @@ -44,6 +43,7 @@ use crate::projection::{ }; use crate::statistics::{ChildStats, StatisticsArgs, StatisticsContext}; use crate::stream::EmptyRecordBatchStream; +use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions, validate_child_count}; use crate::{ DisplayFormatType, ExecutionPlan, metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RatioMetrics}, @@ -552,27 +552,46 @@ impl ExecutionPlan for FilterExec { vec![true] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { - check_if_same_properties!(self, children); - let new_input = children.swap_remove(0); - FilterExecBuilder::from(&*self) - .with_input(new_input) - .build() - .map(|e| Arc::new(e) as _) + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let new_input = children.swap_remove(0); + FilterExecBuilder::from(&*self) + .with_input(new_input) + .build() + .map(|e| Arc::new(e) as _) + } + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -820,8 +839,7 @@ impl ExecutionPlan for FilterExec { self.input .with_preserve_order(preserve_order) .and_then(|new_input| { - Arc::new(self.clone()) - .with_new_children(vec![new_input]) + replace_children_if_necessary(Arc::new(self.clone()), vec![new_input]) .ok() }) } diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index 524887bf03cc..8a477c1021d1 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -34,9 +34,10 @@ use crate::projection::{ use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::EmptyRecordBatchStream; use crate::{ - ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, - ExecutionPlanProperties, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, Statistics, check_if_same_properties, handle_state, + ChildrenPropertiesMode, ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, + ExecutionPlan, ExecutionPlanProperties, PlanProperties, RecordBatchStream, + ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, handle_state, + validate_child_count, }; use arrow::array::{RecordBatch, RecordBatchOptions}; @@ -276,32 +277,51 @@ impl ExecutionPlan for CrossJoinExec { Ok(TreeNodeRecursion::Continue) } + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + + Ok(Arc::new(Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + left_fut: Default::default(), + cache: Arc::clone(&self.cache), + schema: Arc::clone(&self.schema), + })) + } + ChildrenPropertiesMode::Recompute => Ok(Arc::new(CrossJoinExec::new( + Arc::clone(&children[0]), + Arc::clone(&children[1]), + ))), + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new(CrossJoinExec::new( - Arc::clone(&children[0]), - Arc::clone(&children[1]), - ))) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - let left = children.swap_remove(0); - let right = children.swap_remove(0); - - Ok(Arc::new(Self { - left, - right, - metrics: ExecutionPlanMetricsSet::new(), - left_fut: Default::default(), - cache: Arc::clone(&self.cache), - schema: Arc::clone(&self.schema), - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn reset_state(self: Arc) -> Result> { diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index cd9048b58c2b..08d209003ad9 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -22,7 +22,6 @@ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, OnceLock}; use std::vec; -use crate::ExecutionPlanProperties; use crate::execution_plan::{ EmissionType, boundedness_from_children, has_same_children_properties, plan_contains_expression_id, stub_properties, @@ -53,6 +52,10 @@ use crate::projection::{ }; use crate::repartition::REPARTITION_RANDOM_STATE; use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::{ + ChildrenPropertiesMode, ExecutionPlanProperties, ReplaceChildrenOptions, + validate_child_count, +}; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, Partitioning, PlanProperties, @@ -1375,11 +1378,32 @@ impl ExecutionPlan for HashJoinExec { /// This method is called during query optimization when the optimizer creates new /// plan nodes. Importantly, it creates a fresh bounds_accumulator via `try_new` /// rather than cloning the existing one because partitioning may have changed. + fn replace_children( + self: Arc, + children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => { + self.builder().with_new_children(children)?.build_exec() + } + ChildrenPropertiesMode::Recompute => self + .builder() + .recompute_properties() + .with_new_children(children)? + .build_exec(), + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - self.builder().with_new_children(children)?.build_exec() + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn reset_state(self: Arc) -> Result> { @@ -2451,6 +2475,7 @@ mod tests { use crate::filter::FilterExecBuilder; use crate::joins::hash_join::stream::lookup_join_hashmap; use crate::test::{TestMemoryExec, assert_join_metrics}; + use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions}; use crate::{ common, expressions::Column, repartition::RepartitionExec, test::build_table_i32, test::exec::MockExec, @@ -2523,13 +2548,24 @@ mod tests { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index 7069a8b44805..eb1df638c7dc 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -44,9 +44,9 @@ use crate::projection::{ }; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ - DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - PlanProperties, RecordBatchStream, SendableRecordBatchStream, - check_if_same_properties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, PlanProperties, RecordBatchStream, ReplaceChildrenOptions, + SendableRecordBatchStream, validate_child_count, }; use arrow::array::{ @@ -574,43 +574,61 @@ impl ExecutionPlan for NestedLoopJoinExec { ) } + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + Ok(Arc::new(Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + build_side_data: Default::default(), + left_spill_data: Arc::new(OnceAsync::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(), + })) + } + ChildrenPropertiesMode::Recompute => 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 with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - 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()?, - )) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - let left = children.swap_remove(0); - let right = children.swap_remove(0); - - Ok(Arc::new(Self { - left, - right, - metrics: ExecutionPlanMetricsSet::new(), - build_side_data: Default::default(), - left_spill_data: Arc::new(OnceAsync::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(), - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( 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 b60ec1c784de..c42ec67ef80d 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs @@ -53,7 +53,8 @@ use crate::joins::piecewise_merge_join::utils::{ use crate::joins::utils::asymmetric_join_output_partitioning; use crate::metrics::MetricsSet; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlanProperties, check_if_same_properties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlanProperties, + ReplaceChildrenOptions, validate_child_count, }; use crate::{ ExecutionPlan, PlanProperties, @@ -517,56 +518,80 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { } } + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => { + let buffered = children.swap_remove(0); + let streamed = children.swap_remove(0); + Ok(Arc::new(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(), + })) + } + ChildrenPropertiesMode::Recompute => match &children[..] { + [left, right] => Ok(Arc::new(PiecewiseMergeJoinExec::try_new( + Arc::clone(left), + Arc::clone(right), + self.on.clone(), + self.operator, + self.join_type, + self.num_partitions, + )?)), + _ => internal_err!( + "PiecewiseMergeJoin should have 2 children, found {}", + children.len() + ), + }, + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - match &children[..] { - [left, right] => Ok(Arc::new(PiecewiseMergeJoinExec::try_new( - Arc::clone(left), - Arc::clone(right), - self.on.clone(), - self.operator, - self.join_type, - self.num_partitions, - )?)), - _ => internal_err!( - "PiecewiseMergeJoin should have 2 children, found {}", - children.len() - ), - } + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - let buffered = children.swap_remove(0); - let streamed = children.swap_remove(0); - Ok(Arc::new(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(), - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn reset_state(self: Arc) -> Result> { let buffered = Arc::clone(&self.buffered); let streamed = Arc::clone(&self.streamed); - self.with_new_children_and_same_properties(vec![buffered, streamed]) + self.replace_children( + vec![buffered, streamed], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( 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 55a4b2136c4f..b48905500d54 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -40,9 +40,9 @@ use crate::projection::{ use crate::spill::spill_manager::SpillManager; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ - DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - InputDistributionRequirements, PlanProperties, SendableRecordBatchStream, Statistics, - check_if_same_properties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, InputDistributionRequirements, PlanProperties, + ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, validate_child_count, }; use arrow::compute::SortOptions; @@ -449,37 +449,56 @@ impl ExecutionPlan for SortMergeJoinExec { crate::apply_expression_roots(join_keys.chain(filter), f) } + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + Ok(Arc::new(Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + ChildrenPropertiesMode::Recompute => match &children[..] { + [left, right] => Ok(Arc::new(SortMergeJoinExec::try_new( + Arc::clone(left), + Arc::clone(right), + self.on.clone(), + self.filter.clone(), + self.join_type, + self.sort_options.clone(), + self.null_equality, + )?)), + _ => internal_err!("SortMergeJoin wrong number of children"), + }, + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - match &children[..] { - [left, right] => Ok(Arc::new(SortMergeJoinExec::try_new( - Arc::clone(left), - Arc::clone(right), - self.on.clone(), - self.filter.clone(), - self.join_type, - self.sort_options.clone(), - self.null_equality, - )?)), - _ => internal_err!("SortMergeJoin wrong number of children"), - } + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - let left = children.swap_remove(0); - let right = children.swap_remove(0); - Ok(Arc::new(Self { - left, - right, - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 589042a82818..0c6e84b36cc5 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -31,7 +31,6 @@ 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::{boundedness_from_children, emission_type_from_children}; use crate::joins::stream_join_utils::{ @@ -50,6 +49,7 @@ use crate::projection::{ JoinData, ProjectionExec, try_pushdown_through_join_with_column_indices, }; use crate::stream::EmptyRecordBatchStream; +use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions, validate_child_count}; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, InputDistributionRequirements, PlanProperties, RecordBatchStream, @@ -462,36 +462,57 @@ impl ExecutionPlan for SymmetricHashJoinExec { crate::apply_expression_roots(join_keys.chain(filter), f) } + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + Ok(Arc::new(Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + ChildrenPropertiesMode::Recompute => { + Ok(Arc::new(SymmetricHashJoinExec::try_new( + Arc::clone(&children[0]), + Arc::clone(&children[1]), + self.on.clone(), + self.filter.clone(), + &self.join_type, + self.null_equality, + self.left_sort_exprs.clone(), + self.right_sort_exprs.clone(), + self.mode, + )?)) + } + } + } + fn with_new_children( 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]), - self.on.clone(), - self.filter.clone(), - &self.join_type, - self.null_equality, - self.left_sort_exprs.clone(), - self.right_sort_exprs.clone(), - self.mode, - )?)) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - let left = children.swap_remove(0); - let right = children.swap_remove(0); - Ok(Arc::new(Self { - left, - right, - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn metrics(&self) -> Option { diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 941b4e561bb1..9e50a93b2163 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -44,10 +44,12 @@ pub use crate::display::{DefaultDisplay, DisplayAs, DisplayFormatType, VerboseDi pub use crate::distribution_requirements::{ ChildSatisfactionOptions, InputDistributionRequirements, }; +#[expect(deprecated)] pub use crate::execution_plan::{ - AsPhysicalExprRef, ExecutionPlan, ExecutionPlanProperties, PlanProperties, - apply_expression_roots, collect, collect_partitioned, displayable, - execute_input_stream, execute_stream, execute_stream_partitioned, get_plan_string, + AsPhysicalExprRef, ChildrenPropertiesMode, ExecutionPlan, ExecutionPlanProperties, + PlanProperties, ReplaceChildrenOptions, apply_expression_roots, collect, + collect_partitioned, displayable, execute_input_stream, execute_stream, + execute_stream_partitioned, get_plan_string, replace_children_if_necessary, with_new_children_if_necessary, }; pub use crate::metrics::Metric; diff --git a/datafusion/physical-plan/src/limit.rs b/datafusion/physical-plan/src/limit.rs index 9dbdf17dbcbb..dd62c93d1cfe 100644 --- a/datafusion/physical-plan/src/limit.rs +++ b/datafusion/physical-plan/src/limit.rs @@ -29,8 +29,8 @@ use super::{ use crate::execution_plan::{Boundedness, CardinalityEffect}; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ - DisplayFormatType, Distribution, ExecutionPlan, Partitioning, - check_if_same_properties, + ChildrenPropertiesMode, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, + ReplaceChildrenOptions, validate_child_count, }; use arrow::datatypes::SchemaRef; @@ -175,26 +175,45 @@ impl ExecutionPlan for GlobalLimitExec { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let mut new_limit = + GlobalLimitExec::new(children.swap_remove(0), self.skip, self.fetch); + new_limit.set_required_ordering(self.required_ordering.clone()); + Ok(Arc::new(new_limit)) + } + } + } + + fn with_new_children( + self: Arc, + children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - let mut new_limit = - GlobalLimitExec::new(children.swap_remove(0), self.skip, self.fetch); - new_limit.set_required_ordering(self.required_ordering.clone()); - Ok(Arc::new(new_limit)) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -421,25 +440,45 @@ impl ExecutionPlan for LocalLimitExec { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { - check_if_same_properties!(self, children); - let mut new_limit = LocalLimitExec::new(children.swap_remove(0), self.fetch); - new_limit.set_required_ordering(self.required_ordering.clone()); - Ok(Arc::new(new_limit)) + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let mut new_limit = + LocalLimitExec::new(children.swap_remove(0), self.fetch); + new_limit.set_required_ordering(self.required_ordering.clone()); + Ok(Arc::new(new_limit)) + } + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -873,7 +912,7 @@ mod tests { } #[test] - fn with_new_children_preserves_required_ordering() -> Result<()> { + fn replace_children_preserves_required_ordering() -> Result<()> { let source = test::scan_partitioned(1); let schema = source.schema(); let ordering = LexOrdering::new(vec![PhysicalSortExpr { @@ -886,15 +925,19 @@ mod tests { let mut global = GlobalLimitExec::new(Arc::clone(&source), 0, Some(10)); global.set_required_ordering(ordering.clone()); - let rebuilt = - Arc::new(global).with_new_children(vec![test::scan_partitioned(1)])?; + let rebuilt = Arc::new(global).replace_children( + vec![test::scan_partitioned(1)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; let rebuilt = rebuilt.downcast_ref::().unwrap(); assert_eq!(rebuilt.required_ordering(), &ordering); let mut local = LocalLimitExec::new(source, 10); local.set_required_ordering(ordering.clone()); - let rebuilt = - Arc::new(local).with_new_children(vec![test::scan_partitioned(1)])?; + let rebuilt = Arc::new(local).replace_children( + vec![test::scan_partitioned(1)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; let rebuilt = rebuilt.downcast_ref::().unwrap(); assert_eq!(rebuilt.required_ordering(), &ordering); diff --git a/datafusion/physical-plan/src/memory.rs b/datafusion/physical-plan/src/memory.rs index eb141b8c70d5..efe42c7ebc5f 100644 --- a/datafusion/physical-plan/src/memory.rs +++ b/datafusion/physical-plan/src/memory.rs @@ -26,8 +26,8 @@ use crate::coop::cooperative; use crate::execution_plan::{Boundedness, EmissionType, SchedulingType}; use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, - RecordBatchStream, SendableRecordBatchStream, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, + PlanProperties, RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, }; use arrow::array::RecordBatch; @@ -319,9 +319,10 @@ impl ExecutionPlan for LazyMemoryExec { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { assert_or_internal_err!( children.is_empty(), @@ -330,6 +331,16 @@ impl ExecutionPlan for LazyMemoryExec { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index ec54201e7b3d..16b89e9eca92 100644 --- a/datafusion/physical-plan/src/operator_statistics/mod.rs +++ b/datafusion/physical-plan/src/operator_statistics/mod.rs @@ -1027,7 +1027,10 @@ mod tests { use crate::filter::FilterExec; use crate::projection::ProjectionExec; use crate::statistics::StatisticsArgs; - use crate::{DisplayAs, DisplayFormatType, PlanProperties}; + use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, PlanProperties, + ReplaceChildrenOptions, + }; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::stats::Precision; use datafusion_common::{ColumnStatistics, ScalarValue}; @@ -1107,13 +1110,24 @@ mod tests { vec![] } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn properties(&self) -> &Arc { &self.cache } @@ -1214,15 +1228,26 @@ mod tests { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(Arc::new(CustomExec { input: Arc::clone(&children[0]), })) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn properties(&self) -> &Arc { self.input.properties() } diff --git a/datafusion/physical-plan/src/placeholder_row.rs b/datafusion/physical-plan/src/placeholder_row.rs index de07529bab70..67c063b65cbc 100644 --- a/datafusion/physical-plan/src/placeholder_row.rs +++ b/datafusion/physical-plan/src/placeholder_row.rs @@ -23,8 +23,9 @@ use crate::coop::cooperative; use crate::execution_plan::{Boundedness, EmissionType, SchedulingType}; use crate::memory::MemoryStream; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, - SendableRecordBatchStream, Statistics, common, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, + PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, + common, }; use arrow::array::{ArrayRef, NullArray, RecordBatch, RecordBatchOptions}; @@ -145,13 +146,24 @@ impl ExecutionPlan for PlaceholderRowExec { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, @@ -250,16 +262,15 @@ impl PlaceholderRowExec { #[cfg(test)] mod tests { use super::*; - use crate::test; - use crate::with_new_children_if_necessary; + use crate::{execution_plan::replace_children_if_necessary, test}; #[test] - fn with_new_children() -> Result<()> { + fn replace_children() -> Result<()> { let schema = test::aggr_test_schema(); let placeholder = Arc::new(PlaceholderRowExec::new(schema)); - let placeholder_2 = with_new_children_if_necessary( + let placeholder_2 = replace_children_if_necessary( Arc::clone(&placeholder) as Arc, vec![], )?; @@ -267,7 +278,7 @@ mod tests { let too_many_kids = vec![placeholder_2]; assert!( - with_new_children_if_necessary(placeholder, too_many_kids).is_err(), + replace_children_if_necessary(placeholder, too_many_kids).is_err(), "expected error when providing list of kids" ); Ok(()) diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index ecdd78cc2acc..cf362cdee55d 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -27,14 +27,17 @@ use super::{ SendableRecordBatchStream, SortOrderPushdownResult, Statistics, }; use crate::column_rewriter::PhysicalColumnRewriter; -use crate::execution_plan::CardinalityEffect; +use crate::execution_plan::{CardinalityEffect, replace_children_if_necessary}; use crate::filter_pushdown::{ ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, FilterRemapper, PushedDownPredicate, }; use crate::joins::utils::{ColumnIndex, JoinFilter, JoinOn, JoinOnRef}; use crate::statistics::{ChildStats, StatisticsArgs}; -use crate::{DisplayFormatType, ExecutionPlan, PhysicalExpr, check_if_same_properties}; +use crate::{ + ChildrenPropertiesMode, DisplayFormatType, ExecutionPlan, PhysicalExpr, + ReplaceChildrenOptions, validate_child_count, +}; use std::collections::HashMap; use std::pin::Pin; use std::sync::Arc; @@ -337,27 +340,44 @@ impl ExecutionPlan for ProjectionExec { crate::apply_expression_roots(self.projector.projection().as_ref().iter(), f) } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { - check_if_same_properties!(self, children); - ProjectionExec::try_from_projector( - self.projector.clone(), - children.swap_remove(0), + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => ProjectionExec::try_from_projector( + self.projector.clone(), + children.swap_remove(0), + ) + .map(|p| Arc::new(p) as _), + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), ) - .map(|p| Arc::new(p) as _) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -510,11 +530,13 @@ impl ExecutionPlan for ProjectionExec { // Recursively push down to child node match child.try_pushdown_sort(&child_order)? { SortOrderPushdownResult::Exact { inner } => { - let new_exec = Arc::new(self.clone()).with_new_children(vec![inner])?; + let new_exec = + replace_children_if_necessary(Arc::new(self.clone()), vec![inner])?; Ok(SortOrderPushdownResult::Exact { inner: new_exec }) } SortOrderPushdownResult::Inexact { inner } => { - let new_exec = Arc::new(self.clone()).with_new_children(vec![inner])?; + let new_exec = + replace_children_if_necessary(Arc::new(self.clone()), vec![inner])?; Ok(SortOrderPushdownResult::Inexact { inner: new_exec }) } SortOrderPushdownResult::Unsupported => { @@ -530,8 +552,7 @@ impl ExecutionPlan for ProjectionExec { self.input .with_preserve_order(preserve_order) .and_then(|new_input| { - Arc::new(self.clone()) - .with_new_children(vec![new_input]) + replace_children_if_necessary(Arc::new(self.clone()), vec![new_input]) .ok() }) } diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index 4c6f0493adf4..0a56488de84d 100644 --- a/datafusion/physical-plan/src/recursive_query.rs +++ b/datafusion/physical-plan/src/recursive_query.rs @@ -30,8 +30,8 @@ use crate::metrics::{ BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RecordOutput, }; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, }; use arrow::array::{BooleanArray, BooleanBuilder}; use arrow::compute::filter_record_batch; @@ -183,9 +183,10 @@ impl ExecutionPlan for RecursiveQueryExec { ]) } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { RecursiveQueryExec::try_new( self.name.clone(), @@ -197,6 +198,16 @@ impl ExecutionPlan for RecursiveQueryExec { .map(|e| Arc::new(e) as _) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 8e7bec5dee32..063954a72a09 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! This file implements the [`RepartitionExec`] operator, which maps N input +//! This file implements the [`RepartitionExec`] operator, which maps N input //! partitions to M output partitions based on a partitioning scheme, optionally //! maintaining the order of the input rows in the output. @@ -43,8 +43,8 @@ use crate::spill::spill_pool::{self, SpillPoolSink, SpillPoolWriter}; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::{EmptyRecordBatchStream, RecordBatchStreamAdapter}; use crate::{ - DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, Statistics, - check_if_same_properties, + ChildrenPropertiesMode, DisplayFormatType, ExecutionPlan, Partitioning, + PlanProperties, ReplaceChildrenOptions, Statistics, validate_child_count, }; use arrow::array::{Array, PrimitiveArray, RecordBatch, RecordBatchOptions, UInt64Array}; @@ -1564,31 +1564,50 @@ impl ExecutionPlan for RepartitionExec { } } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { - check_if_same_properties!(self, children); - let mut repartition = RepartitionExec::try_new( - children.swap_remove(0), - self.partitioning().clone(), - )?; - if self.preserve_order { - repartition = repartition.with_preserve_order(); + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + state: Default::default(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let mut repartition = RepartitionExec::try_new( + children.swap_remove(0), + self.partitioning().clone(), + )?; + if self.preserve_order { + repartition = repartition.with_preserve_order(); + } + Ok(Arc::new(repartition)) + } } - Ok(Arc::new(repartition)) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - state: Default::default(), - ..Self::clone(&*self) - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn benefits_from_input_partitioning(&self) -> Vec { @@ -3227,13 +3246,24 @@ mod tests { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/physical-plan/src/scalar_subquery.rs b/datafusion/physical-plan/src/scalar_subquery.rs index ee3c2e5d077f..f2b7c5e0b53e 100644 --- a/datafusion/physical-plan/src/scalar_subquery.rs +++ b/datafusion/physical-plan/src/scalar_subquery.rs @@ -37,7 +37,10 @@ use crate::execution_plan::{CardinalityEffect, ExecutionPlan, PlanProperties}; use crate::joins::utils::{OnceAsync, OnceFut}; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::RecordBatchStreamAdapter; -use crate::{DisplayAs, DisplayFormatType, SendableRecordBatchStream}; +use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ReplaceChildrenOptions, + SendableRecordBatchStream, +}; use futures::StreamExt; use futures::TryStreamExt; @@ -164,9 +167,10 @@ impl ExecutionPlan for ScalarSubqueryExec { children } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { // First child is the main input, the rest are subquery plans. let input = children.remove(0); @@ -186,6 +190,16 @@ impl ExecutionPlan for ScalarSubqueryExec { ))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn reset_state(self: Arc) -> Result> { self.results.clear(); Ok(Arc::new(ScalarSubqueryExec { @@ -452,9 +466,10 @@ mod tests { vec![&self.inner] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(Arc::new(Self::new( children.remove(0), @@ -469,6 +484,16 @@ mod tests { Ok(TreeNodeRecursion::Continue) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/sorts/partial_sort.rs b/datafusion/physical-plan/src/sorts/partial_sort.rs index 5f15f8b6cb59..478ac14e119d 100644 --- a/datafusion/physical-plan/src/sorts/partial_sort.rs +++ b/datafusion/physical-plan/src/sorts/partial_sort.rs @@ -61,9 +61,9 @@ use crate::sorts::sort::sort_batch; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::EmptyRecordBatchStream; use crate::{ - DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - Partitioning, PlanProperties, SendableRecordBatchStream, Statistics, - check_if_same_properties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, Partitioning, PlanProperties, ReplaceChildrenOptions, + SendableRecordBatchStream, Statistics, validate_child_count, }; use arrow::compute::concat_batches; @@ -385,31 +385,50 @@ impl ExecutionPlan for PartialSortExec { ) } + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics_set: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let new_partial_sort = PartialSortExec::new( + self.expr.clone(), + Arc::clone(&children[0]), + self.common_prefix_length, + ) + .with_fetch(self.fetch) + .with_preserve_partitioning(self.preserve_partitioning); + + Ok(Arc::new(new_partial_sort)) + } + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - let new_partial_sort = PartialSortExec::new( - self.expr.clone(), - Arc::clone(&children[0]), - self.common_prefix_length, + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), ) - .with_fetch(self.fetch) - .with_preserve_partitioning(self.preserve_partitioning); - - Ok(Arc::new(new_partial_sort)) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics_set: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs b/datafusion/physical-plan/src/sorts/partitioned_topk.rs index 78dd9b9696d4..41ccfab6833d 100644 --- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs +++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs @@ -48,6 +48,7 @@ use futures::TryStreamExt; use crate::execution_plan::{Boundedness, EmissionType}; use crate::metrics::ExecutionPlanMetricsSet; use crate::topk::{PartitionedTopK, PartitionedTopKRank, build_sort_fields}; +use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions}; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, PlanProperties, SendableRecordBatchStream, stream::RecordBatchStreamAdapter, @@ -366,9 +367,10 @@ impl ExecutionPlan for PartitionedTopKExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { assert_eq!(children.len(), 1); Ok(Arc::new(PartitionedTopKExec::try_new( @@ -390,6 +392,16 @@ impl ExecutionPlan for PartitionedTopKExec { ) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 5c6b86acc59c..6c782f513448 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -28,6 +28,7 @@ use parking_lot::RwLock; use crate::common::spawn_buffered; use crate::execution_plan::{ Boundedness, CardinalityEffect, EmissionType, has_same_children_properties, + replace_children_if_necessary, }; use crate::expressions::PhysicalSortExpr; use crate::filter::FilterExec; @@ -51,9 +52,9 @@ use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; use crate::topk::TopK; use crate::topk::TopKDynamicFilters; use crate::{ - DisplayAs, DisplayFormatType, Distribution, EmptyRecordBatchStream, ExecutionPlan, - ExecutionPlanProperties, Partitioning, PlanProperties, SendableRecordBatchStream, - Statistics, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, + EmptyRecordBatchStream, ExecutionPlan, ExecutionPlanProperties, Partitioning, + PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, }; use arrow::array::{RecordBatch, RecordBatchOptions}; @@ -1307,16 +1308,17 @@ impl ExecutionPlan for SortExec { vec![false] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { let mut new_sort = self.cloned(); assert_eq!(children.len(), 1, "SortExec should have exactly one child"); new_sort.input = Arc::clone(&children[0]); - if !has_same_children_properties(self.as_ref(), &children)? { - // Recompute the properties based on the new input since they may have changed + if options.children_properties == ChildrenPropertiesMode::Recompute { + // 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(), @@ -1332,12 +1334,28 @@ impl ExecutionPlan for SortExec { Ok(Arc::new(new_sort)) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + match has_same_children_properties(self.as_ref(), &children)? { + true => self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ), + false => self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ), + } + } + fn reset_state(self: Arc) -> Result> { let children = self.children().into_iter().cloned().collect(); - let new_sort = self.with_new_children(children)?; + let new_sort = replace_children_if_necessary(self, children)?; let mut new_sort = new_sort .downcast_ref::() - .expect("cloned 1 lines above this line, we know the type") + .expect("rebuilt SortExec with new children") .clone(); // Our dynamic filter and execution metrics are the state we need to reset. new_sort.filter = Some(new_sort.create_filter()); @@ -1781,9 +1799,10 @@ mod tests { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } @@ -1795,6 +1814,16 @@ mod tests { Ok(TreeNodeRecursion::Continue) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index ac6f5d18cd2f..ad17f2c2136a 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -26,9 +26,9 @@ use crate::projection::{ProjectionExec, make_with_child, update_ordering}; use crate::sorts::streaming_merge::StreamingMergeBuilder; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ - DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - Partitioning, PlanProperties, SendableRecordBatchStream, Statistics, - check_if_same_properties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, Partitioning, PlanProperties, ReplaceChildrenOptions, + SendableRecordBatchStream, Statistics, validate_child_count, }; use datafusion_common::tree_node::TreeNodeRecursion; @@ -38,7 +38,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::{CardinalityEffect, EvaluationType, SchedulingType}; +use crate::execution_plan::{ + CardinalityEffect, EvaluationType, SchedulingType, replace_children_if_necessary, +}; use log::{debug, trace}; /// Sort preserving merge execution plan @@ -251,8 +253,7 @@ impl ExecutionPlan for SortPreservingMergeExec { self.input .with_preserve_order(preserve_order) .and_then(|new_input| { - Arc::new(self.clone()) - .with_new_children(vec![new_input]) + replace_children_if_necessary(Arc::new(self.clone()), vec![new_input]) .ok() }) } @@ -293,26 +294,43 @@ impl ExecutionPlan for SortPreservingMergeExec { ) } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new( - SortPreservingMergeExec::new(self.expr.clone(), children.swap_remove(0)) - .with_fetch(self.fetch), - )) + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => Ok(Arc::new( + SortPreservingMergeExec::new(self.expr.clone(), children.swap_remove(0)) + .with_fetch(self.fetch), + )), + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -1591,12 +1609,23 @@ mod tests { ) -> Result { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/streaming.rs b/datafusion/physical-plan/src/streaming.rs index a82f8d9441e9..7b0058e79887 100644 --- a/datafusion/physical-plan/src/streaming.rs +++ b/datafusion/physical-plan/src/streaming.rs @@ -30,7 +30,10 @@ use crate::projection::{ ProjectionExec, all_alias_free_columns, new_projections_for_columns, update_ordering, }; use crate::stream::RecordBatchStreamAdapter; -use crate::{ExecutionPlan, Partitioning, SendableRecordBatchStream}; +use crate::{ + ChildrenPropertiesMode, ExecutionPlan, Partitioning, ReplaceChildrenOptions, + SendableRecordBatchStream, +}; use arrow::datatypes::{Schema, SchemaRef}; use datafusion_common::tree_node::TreeNodeRecursion; @@ -280,9 +283,10 @@ impl ExecutionPlan for StreamingTableExec { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { if children.is_empty() { Ok(self) @@ -291,6 +295,16 @@ impl ExecutionPlan for StreamingTableExec { } } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/test.rs b/datafusion/physical-plan/src/test.rs index 68e6ff7eca48..b38a46d16075 100644 --- a/datafusion/physical-plan/src/test.rs +++ b/datafusion/physical-plan/src/test.rs @@ -24,7 +24,6 @@ use std::pin::Pin; use std::sync::Arc; use std::task::Context; -use crate::ExecutionPlan; use crate::common; use crate::execution_plan::{Boundedness, EmissionType}; use crate::memory::MemoryStream; @@ -32,6 +31,7 @@ use crate::metrics::MetricsSet; use crate::statistics::StatisticsArgs; use crate::stream::RecordBatchStreamAdapter; use crate::streaming::PartitionStream; +use crate::{ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions}; use crate::{DisplayAs, DisplayFormatType, PlanProperties}; use arrow::array::{Array, ArrayRef, Int32Array, RecordBatch}; @@ -148,13 +148,24 @@ impl ExecutionPlan for TestMemoryExec { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn repartitioned( &self, _target_partitions: usize, diff --git a/datafusion/physical-plan/src/test/exec.rs b/datafusion/physical-plan/src/test/exec.rs index 9a3f05a6e02a..1e2005e908fb 100644 --- a/datafusion/physical-plan/src/test/exec.rs +++ b/datafusion/physical-plan/src/test/exec.rs @@ -17,6 +17,7 @@ //! Simple iterator over batches for use in testing +use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions}; use crate::{ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, common, @@ -218,13 +219,24 @@ impl ExecutionPlan for MockExec { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + /// Returns a stream which yields data fn execute( &self, @@ -449,9 +461,10 @@ impl ExecutionPlan for BarrierExec { unimplemented!() } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unimplemented!() } @@ -463,6 +476,16 @@ impl ExecutionPlan for BarrierExec { Ok(TreeNodeRecursion::Continue) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + /// Returns a stream which yields data fn execute( &self, @@ -592,9 +615,10 @@ impl ExecutionPlan for ErrorExec { unimplemented!() } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unimplemented!() } @@ -606,6 +630,16 @@ impl ExecutionPlan for ErrorExec { Ok(TreeNodeRecursion::Continue) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + /// Returns a stream which yields data fn execute( &self, @@ -692,13 +726,24 @@ impl ExecutionPlan for StatisticsExec { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -796,9 +841,10 @@ impl ExecutionPlan for BlockingExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { internal_err!("Children cannot be replaced in {self:?}") } @@ -810,6 +856,16 @@ impl ExecutionPlan for BlockingExec { Ok(TreeNodeRecursion::Continue) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -952,13 +1008,24 @@ impl ExecutionPlan for PanicExec { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { internal_err!("Children cannot be replaced in {:?}", self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/tree_node.rs b/datafusion/physical-plan/src/tree_node.rs index aa4f144f9189..dcdceff8693e 100644 --- a/datafusion/physical-plan/src/tree_node.rs +++ b/datafusion/physical-plan/src/tree_node.rs @@ -20,7 +20,8 @@ use std::fmt::{self, Display, Formatter}; use std::sync::Arc; -use crate::{ExecutionPlan, displayable, with_new_children_if_necessary}; +use crate::execution_plan::replace_children_if_necessary; +use crate::{ExecutionPlan, displayable}; use datafusion_common::Result; use datafusion_common::tree_node::{ConcreteTreeNode, DynTreeNode}; @@ -35,7 +36,7 @@ impl DynTreeNode for dyn ExecutionPlan { arc_self: Arc, new_children: Vec>, ) -> Result> { - with_new_children_if_necessary(arc_self, new_children) + replace_children_if_necessary(arc_self, new_children) } } @@ -73,7 +74,7 @@ impl PlanContext { /// if the `PlanContext.children` have been changed. pub fn update_plan_from_children(mut self) -> Result { let children_plans = self.children.iter().map(|c| Arc::clone(&c.plan)).collect(); - self.plan = with_new_children_if_necessary(self.plan, children_plans)?; + self.plan = replace_children_if_necessary(self.plan, children_plans)?; Ok(self) } diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index fb62deecc33d..c1cc5da31aba 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -31,7 +31,6 @@ use super::{ PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, metrics::{ExecutionPlanMetricsSet, MetricsSet}, }; -use crate::check_if_same_properties; use crate::execution_plan::{ CardinalityEffect, InvariantLevel, boundedness_from_children, check_default_invariants, emission_type_from_children, @@ -45,6 +44,7 @@ use crate::metrics::BaselineMetrics; use crate::projection::{ProjectionExec, ProjectionExpr, make_with_child}; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::ObservedStream; +use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions, validate_child_count}; use arrow::datatypes::{Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; @@ -330,23 +330,40 @@ impl ExecutionPlan for UnionExec { Ok(TreeNodeRecursion::Continue) } + fn replace_children( + self: Arc, + children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + inputs: children, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => UnionExec::try_new(children), + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - UnionExec::try_new(children) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn with_new_children_and_same_properties( self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - inputs: children, - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -716,28 +733,47 @@ impl ExecutionPlan for InterleaveExec { Ok(TreeNodeRecursion::Continue) } + fn replace_children( + self: Arc, + children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + inputs: children, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + // New children are no longer interleavable, which might be a bug of optimization rewrite. + assert_or_internal_err!( + can_interleave(children.iter()), + "Can not create InterleaveExec: new children can not be interleaved" + ); + Ok(Arc::new(InterleaveExec::try_new(children)?)) + } + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - // New children are no longer interleavable, which might be a bug of optimization rewrite. - assert_or_internal_err!( - 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)?)) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn with_new_children_and_same_properties( self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - inputs: children, - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index 1877f668c525..3fa274b27a7b 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -28,8 +28,9 @@ use super::metrics::{ use super::{DisplayAs, ExecutionPlanProperties, PlanProperties}; use crate::stream::EmptyRecordBatchStream; use crate::{ - DisplayFormatType, Distribution, ExecutionPlan, RecordBatchStream, - SendableRecordBatchStream, check_if_same_properties, + ChildrenPropertiesMode, DisplayFormatType, Distribution, ExecutionPlan, + RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, + validate_child_count, }; use arrow::array::{ @@ -235,29 +236,46 @@ impl ExecutionPlan for UnnestExec { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new(UnnestExec::new( - children.swap_remove(0), - self.list_column_indices.clone(), - self.struct_column_indices.clone(), - Arc::clone(&self.schema), - self.options.clone(), - )?)) + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => Ok(Arc::new(UnnestExec::new( + children.swap_remove(0), + self.list_column_indices.clone(), + self.struct_column_indices.clone(), + Arc::clone(&self.schema), + self.options.clone(), + )?)), + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn required_input_distribution(&self) -> Vec { 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 c6a417cd4453..d4c98009ba70 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -35,10 +35,10 @@ use crate::windows::{ window_equivalence_properties, }; use crate::{ - ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, - ExecutionPlanProperties, InputDistributionRequirements, InputOrderMode, - PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, WindowExpr, - check_if_same_properties, + ChildrenPropertiesMode, ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, + ExecutionPlan, ExecutionPlanProperties, InputDistributionRequirements, + InputOrderMode, PlanProperties, RecordBatchStream, ReplaceChildrenOptions, + SendableRecordBatchStream, Statistics, WindowExpr, validate_child_count, }; use arrow::compute::take_record_batch; @@ -463,30 +463,49 @@ impl ExecutionPlan for BoundedWindowAggExec { vec![true] } + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let new = BoundedWindowAggExec::try_new( + self.window_expr.clone(), + Arc::clone(&children[0]), + self.input_order_mode.clone(), + self.can_repartition, + )? + .with_state_observer(self.state_observer.clone())?; + Ok(Arc::new(new)) + } + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - let new = BoundedWindowAggExec::try_new( - self.window_expr.clone(), - Arc::clone(&children[0]), - self.input_order_mode.clone(), - self.can_repartition, - )? - .with_state_observer(self.state_observer.clone())?; - Ok(Arc::new(new)) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs index bbf9a14fd5ea..d794e7df9d0a 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -33,10 +33,10 @@ use crate::windows::{ window_equivalence_properties, }; use crate::{ - ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, - ExecutionPlanProperties, InputDistributionRequirements, PhysicalExpr, PlanProperties, - RecordBatchStream, SendableRecordBatchStream, Statistics, WindowExpr, - check_if_same_properties, + ChildrenPropertiesMode, ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, + ExecutionPlan, ExecutionPlanProperties, InputDistributionRequirements, PhysicalExpr, + PlanProperties, RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, + Statistics, WindowExpr, validate_child_count, }; use arrow::array::ArrayRef; @@ -262,27 +262,44 @@ impl ExecutionPlan for WindowAggExec { } } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new(WindowAggExec::try_new( - self.window_expr.clone(), - children.swap_remove(0), - true, - )?)) + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => Ok(Arc::new(WindowAggExec::try_new( + self.window_expr.clone(), + children.swap_remove(0), + true, + )?)), + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( diff --git a/datafusion/physical-plan/src/work_table.rs b/datafusion/physical-plan/src/work_table.rs index 83cd0a15a6d2..b5d6fd47bc46 100644 --- a/datafusion/physical-plan/src/work_table.rs +++ b/datafusion/physical-plan/src/work_table.rs @@ -25,8 +25,8 @@ use crate::execution_plan::{Boundedness, EmissionType, SchedulingType}; use crate::memory::MemoryStream; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, - SendableRecordBatchStream, Statistics, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, }; use crate::statistics::StatisticsArgs; @@ -194,13 +194,24 @@ impl ExecutionPlan for WorkTableExec { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(Arc::clone(&self) as Arc) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + /// Stream the batches that were written to the work table. fn execute( &self, diff --git a/datafusion/proto/tests/cases/plans/dispatch.rs b/datafusion/proto/tests/cases/plans/dispatch.rs index af9d8a62d32f..75f299107e35 100644 --- a/datafusion/proto/tests/cases/plans/dispatch.rs +++ b/datafusion/proto/tests/cases/plans/dispatch.rs @@ -27,8 +27,8 @@ use datafusion::physical_plan::expressions::{PhysicalSortExpr, col}; use datafusion::physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PhysicalExpr, PlanProperties, - SendableRecordBatchStream, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PhysicalExpr, + PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, }; use datafusion::prelude::SessionContext; use datafusion_common::tree_node::TreeNodeRecursion; @@ -81,14 +81,28 @@ impl ExecutionPlan for DowncastDelegatingExec { self.inner.apply_expressions(f) } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { - let inner = Arc::clone(&self.inner).with_new_children(children)?; + let inner = Arc::clone(&self.inner).replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; Ok(Arc::new(Self::new(inner))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn downcast_delegate(&self) -> Option<&dyn ExecutionPlan> { Some(self.inner.as_ref()) } diff --git a/datafusion/proto/tests/cases/plans/dynamic_filters.rs b/datafusion/proto/tests/cases/plans/dynamic_filters.rs index cda649b4c57b..ee0ff9d8b1fa 100644 --- a/datafusion/proto/tests/cases/plans/dynamic_filters.rs +++ b/datafusion/proto/tests/cases/plans/dynamic_filters.rs @@ -47,8 +47,8 @@ use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; use datafusion::physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PhysicalExpr, PlanProperties, - SendableRecordBatchStream, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PhysicalExpr, + PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, }; use datafusion::prelude::SessionContext; use datafusion_common::config::{ConfigOptions, TableParquetOptions}; @@ -789,13 +789,24 @@ impl ExecutionPlan for CustomExecWithExprs { datafusion_physical_plan::apply_expression_roots(&self.exprs, f) } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unreachable!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/proto/tests/cases/plans/limits.rs b/datafusion/proto/tests/cases/plans/limits.rs index e1c33ff94923..a832d46d5315 100644 --- a/datafusion/proto/tests/cases/plans/limits.rs +++ b/datafusion/proto/tests/cases/plans/limits.rs @@ -38,6 +38,7 @@ use datafusion::physical_plan::coop::CooperativeExec; use datafusion::physical_plan::empty::EmptyExec; use datafusion::physical_plan::expressions::{PhysicalSortExpr, col}; use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; +use datafusion::physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}; use datafusion::prelude::SessionContext; use datafusion_common::Result; use datafusion_common::config::ConfigOptions; @@ -148,7 +149,10 @@ fn roundtrip_limit_required_ordering_reaches_data_source() -> Result<()> { roundtrip_test_and_return(Arc::new(limit), &ctx, &codec, &proto_converter)?; // Child replacement must not erase the decoded ordering before pushdown. - let rebuilt = decoded.with_new_children(vec![make_scan()])?; + let rebuilt = decoded.replace_children( + vec![make_scan()], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; let optimized = LimitPushdown::new().optimize(rebuilt, &ConfigOptions::default())?; diff --git a/docs/source/library-user-guide/custom-table-providers.md b/docs/source/library-user-guide/custom-table-providers.md index c6a316aa74b9..c094f8bf7eb1 100644 --- a/docs/source/library-user-guide/custom-table-providers.md +++ b/docs/source/library-user-guide/custom-table-providers.md @@ -247,14 +247,22 @@ impl ExecutionPlan for MyExecPlan { vec![] // Leaf node -- no children } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { assert!(children.is_empty()); Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute)) + } + fn execute( &self, partition: usize, @@ -655,7 +663,7 @@ and reading files that cannot possibly match the query. # use datafusion::execution::context::TaskContext; # use datafusion::logical_expr::{Expr, TableProviderFilterPushDown}; # use datafusion::physical_expr::EquivalenceProperties; -# use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PhysicalExpr, PlanProperties}; +# use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PhysicalExpr, PlanProperties, ChildrenPropertiesMode, ReplaceChildrenOptions}; # use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; # /// A table provider backed by date-partitioned directories. @@ -764,7 +772,15 @@ impl DatePartitionedTable { # fn name(&self) -> &str { "DatePartitionedExec" } # fn properties(&self) -> &Arc { &self.properties } # fn children(&self) -> Vec<&Arc> { vec![] } -# fn with_new_children(self: Arc, _: Vec>) -> Result> { Ok(self) } +# fn replace_children(self: Arc, _: Vec>, _: ReplaceChildrenOptions) -> Result> { Ok(self) } +# +# fn with_new_children( +# self: Arc, +# children: Vec>, +# ) -> Result> { +# self.replace_children(children, ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute)) +# } +# # fn execute(&self, _: usize, _: Arc) -> Result { todo!() } # fn apply_expressions(&self, _f: &mut dyn FnMut(&Arc) -> Result) -> Result { Ok(TreeNodeRecursion::Continue) } # } @@ -801,10 +817,8 @@ use datafusion::physical_expr::EquivalenceProperties; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ -# DisplayAs, DisplayFormatType, - ExecutionPlan, Partitioning, -# PhysicalExpr, - PlanProperties, +# DisplayAs, DisplayFormatType, PhysicalExpr, + ChildrenPropertiesMode, ReplaceChildrenOptions, ExecutionPlan, Partitioning, PlanProperties, }; use futures::stream; @@ -874,13 +888,21 @@ impl ExecutionPlan for CountingExec { fn properties(&self) -> &Arc { &self.properties } fn children(&self) -> Vec<&Arc> { vec![] } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute)) + } + fn execute( &self, partition: usize, diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md index 7d0c19c846ca..d64f287ea0b5 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -576,6 +576,78 @@ See [PR #22733](https://github.com/apache/datafusion/pull/22733) for details, including the per-variant size breakdown and benchmark results. +### `ExecutionPlan::with_new_children` and `ExecutionPlan::with_new_children_and_same_properties` deprecated + +`with_new_children` and `with_new_children_and_same_properties` have been +deprecated. These methods are used to replace the child plans of an +`ExecutionPlan` while leaving the plan otherwise identical. + +`with_new_children_if_necessary` has also been deprecated in favor of +`replace_children_if_necessary` for consistency in naming. + +As noted [here](https://github.com/apache/datafusion/pull/23332#discussion_r3554897693), +while the addition of `with_new_children_and_same_properties` has the benefit +of skipping potentially expensive computation in the case that replacement children +have the same properties as the original children, it widens the API surface area +of `ExecutionPlan` in a way that could be confusing for users. + +Thus, to rectify this, we unify these methods by introducing `replace_children`. +`replace_children` solves this problem by taking `ReplaceChildrenOptions`, +which includes a `ChildrenPropertiesMode`. The mode has two variants, +`Keep` and `Recompute`, which tell `replace_children` whether plan +properties can be reused or need to be recomputed. + +This method is called from `replace_children_if_necessary`, which is the +standard entry point that should be used for replacing the children of a node. + +**Migration guide:** + +To migrate from `with_new_children` and `with_new_children_and_same_properties` +to `replace_children`, it is recommended to implement `replace_children` with +a `match` statement matching on the `ChildrenPropertiesMode`. In the case that +the properties match the children, `ChildrenPropertiesMode::Keep`, +follow the body of `with_new_children_and_same_properties`. In the case that +the properties do not match the children, `ChildrenPropertiesMode::Recompute`, +follow the body of `with_new_children`. + +For example, take a look at the implementation for `FilterExec`: + +``` + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let new_input = children.swap_remove(0); + FilterExecBuilder::from(&*self) + .with_input(new_input) + .build() + .map(|e| Arc::new(e) as _) + } + } + } +``` + +In the case that the options indicate the properties are the same, we can simply +swap the children without having to recompute the properties. In the other case, +we create a new node from scratch. + +To ensure that this works correctly, it is recommended that users also look +through their codebase and ensure that they use `replace_children_if_necessary` +for these changes — `replace_children_if_necessary` should be preferred over +manual use of `replace_children`, since `replace_children_if_necessary` will +call `replace_children` with the correct options filled in. + +See [PR #23903](https://github.com/apache/datafusion/pull/23903) for details. + ### `ListingOptions::target_partitions` and `collect_stat` removed The `target_partitions` and `collect_stat` fields on