nodedb_sql/planner/
aggregate_order.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Records the SELECT-list interleaving of GROUP BY keys and aggregate
4//! expressions so the output-schema builder emits aggregate result columns in
5//! the user's SELECT order rather than a hardcoded group-keys-first order.
6
7use sqlparser::ast;
8
9use crate::error::Result;
10use crate::functions::registry::FunctionRegistry;
11use crate::parser::normalize::normalize_ident;
12use crate::planner::aggregate::{
13    expr_column_name, extract_aggregates_from_projection, function_args_exprs, key_column_name,
14    normalize_function_name,
15};
16use crate::resolver::expr::convert_expr;
17use crate::types::query::AggOutputSlot;
18use crate::types_expr::SqlExpr;
19
20/// Walk `projection` in output order and classify each item as a GROUP BY key
21/// slot or an aggregate slot, producing the interleaving that `output_schema`
22/// replays to emit columns in SELECT-list order.
23///
24/// `group_by` is the canonical key list; [`AggOutputSlot::GroupKey`] indexes
25/// it. Real aggregates are numbered in projection order from 0;
26/// `GROUPING(col)` pseudo-aggregates are numbered after all real aggregates,
27/// mirroring how the planner appends them to the aggregate list. Items that
28/// are neither a bare group key nor an aggregate (a bare non-key expression)
29/// produce no slot — they have no output column in the aggregate result today.
30pub fn compute_output_order(
31    projection: &[ast::SelectItem],
32    group_by: &[SqlExpr],
33    functions: &FunctionRegistry,
34) -> Result<Vec<AggOutputSlot>> {
35    let real_agg_count = extract_aggregates_from_projection(projection, functions)?.len();
36    let mut order = Vec::new();
37    let mut agg_cursor = 0usize;
38    let mut grouping_cursor = 0usize;
39    for item in projection {
40        let expr = match item {
41            ast::SelectItem::UnnamedExpr(expr) => expr,
42            ast::SelectItem::ExprWithAlias { expr, .. } => expr,
43            _ => continue,
44        };
45        // Bare column that names one of the GROUP BY keys.
46        if let Some(name) = expr_column_name(expr)
47            && let Some(index) = group_by
48                .iter()
49                .position(|key| key_column_name(key) == Some(name.as_str()))
50        {
51            order.push(AggOutputSlot::GroupKey(index));
52            continue;
53        }
54        // Computed expression that structurally matches a GROUP BY key
55        // (e.g. `UPPER(label)` selected AND grouped). It is neither a bare
56        // column nor an aggregate, so it must be classified before the
57        // grouping / aggregate checks below. `SqlExpr` has no `PartialEq`
58        // (it can carry a subquery plan), so compare the canonical `Debug`
59        // rendering — both the projection expr and the GROUP BY key pass
60        // through the same `convert_expr`, so equal expressions render
61        // identically.
62        if let Ok(converted) = convert_expr(expr) {
63            let rendered = format!("{converted:?}");
64            if let Some(index) = group_by
65                .iter()
66                .position(|key| format!("{key:?}") == rendered)
67            {
68                order.push(AggOutputSlot::GroupKey(index));
69                continue;
70            }
71        }
72        // `GROUPING(col)` pseudo-aggregates: not in the function registry, so
73        // `contains_aggregate` misses them, and the planner appends them after
74        // every real aggregate.
75        let grouping_count = count_grouping_calls(expr);
76        if grouping_count > 0 {
77            for _ in 0..grouping_count {
78                order.push(AggOutputSlot::Aggregate(real_agg_count + grouping_cursor));
79                grouping_cursor += 1;
80            }
81            continue;
82        }
83        // Ordinary aggregate expression(s).
84        if crate::aggregate_walk::contains_aggregate(expr, functions) {
85            let alias = match item {
86                ast::SelectItem::ExprWithAlias { alias, .. } => normalize_ident(alias),
87                _ => format!("{expr}").to_lowercase(),
88            };
89            let produced =
90                crate::aggregate_walk::extract_aggregates(expr, &alias, functions)?.len();
91            for _ in 0..produced {
92                order.push(AggOutputSlot::Aggregate(agg_cursor));
93                agg_cursor += 1;
94            }
95        }
96    }
97    Ok(order)
98}
99
100/// Count the `GROUPING(col)` calls reachable in `expr`, mirroring the planner's
101/// grouping extraction: each argument of a `GROUPING(...)` call yields one
102/// pseudo-aggregate, and binary operators are traversed on both sides.
103fn count_grouping_calls(expr: &ast::Expr) -> usize {
104    match expr {
105        ast::Expr::Function(f) if normalize_function_name(f).eq_ignore_ascii_case("grouping") => {
106            function_args_exprs(f).len()
107        }
108        ast::Expr::BinaryOp { left, right, .. } => {
109            count_grouping_calls(left) + count_grouping_calls(right)
110        }
111        _ => 0,
112    }
113}
nodedb_sql/planner/aggregate_order.rs

nodedb_sql/planner/
aggregate_order.rs
