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

uqa_sql/semantics/
aggregates.rs
