Migrate usages away from with_new_children.

Author: JSOD11

datafusion/catalog/src/memory/table.rs

@@ -574,7 +574,7 @@ impl ExecutionPlan for DmlResultExec {
     fn replace_children(
         self: Arc<Self>,
         _: Vec<Arc<dyn ExecutionPlan>>,
-        _: datafusion_physical_plan::ChildrenPropertiesHint,
+        _: ChildrenPropertiesHint,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         Ok(self)
     }

datafusion/core/src/physical_planner.rs

@@ -4981,10 +4981,16 @@ digraph {
         // ok plan
         let ok_node: Arc<dyn ExecutionPlan> = 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)],
+                    ChildrenPropertiesHint::Recompute,
+                )?,
+                Arc::clone(&child),
+            ],
+            ChildrenPropertiesHint::Recompute,
+        )?;
 
         // Test: check should pass with same schema
         let equal_schema = ok_plan.schema();
@@ -5016,10 +5022,16 @@ digraph {
 
         // Test: should fail when descendent extension node fails
         let failing_node: Arc<dyn ExecutionPlan> = 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)],
+                    ChildrenPropertiesHint::Recompute,
+                )?,
+                Arc::clone(&child),
+            ],
+            ChildrenPropertiesHint::Recompute,
+        )?;
         let result = OptimizationInvariantChecker::new(&rule)
             .check(&invalid_plan, &ok_plan.schema());
         if cfg!(debug_assertions) {
@@ -5105,10 +5117,16 @@ digraph {
         let failing_node: Arc<dyn ExecutionPlan> = Arc::new(ExecutableInvariantFails);
         let ok_node: Arc<dyn ExecutionPlan> = 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)],
+                    ChildrenPropertiesHint::Recompute,
+                )?,
+                Arc::clone(&child),
+            ],
+            ChildrenPropertiesHint::Recompute,
+        )?;
         let expected_err = InvariantChecker(InvariantLevel::Executable)
             .check(&plan)
             .unwrap_err();

datafusion/ffi/src/execution_plan.rs

@@ -151,7 +151,9 @@ 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, ChildrenPropertiesHint::Recompute)
+    );
 
     FFI_Result::Ok(FFI_ExecutionPlan::new(new_plan, runtime))
 }
@@ -268,7 +270,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::<ForeignExecutionPlan>() {
                 updated_children = true;
@@ -281,7 +283,9 @@ fn pass_runtime_to_children(
         })
         .collect::<Result<Vec<_>>>()?;
     if updated_children {
-        Arc::clone(plan).with_new_children(children).map(Some)
+        Arc::clone(plan)
+            .replace_children(children, ChildrenPropertiesHint::Recompute)
+            .map(Some)
     } else {
         Ok(None)
     }
@@ -636,7 +640,8 @@ 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], ChildrenPropertiesHint::Recompute)?;
         assert_eq!(parent_foreign.children().len(), 1);
 
         // Version 2: Adding child to the local plan
@@ -646,7 +651,8 @@ pub mod tests {
         let child_foreign = <Arc<dyn ExecutionPlan>>::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], ChildrenPropertiesHint::Recompute)?;
         let mut parent_local = FFI_ExecutionPlan::new(parent_plan, None);
         parent_local.library_marker_id = crate::mock_foreign_marker_id;
         let parent_foreign = <Arc<dyn ExecutionPlan>>::try_from(&parent_local)?;

datafusion/ffi/tests/ffi_execution_plan.rs

@@ -25,7 +25,7 @@ mod tests {
     use datafusion_ffi::execution_plan::ForeignExecutionPlan;
     use datafusion_ffi::execution_plan::{ExecutionPlanPrivateData, tests::EmptyExec};
     use datafusion_ffi::tests::utils::get_module;
-    use datafusion_physical_plan::ExecutionPlan;
+    use datafusion_physical_plan::{ChildrenPropertiesHint, ExecutionPlan};
     use std::sync::Arc;
 
     #[test]
@@ -92,7 +92,8 @@ 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], ChildrenPropertiesHint::Recompute)?;
 
         unsafe {
             // Originally the runtime is not set. We go through the unsafe casting
@@ -107,7 +108,8 @@ 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], ChildrenPropertiesHint::Recompute)?;
 
         // Adding the grandchild beneath this FFI plan should get the runtime passed down.
         let runtime = tokio::runtime::Builder::new_current_thread()

datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs

@@ -760,16 +760,19 @@ 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 = with_new_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);
     }
     // Build parent with an unordered child via CoalescePartitionsExec.
     let unordered_child: Arc<dyn ExecutionPlan> =
         Arc::new(CoalescePartitionsExec::new(Arc::clone(ordered_child)));
