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

uqa_sql/semantics/
aggregates.rs
