Skip to main content

uqa_sql/plan/
scalar.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! SQL AST to executable scalar IR lowering and aggregate classification.
8
9use super::{
10    AggregateClassifier, Expr, FrameBound, OrderBy, QueryPlan, ScalarExpr, ScalarFrameBound,
11    ScalarOrder, ScalarWindowFrame, ScalarWindowSpec, WindowSpec,
12};
13use crate::schema::retention::CatalogRetentionError;
14use resources::{Control, Lowering, Result};
15use source::{Node, Source};
16use uqa_core::{
17    memory::{Budgeted, MemoryBudget},
18    CancellationToken,
19};
20
21mod binding;
22mod copy;
23mod resources;
24mod source;
25mod window;
26
27impl super::ExpressionPlan {
28    /// Lower a borrowed, validated column expression directly into admitted scalar IR. Destination strings, value payloads, bindings, vector capacities and boxes acquire the supplied allowance before allocation, and the result retains those leases. Both the retained definition's original cancellation and the invoking reader's cancellation remain active during lowering. This controls AST-to-IR production only; subsequent type binding and evaluation require their own resource contracts. Query children violate the validated column-expression invariant.
29    pub fn lower_column_budgeted(
30        expression: &Expr,
31        budget: &MemoryBudget,
32        original: &CancellationToken,
33        invoking: &CancellationToken,
34    ) -> Result<Budgeted<ScalarExpr>> {
35        let mut lowering = Lowering {
36            control: Some(Control::new(budget, original, invoking)),
37        };
38        let scalar = lowering.expression(
39            Source::Borrowed(expression),
40            &super::NoRegisteredAggregates,
41            &mut Vec::new(),
42        )?;
43        lowering.finish(scalar)
44    }
45}
46
47pub(super) fn lower_scalar_expression(
48    expression: Expr,
49    aggregates: &dyn AggregateClassifier,
50    subqueries: &mut Vec<QueryPlan>,
51) -> ScalarExpr {
52    Lowering { control: None }
53        .expression(Source::Owned(expression), aggregates, subqueries)
54        .expect("owned lowering has no admission failure")
55}
56
57pub(super) fn lower_window_spec(
58    specification: WindowSpec,
59    aggregates: &dyn AggregateClassifier,
60    subqueries: &mut Vec<QueryPlan>,
61) -> ScalarWindowSpec {
62    Lowering { control: None }
63        .window(Source::Owned(specification), aggregates, subqueries)
64        .expect("owned lowering has no admission failure")
65}
66
67impl Lowering<'_> {
68    #[expect(
69        clippy::too_many_lines,
70        reason = "plan lowering preserves exhaustive variants and structural identities"
71    )]
72    fn expression(
73        &mut self,
74        expression: Source<'_, Expr>,
75        aggregates: &dyn AggregateClassifier,
76        subqueries: &mut Vec<QueryPlan>,
77    ) -> Result<ScalarExpr> {
78        self.check()?;
79        Ok(match expression.node() {
80            Node::Star => ScalarExpr::Star,
81            Node::QualifiedStar(name) => ScalarExpr::QualifiedStar(self.text(name)?),
82            Node::Default => ScalarExpr::Default,
83            Node::Column(name) => ScalarExpr::Column(self.text(name)?),
84            Node::QualifiedColumn { qualifier, column } => ScalarExpr::QualifiedColumn {
85                qualifier: self.text(qualifier)?,
86                column: self.text(column)?,
87            },
88            Node::InternalColumn(column) => ScalarExpr::InternalColumn(column),
89            Node::Literal(value) => ScalarExpr::Literal(self.value(value)?),
90            Node::TypedLiteral {
91                value,
92                ty,
93                collation_oid,
94                composite_source,
95            } => ScalarExpr::TypedLiteral {
96                collation_oid,
97                composite_source: self.composite_source(composite_source)?,
98                value: self.value(value)?,
99                ty: self.text(ty)?,
100                bound_type: None,
101                parameter_index: None,
102            },
103            Node::Param(index) => ScalarExpr::Param(index),
104            Node::Func {
105                name,
106                binding,
107                args,
108                distinct,
109                order_by,
110                order_syntax,
111                filter,
112            } => ScalarExpr::Func {
113                name: self.text(name)?,
114                binding: binding.map(|binding| self.binding(binding)).transpose()?,
115                args: self.map(args, |this, argument| {
116                    this.expression(argument, aggregates, subqueries)
117                })?,
118                distinct,
119                order_by: self.map(order_by, |this, order| {
120                    this.order(order, aggregates, subqueries)
121                })?,
122                order_syntax,
123                filter: filter
124                    .map(|filter| self.child(filter, aggregates, subqueries))
125                    .transpose()?,
126            },
127            Node::Array(items) => ScalarExpr::Array(self.map(items, |this, item| {
128                this.expression(item, aggregates, subqueries)
129            })?),
130            Node::CompositeRow { items, binding } => ScalarExpr::CompositeRow {
131                items: self.map(items, |this, item| {
132                    this.expression(item, aggregates, subqueries)
133                })?,
134                binding: self.composite_binding(binding)?,
135                bound_type: None,
136            },
137            Node::Row(items) => ScalarExpr::Row(self.map(items, |this, item| {
138                this.expression(item, aggregates, subqueries)
139            })?),
140            Node::Binary { op, lhs, rhs } => ScalarExpr::Binary {
141                op,
142                lhs: self.child(lhs, aggregates, subqueries)?,
143                rhs: self.child(rhs, aggregates, subqueries)?,
144            },
145            Node::UnaryMinus(expression) => {
146                ScalarExpr::UnaryMinus(self.child(expression, aggregates, subqueries)?)
147            }
148            Node::Not(expression) => {
149                ScalarExpr::Not(self.child(expression, aggregates, subqueries)?)
150            }
151            Node::And(items) => ScalarExpr::And(self.map(items, |this, item| {
152                this.expression(item, aggregates, subqueries)
153            })?),
154            Node::Or(items) => ScalarExpr::Or(self.map(items, |this, item| {
155                this.expression(item, aggregates, subqueries)
156            })?),
157            Node::IsNull { expr, negated } => ScalarExpr::IsNull {
158                expr: self.child(expr, aggregates, subqueries)?,
159                negated,
160            },
161            Node::Between { expr, low, high } => ScalarExpr::Between {
162                expr: self.child(expr, aggregates, subqueries)?,
163                low: self.child(low, aggregates, subqueries)?,
164                high: self.child(high, aggregates, subqueries)?,
165            },
166            Node::InList {
167                expr,
168                list,
169                negated,
170            } => ScalarExpr::InList {
171                expr: self.child(expr, aggregates, subqueries)?,
172                list: self.map(list, |this, item| {
173                    this.expression(item, aggregates, subqueries)
174                })?,
175                negated,
176            },
177            Node::WindowCall {
178                name,
179                args,
180                spec,
181                filter,
182                modifiers,
183            } => ScalarExpr::WindowCall {
184                modifiers,
185                name: self.text(name)?,
186                args: self.map(args, |this, argument| {
187                    this.expression(argument, aggregates, subqueries)
188                })?,
189                spec: self.window(spec, aggregates, subqueries)?,
190                filter: filter
191                    .map(|filter| self.child(filter, aggregates, subqueries))
192                    .transpose()?,
193            },
194            Node::Case {
195                base,
196                when,
197                else_branch,
198            } => ScalarExpr::Case {
199                base: base
200                    .map(|base| self.child(base, aggregates, subqueries))
201                    .transpose()?,
202                when: self.map(when, |this, pair| {
203                    let (condition, result) = pair.pair();
204                    Ok((
205                        this.expression(condition, aggregates, subqueries)?,
206                        this.expression(result, aggregates, subqueries)?,
207                    ))
208                })?,
209                else_branch: else_branch
210                    .map(|branch| self.child(branch, aggregates, subqueries))
211                    .transpose()?,
212            },
213            Node::Cast { expr, ty, implicit } => ScalarExpr::Cast {
214                implicit,
215                expr: self.child(expr, aggregates, subqueries)?,
216                ty: self.text(ty)?,
217            },
218            Node::ScalarSubquery(query) => {
219                ScalarExpr::ScalarSubquery(self.query(query, aggregates, subqueries)?)
220            }
221            Node::Exists { body, negated } => ScalarExpr::Exists {
222                subquery: self.query(body, aggregates, subqueries)?,
223                negated,
224            },
225            Node::InSubquery {
226                expr,
227                body,
228                negated,
229            } => {
230                let expr = self.child(expr, aggregates, subqueries)?;
231                ScalarExpr::InSubquery {
232                    expr,
233                    subquery: self.query(body, aggregates, subqueries)?,
234                    negated,
235                }
236            }
237        })
238    }
239
240    fn child(
241        &mut self,
242        expression: Source<'_, Box<Expr>>,
243        aggregates: &dyn AggregateClassifier,
244        subqueries: &mut Vec<QueryPlan>,
245    ) -> Result<Box<ScalarExpr>> {
246        self.boxed(|this| this.expression(expression.unbox(), aggregates, subqueries))
247    }
248
249    fn query(
250        &self,
251        query: Source<'_, Box<crate::ast::SelectStmt>>,
252        aggregates: &dyn AggregateClassifier,
253        subqueries: &mut Vec<QueryPlan>,
254    ) -> Result<usize> {
255        self.check()?;
256        let Source::Owned(query) = query else {
257            return Err(CatalogRetentionError::UnexpectedSubquery);
258        };
259        let id = subqueries.len();
260        subqueries.push(QueryPlan::lower_with(*query, aggregates));
261        Ok(id)
262    }
263}
264
265pub(crate) fn is_builtin_aggregate(name: &str) -> bool {
266    crate::ast::is_builtin_aggregate_function(name)
267}
268
269#[cfg(test)]
270mod tests;