nodedb_sql/planner/
aggregate.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! GROUP BY and aggregate planning over a single collection.
4
5use sqlparser::ast;
6
7use crate::engine_rules::{self, AggregateParams, TimeBucketKey};
8use crate::error::Result;
9use crate::functions::registry::{FunctionRegistry, SearchTrigger};
10use crate::parser::normalize::normalize_ident;
11use crate::planner::aggregate_clauses::plan_aggregate_clauses;
12use crate::resolver::ColumnScope;
13use crate::resolver::columns::{ResolvedTable, TableScope};
14use crate::resolver::expr::convert_expr;
15use crate::temporal::TemporalScope;
16use crate::types::*;
17
18/// The inputs of one aggregate query.
19pub struct AggregateQuery<'a> {
20    pub select: &'a ast::Select,
21    pub order_by: Option<&'a ast::OrderBy>,
22    /// The single collection the aggregate reads.
23    pub table: &'a ResolvedTable,
24    /// The WHERE clause, already converted against `scope`.
25    pub filters: &'a [Filter],
26    pub scope: &'a TableScope,
27    pub temporal: &'a TemporalScope,
28    /// Set when semi or anti subquery joins must filter the input rows below
29    /// the aggregate. The engine then plans an aggregate over an input plan
30    /// the caller can wrap.
31    pub input_joins: bool,
32}
33
34/// The aggregate's row limit without a `LIMIT` clause: every group. An
35/// explicit `LIMIT n` overwrites it in `apply_limit`.
36pub const AGGREGATE_NO_LIMIT: usize = usize::MAX;
37
38/// Plan an aggregate query (GROUP BY + aggregate functions).
39pub fn plan_aggregate(query: AggregateQuery<'_>, functions: &FunctionRegistry) -> Result<SqlPlan> {
40    let AggregateQuery {
41        select,
42        order_by,
43        table,
44        filters,
45        scope,
46        temporal,
47        input_joins,
48    } = query;
49    let clauses = plan_aggregate_clauses(select, order_by, &table.info.columns, scope, functions)?;
50    let time_bucket = time_bucket_key(&clauses.group_by, functions);
51
52    let rules = engine_rules::resolve_engine_rules(table.info.engine);
53    let mut plan = rules.plan_aggregate(AggregateParams {
54        collection: table.name.clone(),
55        alias: table.alias.clone(),
56        filters: filters.to_vec(),
57        group_by: clauses.group_by,
58        aggregates: clauses.aggregates,
59        having: clauses.having,
60        limit: AGGREGATE_NO_LIMIT,
61        time_bucket,
62        grouping_sets: clauses.grouping_sets,
63        columns: table.info.columns.clone(),
64        has_auto_tier: table.info.has_auto_tier,
65        bitemporal: table.info.bitemporal,
66        temporal: *temporal,
67        hidden_aggregates: clauses.hidden_aggregates,
68        input_joins,
69    })?;
70
71    // The engine rules build the Aggregate node without a projection in scope,
72    // so the SELECT-list output aliases for the GROUP BY keys are attached here
73    // (Postgres: `SELECT k AS label ... GROUP BY k` yields output column
74    // `label`, not `k`).
75    if let SqlPlan::Aggregate {
76        group_by_aliases,
77        output_order,
78        ..
79    } = &mut plan
80    {
81        *group_by_aliases = clauses.group_by_aliases;
82        *output_order = clauses.output_order;
83    }
84    Ok(plan)
85}
86
87/// The `time_bucket(...)` key among the GROUP BY keys, with its interval.
88fn time_bucket_key(group_by: &[SqlExpr], functions: &FunctionRegistry) -> Option<TimeBucketKey> {
89    group_by
90        .iter()
91        .enumerate()
92        .find_map(|(group_index, key)| match key {
93            SqlExpr::Function { name, args, .. }
94                if functions.search_trigger(name) == SearchTrigger::TimeBucket =>
95            {
96                Some(TimeBucketKey {
97                    group_index,
98                    interval_ms: bucket_interval_ms(args),
99                })
100            }
101            _ => None,
102        })
103}
104
105/// The bucket width a `time_bucket` call names, in milliseconds.
106///
107/// Both argument orders are accepted: `time_bucket('1 hour', ts)`,
108/// `time_bucket(ts, '1 hour')`, and integer seconds `time_bucket(3600, ts)`.
109/// 0 when no argument is a readable positive interval literal.
110fn bucket_interval_ms(args: &[SqlExpr]) -> i64 {
111    args.iter()
112        .find_map(|arg| match arg {
113            SqlExpr::Literal(SqlValue::String(s)) => {
114                Some(parse_interval_to_ms(s)).filter(|ms| *ms > 0)
115            }
116            SqlExpr::Literal(SqlValue::Int(secs)) if *secs > 0 => secs.checked_mul(1000),
117            _ => None,
118        })
119        .unwrap_or(0)
120}
121
122/// Parse an interval string to milliseconds.
123///
124/// Delegates to the canonical `nodedb_types::kv_parsing::parse_interval_to_ms`.
125fn parse_interval_to_ms(s: &str) -> i64 {
126    nodedb_types::kv_parsing::parse_interval_to_ms(s)
127        .map(|ms| ms as i64)
128        .unwrap_or(0)
129}
130
131/// Scan the SELECT projection for `GROUPING(col)` calls and return synthetic
132/// `AggregateExpr` entries so the executor can compute the per-set bitmask value.
133///
134/// For `GROUPING(col)`, the canonical index of `col` in `canonical_keys` is
135/// encoded into the `field` of the resulting `AggregateExpr` (as a decimal string).
136/// The executor reads this index and checks the grouping-set bitmask at run time.
137pub(super) fn extract_grouping_calls(
138    items: &[ast::SelectItem],
139    canonical_keys: &[SqlExpr],
140    scope: &TableScope,
141) -> Result<Vec<AggregateExpr>> {
142    let mut out = Vec::new();
143    for item in items {
144        let (expr, alias): (&ast::Expr, String) = match item {
145            ast::SelectItem::UnnamedExpr(expr) => (expr, format!("{expr}")),
146            ast::SelectItem::ExprWithAlias { expr, alias } => (expr, normalize_ident(alias)),
147            _ => continue,
148        };
149        collect_grouping_from_expr(expr, &alias, canonical_keys, &mut out, scope)?;
150    }
151    Ok(out)
152}
153
154/// Recursively collect `GROUPING(col)` calls from an expression.
155fn collect_grouping_from_expr(
156    expr: &ast::Expr,
157    alias: &str,
158    canonical_keys: &[SqlExpr],
159    out: &mut Vec<AggregateExpr>,
160    scope: &TableScope,
161) -> Result<()> {
162    match expr {
163        ast::Expr::Function(f) => {
164            let name = normalize_function_name(f);
165            if name.eq_ignore_ascii_case("grouping") {
166                // Extract the column argument(s).
167                let args = function_args_exprs(f);
168                for col_expr in &args {
169                    let canonical_idx = crate::planner::grouping_sets::resolve_grouping_col(
170                        col_expr,
171                        canonical_keys,
172                    )?;
173                    // Encode index in the field name; alias is user-visible output name.
174                    out.push(AggregateExpr {
175                        function: "grouping".into(),
176                        args: vec![convert_expr(col_expr, &ColumnScope::Relations(scope))?],
177                        alias: alias.to_string(),
178                        distinct: false,
179                        grouping_col_index: Some(canonical_idx),
180                        filter_text: None,
181                        arg_text: None,
182                    });
183                }
184            }
185        }
186        // Recurse into binary ops and other wrappers.
187        ast::Expr::BinaryOp { left, right, .. } => {
188            collect_grouping_from_expr(left, alias, canonical_keys, out, scope)?;
189            collect_grouping_from_expr(right, alias, canonical_keys, out, scope)?;
190        }
191        _ => {}
192    }
193    Ok(())
194}
195
196/// Extract the positional expression arguments from a function call.
197pub(super) fn function_args_exprs(f: &ast::Function) -> Vec<&ast::Expr> {
198    match &f.args {
199        ast::FunctionArguments::List(list) => list
200            .args
201            .iter()
202            .filter_map(|a| match a {
203                ast::FunctionArg::Unnamed(ast::FunctionArgExpr::Expr(e)) => Some(e),
204                _ => None,
205            })
206            .collect(),
207        _ => Vec::new(),
208    }
209}
210
211/// Return the simple lowercase name of a function (unqualified part only).
212pub(super) fn normalize_function_name(f: &ast::Function) -> String {
213    f.name
214        .0
215        .last()
216        .map(|p| match p {
217            ast::ObjectNamePart::Identifier(ident) => normalize_ident(ident),
218            _ => String::new(),
219        })
220        .unwrap_or_default()
221}
222
223/// Extract aggregate expressions from SELECT projection.
224///
225/// `scope` gates each aggregate argument, so a column no relation declares
226/// raises an undefined-column error instead of planning.
227pub fn extract_aggregates_from_projection(
228    items: &[ast::SelectItem],
229    functions: &FunctionRegistry,
230    scope: &TableScope,
231) -> Result<Vec<AggregateExpr>> {
232    let mut aggregates = Vec::new();
233    for item in items {
234        let (expr, alias): (&ast::Expr, String) = match item {
235            // Lowercase the unparsed expression text so the resulting
236            // alias (which becomes the JSON row key after the
237            // `apply_user_aliases_to_rows` rename) matches the pgwire
238            // lookup key built by `expr_column_names`, which also
239            // lowercases. Without this match the row description
240            // column `count(distinct user_id)` would not resolve to
241            // the JSON value stored under `COUNT(DISTINCT user_id)`,
242            // and the client would see NULL.
243            ast::SelectItem::UnnamedExpr(expr) => (expr, format!("{expr}").to_lowercase()),
244            ast::SelectItem::ExprWithAlias { expr, alias } => (expr, normalize_ident(alias)),
245            _ => continue,
246        };
247        let mut extracted =
248            crate::aggregate_walk::extract_aggregates(expr, &alias, functions, scope)?;
249        aggregates.append(&mut extracted);
250    }
251    Ok(aggregates)
252}
253
254#[cfg(test)]
255mod tests {
256    use super::*;
257
258    #[test]
259    fn parse_intervals() {
260        assert_eq!(parse_interval_to_ms("1h"), 3_600_000);
261        assert_eq!(parse_interval_to_ms("15m"), 900_000);
262        assert_eq!(parse_interval_to_ms("30s"), 30_000);
263        assert_eq!(parse_interval_to_ms("7d"), 604_800_000);
264        // Word-form intervals.
265        assert_eq!(parse_interval_to_ms("1 hour"), 3_600_000);
266        assert_eq!(parse_interval_to_ms("2 hours"), 7_200_000);
267        assert_eq!(parse_interval_to_ms("15 minutes"), 900_000);
268        assert_eq!(parse_interval_to_ms("30 seconds"), 30_000);
269        assert_eq!(parse_interval_to_ms("1 day"), 86_400_000);
270        assert_eq!(parse_interval_to_ms("5 min"), 300_000);
271    }
272
273    fn column(name: &str) -> SqlExpr {
274        SqlExpr::Column {
275            table: None,
276            name: name.into(),
277        }
278    }
279
280    fn time_bucket(interval: SqlValue) -> SqlExpr {
281        SqlExpr::Function {
282            name: "time_bucket".into(),
283            args: vec![SqlExpr::Literal(interval), column("ts")],
284            distinct: false,
285        }
286    }
287
288    #[test]
289    fn time_bucket_key_finds_the_bucket_among_plain_keys() {
290        let functions = FunctionRegistry::new();
291        let keys = vec![
292            column("host"),
293            time_bucket(SqlValue::String("1 hour".into())),
294        ];
295        assert_eq!(
296            time_bucket_key(&keys, &functions),
297            Some(TimeBucketKey {
298                group_index: 1,
299                interval_ms: 3_600_000,
300            })
301        );
302    }
303
304    #[test]
305    fn time_bucket_key_reads_integer_seconds() {
306        let functions = FunctionRegistry::new();
307        let keys = vec![time_bucket(SqlValue::Int(300))];
308        assert_eq!(
309            time_bucket_key(&keys, &functions).map(|k| k.interval_ms),
310            Some(300_000)
311        );
312    }
313
314    #[test]
315    fn time_bucket_key_is_absent_for_plain_columns() {
316        let functions = FunctionRegistry::new();
317        assert_eq!(time_bucket_key(&[column("qtype")], &functions), None);
318    }
319
320    #[test]
321    fn unreadable_interval_is_zero() {
322        let functions = FunctionRegistry::new();
323        let keys = vec![time_bucket(SqlValue::String("soon".into()))];
324        assert_eq!(
325            time_bucket_key(&keys, &functions).map(|k| k.interval_ms),
326            Some(0)
327        );
328    }
329}
nodedb_sql/planner/aggregate.rs

nodedb_sql/planner/
aggregate.rs
