diff --git a/datafusion/physical-optimizer/src/hash_join_buffering.rs b/datafusion/physical-optimizer/src/hash_join_buffering.rs index 7a198cac13fc9..44ebddefffd51 100644 --- a/datafusion/physical-optimizer/src/hash_join_buffering.rs +++ b/datafusion/physical-optimizer/src/hash_join_buffering.rs @@ -19,9 +19,9 @@ use crate::PhysicalOptimizerRule; use datafusion_common::JoinSide; 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::joins::HashJoinExec; +use datafusion_physical_plan::{ExecutionPlan, with_new_children_if_necessary}; use std::sync::Arc; /// Looks for all the [HashJoinExec]s in the plan and places a [BufferExec] node with the @@ -74,19 +74,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), - ])? + with_new_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)), - ])? + with_new_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/output_requirements.rs b/datafusion/physical-optimizer/src/output_requirements.rs index b9d0d06da1dda..0583e5ee885a5 100644 --- a/datafusion/physical-optimizer/src/output_requirements.rs +++ b/datafusion/physical-optimizer/src/output_requirements.rs @@ -42,6 +42,7 @@ use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeE use datafusion_physical_plan::{ ChildStats, DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, SendableRecordBatchStream, StatisticsArgs, + with_new_children_if_necessary, }; /// This rule either adds or removes [`OutputRequirements`]s to/from the physical @@ -460,7 +461,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((with_new_children_if_necessary(plan, children)?, true)); } } Ok((plan, false)) diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 3b9d5d258a838..ba27399f9ae41 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -247,6 +247,21 @@ 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 if possible as indicated by the hint. + fn replace_children( + self: Arc, + children: Vec>, + hint: ChildrenPropertiesHint, + ) -> Result> { + match hint { + ChildrenPropertiesHint::SameProperties => { + self.with_new_children_and_same_properties(children) + } + ChildrenPropertiesHint::Recompute => self.with_new_children(children), + } + } + /// Returns a new `ExecutionPlan` where all existing children were replaced /// by the `children`, in order fn with_new_children( @@ -306,7 +321,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. @@ -427,7 +442,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( @@ -659,7 +674,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. @@ -846,6 +861,18 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { } } +/// A hint from `with_new_children_if_necessary` to `replace_children` indicating +/// whether the properties of the new children must be recomputed. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ChildrenPropertiesHint { + /// The properties of the new children are identical to the properties of the + /// existing children, so we can skip recomputation of plan properties. + SameProperties, + /// The properties of the new children are different from the properties of the + /// existing children, so we must recompute the properties from scratch. + Recompute, +} + impl dyn ExecutionPlan { /// Returns `true` if the plan is of type `T`. /// @@ -1176,12 +1203,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] @@ -1396,11 +1421,12 @@ 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, ChildrenPropertiesHint::SameProperties); } } // Layer 3: full recompute. - plan.with_new_children(children) + plan.replace_children(children, ChildrenPropertiesHint::Recompute) } /// Return a [`DisplayableExecutionPlan`] wrapper around an @@ -2209,9 +2235,18 @@ 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 + // Using the same child means we return the original plan Arc verbatim, so even when + // no override exists for `with_new_children_and_same_properties` we do not recompute. + let out = with_new_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 // enters the "same properties" branch and calls the trait method, - // whose default forwards to `with_new_children`. + // whose default forwards to `with_new_children`, causing recomputation. let out = with_new_children_if_necessary( Arc::clone(&parent_dyn), vec![Arc::clone(&leaf_b)], diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 8cba650b79770..a8dd87332bb71 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -45,9 +45,9 @@ pub use crate::distribution_requirements::{ ChildSatisfactionOptions, InputDistributionRequirements, }; pub use crate::execution_plan::{ - ExecutionPlan, ExecutionPlanProperties, PlanProperties, collect, collect_partitioned, - displayable, execute_input_stream, execute_stream, execute_stream_partitioned, - get_plan_string, with_new_children_if_necessary, + ChildrenPropertiesHint, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + collect, collect_partitioned, displayable, execute_input_stream, execute_stream, + execute_stream_partitioned, get_plan_string, with_new_children_if_necessary, }; pub use crate::metrics::Metric; pub use crate::ordering::InputOrderMode; diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 2add3e1eb82f0..0a6571ffaf74d 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, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, Partitioning, PlanProperties, SendableRecordBatchStream, + Statistics, }; use datafusion_common::{Result, assert_eq_or_internal_err, internal_err}; @@ -281,26 +281,36 @@ impl ExecutionPlan for SortPreservingMergeExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + hint: ChildrenPropertiesHint, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new( - SortPreservingMergeExec::new(self.expr.clone(), children.swap_remove(0)) - .with_fetch(self.fetch), - )) + match hint { + ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesHint::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, ChildrenPropertiesHint::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, ChildrenPropertiesHint::SameProperties) } fn execute( diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index fbe849229941a..6086515ff1834 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -28,8 +28,8 @@ use super::metrics::{ use super::{DisplayAs, ExecutionPlanProperties, PlanProperties}; use crate::stream::EmptyRecordBatchStream; use crate::{ - DisplayFormatType, Distribution, ExecutionPlan, RecordBatchStream, - SendableRecordBatchStream, check_if_same_properties, + ChildrenPropertiesHint, DisplayFormatType, Distribution, ExecutionPlan, + RecordBatchStream, SendableRecordBatchStream, }; use arrow::array::{ @@ -227,29 +227,39 @@ impl ExecutionPlan for UnnestExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + hint: ChildrenPropertiesHint, ) -> 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(), - )?)) + match hint { + ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesHint::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, ChildrenPropertiesHint::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, ChildrenPropertiesHint::SameProperties) } fn required_input_distribution(&self) -> Vec {