feat: implement replace_children for custom ExecutionPlan nodes

benchmarks/benches/broadcast_cache_scenarios.rs

@@ -1,3 +1,4 @@
+use datafusion::physical_plan::ReplaceChildrenOptions;
 use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
 use datafusion::arrow::array::UInt8Array;
 use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
@@ -148,6 +149,14 @@ impl ExecutionPlan for SyntheticExec {
         unimplemented!()
     }
 
+    fn replace_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        _options: ReplaceChildrenOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        self.with_new_children(children)
+    }
+
     fn execute(
         &self,
         partition: usize,

examples/custom_execution_plan.rs

@@ -1,3 +1,4 @@
+use datafusion::physical_plan::ReplaceChildrenOptions;
 //! This example demonstrates how to create a custom execution plan that works with
 //! Distributed DataFusion. It implements a `numbers(start, end)` table function that
 //! generates a sequence of numbers and can be distributed across multiple workers.
@@ -187,6 +188,14 @@ impl ExecutionPlan for NumbersExec {
         Ok(self)
     }
 
+    fn replace_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        _options: ReplaceChildrenOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        self.with_new_children(children)
+    }
+
     fn execute(
         &self,
         _partition: usize,

examples/custom_worker_url_routing.rs

@@ -1,3 +1,4 @@
+use datafusion::physical_plan::ReplaceChildrenOptions;
 //! Demonstrates **custom task routing** for **cache affinity**: consistently routing each parquet
 //! file to the *same* worker so that worker can serve it from an in-memory cache on repeat queries.
 //!
@@ -108,6 +109,14 @@ impl ExecutionPlan for CacheExec {
         Ok(CacheExec::new(children.remove(0)))
     }
 
+    fn replace_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        _options: ReplaceChildrenOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        self.with_new_children(children)
+    }
+
     fn execute(
         &self,
         partition: usize,

examples/work_unit_feed.rs

@@ -1,3 +1,4 @@
+use datafusion::physical_plan::ReplaceChildrenOptions;
 //! Demonstrates **work unit feeds**: a distributed leaf node whose work is discovered on the
 //! coordinator *at runtime* and streamed to the workers while the query runs, instead of being
 //! known at planning time (think a paginated API, a queue, or a catalog handing out keys).
@@ -163,6 +164,14 @@ impl ExecutionPlan for RemoteScanExec {
         Ok(self)
     }
 
+    fn replace_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        _options: ReplaceChildrenOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        self.with_new_children(children)
+    }
+
     fn execute(
         &self,
         partition: usize,

src/coordinator/distributed.rs

@@ -13,7 +13,7 @@ use datafusion::physical_expr::PhysicalExpr;
 use datafusion::physical_expr_common::metrics::MetricsSet;
 use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
 use datafusion::physical_plan::stream::RecordBatchReceiverStreamBuilder;
-use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
+use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, ReplaceChildrenOptions};
 use futures::StreamExt;
 use std::fmt::Formatter;
 use std::sync::{Arc, Mutex};
@@ -176,6 +176,15 @@ impl ExecutionPlan for DistributedExec {
         }))
     }
 
+    fn replace_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        _options: ReplaceChildrenOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        // Prefer replace_children over deprecated with_new_children (#657).
+        self.with_new_children(children)
+    }
+
     fn execute(
         &self,
         partition: usize,

src/execution_plans/broadcast.rs

@@ -9,8 +9,7 @@ use datafusion::execution::{SendableRecordBatchStream, TaskContext};
 use datafusion::physical_expr::PhysicalExpr;
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion::physical_plan::{
-    DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, internal_err,
-};
+    DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, internal_err, ReplaceChildrenOptions};
 use futures::{Stream, StreamExt};
 use std::fmt::Formatter;
 use std::pin::Pin;
@@ -155,6 +154,15 @@ impl ExecutionPlan for BroadcastExec {
         )))
     }
 
