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            // P88 DISTINCT aggregates: the parser encodes `COUNT(DISTINCT x)`
17            // as function name `COUNT_DISTINCT`; split back into the base
18            // function + a per-function distinct flag.
19            let mut funcs: Vec<akar_function::AggregateFunction> = Vec::with_capacity(a.aggregates.len());
20            let mut distinct_flags: Vec<bool> = Vec::with_capacity(a.aggregates.len());
21            for (n, args) in &a.aggregates {
22                let (base, distinct) = split_distinct_name(n);
23                distinct_flags.push(distinct);
24                // Detect COUNT(*) from Star arg: override name to get CountStar
25                let effective_name = if base == "COUNT" && args.iter().any(|e| matches!(e, Expression::Star)) {
26                    "COUNT(*)"
27                } else {
28                    base
29                };
30                funcs.push(crate::physical::order_aggregate::parse_aggregate_function(
31                    effective_name,
32                ));
33            }
34            let agg_expressions: Vec<Vec<Expression>> = a.aggregates.iter().map(|(_, args)| args.clone()).collect();
35
36            // Resolve GROUP BY expressions to actual column indices using input field_names
37            let field_names = current_input.first().map(|c| c.field_names.as_slice()).unwrap_or(&[]);
38            let group_by_cols = if a.group_by.is_empty() {
39                Vec::new()
40            } else {
41                resolve_group_by_indices(&a.group_by, field_names)
42            };
43
44            let shared_state = std::sync::Arc::new(crate::physical::order_aggregate::SharedAggregateState::new(
45                funcs,
46                group_by_cols,
47                agg_expressions,
48                distinct_flags,
49            ));
50
51            let agg_scan = crate::physical::order_aggregate::PhysicalAggregateScan {
52                shared_state: shared_state.clone(),
53            };
54            let agg_finalize = crate::physical::order_aggregate::PhysicalAggregateFinalize { shared_state };
55
56            // Phase 1: Scan and accumulate (returns empty chunk in sequential push-down)
57            let _ = agg_scan.execute(current_input)?;
58
59            // Phase 2: Finalize and yield grouped chunks
60            let result = agg_finalize.execute(vec![])?;
61
62            // P52.56: the aggregate output chunks carried no field_names, so
63            // result columns were positional-only and an alias (`AS cnt`) never
64            // reached the result. Propagate group-by variable names + aggregate
65            // function names onto the output chunks. A Projection above the
66            // aggregate still resolves positionally when a name doesn't match,
67            // so this is safe.
68            let names = aggregate_field_names(a);
69            let result: Vec<DataChunk> = result
70                .into_iter()
71                .map(|chunk| chunk.with_names(names.clone()))
72                .collect();
73
74            Ok(result)
75        }
76        LogicalOperator::CountRelTable(crt) => {
77            let physical = PhysicalCountRelTable {
78                table_name: crt.table_name.clone(),
79                table_id: crt.table_id,
80                table_catalog: ctx.table_catalog.clone(),
81            };
82            let result = physical.execute(vec![])?;
83            Ok(result)
84        }
85        _ => Err(format!("Not an aggregate operator: {:?}", op).into()),
86    }
87}
88
89/// Split an aggregate name that may carry the parser's DISTINCT encoding
90/// (P88): `COUNT_DISTINCT` → (`COUNT`, true). Aggregate names reach the
91/// processor uppercased by aggregate_detection.
92fn split_distinct_name(name: &str) -> (&str, bool) {
93    match name.strip_suffix("_DISTINCT") {
94        Some(base) => (base, true),
95        None => (name, false),
96    }
97}
98
99/// Build output field names for an aggregate result: the group-by variable
100/// names followed by the aggregate function names (P52.56). Group-by naming
101/// mirrors `expression_field_name` in map_projection.rs so the projection above
102/// the aggregate resolves columns by name (P53.16).
103fn aggregate_field_names(a: &akar_planner::logical_operator::LogicalAggregate) -> Vec<String> {
104    let mut names: Vec<String> = a
105        .group_by
106        .iter()
107        .map(|e| match e {
108            Expression::Variable(v) => v.clone(),
109            Expression::PropertyAccess(obj, prop) => {
110                if let Expression::Variable(var) = &**obj {
111                    format!("{var}.{prop}")
112                } else {
113                    prop.clone()
114                }
115            }
116            other => format!("{other:?}"),
117        })
118        .collect();
119    for (fname, args) in &a.aggregates {
120        let (fname, _) = split_distinct_name(fname);
121        let effective = if fname == "COUNT" && args.iter().any(|e| matches!(e, Expression::Star)) {
122            "COUNT(*)".to_string()
123        } else if args.len() == 1 {
124            match &args[0] {
125                Expression::Variable(v) => format!("{fname}({v})"),
126                Expression::PropertyAccess(obj, prop) => {
127                    if let Expression::Variable(base) = &**obj {
128                        format!("{fname}({base}.{prop})")
129                    } else {
130                        format!("{fname}({prop})")
131                    }
132                }
133                Expression::Star => format!("{fname}(*)"),
134                _ => fname.to_string(),
135            }
136        } else {
137            fname.to_string()
138        };
139        names.push(effective);
140    }
141    names
142}
akar_processor/processor/mapper/map_aggregate.rs

akar_processor/processor/mapper/
map_aggregate.rs
