uqa_sql/semantics/aggregates/
slots.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Assign internal scalar slots to projection and HAVING aggregates.
8
9use super::{exprs_match, is_aggregate};
10use crate::plan::AggregateClassifier;
11use crate::{SQLError, ScalarExpr};
12
13pub fn compile_projection_aggregate_slots(
14    context: &dyn AggregateClassifier,
15    expr: &ScalarExpr,
16    relation: crate::ast::InternalRelationId,
17    cursor: &mut usize,
18) -> Result<ScalarExpr, SQLError> {
19    rewrite_aggregates(context, expr, &mut |_| {
20        let slot = aggregate_slot(relation, *cursor);
21        *cursor += 1;
22        Ok(slot)
23    })
24}
25
26pub fn compile_having_aggregate_slots(
27    context: &dyn AggregateClassifier,
28    expr: &ScalarExpr,
29    relation: crate::ast::InternalRelationId,
30    aggregate_targets: &[ScalarExpr],
31) -> Result<ScalarExpr, SQLError> {
32    rewrite_aggregates(context, expr, &mut |aggregate| {
33        aggregate_targets
34            .iter()
35            .position(|target| exprs_match(target, aggregate))
36            .map(|index| aggregate_slot(relation, index))
37            .ok_or_else(|| {
38                SQLError::Unsupported(
39                    "HAVING references an aggregate that is not in the aggregate plan".into(),
40                )
41            })
42    })
43}
44
45pub fn aggregate_slot_index(
46    column: crate::ast::InternalColumnRef,
47    relation: crate::ast::InternalRelationId,
48) -> Option<usize> {
49    (column.relation() == relation).then(|| column.attribute())
50}
51
52fn aggregate_slot(relation: crate::ast::InternalRelationId, index: usize) -> ScalarExpr {
53    ScalarExpr::InternalColumn(relation.column(index))
54}
55
56#[expect(
57    clippy::too_many_lines,
58    reason = "preserves aggregate NULL and type order"
59)]
60fn rewrite_aggregates(
61    context: &dyn AggregateClassifier,
62    expr: &ScalarExpr,
63    replace: &mut impl FnMut(&ScalarExpr) -> Result<ScalarExpr, SQLError>,
64) -> Result<ScalarExpr, SQLError> {
65    if is_aggregate(context, expr) {
66        return replace(expr);
67    }
68    match expr {
69        ScalarExpr::Func {
70            name,
71            binding,
72            args,
73            distinct,
74            order_by,
75            filter,
76        } => Ok(ScalarExpr::Func {
77            name: name.clone(),
78            binding: binding.clone(),
79            args: args
80                .iter()
81                .map(|arg| rewrite_aggregates(context, arg, replace))
82                .collect::<Result<Vec<_>, _>>()?,
83            distinct: *distinct,
84            order_by: order_by.clone(),
85            filter: filter
86                .as_deref()
87                .map(|filter| rewrite_aggregates(context, filter, replace).map(Box::new))
88                .transpose()?,
89        }),
90        ScalarExpr::Array(items) => Ok(ScalarExpr::Array(
91            items
92                .iter()
93                .map(|item| rewrite_aggregates(context, item, replace))
94                .collect::<Result<Vec<_>, _>>()?,
95        )),
96        ScalarExpr::Row(items) => Ok(ScalarExpr::Row(
97            items
98                .iter()
99                .map(|item| rewrite_aggregates(context, item, replace))
100                .collect::<Result<Vec<_>, _>>()?,
101        )),
102        ScalarExpr::Binary { op, lhs, rhs } => Ok(ScalarExpr::Binary {
103            op: *op,
104            lhs: Box::new(rewrite_aggregates(context, lhs, replace)?),
105            rhs: Box::new(rewrite_aggregates(context, rhs, replace)?),
106        }),
107        ScalarExpr::Not(inner) => Ok(ScalarExpr::Not(Box::new(rewrite_aggregates(
108            context, inner, replace,
109        )?))),
110        ScalarExpr::UnaryMinus(inner) => Ok(ScalarExpr::UnaryMinus(Box::new(rewrite_aggregates(
111            context, inner, replace,
112        )?))),
113        ScalarExpr::And(parts) => Ok(ScalarExpr::And(
114            parts
115                .iter()
116                .map(|part| rewrite_aggregates(context, part, replace))
117                .collect::<Result<Vec<_>, _>>()?,
118        )),
119        ScalarExpr::Or(parts) => Ok(ScalarExpr::Or(
120            parts
121                .iter()
122                .map(|part| rewrite_aggregates(context, part, replace))
123                .collect::<Result<Vec<_>, _>>()?,
124        )),
125        ScalarExpr::IsNull { expr, negated } => Ok(ScalarExpr::IsNull {
126            expr: Box::new(rewrite_aggregates(context, expr, replace)?),
127            negated: *negated,
128        }),
129        ScalarExpr::Between { expr, low, high } => Ok(ScalarExpr::Between {
130            expr: Box::new(rewrite_aggregates(context, expr, replace)?),
131            low: Box::new(rewrite_aggregates(context, low, replace)?),
132            high: Box::new(rewrite_aggregates(context, high, replace)?),
133        }),
134        ScalarExpr::InList {
135            expr,
136            list,
137            negated,
138        } => Ok(ScalarExpr::InList {
139            expr: Box::new(rewrite_aggregates(context, expr, replace)?),
140            list: list
141                .iter()
142                .map(|item| rewrite_aggregates(context, item, replace))
143                .collect::<Result<Vec<_>, _>>()?,
144            negated: *negated,
145        }),
146        ScalarExpr::Case {
147            base,
148            when,
149            else_branch,
150        } => Ok(ScalarExpr::Case {
151            base: base
152                .as_deref()
153                .map(|base| rewrite_aggregates(context, base, replace).map(Box::new))
154                .transpose()?,
155            when: when
156                .iter()
157                .map(|(condition, result)| {
158                    Ok((
159                        rewrite_aggregates(context, condition, replace)?,
160                        rewrite_aggregates(context, result, replace)?,
161                    ))
162                })
163                .collect::<Result<Vec<_>, SQLError>>()?,
164            else_branch: else_branch
165                .as_deref()
166                .map(|branch| rewrite_aggregates(context, branch, replace).map(Box::new))
167                .transpose()?,
168        }),
169        ScalarExpr::Cast { expr, ty } => Ok(ScalarExpr::Cast {
170            expr: Box::new(rewrite_aggregates(context, expr, replace)?),
171            ty: ty.clone(),
172        }),
173        ScalarExpr::InSubquery {
174            expr,
175            subquery,
176            negated,
177        } => Ok(ScalarExpr::InSubquery {
178            expr: Box::new(rewrite_aggregates(context, expr, replace)?),
179            subquery: *subquery,
180            negated: *negated,
181        }),
182        other => Ok(other.clone()),
183    }
184}
uqa_sql/semantics/aggregates/slots.rs

uqa_sql/semantics/aggregates/
slots.rs