+    fn replace_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        _options: ReplaceChildrenOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        // Prefer replace_children over deprecated with_new_children (#657).
+        self.with_new_children(children)
+    }
+
     fn execute(
         &self,
         partition: usize,

src/execution_plans/children_isolator_union.rs

@@ -12,8 +12,7 @@ use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSe
 use datafusion::physical_plan::union::UnionExec;
 use datafusion::physical_plan::{
     DisplayAs, DisplayFormatType, EmptyRecordBatchStream, ExecutionPlan, ExecutionPlanProperties,
-    Partitioning, PlanProperties,
-};
+    Partitioning, PlanProperties, ReplaceChildrenOptions};
 use futures::{Stream, StreamExt};
 use itertools::Itertools;
 use std::fmt::Formatter;
@@ -295,6 +294,15 @@ impl ExecutionPlan for ChildrenIsolatorUnionExec {
         )?))
     }
 
+    fn replace_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        _options: ReplaceChildrenOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        // Prefer replace_children over deprecated with_new_children (#657).
+        self.with_new_children(children)
+    }
+
     fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
         self.children.iter().collect()
     }

src/execution_plans/distributed_leaf.rs

@@ -5,8 +5,7 @@ use datafusion::execution::{SendableRecordBatchStream, TaskContext};
 use datafusion::physical_expr::PhysicalExpr;
 use datafusion::physical_expr_common::metrics::MetricsSet;
 use datafusion::physical_plan::{
-    DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, StatisticsArgs,
-};
+    DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, StatisticsArgs, ReplaceChildrenOptions};
 use std::fmt::Formatter;
 use std::sync::Arc;
 
@@ -177,6 +176,15 @@ impl ExecutionPlan for DistributedLeafExec {
         Ok(self)
     }
 
+    fn replace_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        _options: ReplaceChildrenOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        // Prefer replace_children over deprecated with_new_children (#657).
+        self.with_new_children(children)
+    }
+
     fn execute(
         &self,
         partition: usize,

src/execution_plans/metrics.rs

@@ -77,6 +77,15 @@ impl ExecutionPlan for MetricsWrapperExec {
         }))
     }
 
+    fn replace_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        _options: ReplaceChildrenOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        // Prefer replace_children over deprecated with_new_children (#657).
+        self.with_new_children(children)
+    }
+
     fn execute(
         &self,
         _partition: usize,

src/execution_plans/network_broadcast.rs

@@ -12,8 +12,7 @@ use datafusion::physical_expr_common::metrics::MetricsSet;
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion::physical_plan::{
     DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, Statistics,
-    StatisticsArgs,
-};
+    StatisticsArgs, ReplaceChildrenOptions};
 use std::fmt::Formatter;
 use std::sync::Arc;
 use uuid::Uuid;
@@ -242,6 +241,15 @@ impl ExecutionPlan for NetworkBroadcastExec {
         Ok(Arc::new(self_clone))
     }
 
+    fn replace_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        _options: ReplaceChildrenOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        // Prefer replace_children over deprecated with_new_children (#657).
+        self.with_new_children(children)
+    }
+
     fn execute(
         &self,
         partition: usize,

src/execution_plans/network_coalesce.rs

@@ -264,6 +264,15 @@ impl ExecutionPlan for NetworkCoalesceExec {
         Ok(Arc::new(self_clone))
     }
 
+    fn replace_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        _options: ReplaceChildrenOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        // Prefer replace_children over deprecated with_new_children (#657).
+        self.with_new_children(children)
+    }
+
     fn execute(
         &self,
         partition: usize,

src/execution_plans/network_shuffle.rs

@@ -13,8 +13,7 @@ use datafusion::physical_expr_common::metrics::MetricsSet;
 use datafusion::physical_plan::repartition::RepartitionExec;
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion::physical_plan::{
-    DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, Statistics, StatisticsArgs,
-};
+    DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, Statistics, StatisticsArgs, ReplaceChildrenOptions};
 use std::fmt::Formatter;
 use std::sync::Arc;
 use uuid::Uuid;
@@ -218,6 +217,15 @@ impl ExecutionPlan for NetworkShuffleExec {
         Ok(Arc::new(self_clone))
     }
 
