Implement fast path of with_new_children() in ExecutionPlan (#2168)

* Implement fast path of with_new_children() in ExecutionPlan

* resolve review comments

* refine comments

Author: mingmwang

ballista/rust/core/src/execution_plans/distributed_query.rs

@@ -148,7 +148,7 @@ impl<T: 'static + AsLogicalPlan> ExecutionPlan for DistributedQueryExec<T> {
     }
 
     fn with_new_children(
-        &self,
+        self: Arc<Self>,
         _children: Vec<Arc<dyn ExecutionPlan>>,
     ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
         Ok(Arc::new(DistributedQueryExec {

ballista/rust/core/src/execution_plans/shuffle_reader.rs

@@ -95,7 +95,7 @@ impl ExecutionPlan for ShuffleReaderExec {
     }
 
     fn with_new_children(
-        &self,
+        self: Arc<Self>,
         _children: Vec<Arc<dyn ExecutionPlan>>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         Err(DataFusionError::Plan(

ballista/rust/core/src/execution_plans/shuffle_writer.rs

@@ -342,10 +342,9 @@ impl ExecutionPlan for ShuffleWriterExec {
     }
 
     fn with_new_children(
-        &self,
+        self: Arc<Self>,
         children: Vec<Arc<dyn ExecutionPlan>>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        assert!(children.len() == 1);
         Ok(Arc::new(ShuffleWriterExec::try_new(
             self.job_id.clone(),
             self.stage_id,

ballista/rust/core/src/execution_plans/unresolved_shuffle.rs

@@ -92,7 +92,7 @@ impl ExecutionPlan for UnresolvedShuffleExec {
     }
 
     fn with_new_children(
-        &self,
+        self: Arc<Self>,
         _children: Vec<Arc<dyn ExecutionPlan>>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         Err(DataFusionError::Plan(

ballista/rust/core/src/serde/mod.rs

@@ -505,18 +505,13 @@ mod tests {
         }
 
         fn with_new_children(
-            &self,
+            self: Arc<Self>,
             children: Vec<Arc<dyn ExecutionPlan>>,
         ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
-            match children.len() {
-                1 => Ok(Arc::new(TopKExec {
-                    input: children[0].clone(),
-                    k: self.k,
-                })),
-                _ => Err(DataFusionError::Internal(
-                    "TopKExec wrong number of children".to_string(),
-                )),
-            }
+            Ok(Arc::new(TopKExec {
+                input: children[0].clone(),
+                k: self.k,
+            }))
         }
 
         /// Execute one partition and return an iterator over RecordBatch

ballista/rust/executor/src/collect.rs

@@ -72,7 +72,7 @@ impl ExecutionPlan for CollectExec {
     }
 
     fn with_new_children(
-        &self,
+        self: Arc<Self>,
         _children: Vec<Arc<dyn ExecutionPlan>>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         unimplemented!()

ballista/rust/scheduler/src/planner.rs

@@ -30,7 +30,9 @@ use ballista_core::{
 use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion::physical_plan::repartition::RepartitionExec;
 use datafusion::physical_plan::windows::WindowAggExec;
-use datafusion::physical_plan::{ExecutionPlan, Partitioning};
+use datafusion::physical_plan::{
+    with_new_children_if_necessary, ExecutionPlan, Partitioning,
+};
 use futures::future::BoxFuture;
 use futures::FutureExt;
 use log::info;
@@ -99,7 +101,7 @@ impl DistributedPlanner {
                 stages.append(&mut child_stages);
             }
 
-            if let Some(coalesce) = execution_plan
+            if let Some(_coalesce) = execution_plan
                 .as_any()
                 .downcast_ref::<CoalescePartitionsExec>()
             {
@@ -122,7 +124,10 @@ impl DistributedPlanner {
                 ));
                 stages.push(shuffle_writer);
                 Ok((
-                    coalesce.with_new_children(vec![unresolved_shuffle])?,
+                    with_new_children_if_necessary(
+                        execution_plan,
+                        vec![unresolved_shuffle],
+                    )?,
                     stages,
                 ))
             } else if let Some(repart) =
@@ -163,7 +168,10 @@ impl DistributedPlanner {
                     window
                 )))
             } else {
-                Ok((execution_plan.with_new_children(children)?, stages))
+                Ok((
+                    with_new_children_if_necessary(execution_plan, children)?,
+                    stages,
+                ))
             }
         }
         .boxed()
@@ -197,7 +205,7 @@ pub fn find_unresolved_shuffles(
 }
 
 pub fn remove_unresolved_shuffles(
-    stage: &dyn ExecutionPlan,
+    stage: Arc<dyn ExecutionPlan>,
     partition_locations: &HashMap<usize, HashMap<usize, Vec<PartitionLocation>>>,
 ) -> Result<Arc<dyn ExecutionPlan>> {
     let mut new_children: Vec<Arc<dyn ExecutionPlan>> = vec![];
@@ -240,13 +248,10 @@ pub fn remove_unresolved_shuffles(
                 unresolved_shuffle.schema().clone(),
             )?))
         } else {
-            new_children.push(remove_unresolved_shuffles(
-                child.as_ref(),
-                partition_locations,
-            )?);
+            new_children.push(remove_unresolved_shuffles(child, partition_locations)?);
         }
     }
-    Ok(stage.with_new_children(new_children)?)
+    Ok(with_new_children_if_necessary(stage, new_children)?)
 }
 
 fn create_shuffle_writer(

ballista/rust/scheduler/src/scheduler_server/mod.rs

@@ -241,7 +241,7 @@ impl Default for SessionContextRegistry {
 }
 
 impl SessionContextRegistry {
-    /// Create the registry that object stores can registered into.
+    /// Create the registry that session contexts can registered into.
     /// ['LocalFileSystem'] store is registered in by default to support read local files natively.
     pub fn new() -> Self {
         Self {

ballista/rust/scheduler/src/scheduler_server/query_stage_scheduler.rs

@@ -304,8 +304,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> QueryStageSchedul
                 }
             }
 
-            let plan =
-                remove_unresolved_shuffles(stage_plan.as_ref(), &partition_locations)?;
+            let plan = remove_unresolved_shuffles(stage_plan, &partition_locations)?;
             self.state.save_stage_plan(job_id, stage_id, plan).await?;
         }
 

datafusion-examples/examples/custom_datasource.rs

@@ -21,7 +21,7 @@ use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::dataframe::DataFrame;
 use datafusion::datasource::TableProvider;
-use datafusion::error::{DataFusionError, Result};
+use datafusion::error::Result;
 use datafusion::execution::context::TaskContext;
 use datafusion::logical_plan::{Expr, LogicalPlanBuilder};
 use datafusion::physical_plan::expressions::PhysicalSortExpr;
@@ -219,17 +219,10 @@ impl ExecutionPlan for CustomExec {
     }
 
     fn with_new_children(
-        &self,
-        children: Vec<Arc<dyn ExecutionPlan>>,
+        self: Arc<Self>,
+        _: Vec<Arc<dyn ExecutionPlan>>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        if children.is_empty() {
-            Ok(Arc::new(self.clone()))
-        } else {
-            Err(DataFusionError::Internal(format!(
-                "Children cannot be replaced in {:?}",
-                self
-            )))
-        }
+        Ok(self)
     }
 
     async fn execute(

datafusion/core/src/physical_optimizer/coalesce_batches.rs

@@ -19,6 +19,7 @@
 //! in bigger batches to avoid overhead with small batches
 
 use super::optimizer::PhysicalOptimizerRule;
+use crate::physical_plan::with_new_children_if_necessary;
 use crate::{
     error::Result,
     physical_plan::{
@@ -69,7 +70,7 @@ impl PhysicalOptimizerRule for CoalesceBatches {
             // leaf node, children cannot be replaced
             Ok(plan.clone())
         } else {
-            let plan = plan.with_new_children(children)?;
+            let plan = with_new_children_if_necessary(plan, children)?;
             Ok(if wrap_in_coalesce {
                 // TODO we should add specific configuration settings for coalescing batches and
                 // we should do that once https://issues.apache.org/jira/browse/ARROW-11059 is

datafusion/core/src/physical_optimizer/merge_exec.rs

@@ -19,6 +19,7 @@
 //! with more than one partition, to coalesce them into one partition
 //! when the node needs a single partition
 use super::optimizer::PhysicalOptimizerRule;
+use crate::physical_plan::with_new_children_if_necessary;
 use crate::{
     error::Result,
     physical_plan::{coalesce_partitions::CoalescePartitionsExec, Distribution},
@@ -52,9 +53,14 @@ impl PhysicalOptimizerRule for AddCoalescePartitionsExec {
                 .map(|child| self.optimize(child.clone(), config))
                 .collect::<Result<Vec<_>>>()?;
             match plan.required_child_distribution() {
-                Distribution::UnspecifiedDistribution => plan.with_new_children(children),
-                Distribution::HashPartitioned(_) => plan.with_new_children(children),
-                Distribution::SinglePartition => plan.with_new_children(
+                Distribution::UnspecifiedDistribution => {
+                    with_new_children_if_necessary(plan, children)
+                }
+                Distribution::HashPartitioned(_) => {
+                    with_new_children_if_necessary(plan, children)
+                }
+                Distribution::SinglePartition => with_new_children_if_necessary(
+                    plan,
                     children
                         .iter()
                         .map(|child| {

datafusion/core/src/physical_optimizer/repartition.rs

@@ -20,7 +20,9 @@ use std::sync::Arc;
 
 use super::optimizer::PhysicalOptimizerRule;
 use crate::physical_plan::Partitioning::*;
-use crate::physical_plan::{repartition::RepartitionExec, ExecutionPlan};
+use crate::physical_plan::{
+    repartition::RepartitionExec, with_new_children_if_necessary, ExecutionPlan,
+};
 use crate::{error::Result, execution::context::SessionConfig};
 
 /// Optimizer that introduces repartition to introduce more
@@ -191,7 +193,7 @@ fn optimize_partitions(
                 )
             })
             .collect::<Result<_>>()?;
-        plan.with_new_children(children)?
+        with_new_children_if_necessary(plan, children)?
     };
 
     // decide if we should bother trying to repartition the output of this plan

datafusion/core/src/physical_optimizer/utils.rs

@@ -21,7 +21,7 @@ use super::optimizer::PhysicalOptimizerRule;
 use crate::execution::context::SessionConfig;
 
 use crate::error::Result;
-use crate::physical_plan::ExecutionPlan;
+use crate::physical_plan::{with_new_children_if_necessary, ExecutionPlan};
 use std::sync::Arc;
 
 /// Convenience rule for writing optimizers: recursively invoke
@@ -42,6 +42,6 @@ pub fn optimize_children(
     if children.is_empty() {
         Ok(Arc::clone(&plan))
     } else {
-        plan.with_new_children(children)
+        with_new_children_if_necessary(plan, children)
     }
 }

datafusion/core/src/physical_plan/analyze.rs

@@ -92,21 +92,14 @@ impl ExecutionPlan for AnalyzeExec {
     }
 
     fn with_new_children(
-        &self,
+        self: Arc<Self>,
         mut children: Vec<Arc<dyn ExecutionPlan>>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        if children.len() == 1 {
-            Ok(Arc::new(Self::new(
-                self.verbose,
-                children.pop().unwrap(),
-                self.schema.clone(),
-            )))
-        } else {
-            Err(DataFusionError::Internal(format!(
-                "Invalid child count for AnalyzeExec. Expected 1 got {}",
-                children.len()
-            )))
-        }
+        Ok(Arc::new(Self::new(
+            self.verbose,
+            children.pop().unwrap(),
+            self.schema.clone(),
+        )))
     }
 
     async fn execute(

datafusion/core/src/physical_plan/coalesce_batches.rs

@@ -23,7 +23,7 @@ use std::pin::Pin;
 use std::sync::Arc;
 use std::task::{Context, Poll};
 
-use crate::error::{DataFusionError, Result};
+use crate::error::Result;
 use crate::physical_plan::{
     DisplayFormatType, ExecutionPlan, Partitioning, RecordBatchStream,
     SendableRecordBatchStream,
@@ -107,18 +107,13 @@ impl ExecutionPlan for CoalesceBatchesExec {
     }
 
     fn with_new_children(
-        &self,
+        self: Arc<Self>,
         children: Vec<Arc<dyn ExecutionPlan>>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        match children.len() {
-            1 => Ok(Arc::new(CoalesceBatchesExec::new(
-                children[0].clone(),
-                self.target_batch_size,
-            ))),
-            _ => Err(DataFusionError::Internal(
-                "CoalesceBatchesExec wrong number of children".to_string(),
-            )),
-        }
+        Ok(Arc::new(CoalesceBatchesExec::new(
+            children[0].clone(),
+            self.target_batch_size,
+        )))
     }
 
     async fn execute(

datafusion/core/src/physical_plan/coalesce_partitions.rs

@@ -96,15 +96,10 @@ impl ExecutionPlan for CoalescePartitionsExec {
     }
 
     fn with_new_children(
-        &self,
+        self: Arc<Self>,
         children: Vec<Arc<dyn ExecutionPlan>>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        match children.len() {
-            1 => Ok(Arc::new(CoalescePartitionsExec::new(children[0].clone()))),
-            _ => Err(DataFusionError::Internal(
-                "CoalescePartitionsExec wrong number of children".to_string(),
-            )),
-        }
+        Ok(Arc::new(CoalescePartitionsExec::new(children[0].clone())))
     }
 
     async fn execute(

datafusion/core/src/physical_plan/cross_join.rs

@@ -32,10 +32,7 @@ use super::{
     coalesce_partitions::CoalescePartitionsExec, join_utils::check_join_is_valid,
     ColumnStatistics, Statistics,
 };
-use crate::{
-    error::{DataFusionError, Result},
-    scalar::ScalarValue,
-};
+use crate::{error::Result, scalar::ScalarValue};
 use async_trait::async_trait;
 use std::time::Instant;
 
@@ -120,18 +117,13 @@ impl ExecutionPlan for CrossJoinExec {
     }
 
     fn with_new_children(
-        &self,
+        self: Arc<Self>,
         children: Vec<Arc<dyn ExecutionPlan>>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        match children.len() {
-            2 => Ok(Arc::new(CrossJoinExec::try_new(
-                children[0].clone(),
-                children[1].clone(),
-            )?)),
-            _ => Err(DataFusionError::Internal(
-                "CrossJoinExec wrong number of children".to_string(),
-            )),
-        }
+        Ok(Arc::new(CrossJoinExec::try_new(
+            children[0].clone(),
+            children[1].clone(),
+        )?))
     }
 
     fn output_partitioning(&self) -> Partitioning {