feat: replace_children

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::<BufferExec>() {
                         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::<BufferExec>() {
                         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)),
+                        ],
+                    )?
                 },
             ))
         })

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))

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<dyn ExecutionPlan>>;
 
+    /// 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<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        hint: ChildrenPropertiesHint,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        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<TaskContext>,
     ///     ) -> Result<SendableRecordBatchStream> {
-    ///         // 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<dyn ExecutionPlan> = 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)],

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;

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<Self>,
         mut children: Vec<Arc<dyn ExecutionPlan>>,
+        hint: ChildrenPropertiesHint,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        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<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        self.replace_children(children, ChildrenPropertiesHint::Recompute)
     }
 
     fn with_new_children_and_same_properties(
         self: Arc<Self>,
-        mut children: Vec<Arc<dyn ExecutionPlan>>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        Ok(Arc::new(Self {
-            input: children.swap_remove(0),
-            metrics: ExecutionPlanMetricsSet::new(),
-            ..Self::clone(&*self)
-        }))
+        self.replace_children(children, ChildrenPropertiesHint::SameProperties)
     }
 
     fn execute(

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<Self>,
         mut children: Vec<Arc<dyn ExecutionPlan>>,
+        hint: ChildrenPropertiesHint,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        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<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        self.replace_children(children, ChildrenPropertiesHint::Recompute)
     }
 
     fn with_new_children_and_same_properties(
         self: Arc<Self>,
-        mut children: Vec<Arc<dyn ExecutionPlan>>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        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<Distribution> {