From e4a6bcf7e04c483059179e22d7736a519abbb67d Mon Sep 17 00:00:00 2001 From: Piotr Findeisen Date: Mon, 21 Oct 2024 11:29:33 +0200 Subject: [PATCH] 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. --- .../src/aggregate_statistics.rs | 22 ++++++------------- 1 file changed, 7 insertions(+), 15 deletions(-) diff --git a/datafusion/physical-optimizer/src/aggregate_statistics.rs b/datafusion/physical-optimizer/src/aggregate_statistics.rs index fd21362fd3eb9..5cee328cfbae2 100644 --- a/datafusion/physical-optimizer/src/aggregate_statistics.rs +++ b/datafusion/physical-optimizer/src/aggregate_statistics.rs @@ -115,22 +115,14 @@ fn take_optimizable(node: &dyn ExecutionPlan) -> Option> 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::() + let child = final_agg_exec.input(); + if let Some(partial_agg_exec) = child.as_any().downcast_ref::() + { + 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)); } } }