Skip to content

Commit

Permalink
Prevent take_optimizable from discarding arbitrary plan node
Browse files Browse the repository at this point in the history
`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.
  • Loading branch information
findepi committed Oct 21, 2024
1 parent 1e046f5 commit e4a6bcf
Showing 1 changed file with 7 additions and 15 deletions.
22 changes: 7 additions & 15 deletions datafusion/physical-optimizer/src/aggregate_statistics.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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));
}
}
}
Expand Down

0 comments on commit e4a6bcf

Please sign in to comment.