uqa_sql/semantics/
aggregates.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Aggregate discovery and expression equivalence for SQL grouping.
8
9use crate::plan::{AggregateClassifier, ProjectionPlan, QueryBlockPlan};
10use crate::ScalarExpr;
11use uqa_core::Value;
12
13#[cfg(test)]
14mod tests;
15
16pub fn exprs_match(lhs: &ScalarExpr, rhs: &ScalarExpr) -> bool {
17    match (lhs, rhs) {
18        (ScalarExpr::Star, ScalarExpr::Star) => true,
19        (ScalarExpr::Column(a), ScalarExpr::Column(b)) => a == b,
20        (
21            ScalarExpr::QualifiedColumn {
22                qualifier: aq,
23                column: ac,
24                ..
25            },
26            ScalarExpr::QualifiedColumn {
27                qualifier: bq,
28                column: bc,
29                ..
30            },
31        ) => aq == bq && ac == bc,
32        (ScalarExpr::Column(c), ScalarExpr::QualifiedColumn { column, .. })
33        | (ScalarExpr::QualifiedColumn { column, .. }, ScalarExpr::Column(c)) => c == column,
34        (ScalarExpr::Literal(a), ScalarExpr::Literal(b)) => literals_equal(a, b),
35        (
36            ScalarExpr::TypedLiteral {
37                value: av,
38                ty: at,
39                bound_type: ab,
40                parameter_index: ap,
41            },
42            ScalarExpr::TypedLiteral {
43                value: bv,
44                ty: bt,
45                bound_type: bb,
46                parameter_index: bp,
47            },
48        ) => at == bt && ab == bb && ap == bp && literals_equal(av, bv),
49        (ScalarExpr::Param(a), ScalarExpr::Param(b)) => a == b,
50        (ScalarExpr::Position(a), ScalarExpr::Position(b)) => a == b,
51        (ScalarExpr::InternalColumn(a), ScalarExpr::InternalColumn(b)) => a == b,
52        (
53            ScalarExpr::Func {
54                name: an,
55                binding: ab,
56                args: aa,
57                distinct: ad,
58                order_by: ao,
59                filter: af,
60            },
61            ScalarExpr::Func {
62                name: bn,
63                binding: bb,
64                args: ba,
65                distinct: bd,
66                order_by: bo,
67                filter: bf,
68            },
69        ) => {
70            an.eq_ignore_ascii_case(bn)
71                && ab == bb
72                && ad == bd
73                && aa.len() == ba.len()
74                && aa.iter().zip(ba.iter()).all(|(x, y)| exprs_match(x, y))
75                && ao.len() == bo.len()
76                && ao.iter().zip(bo.iter()).all(|(x, y)| {
77                    x.descending == y.descending
78                        && x.nulls == y.nulls
79                        && exprs_match(&x.expr, &y.expr)
80                })
81                && match (af.as_deref(), bf.as_deref()) {
82                    (None, None) => true,
83                    (Some(x), Some(y)) => exprs_match(x, y),
84                    _ => false,
85                }
86        }
87        (
88            ScalarExpr::Binary {
89                op: ao,
90                lhs: al,
91                rhs: ar,
92            },
93            ScalarExpr::Binary {
94                op: bo,
95                lhs: bl,
96                rhs: br,
97            },
98        ) => ao == bo && exprs_match(al, bl) && exprs_match(ar, br),
99        (ScalarExpr::And(a), ScalarExpr::And(b)) | (ScalarExpr::Or(a), ScalarExpr::Or(b)) => {
100            a.len() == b.len() && a.iter().zip(b.iter()).all(|(x, y)| exprs_match(x, y))
101        }
102        (ScalarExpr::Not(a), ScalarExpr::Not(b))
103        | (ScalarExpr::UnaryMinus(a), ScalarExpr::UnaryMinus(b)) => exprs_match(a, b),
104        (ScalarExpr::Cast { expr: a, ty: at }, ScalarExpr::Cast { expr: b, ty: bt }) => {
105            at == bt && exprs_match(a, b)
106        }
107        _ => false,
108    }
109}
110
111pub fn literals_equal(a: &Value, b: &Value) -> bool {
112    a.has_same_representation(b)
113}
114
115pub fn has_aggregate(aggregates: &dyn AggregateClassifier, projections: &[ProjectionPlan]) -> bool {
116    projections
117        .iter()
118        .any(|p| contains_aggregate(aggregates, &p.expr))
119}
120
121pub fn is_aggregate(aggregates: &dyn AggregateClassifier, expr: &ScalarExpr) -> bool {
122    is_builtin_aggregate(expr)
123        || matches!(expr, ScalarExpr::Func { name, .. } if aggregates.is_registered_aggregate(name))
124}
125
126pub use super::is_builtin_aggregate;
127
128pub fn aggregate_exprs<'a>(
129    aggregates: &dyn AggregateClassifier,
130    projections: &'a [ProjectionPlan],
131) -> Vec<&'a ScalarExpr> {
132    let mut out = Vec::new();
133    for projection in projections {
134        collect_aggregate_exprs(aggregates, &projection.expr, &mut out);
135    }
136    out
137}
138
139/// Aggregate states needed by a query block, in accumulator order.
140///
141/// Projection aggregates remain first (including repeated expressions) because
142/// projection rewriting consumes them positionally. Aggregates referenced only
143/// by HAVING are appended once as hidden targets, matching `PostgreSQL`'s rule
144/// that HAVING need not expose an aggregate in the SELECT list.
145pub fn aggregate_targets<'a>(
146    aggregates: &dyn AggregateClassifier,
147    statement: &'a QueryBlockPlan,
148) -> Vec<&'a ScalarExpr> {
149    let mut targets = aggregate_exprs(aggregates, &statement.projections);
150    if let Some(having) = statement.having.as_ref() {
151        let mut hidden = Vec::new();
152        collect_aggregate_exprs(aggregates, having, &mut hidden);
153        for aggregate in hidden {
154            if !targets
155                .iter()
156                .any(|existing| exprs_match(existing, aggregate))
157            {
158                targets.push(aggregate);
159            }
160        }
161    }
162    targets
163}
164
165pub fn collect_aggregate_exprs<'a>(
166    aggregates: &dyn AggregateClassifier,
167    expr: &'a ScalarExpr,
168    out: &mut Vec<&'a ScalarExpr>,
169) {
170    if is_aggregate(aggregates, expr) {
171        out.push(expr);
172        return;
173    }
174    match expr {
175        ScalarExpr::Func { args, filter, .. } => {
176            for arg in args {
177                collect_aggregate_exprs(aggregates, arg, out);
178            }
179            if let Some(filter) = filter.as_deref() {
180                collect_aggregate_exprs(aggregates, filter, out);
181            }
182        }
183        ScalarExpr::Array(items)
184        | ScalarExpr::Row(items)
185        | ScalarExpr::And(items)
186        | ScalarExpr::Or(items) => {
187            for item in items {
188                collect_aggregate_exprs(aggregates, item, out);
189            }
190        }
191        ScalarExpr::Binary { lhs, rhs, .. } => {
192            collect_aggregate_exprs(aggregates, lhs, out);
193            collect_aggregate_exprs(aggregates, rhs, out);
194        }
195        ScalarExpr::Not(inner)
196        | ScalarExpr::UnaryMinus(inner)
197        | ScalarExpr::Cast { expr: inner, .. } => {
198            collect_aggregate_exprs(aggregates, inner, out);
199        }
200        ScalarExpr::IsNull { expr, .. } => collect_aggregate_exprs(aggregates, expr, out),
201        ScalarExpr::Between { expr, low, high } => {
202            collect_aggregate_exprs(aggregates, expr, out);
203            collect_aggregate_exprs(aggregates, low, out);
204            collect_aggregate_exprs(aggregates, high, out);
205        }
206        ScalarExpr::InList { expr, list, .. } => {
207            collect_aggregate_exprs(aggregates, expr, out);
208            for item in list {
209                collect_aggregate_exprs(aggregates, item, out);
210            }
211        }
212        ScalarExpr::Case {
213            base,
214            when,
215            else_branch,
216        } => {
217            if let Some(base) = base.as_deref() {
218                collect_aggregate_exprs(aggregates, base, out);
219            }
220            for (condition, result) in when {
221                collect_aggregate_exprs(aggregates, condition, out);
222                collect_aggregate_exprs(aggregates, result, out);
223            }
224            if let Some(else_branch) = else_branch.as_deref() {
225                collect_aggregate_exprs(aggregates, else_branch, out);
226            }
227        }
228        ScalarExpr::InSubquery { expr, .. } => collect_aggregate_exprs(aggregates, expr, out),
229        ScalarExpr::Default
230        | ScalarExpr::Star
231        | ScalarExpr::QualifiedStar(_)
232        | ScalarExpr::Column(_)
233        | ScalarExpr::Position(_)
234        | ScalarExpr::InternalColumn(_)
235        | ScalarExpr::QualifiedColumn { .. }
236        | ScalarExpr::Literal(_)
237        | ScalarExpr::TypedLiteral { .. }
238        | ScalarExpr::Param(_)
239        | ScalarExpr::WindowCall { .. }
240        | ScalarExpr::ScalarSubquery(_)
241        | ScalarExpr::Exists { .. } => {}
242    }
243}
244
245pub fn contains_aggregate(aggregates: &dyn AggregateClassifier, expr: &ScalarExpr) -> bool {
246    let mut found = Vec::new();
247    collect_aggregate_exprs(aggregates, expr, &mut found);
248    !found.is_empty()
249}
250
251/// Collect the top-level column names an expression reads. Returns
252/// `false` when the expression can reach arbitrary fields (`*`,
253/// subqueries, window calls), in which case callers must materialise
254/// whole documents.
255pub fn expr_references_columns(expr: &ScalarExpr) -> bool {
256    match expr {
257        ScalarExpr::Star
258        | ScalarExpr::QualifiedStar(_)
259        | ScalarExpr::Column(_)
260        | ScalarExpr::Position(_)
261        | ScalarExpr::InternalColumn(_)
262        | ScalarExpr::QualifiedColumn { .. } => true,
263        ScalarExpr::Func { args, filter, .. } => {
264            args.iter().any(expr_references_columns)
265                || filter.as_deref().is_some_and(expr_references_columns)
266        }
267        ScalarExpr::Array(items)
268        | ScalarExpr::Row(items)
269        | ScalarExpr::And(items)
270        | ScalarExpr::Or(items) => items.iter().any(expr_references_columns),
271        ScalarExpr::Binary { lhs, rhs, .. } => {
272            expr_references_columns(lhs) || expr_references_columns(rhs)
273        }
274        ScalarExpr::Not(inner)
275        | ScalarExpr::UnaryMinus(inner)
276        | ScalarExpr::Cast { expr: inner, .. } => expr_references_columns(inner),
277        ScalarExpr::IsNull { expr, .. } => expr_references_columns(expr),
278        ScalarExpr::Between { expr, low, high } => {
279            expr_references_columns(expr)
280                || expr_references_columns(low)
281                || expr_references_columns(high)
282        }
283        ScalarExpr::InList { expr, list, .. } => {
284            expr_references_columns(expr) || list.iter().any(expr_references_columns)
285        }
286        ScalarExpr::WindowCall { args, .. } => args.iter().any(expr_references_columns),
287        ScalarExpr::Case {
288            base,
289            when,
290            else_branch,
291        } => {
292            base.as_deref().is_some_and(expr_references_columns)
293                || when.iter().any(|(condition, result)| {
294                    expr_references_columns(condition) || expr_references_columns(result)
295                })
296                || else_branch.as_deref().is_some_and(expr_references_columns)
297        }
298        ScalarExpr::InSubquery { expr, .. } => expr_references_columns(expr),
299        ScalarExpr::ScalarSubquery(_) | ScalarExpr::Exists { .. } => true,
300        ScalarExpr::Default
301        | ScalarExpr::Literal(_)
302        | ScalarExpr::TypedLiteral { .. }
303        | ScalarExpr::Param(_) => false,
304    }
305}
306
307mod slots;
308pub use slots::{
309    aggregate_slot_index, compile_group_slots, compile_having_aggregate_slots,
310    compile_projection_aggregate_slots, select_grouping_set,
311};
uqa_sql/semantics/aggregates.rs

uqa_sql/semantics/
aggregates.rs
