Skip to main content

akar_processor/processor/mapper/
map_aggregate.rs

1use super::ExecutionContext;
2use crate::physical::order_aggregate::resolve_group_by_indices;
3use crate::physical_operator::*;
4use akar_common::error::ProcessorError;
5use akar_common::vector::DataChunk;
6use akar_parser::ast::Expression;
7use akar_planner::logical_operator::LogicalOperator;
8
9pub fn map_and_execute_aggregate(
10    op: &LogicalOperator,
11    current_input: Vec<DataChunk>,
12    ctx: &mut ExecutionContext,
13) -> Result<Vec<DataChunk>, ProcessorError> {
14    match op {
15        LogicalOperator::Aggregate(a) => {
16            let funcs: Vec<akar_function::AggregateFunction> = a
17                .aggregates
18                .iter()
19                .map(|(n, args)| {
20                    // Detect COUNT(*) from Star arg: override name to get CountStar
21                    let effective_name = if n == "COUNT" && args.iter().any(|e| matches!(e, Expression::Star)) {
22                        "COUNT(*)"
23                    } else {
24                        n
25                    };
26                    crate::physical::order_aggregate::parse_aggregate_function(effective_name)
27                })
28                .collect();
29            let agg_expressions: Vec<Vec<Expression>> = a.aggregates.iter().map(|(_, args)| args.clone()).collect();
30
31            // Resolve GROUP BY expressions to actual column indices using input field_names
32            let field_names = current_input.first().map(|c| c.field_names.as_slice()).unwrap_or(&[]);
33            let group_by_cols = if a.group_by.is_empty() {
34                Vec::new()
35            } else {
36                resolve_group_by_indices(&a.group_by, field_names)
37            };
38
39            let shared_state = std::sync::Arc::new(crate::physical::order_aggregate::SharedAggregateState::new(
40                funcs,
41                group_by_cols,
42                agg_expressions,
43            ));
44
45            let agg_scan = crate::physical::order_aggregate::PhysicalAggregateScan {
46                shared_state: shared_state.clone(),
47            };
48            let agg_finalize = crate::physical::order_aggregate::PhysicalAggregateFinalize { shared_state };
49
50            // Phase 1: Scan and accumulate (returns empty chunk in sequential push-down)
51            let _ = agg_scan.execute(current_input)?;
52
53            // Phase 2: Finalize and yield grouped chunks
54            let result = agg_finalize.execute(vec![])?;
55
56            // P52.56: the aggregate output chunks carried no field_names, so
57            // result columns were positional-only and an alias (`AS cnt`) never
58            // reached the result. Propagate group-by variable names + aggregate
59            // function names onto the output chunks. A Projection above the
60            // aggregate still resolves positionally when a name doesn't match,
61            // so this is safe.
62            let names = aggregate_field_names(a);
63            let result: Vec<DataChunk> = result
64                .into_iter()
65                .map(|chunk| chunk.with_names(names.clone()))
66                .collect();
67
68            Ok(result)
69        }
70        LogicalOperator::CountRelTable(crt) => {
71            let physical = PhysicalCountRelTable {
72                table_name: crt.table_name.clone(),
73                table_id: crt.table_id,
74                table_catalog: ctx.table_catalog.clone(),
75            };
76            let result = physical.execute(vec![])?;
77            Ok(result)
78        }
79        _ => Err(format!("Not an aggregate operator: {:?}", op).into()),
80    }
81}
82
83/// Build output field names for an aggregate result: the group-by variable
84/// names followed by the aggregate function names (P52.56). Group-by naming
85/// mirrors `expression_field_name` in map_projection.rs so the projection above
86/// the aggregate resolves columns by name (P53.16).
87fn aggregate_field_names(a: &akar_planner::logical_operator::LogicalAggregate) -> Vec<String> {
88    let mut names: Vec<String> = a
89        .group_by
90        .iter()
91        .map(|e| match e {
92            Expression::Variable(v) => v.clone(),
93            Expression::PropertyAccess(obj, prop) => {
94                if let Expression::Variable(var) = &**obj {
95                    format!("{var}.{prop}")
96                } else {
97                    prop.clone()
98                }
99            }
100            other => format!("{other:?}"),
101        })
102        .collect();
103    for (fname, args) in &a.aggregates {
104        let effective = if fname == "COUNT" && args.iter().any(|e| matches!(e, Expression::Star)) {
105            "COUNT(*)".to_string()
106        } else if args.len() == 1 {
107            match &args[0] {
108                Expression::Variable(v) => format!("{fname}({v})"),
109                Expression::PropertyAccess(obj, prop) => {
110                    if let Expression::Variable(base) = &**obj {
111                        format!("{fname}({base}.{prop})")
112                    } else {
113                        format!("{fname}({prop})")
114                    }
115                }
116                Expression::Star => format!("{fname}(*)"),
117                _ => fname.clone(),
118            }
119        } else {
120            fname.clone()
121        };
122        names.push(effective);
123    }
124    names
125}