uqa_sql/semantics/
grouping_sets.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! PostgreSQL-compatible grouping-set preparation after input schema binding.
8
9use std::collections::HashSet;
10
11use crate::ast::ColumnType;
12use crate::{RowSchema, ScalarExpr, ScalarFrameBound};
13use uqa_core::Value;
14
15use crate::{plan::QueryBlockPlan, FunctionTypeResolver, SQLError, SQLParam};
16
17mod expressions;
18mod names;
19pub(crate) use names::resolve_grouping_expression_reference;
20pub use names::{bind_grouping_names, resolve_grouping_expression};
21mod validation;
22pub use validation::validate_grouped_expressions;
23#[cfg(test)]
24mod tests;
25
26/// Resolve grouping names against the input before expanding output aliases and deduplicating `GROUP BY DISTINCT` sets. Equivalent qualified and unqualified columns share one identity without conflating expressions that `PostgreSQL` resolves to different operator inputs.
27pub fn prepare_grouping_sets(
28    engine: &dyn crate::routines::RoutineResolution,
29    statement: &QueryBlockPlan,
30    schema: &RowSchema,
31    params: &[SQLParam],
32) -> Result<Option<QueryBlockPlan>, SQLError> {
33    if statement.group_by.is_empty() && statement.grouping_sets.is_empty() {
34        return Ok(None);
35    }
36
37    let mut prepared = statement.clone();
38    // A block's operator holds the block's own columns; a lateral block binds its output names against the scope of its outer row before it builds the operator.
39    let mut changed = bind_grouping_names(engine, &mut prepared, schema, None, params)?;
40    changed |= expressions::bind_grouping_expressions(engine, &mut prepared, schema, params)?;
41    if !prepared.group_distinct {
42        return Ok(changed.then_some(prepared));
43    }
44    prepared.group_distinct = false;
45    let mut seen = HashSet::with_capacity(prepared.grouping_sets.len());
46    let mut distinct = Vec::with_capacity(prepared.grouping_sets.len());
47    for grouping_set in std::mem::take(&mut prepared.grouping_sets) {
48        let identity = grouping_set_identity(engine, &grouping_set, schema, params)?;
49        if seen.insert(identity) {
50            distinct.push(grouping_set);
51        }
52    }
53    prepared.grouping_sets = distinct;
54    Ok(Some(prepared))
55}
56
57fn grouping_set_identity(
58    engine: &dyn FunctionTypeResolver,
59    grouping_set: &[ScalarExpr],
60    schema: &RowSchema,
61    params: &[SQLParam],
62) -> Result<Vec<Vec<u8>>, SQLError> {
63    let mut identity = grouping_set
64        .iter()
65        .map(|expression| expression_identity(engine, expression, schema, params))
66        .collect::<Result<Vec<_>, _>>()?;
67    identity.sort_unstable();
68    identity.dedup();
69    Ok(identity)
70}
71
72pub(crate) fn expression_identity(
73    engine: &dyn FunctionTypeResolver,
74    expression: &ScalarExpr,
75    schema: &RowSchema,
76    params: &[SQLParam],
77) -> Result<Vec<u8>, SQLError> {
78    let expression =
79        crate::bind_type_introspection_with_resolver(expression.clone(), schema, params, engine);
80    let expression = normalize_expression(engine, expression, schema, params)?;
81    serde_json::to_vec(&expression).map_err(|error| {
82        SQLError::Internal(format!(
83            "serialize GROUP BY DISTINCT expression identity: {error}"
84        ))
85    })
86}
87
88#[expect(
89    clippy::too_many_lines,
90    reason = "preserves SELECT schema and row identity"
91)]
92fn normalize_expression(
93    engine: &dyn FunctionTypeResolver,
94    expression: ScalarExpr,
95    schema: &RowSchema,
96    params: &[SQLParam],
97) -> Result<ScalarExpr, SQLError> {
98    Ok(match expression {
99        ScalarExpr::Column(column) => schema
100            .unqualified_position(&column)
101            .map_or(ScalarExpr::Column(column), ScalarExpr::Position),
102        ScalarExpr::QualifiedColumn { qualifier, column } => {
103            schema.qualified_position(&qualifier, &column).map_or(
104                ScalarExpr::QualifiedColumn { qualifier, column },
105                ScalarExpr::Position,
106            )
107        }
108        ScalarExpr::Func {
109            order_syntax,
110            name,
111            binding,
112            args,
113            distinct,
114            order_by,
115            filter,
116        } => {
117            let name = canonical_function_name(name);
118            let argument_types = args
119                .iter()
120                .map(|argument| expression_type(engine, argument, schema, params))
121                .collect::<Result<Vec<_>, _>>()?;
122            let targets = crate::builtin_function_argument_targets(&name, &argument_types);
123            ScalarExpr::Func {
124                order_syntax,
125                name,
126                binding,
127                args: args
128                    .into_iter()
129                    .zip(targets)
130                    .map(|(argument, target)| {
131                        normalize_unknown_literal(engine, argument, target.as_ref(), schema, params)
132                    })
133                    .collect::<Result<Vec<_>, _>>()?,
134                distinct,
135                order_by: order_by
136                    .into_iter()
137                    .map(|mut order| {
138                        order.expr = normalize_expression(engine, order.expr, schema, params)?;
139                        Ok(order)
140                    })
141                    .collect::<Result<Vec<_>, SQLError>>()?,
142                filter: filter
143                    .map(|expression| {
144                        normalize_expression(engine, *expression, schema, params).map(Box::new)
145                    })
146                    .transpose()?,
147            }
148        }
149        ScalarExpr::Array(items) => {
150            ScalarExpr::Array(normalize_items(engine, items, schema, params)?)
151        }
152        ScalarExpr::Row(items) => ScalarExpr::Row(normalize_items(engine, items, schema, params)?),
153        ScalarExpr::CompositeRow {
154            items,
155            binding,
156            bound_type,
157        } => ScalarExpr::CompositeRow {
158            items: normalize_items(engine, items, schema, params)?,
159            binding: binding.clone(),
160            bound_type: bound_type.clone(),
161        },
162        ScalarExpr::Binary { op, lhs, rhs } => {
163            let left_type = expression_type(engine, &lhs, schema, params)?;
164            let right_type = expression_type(engine, &rhs, schema, params)?;
165            ScalarExpr::Binary {
166                op,
167                lhs: Box::new(normalize_unknown_literal(
168                    engine,
169                    *lhs,
170                    left_type.is_none().then_some(right_type.as_ref()).flatten(),
171                    schema,
172                    params,
173                )?),
174                rhs: Box::new(normalize_unknown_literal(
175                    engine,
176                    *rhs,
177                    right_type.is_none().then_some(left_type.as_ref()).flatten(),
178                    schema,
179                    params,
180                )?),
181            }
182        }
183        ScalarExpr::UnaryMinus(expression) => {
184            let expression = normalize_expression(engine, *expression, schema, params)?;
185            if let ScalarExpr::Literal(
186                value @ (Value::Int(_) | Value::Float(_) | Value::Decimal(_)),
187            ) = &expression
188            {
189                ScalarExpr::Literal(crate::expr::negate_value(value, None)?)
190            } else {
191                ScalarExpr::UnaryMinus(Box::new(expression))
192            }
193        }
194        ScalarExpr::Not(expression) => ScalarExpr::Not(Box::new(normalize_expression(
195            engine,
196            *expression,
197            schema,
198            params,
199        )?)),
200        ScalarExpr::And(items) => ScalarExpr::And(normalize_items(engine, items, schema, params)?),
201        ScalarExpr::Or(items) => ScalarExpr::Or(normalize_items(engine, items, schema, params)?),
202        ScalarExpr::IsNull { expr, negated } => ScalarExpr::IsNull {
203            expr: Box::new(normalize_expression(engine, *expr, schema, params)?),
204            negated,
205        },
206        ScalarExpr::Between { expr, low, high } => {
207            let target = expression_type(engine, &expr, schema, params)?;
208            ScalarExpr::Between {
209                expr: Box::new(normalize_expression(engine, *expr, schema, params)?),
210                low: Box::new(normalize_unknown_literal(
211                    engine,
212                    *low,
213                    target.as_ref(),
214                    schema,
215                    params,
216                )?),
217                high: Box::new(normalize_unknown_literal(
218                    engine,
219                    *high,
220                    target.as_ref(),
221                    schema,
222                    params,
223                )?),
224            }
225        }
226        ScalarExpr::InList {
227            expr,
228            list,
229            negated,
230        } => {
231            let target = expression_type(engine, &expr, schema, params)?;
232            ScalarExpr::InList {
233                expr: Box::new(normalize_expression(engine, *expr, schema, params)?),
234                list: list
235                    .into_iter()
236                    .map(|item| {
237                        normalize_unknown_literal(engine, item, target.as_ref(), schema, params)
238                    })
239                    .collect::<Result<Vec<_>, _>>()?,
240                negated,
241            }
242        }
243        ScalarExpr::WindowCall {
244            name,
245            args,
246            mut spec,
247            filter,
248            modifiers,
249        } => {
250            spec.partition_by = normalize_items(engine, spec.partition_by, schema, params)?;
251            for order in &mut spec.order_by {
252                order.expr = normalize_expression(engine, order.expr.clone(), schema, params)?;
253            }
254            if let Some(frame) = &mut spec.frame {
255                normalize_frame_bound(engine, &mut frame.start, schema, params)?;
256                normalize_frame_bound(engine, &mut frame.end, schema, params)?;
257            }
258            ScalarExpr::WindowCall {
259                modifiers,
260                name: canonical_function_name(name),
261                args: normalize_items(engine, args, schema, params)?,
262                spec,
263                filter: filter
264                    .map(|expression| {
265                        normalize_expression(engine, *expression, schema, params).map(Box::new)
266                    })
267                    .transpose()?,
268            }
269        }
270        ScalarExpr::Case {
271            base,
272            when,
273            else_branch,
274        } => ScalarExpr::Case {
275            base: base
276                .map(|expression| {
277                    normalize_expression(engine, *expression, schema, params).map(Box::new)
278                })
279                .transpose()?,
280            when: when
281                .into_iter()
282                .map(|(condition, result)| {
283                    Ok((
284                        normalize_expression(engine, condition, schema, params)?,
285                        normalize_expression(engine, result, schema, params)?,
286                    ))
287                })
288                .collect::<Result<Vec<_>, SQLError>>()?,
289            else_branch: else_branch
290                .map(|expression| {
291                    normalize_expression(engine, *expression, schema, params).map(Box::new)
292                })
293                .transpose()?,
294        },
295        ScalarExpr::Collated {
296            expr,
297            ty,
298            collation,
299        } => {
300            let mut expression = normalize_expression(engine, *expr, schema, params)?;
301            match (&mut expression, collation.implicit_oid()) {
302                (
303                    ScalarExpr::TypedLiteral {
304                        bound_type,
305                        collation_oid,
306                        ..
307                    },
308                    Some(oid),
309                ) if bound_type.as_ref() == Some(&ty) => {
310                    *collation_oid = Some(oid);
311                    normalize_expression(engine, expression, schema, params)?
312                }
313                _ => ScalarExpr::Collated {
314                    expr: Box::new(expression),
315                    ty,
316                    collation,
317                },
318            }
319        }
320        ScalarExpr::Cast { expr, ty, .. } => {
321            let source_type = expression_type(engine, &expr, schema, params)?;
322            let target_type = crate::type_resolution::resolve_declared_column_type(
323                engine,
324                &ColumnType::Named(ty),
325            )?;
326            let expression = normalize_expression(engine, *expr, schema, params)?;
327            if source_type.as_ref() == Some(&target_type) {
328                expression
329            } else if input_requires_catalog(&target_type) {
330                ScalarExpr::Cast {
331                    implicit: false,
332                    expr: Box::new(expression),
333                    ty: target_type.catalog_name(),
334                }
335            } else if let ScalarExpr::Literal(value @ Value::Str(_)) = &expression {
336                let input_type = if matches!(
337                    target_type.without_temporal_modifiers(),
338                    ColumnType::Interval
339                ) {
340                    target_type.clone()
341                } else {
342                    target_type.without_type_modifiers()
343                };
344                let value = crate::expr::cast_value(value, &input_type.sql_name())?;
345                let input = normalize_expression(
346                    engine,
347                    ScalarExpr::TypedLiteral {
348                        collation_oid: None,
349                        composite_source: None,
350                        value,
351                        ty: input_type.catalog_name(),
352                        bound_type: Some(input_type.clone()),
353                        parameter_index: None,
354                    },
355                    schema,
356                    params,
357                )?;
358                if input_type == target_type {
359                    input
360                } else {
361                    ScalarExpr::Cast {
362                        implicit: false,
363                        expr: Box::new(input),
364                        ty: target_type.catalog_name(),
365                    }
366                }
367            } else if let ScalarExpr::Literal(Value::Null) = expression {
368                ScalarExpr::TypedLiteral {
369                    collation_oid: None,
370                    composite_source: None,
371                    value: Value::Null,
372                    ty: target_type.catalog_name(),
373                    bound_type: Some(target_type),
374                    parameter_index: None,
375                }
376            } else {
377                ScalarExpr::Cast {
378                    implicit: false,
379                    expr: Box::new(expression),
380                    ty: target_type.catalog_name(),
381                }
382            }
383        }
384        ScalarExpr::InSubquery {
385            expr,
386            subquery,
387            negated,
388        } => ScalarExpr::InSubquery {
389            expr: Box::new(normalize_expression(engine, *expr, schema, params)?),
390            subquery,
391            negated,
392        },
393        ScalarExpr::TypedLiteral {
394            value,
395            ty,
396            bound_type,
397            parameter_index,
398            collation_oid,
399            composite_source,
400        } => {
401            let literal = ScalarExpr::Literal(value.clone());
402            let declared = match bound_type {
403                Some(ty) => ty,
404                None => crate::type_resolution::resolve_declared_column_type(
405                    engine,
406                    &ColumnType::Named(ty),
407                )?,
408            };
409            let collation_oid = collation_oid.filter(|oid| {
410                i64::from(*oid) != crate::catalog::type_metadata::pg_type_collation_oid(&declared)
411            });
412            if parameter_index.is_none()
413                && composite_source.is_none()
414                && collation_oid.is_none()
415                && expression_type(engine, &literal, schema, params)?.as_ref() == Some(&declared)
416            {
417                literal
418            } else {
419                ScalarExpr::TypedLiteral {
420                    collation_oid,
421                    composite_source,
422                    value,
423                    ty: declared.catalog_name(),
424                    bound_type: Some(declared),
425                    parameter_index,
426                }
427            }
428        }
429        expression @ (ScalarExpr::Star
430        | ScalarExpr::QualifiedStar(_)
431        | ScalarExpr::Default
432        | ScalarExpr::Position(_)
433        | ScalarExpr::InternalColumn(_)
434        | ScalarExpr::Literal(_)
435        | ScalarExpr::Param(_)
436        | ScalarExpr::ScalarSubquery(_)
437        | ScalarExpr::Exists { .. }) => expression,
438    })
439}
440
441fn input_requires_catalog(ty: &ColumnType) -> bool {
442    match ty {
443        ColumnType::Named(_)
444        | ColumnType::Domain { .. }
445        | ColumnType::Enum(_)
446        | ColumnType::Composite(_)
447        | ColumnType::Regproc
448        | ColumnType::Regprocedure
449        | ColumnType::Regclass
450        | ColumnType::Regcollation
451        | ColumnType::Regnamespace
452        | ColumnType::Regrole
453        | ColumnType::Regtype
454        | ColumnType::Record
455        | ColumnType::AnyArray => true,
456        ColumnType::Array(element) => input_requires_catalog(element),
457        _ => false,
458    }
459}
460
461fn normalize_items(
462    engine: &dyn FunctionTypeResolver,
463    items: Vec<ScalarExpr>,
464    schema: &RowSchema,
465    params: &[SQLParam],
466) -> Result<Vec<ScalarExpr>, SQLError> {
467    items
468        .into_iter()
469        .map(|item| normalize_expression(engine, item, schema, params))
470        .collect()
471}
472
473fn normalize_unknown_literal(
474    engine: &dyn FunctionTypeResolver,
475    expression: ScalarExpr,
476    target: Option<&ColumnType>,
477    schema: &RowSchema,
478    params: &[SQLParam],
479) -> Result<ScalarExpr, SQLError> {
480    if matches!(expression, ScalarExpr::Literal(Value::Null)) {
481        if let Some(target) = target {
482            return normalize_expression(
483                engine,
484                ScalarExpr::Cast {
485                    implicit: true,
486                    expr: Box::new(expression),
487                    ty: target.catalog_name(),
488                },
489                schema,
490                params,
491            );
492        }
493    }
494    normalize_expression(engine, expression, schema, params)
495}
496
497fn normalize_frame_bound(
498    engine: &dyn FunctionTypeResolver,
499    bound: &mut ScalarFrameBound,
500    schema: &RowSchema,
501    params: &[SQLParam],
502) -> Result<(), SQLError> {
503    match bound {
504        ScalarFrameBound::Preceding(expression) | ScalarFrameBound::Following(expression) => {
505            **expression = normalize_expression(engine, (**expression).clone(), schema, params)?;
506        }
507        ScalarFrameBound::UnboundedPreceding
508        | ScalarFrameBound::UnboundedFollowing
509        | ScalarFrameBound::CurrentRow => {}
510    }
511    Ok(())
512}
513
514fn expression_type(
515    engine: &dyn FunctionTypeResolver,
516    expression: &ScalarExpr,
517    schema: &RowSchema,
518    params: &[SQLParam],
519) -> Result<Option<ColumnType>, SQLError> {
520    crate::scalar_type_with_resolver(expression, schema, params, engine)
521}
522
523fn canonical_function_name(name: String) -> String {
524    let lower = name.to_ascii_lowercase();
525    match lower.strip_prefix("pg_catalog.") {
526        Some(unqualified) => unqualified.to_owned(),
527        None => lower,
528    }
529}
uqa_sql/semantics/grouping_sets.rs

uqa_sql/semantics/
grouping_sets.rs
