[GitHub] [arrow] andygrove commented on a change in pull request #8034: ARROW-9464: [Rust] [DataFusion] Physical plan optimization rule to insert MergeExec when needed

[GitBox <gi...@apache.org>]

andygrove commented on a change in pull request #8034:
URL: https://github.com/apache/arrow/pull/8034#discussion_r476475961



##########
File path: rust/datafusion/src/execution/physical_plan/planner.rs
##########
@@ -61,6 +61,55 @@ impl PhysicalPlanner for DefaultPhysicalPlanner {
         &self,
         logical_plan: &LogicalPlan,
         ctx_state: &ExecutionContextState,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        let plan = self.create_initial_plan(logical_plan, ctx_state)?;
+        self.optimize_plan(plan, ctx_state)
+    }
+}
+
+impl DefaultPhysicalPlanner {
+    /// Create a physical plan from a logical plan
+    fn optimize_plan(
+        &self,
+        plan: Arc<dyn ExecutionPlan>,
+        ctx_state: &ExecutionContextState,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        let children = plan
+            .children()
+            .iter()
+            .map(|child| self.optimize_plan(child.clone(), ctx_state))
+            .collect::<Result<Vec<_>>>()?;
+
+        if children.len() == 0 {
+            // leaf node, children cannot be replaced
+            Ok(plan.clone())
+        } else {
+            match plan.required_child_distribution() {
+                Distribution::UnspecifiedDistribution => plan.with_new_children(children),
+                Distribution::SinglePartition => plan.with_new_children(

Review comment:
       I can't take credit for the cleverness unfortunately, since I stole all of this from Apache Spark. 




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org