1#include "duckdb/optimizer/common_aggregate_optimizer.hpp"
2
3#include "duckdb/planner/expression/bound_columnref_expression.hpp"
4#include "duckdb/planner/operator/logical_aggregate.hpp"
5#include "duckdb/parser/expression_map.hpp"
6#include "duckdb/planner/column_binding_map.hpp"
7
8namespace duckdb {
9
10void CommonAggregateOptimizer::VisitOperator(LogicalOperator &op) {
11 LogicalOperatorVisitor::VisitOperator(op);
12 switch (op.type) {
13 case LogicalOperatorType::LOGICAL_AGGREGATE_AND_GROUP_BY:
14 ExtractCommonAggregates(aggr&: op.Cast<LogicalAggregate>());
15 break;
16 default:
17 break;
18 }
19}
20
21unique_ptr<Expression> CommonAggregateOptimizer::VisitReplace(BoundColumnRefExpression &expr,
22 unique_ptr<Expression> *expr_ptr) {
23 // check if this column ref points to an aggregate that was remapped; if it does we remap it
24 auto entry = aggregate_map.find(x: expr.binding);
25 if (entry != aggregate_map.end()) {
26 expr.binding = entry->second;
27 }
28 return nullptr;
29}
30
31void CommonAggregateOptimizer::ExtractCommonAggregates(LogicalAggregate &aggr) {
32 expression_map_t<idx_t> aggregate_remap;
33 idx_t total_erased = 0;
34 for (idx_t i = 0; i < aggr.expressions.size(); i++) {
35 idx_t original_index = i + total_erased;
36 auto entry = aggregate_remap.find(x: *aggr.expressions[i]);
37 if (entry == aggregate_remap.end()) {
38 // aggregate does not exist yet: add it to the map
39 aggregate_remap[*aggr.expressions[i]] = i;
40 if (i != original_index) {
41 // this aggregate is not erased, however an agregate BEFORE it has been erased
42 // so we need to remap this aggregaet
43 ColumnBinding original_binding(aggr.aggregate_index, original_index);
44 ColumnBinding new_binding(aggr.aggregate_index, i);
45 aggregate_map[original_binding] = new_binding;
46 }
47 } else {
48 // aggregate already exists! we can remove this entry
49 total_erased++;
50 aggr.expressions.erase(position: aggr.expressions.begin() + i);
51 i--;
52 // we need to remap any references to this aggregate so they point to the other aggregate
53 ColumnBinding original_binding(aggr.aggregate_index, original_index);
54 ColumnBinding new_binding(aggr.aggregate_index, entry->second);
55 aggregate_map[original_binding] = new_binding;
56 }
57 }
58}
59
60} // namespace duckdb
61