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_selected(expr, &mut |expression| {
20        if !is_aggregate(context, expression) {
21            return Ok(None);
22        }
23        let slot = aggregate_slot(relation, *cursor);
24        *cursor += 1;
25        Ok(Some(slot))
26    })
27}
28
29pub fn compile_having_aggregate_slots(
30    context: &dyn AggregateClassifier,
31    expr: &ScalarExpr,
32    relation: crate::ast::InternalRelationId,
33    aggregate_targets: &[ScalarExpr],
34) -> Result<ScalarExpr, SQLError> {
35    rewrite_selected(expr, &mut |aggregate| {
36        if !is_aggregate(context, aggregate) {
37            return Ok(None);
38        }
39        aggregate_targets
40            .iter()
41            .position(|target| exprs_match(target, aggregate))
42            .map(|index| Some(aggregate_slot(relation, index)))
43            .ok_or_else(|| {
44                SQLError::Unsupported(
45                    "HAVING references an aggregate that is not in the aggregate plan".into(),
46                )
47            })
48    })
49}
50
51/// Group values follow aggregate finalizers in the retained scalar slot namespace. Match a complete grouped expression before descending into its inputs.
52pub fn compile_group_slots(
53    expression: &ScalarExpr,
54    groups: &[ScalarExpr],
55    relation: crate::ast::InternalRelationId,
56    first_slot: usize,
57) -> Result<ScalarExpr, SQLError> {
58    rewrite_selected(expression, &mut |expression| {
59        groups
60            .iter()
61            .position(|group| exprs_match(group, expression))
62            .map(|index| {
63                first_slot
64                    .checked_add(index)
65                    .map(|index| aggregate_slot(relation, index))
66                    .ok_or_else(|| SQLError::Internal("aggregate scalar slot overflow".into()))
67            })
68            .transpose()
69    })
70}
71
72/// An absent grouping-set key is NULL as a complete expression, even when another selected key exposes one of its inputs. Aggregate arguments still read the original input rows.
73pub fn select_grouping_set(
74    context: &dyn AggregateClassifier,
75    statement: &crate::plan::QueryBlockPlan,
76    selected: &[ScalarExpr],
77) -> Result<crate::plan::QueryBlockPlan, SQLError> {
78    let groups = statement
79        .group_by
80        .iter()
81        .chain(statement.grouping_sets.iter().flatten())
82        .collect::<Vec<_>>();
83    let mut active = statement.clone();
84    active.group_by = selected.to_vec();
85    active.grouping_sets.clear();
86    let mut rewrite = |expression: &ScalarExpr| {
87        if is_aggregate(context, expression) {
88            return Ok(Some(expression.clone()));
89        }
90        Ok(groups
91            .iter()
92            .any(|group| exprs_match(group, expression))
93            .then(|| {
94                if selected.iter().any(|group| exprs_match(group, expression)) {
95                    expression.clone()
96                } else {
97                    ScalarExpr::Literal(uqa_core::Value::Null)
98                }
99            }))
100    };
101    for projection in &mut active.projections {
102        projection.expr = rewrite_selected(&projection.expr, &mut rewrite)?;
103    }
104    if let Some(having) = &mut active.having {
105        *having = rewrite_selected(having, &mut rewrite)?;
106    }
107    Ok(active)
108}
109
110pub fn aggregate_slot_index(
111    column: crate::ast::InternalColumnRef,
112    relation: crate::ast::InternalRelationId,
113) -> Option<usize> {
114    (column.relation() == relation).then(|| column.attribute())
115}
116
117fn aggregate_slot(relation: crate::ast::InternalRelationId, index: usize) -> ScalarExpr {
118    ScalarExpr::InternalColumn(relation.column(index))
119}
120
121#[expect(
122    clippy::too_many_lines,
123    reason = "preserves aggregate NULL and type order"
124)]
125fn rewrite_selected(
126    expr: &ScalarExpr,
127    replace: &mut impl FnMut(&ScalarExpr) -> Result<Option<ScalarExpr>, SQLError>,
128) -> Result<ScalarExpr, SQLError> {
129    if let Some(expression) = replace(expr)? {
130        return Ok(expression);
131    }
132    match expr {
133        ScalarExpr::Func {
134            name,
135            binding,
136            args,
137            distinct,
138            order_by,
139            filter,
140        } => Ok(ScalarExpr::Func {
141            name: name.clone(),
142            binding: binding.clone(),
143            args: args
144                .iter()
145                .map(|arg| rewrite_selected(arg, replace))
146                .collect::<Result<Vec<_>, _>>()?,
147            distinct: *distinct,
148            order_by: order_by.clone(),
149            filter: filter
150                .as_deref()
151                .map(|filter| rewrite_selected(filter, replace).map(Box::new))
152                .transpose()?,
153        }),
154        ScalarExpr::Array(items) => Ok(ScalarExpr::Array(
155            items
156                .iter()
157                .map(|item| rewrite_selected(item, replace))
158                .collect::<Result<Vec<_>, _>>()?,
159        )),
160        ScalarExpr::Row(items) => Ok(ScalarExpr::Row(
161            items
162                .iter()
163                .map(|item| rewrite_selected(item, replace))
164                .collect::<Result<Vec<_>, _>>()?,
165        )),
166        ScalarExpr::Binary { op, lhs, rhs } => Ok(ScalarExpr::Binary {
167            op: *op,
168            lhs: Box::new(rewrite_selected(lhs, replace)?),
169            rhs: Box::new(rewrite_selected(rhs, replace)?),
170        }),
171        ScalarExpr::Not(inner) => Ok(ScalarExpr::Not(Box::new(rewrite_selected(inner, replace)?))),
172        ScalarExpr::UnaryMinus(inner) => Ok(ScalarExpr::UnaryMinus(Box::new(rewrite_selected(
173            inner, replace,
174        )?))),
175        ScalarExpr::And(parts) => Ok(ScalarExpr::And(
176            parts
177                .iter()
178                .map(|part| rewrite_selected(part, replace))
179                .collect::<Result<Vec<_>, _>>()?,
180        )),
181        ScalarExpr::Or(parts) => Ok(ScalarExpr::Or(
182            parts
183                .iter()
184                .map(|part| rewrite_selected(part, replace))
185                .collect::<Result<Vec<_>, _>>()?,
186        )),
187        ScalarExpr::IsNull { expr, negated } => Ok(ScalarExpr::IsNull {
188            expr: Box::new(rewrite_selected(expr, replace)?),
189            negated: *negated,
190        }),
191        ScalarExpr::Between { expr, low, high } => Ok(ScalarExpr::Between {
192            expr: Box::new(rewrite_selected(expr, replace)?),
193            low: Box::new(rewrite_selected(low, replace)?),
194            high: Box::new(rewrite_selected(high, replace)?),
195        }),
196        ScalarExpr::InList {
197            expr,
198            list,
199            negated,
200        } => Ok(ScalarExpr::InList {
201            expr: Box::new(rewrite_selected(expr, replace)?),
202            list: list
203                .iter()
204                .map(|item| rewrite_selected(item, replace))
205                .collect::<Result<Vec<_>, _>>()?,
206            negated: *negated,
207        }),
208        ScalarExpr::Case {
209            base,
210            when,
211            else_branch,
212        } => Ok(ScalarExpr::Case {
213            base: base
214                .as_deref()
215                .map(|base| rewrite_selected(base, replace).map(Box::new))
216                .transpose()?,
217            when: when
218                .iter()
219                .map(|(condition, result)| {
220                    Ok((
221                        rewrite_selected(condition, replace)?,
222                        rewrite_selected(result, replace)?,
223                    ))
224                })
225                .collect::<Result<Vec<_>, SQLError>>()?,
226            else_branch: else_branch
227                .as_deref()
228                .map(|branch| rewrite_selected(branch, replace).map(Box::new))
229                .transpose()?,
230        }),
231        ScalarExpr::Cast { expr, ty } => Ok(ScalarExpr::Cast {
232            expr: Box::new(rewrite_selected(expr, replace)?),
233            ty: ty.clone(),
234        }),
235        ScalarExpr::InSubquery {
236            expr,
237            subquery,
238            negated,
239        } => Ok(ScalarExpr::InSubquery {
240            expr: Box::new(rewrite_selected(expr, replace)?),
241            subquery: *subquery,
242            negated: *negated,
243        }),
244        other => Ok(other.clone()),
245    }
246}
uqa_sql/semantics/aggregates/slots.rs

uqa_sql/semantics/aggregates/
slots.rs
