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

uqa_sql/semantics/
aggregates.rs
