Skip to main content

nodedb_sql/planner/
agg_bind.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Binding aggregate calls in HAVING and post-aggregate ORDER BY.
4//!
5//! Both clauses are evaluated against finalized group rows, so an aggregate
6//! either mentions must actually be computed and must be addressed by the
7//! column name it lands under. Left as a literal `sum(...)` call, it reaches
8//! the executor as a scalar function over a row that has no such column — the
9//! predicate matches nothing, or the sort key resolves to nothing.
10//!
11//! The two clauses differ in *which* name they need, because they run on
12//! opposite sides of the rename to user aliases:
13//!
14//! - HAVING filters before the rename, so it addresses the canonical key
15//!   (`sum(amount)`).
16//! - ORDER BY sorts after the rename, so it addresses the output name
17//!   (`total`, when the projection said `SUM(amount) AS total`).
18
19use sqlparser::ast;
20
21use crate::aggregate_walk::extract_aggregates;
22use crate::error::{Result, SqlError};
23use crate::functions::registry::{FunctionCategory, FunctionRegistry};
24use crate::parser::normalize::normalize_ident;
25use crate::planner::agg_naming::aggregate_output_key;
26use crate::resolver::columns::TableScope;
27use crate::types::AggregateExpr;
28
29/// Which name an aggregate call is rewritten to.
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub enum BindName {
32    /// The canonical key a group row carries before user aliases are applied.
33    Canonical,
34    /// The output column name the client sees, after that rename.
35    Output,
36}
37
38/// Rewrite every aggregate call in `expr` to a reference to its computed
39/// column, registering aggregates the projection did not already request.
40///
41/// `scope` carries the input relations each aggregate argument resolves
42/// against.
43pub fn bind_aggregate_calls(
44    expr: &ast::Expr,
45    projection: &[ast::SelectItem],
46    aggregates: &mut Vec<AggregateExpr>,
47    functions: &FunctionRegistry,
48    name: BindName,
49    scope: &TableScope,
50) -> Result<ast::Expr> {
51    let resolved = resolve_select_aliases(expr, projection);
52    bind(&resolved, aggregates, functions, name, scope)
53}
54
55/// Substitute SELECT-list output aliases referenced by the expression.
56///
57/// `SELECT SUM(amount) AS total ... HAVING total > 0` names an output column
58/// that does not exist yet when the clause runs. Replacing it with the
59/// underlying expression lets the aggregate binding below address the column
60/// that does exist.
61fn resolve_select_aliases(expr: &ast::Expr, projection: &[ast::SelectItem]) -> ast::Expr {
62    match expr {
63        ast::Expr::Identifier(ident) => {
64            let needle = normalize_ident(ident);
65            projection
66                .iter()
67                .find_map(|item| match item {
68                    ast::SelectItem::ExprWithAlias { expr, alias }
69                        if normalize_ident(alias) == needle =>
70                    {
71                        Some(expr.clone())
72                    }
73                    _ => None,
74                })
75                .unwrap_or_else(|| expr.clone())
76        }
77        ast::Expr::BinaryOp { left, op, right } => ast::Expr::BinaryOp {
78            left: Box::new(resolve_select_aliases(left, projection)),
79            op: op.clone(),
80            right: Box::new(resolve_select_aliases(right, projection)),
81        },
82        ast::Expr::UnaryOp { op, expr } => ast::Expr::UnaryOp {
83            op: *op,
84            expr: Box::new(resolve_select_aliases(expr, projection)),
85        },
86        ast::Expr::Nested(inner) => {
87            ast::Expr::Nested(Box::new(resolve_select_aliases(inner, projection)))
88        }
89        other => other.clone(),
90    }
91}
92
93fn bind(
94    expr: &ast::Expr,
95    aggregates: &mut Vec<AggregateExpr>,
96    functions: &FunctionRegistry,
97    name: BindName,
98    scope: &TableScope,
99) -> Result<ast::Expr> {
100    match expr {
101        ast::Expr::Function(func) if is_aggregate_call(func, functions) => {
102            let column = register_aggregate(expr, aggregates, functions, name, scope)?;
103            Ok(ast::Expr::Identifier(ast::Ident::new(column)))
104        }
105        ast::Expr::BinaryOp { left, op, right } => Ok(ast::Expr::BinaryOp {
106            left: Box::new(bind(left, aggregates, functions, name, scope)?),
107            op: op.clone(),
108            right: Box::new(bind(right, aggregates, functions, name, scope)?),
109        }),
110        ast::Expr::UnaryOp { op, expr } => Ok(ast::Expr::UnaryOp {
111            op: *op,
112            expr: Box::new(bind(expr, aggregates, functions, name, scope)?),
113        }),
114        ast::Expr::Nested(inner) => Ok(ast::Expr::Nested(Box::new(bind(
115            inner, aggregates, functions, name, scope,
116        )?))),
117        ast::Expr::IsNull(inner) => Ok(ast::Expr::IsNull(Box::new(bind(
118            inner, aggregates, functions, name, scope,
119        )?))),
120        ast::Expr::IsNotNull(inner) => Ok(ast::Expr::IsNotNull(Box::new(bind(
121            inner, aggregates, functions, name, scope,
122        )?))),
123        ast::Expr::Between {
124            expr,
125            negated,
126            low,
127            high,
128        } => Ok(ast::Expr::Between {
129            expr: Box::new(bind(expr, aggregates, functions, name, scope)?),
130            negated: *negated,
131            low: Box::new(bind(low, aggregates, functions, name, scope)?),
132            high: Box::new(bind(high, aggregates, functions, name, scope)?),
133        }),
134        other => Ok(other.clone()),
135    }
136}
137
138pub(super) fn is_aggregate_call(func: &ast::Function, functions: &FunctionRegistry) -> bool {
139    let name = func
140        .name
141        .0
142        .iter()
143        .map(|part| match part {
144            ast::ObjectNamePart::Identifier(ident) => normalize_ident(ident),
145            _ => String::new(),
146        })
147        .collect::<Vec<_>>()
148        .join(".");
149    matches!(
150        functions.lookup(&name).map(|m| m.category),
151        Some(FunctionCategory::Aggregate)
152    )
153}
154
155/// Ensure the aggregate `expr` is in the computed list and return the column
156/// name it should be addressed by.
157fn register_aggregate(
158    expr: &ast::Expr,
159    aggregates: &mut Vec<AggregateExpr>,
160    functions: &FunctionRegistry,
161    name: BindName,
162    scope: &TableScope,
163) -> Result<String> {
164    // The alias is replaced below for a newly registered aggregate, so the
165    // placeholder is never observable.
166    let mut extracted = extract_aggregates(expr, "", functions, scope)?;
167    let Some(mut agg) = extracted.pop() else {
168        return Err(SqlError::Unsupported {
169            detail: format!("aggregate `{expr}` could not be extracted"),
170        });
171    };
172    if !extracted.is_empty() {
173        return Err(SqlError::Unsupported {
174            detail: format!("nested aggregates are not supported: `{expr}`"),
175        });
176    }
177
178    let key = aggregate_output_key(&agg);
179
180    // An aggregate the projection already computes needs no second entry. Its
181    // output name is its alias, which is what a post-rename caller must use.
182    if let Some(existing) = aggregates
183        .iter()
184        .find(|existing| aggregate_output_key(existing) == key)
185    {
186        return Ok(match name {
187            BindName::Canonical => key,
188            BindName::Output => existing.alias.clone(),
189        });
190    }
191
192    // Carry the canonical key as the alias so no user-facing rename is
193    // attached: an aggregate introduced only by HAVING or ORDER BY is an input
194    // to filtering or sorting, not an output column the client asked for. Both
195    // names therefore coincide.
196    agg.alias = key.clone();
197    aggregates.push(agg);
198    Ok(key)
199}