-    let without_ordered = Arc::clone(parent).with_new_children(vec![unordered_child])?;
+    let without_ordered =
+        with_new_children_if_necessary(Arc::clone(parent), vec![unordered_child])?;
     Ok(without_ordered.pipeline_behavior() == EmissionType::Final)
 }
 

datafusion/physical-optimizer/src/limit_pushdown.rs

@@ -77,7 +77,9 @@ use datafusion_physical_plan::placeholder_row::PlaceholderRowExec;
 use datafusion_physical_plan::projection::ProjectionExec;
 use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
 use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext};
-use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
+use datafusion_physical_plan::{
+    ExecutionPlan, ExecutionPlanProperties, with_new_children_if_necessary,
+};
 /// This rule inspects [`ExecutionPlan`]'s and pushes down the fetch limit from
 /// the parent to the child if applicable.
 #[derive(Default, Debug)]
@@ -403,7 +405,7 @@ pub(crate) fn pushdown_limits(
         .collect::<Result<_>>()?;
 
     if changed {
-        new_node.data.with_new_children(new_children)
+        with_new_children_if_necessary(new_node.data, new_children)
     } else {
         Ok(new_node.data)
     }

datafusion/physical-optimizer/src/topk_repartition.rs

@@ -55,7 +55,9 @@ use std::sync::Arc;
 use datafusion_physical_plan::coalesce_batches::CoalesceBatchesExec;
 use datafusion_physical_plan::repartition::RepartitionExec;
 use datafusion_physical_plan::sorts::sort::SortExec;
-use datafusion_physical_plan::{ExecutionPlan, Partitioning};
+use datafusion_physical_plan::{
+    ExecutionPlan, Partitioning, with_new_children_if_necessary,
+};
 
 /// A physical optimizer rule that pushes TopK (Sort with fetch) past
 /// hash repartition when the partition key is a prefix of the sort key.
@@ -151,7 +153,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])?
+                with_new_children_if_necessary(parent, vec![new_repartition])?
             } else {
                 new_repartition
             };

datafusion/physical-optimizer/src/window_topn.rs

@@ -58,7 +58,6 @@ use datafusion_common::{Result, ScalarValue};
 use datafusion_expr::Operator;
 use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal};
 use datafusion_physical_expr::window::StandardWindowExpr;
-use datafusion_physical_plan::ExecutionPlan;
 use datafusion_physical_plan::filter::FilterExec;
 use datafusion_physical_plan::projection::ProjectionExec;
 use datafusion_physical_plan::repartition::RepartitionExec;
@@ -67,6 +66,7 @@ use datafusion_physical_plan::sorts::partitioned_topk::{
 };
 use datafusion_physical_plan::sorts::sort::SortExec;
 use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr};
+use datafusion_physical_plan::{ExecutionPlan, with_new_children_if_necessary};
 
 /// Physical optimizer rule that converts per-partition `ROW_NUMBER` and
 /// `RANK` top-K queries into a more efficient plan using
@@ -192,13 +192,13 @@ impl WindowTopN {
         .ok()?;
 
         // Step 8: Rebuild window with new child
-        let mut result = window_exec
-            .with_new_children(vec![Arc::new(partitioned_topk)])
-            .ok()?;
+        let mut result =
+            with_new_children_if_necessary(window_exec, vec![Arc::new(partitioned_topk)])
+                .ok()?;
 
         // Step 9: Rebuild intermediate nodes (ProjectionExec/RepartitionExec)
         for node in intermediates.into_iter().rev() {
-            result = node.with_new_children(vec![result]).ok()?;
+            result = with_new_children_if_necessary(node, vec![result]).ok()?;
         }
 
         Some(result)

datafusion/physical-plan/src/async_func.rs

@@ -297,6 +297,7 @@ impl AsyncFuncExec {
         node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
         ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
+        use datafusion_common::assert_eq_or_internal_err;
         use datafusion_proto_models::protobuf;
         let async_func = crate::expect_plan_variant!(
             node,

datafusion/physical-plan/src/buffer.rs

@@ -28,7 +28,7 @@ use crate::statistics::{ChildStats, StatisticsArgs};
 use crate::stream::RecordBatchStreamAdapter;
 use crate::{
     ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
-    SortOrderPushdownResult, validate_child_count,
+    SortOrderPushdownResult, validate_child_count, with_new_children_if_necessary,
 };
 use arrow::array::RecordBatch;
 use datafusion_common::config::ConfigOptions;
@@ -273,9 +273,10 @@ impl ExecutionPlan for BufferExec {
         projection: &ProjectionExec,
     ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
         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(with_new_children_if_necessary(
+                Arc::new(self.clone()),
+                vec![new_input],
+            )?)),
             None => Ok(None),
         }
     }