nodedb_sql/planner/subquery/
scalar.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Scalar subqueries: `col > (SELECT AGG(...) FROM tbl ...)`.
4
5use sqlparser::ast::{self, Expr, SetExpr};
6
7use crate::error::Result;
8use crate::functions::registry::FunctionRegistry;
9use crate::parser::normalize::normalize_ident;
10use crate::types::*;
11
12use super::extract::SubqueryJoin;
13
14fn canonical_aggregate_key(function: &str, field: &str) -> String {
15    format!("{function}({field})")
16}
17
18/// Result of planning a scalar subquery.
19pub(super) struct ScalarSubqueryResult {
20    pub(super) join: SubqueryJoin,
21    pub(super) replacement_expr: Expr,
22}
23
24/// Plan a scalar subquery (e.g., `(SELECT AVG(amount) FROM orders)`).
25///
26/// Rewrites `col > (SELECT AVG(amount) FROM orders)` as:
27///   cross-join with the aggregate result (1 row), then filter `col > result_col`.
28///
29/// The cross-join produces a cartesian product, but since the aggregate returns
30/// exactly 1 row, every outer row gets paired with that single result row.
31pub(super) fn try_plan_scalar_subquery(
32    subquery: &ast::Query,
33    catalog: &dyn SqlCatalog,
34    functions: &FunctionRegistry,
35    temporal: crate::TemporalScope,
36) -> Result<Option<ScalarSubqueryResult>> {
37    let inner_plan = crate::planner::select::plan_query(subquery, catalog, functions, temporal)?;
38
39    // Extract the result column name from the subquery's SELECT list.
40    let result_col = match extract_scalar_column(subquery) {
41        Some(col) => col,
42        None => return Ok(None),
43    };
44
45    let replacement = Expr::Identifier(ast::Ident::new(&result_col));
46
47    Ok(Some(ScalarSubqueryResult {
48        join: SubqueryJoin {
49            // A cross join pairs every row with every row, so it has no key.
50            on: Vec::new(),
51            inner_plan,
52            join_type: JoinType::Cross,
53        },
54        replacement_expr: replacement,
55    }))
56}
57
58/// Extract the projected column name from a scalar subquery.
59///
60/// Handles aliased aggregates like `SELECT AVG(amount) AS avg_amount`.
61/// For unaliased aggregates, returns the canonical aggregate key emitted by
62/// the aggregate executor (e.g. `avg(amount)`, `count(*)`).
63fn extract_scalar_column(query: &ast::Query) -> Option<String> {
64    let select = match &*query.body {
65        SetExpr::Select(s) => s,
66        _ => return None,
67    };
68    if select.projection.len() != 1 {
69        return None;
70    }
71    match &select.projection[0] {
72        ast::SelectItem::ExprWithAlias { alias, .. } => Some(normalize_ident(alias)),
73        ast::SelectItem::UnnamedExpr(expr) => match expr {
74            Expr::Identifier(ident) => Some(normalize_ident(ident)),
75            Expr::CompoundIdentifier(parts) if parts.len() >= 3 => {
76                // Schema-qualified: return None to propagate "unsupported" through convert_expr.
77                None
78            }
79            Expr::CompoundIdentifier(parts) if parts.len() == 2 => Some(normalize_ident(&parts[1])),
80            Expr::Function(func) => {
81                let func_name = func
82                    .name
83                    .0
84                    .iter()
85                    .map(|p| match p {
86                        ast::ObjectNamePart::Identifier(ident) => normalize_ident(ident),
87                        _ => String::new(),
88                    })
89                    .collect::<Vec<_>>()
90                    .join(".")
91                    .to_lowercase();
92                let arg = match &func.args {
93                    ast::FunctionArguments::List(arg_list) => arg_list
94                        .args
95                        .first()
96                        .and_then(|a| match a {
97                            ast::FunctionArg::Unnamed(ast::FunctionArgExpr::Expr(
98                                Expr::Identifier(ident),
99                            )) => Some(normalize_ident(ident)),
100                            ast::FunctionArg::Unnamed(ast::FunctionArgExpr::Expr(
101                                Expr::CompoundIdentifier(parts),
102                            )) if parts.len() >= 3 => None,
103                            ast::FunctionArg::Unnamed(ast::FunctionArgExpr::Expr(
104                                Expr::CompoundIdentifier(parts),
105                            )) if parts.len() == 2 => Some(normalize_ident(&parts[1])),
106                            ast::FunctionArg::Unnamed(
107                                ast::FunctionArgExpr::Wildcard
108                                | ast::FunctionArgExpr::QualifiedWildcard(_),
109                            ) => Some("all".to_string()),
110                            _ => None,
111                        })
112                        .unwrap_or_else(|| "*".to_string()),
113                    _ => "*".to_string(),
114                };
115                Some(canonical_aggregate_key(&func_name, &arg))
116            }
117            _ => None,
118        },
119        _ => None,
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use super::extract_scalar_column;
126    use crate::parser::statement::parse_sql;
127    use sqlparser::ast::Statement;
128
129    #[test]
130    fn unaliased_scalar_aggregate_uses_canonical_aggregate_key() {
131        let statements = parse_sql("SELECT AVG(amount) FROM orders").unwrap();
132        let Statement::Query(query) = &statements[0] else {
133            panic!("expected query");
134        };
135        assert_eq!(extract_scalar_column(query), Some("avg(amount)".into()));
136    }
137}
nodedb_sql/planner/subquery/scalar.rs

nodedb_sql/planner/subquery/
scalar.rs
