Prevent take_optimizable from discarding arbitrary plan node

`take_optimizable` started from inspecting top level plan node (it
should be final aggregation) and then descended trying to find matching
partial aggregation. When doing so, it would ignore any single-source
nodes it passes by. As a result, `AggregateStatistics` could change the
plan in an undesired manner.

Author: findepi

datafusion/physical-optimizer/src/aggregate_statistics.rs

@@ -115,22 +115,14 @@ fn take_optimizable(node: &dyn ExecutionPlan) -> Option<Arc<dyn ExecutionPlan>>
         if !final_agg_exec.mode().is_first_stage()
             && final_agg_exec.group_expr().is_empty()
         {
-            let mut child = Arc::clone(final_agg_exec.input());
-            loop {
-                if let Some(partial_agg_exec) =
-                    child.as_any().downcast_ref::<AggregateExec>()
+            let child = final_agg_exec.input();
+            if let Some(partial_agg_exec) = child.as_any().downcast_ref::<AggregateExec>()
+            {
+                if partial_agg_exec.mode().is_first_stage()
+                    && partial_agg_exec.group_expr().is_empty()
+                    && partial_agg_exec.filter_expr().iter().all(|e| e.is_none())
                 {
-                    if partial_agg_exec.mode().is_first_stage()
-                        && partial_agg_exec.group_expr().is_empty()
-                        && partial_agg_exec.filter_expr().iter().all(|e| e.is_none())
-                    {
-                        return Some(child);
-                    }
-                }
-                if let [childrens_child] = child.children().as_slice() {
-                    child = Arc::clone(childrens_child);
-                } else {
-                    break;
+                    return Some(Arc::clone(child));
                 }
             }
         }