+    fn replace_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        _options: ReplaceChildrenOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        // Prefer replace_children over deprecated with_new_children (#657).
+        self.with_new_children(children)
+    }
+
     fn execute(
         &self,
         partition: usize,

src/execution_plans/sampler.rs

@@ -17,8 +17,7 @@ use datafusion::physical_expr_common::metrics::{Gauge, MetricValue, MetricsSet};
 use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder, Time};
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion::physical_plan::{
-    DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties,
-};
+    DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, ReplaceChildrenOptions};
 use futures::stream::FusedStream;
 use futures::{Stream, StreamExt, TryFutureExt};
 use std::collections::VecDeque;
@@ -551,6 +550,15 @@ impl ExecutionPlan for SamplerExec {
         Ok(Arc::new(Self::new(require_one_child(children)?)))
     }
 
+    fn replace_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        _options: ReplaceChildrenOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        // Prefer replace_children over deprecated with_new_children (#657).
+        self.with_new_children(children)
+    }
+
     fn execute(
         &self,
         partition: usize,

src/explain_analyze.rs

@@ -12,8 +12,7 @@ use datafusion::physical_expr::PhysicalExpr;
 use datafusion::physical_plan::analyze::AnalyzeExec;
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion::physical_plan::{
-    DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, PlanProperties,
-};
+    DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, PlanProperties, ReplaceChildrenOptions};
 use futures::{StreamExt, stream};
 use std::fmt::Formatter;
 use std::sync::Arc;
@@ -80,6 +79,15 @@ impl ExecutionPlan for DistributedAnalyzeExec {
         }))
     }
 
+    fn replace_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        _options: ReplaceChildrenOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        // Prefer replace_children over deprecated with_new_children (#657).
+        self.with_new_children(children)
+    }
+
     fn execute(
         &self,
         partition: usize,

src/test_utils/mock_exec.rs

@@ -7,7 +7,7 @@ use datafusion::physical_expr::{EquivalenceProperties, Partitioning, PhysicalExp
 use datafusion::physical_plan::common::compute_record_batch_statistics;
 use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
 use datafusion::physical_plan::stream::{RecordBatchReceiverStream, RecordBatchStreamAdapter};
-use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
+use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, ReplaceChildrenOptions};
 use futures::{Stream, stream};
 use std::pin::Pin;
 use std::sync::Arc;
@@ -166,6 +166,15 @@ impl ExecutionPlan for MockExec {
         unimplemented!()
     }
 
+    fn replace_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        _options: ReplaceChildrenOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        // Prefer replace_children over deprecated with_new_children (#657).
+        self.with_new_children(children)
+    }
+
     /// Returns a stream which yields data
     fn execute(
         &self,

src/test_utils/routing.rs

@@ -1,3 +1,4 @@
+use datafusion::physical_plan::ReplaceChildrenOptions;
 use arrow::{
     array::{Int64Array, RecordBatch, StringArray},
     datatypes::{DataType, Field, Schema, SchemaRef},
@@ -197,6 +198,15 @@ impl ExecutionPlan for URLEmitterExec {
         Ok(Arc::new(self.as_ref().clone()))
     }
 
+    fn replace_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        _options: ReplaceChildrenOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        // Prefer replace_children over deprecated with_new_children (#657).
+        self.with_new_children(children)
+    }
+
     fn execute(
         &self,
         partition: usize,

src/test_utils/test_work_unit_feed.rs

@@ -20,7 +20,7 @@ use datafusion::physical_expr_common::metrics::MetricsSet;
 use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
 use datafusion::physical_plan::metrics::{Count, ExecutionPlanMetricsSet, MetricBuilder};
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
-use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
+use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, ReplaceChildrenOptions};
 use datafusion_proto::physical_plan::{PhysicalExtensionCodec, PhysicalProtoConverterExtension};
 use datafusion_proto::protobuf::proto_error;
 use futures::StreamExt;
@@ -403,6 +403,15 @@ impl ExecutionPlan for RowGeneratorExec {
         Ok(Arc::new(self.as_ref().clone()))
     }
 
+    fn replace_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        _options: ReplaceChildrenOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        // Prefer replace_children over deprecated with_new_children (#657).
+        self.with_new_children(children)
+    }
+
     fn execute(
         &self,
         partition: usize,