nodedb_lite/query/query_ops/
aggregate.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Aggregate and PartialAggregate QueryOp implementations for Lite.
3//!
4//! Supports GROUP BY, aggregate functions (COUNT, SUM, AVG, MIN, MAX,
5//! TS_STDDEV, and the DISTINCT forms of COUNT, SUM and AVG), HAVING,
6//! ORDER BY, and GROUPING SETS (ROLLUP/CUBE) expansion. Any other function
7//! is refused.
8
9use std::collections::HashMap;
10
11use nodedb_physical::physical_plan::query::AggregateSpec;
12use nodedb_physical::physical_plan::{PARAMETERIZED_AGGREGATE_FUNCTIONS, SortKeySpec};
13use nodedb_query::EvalError;
14use nodedb_query::expr::GroupKeySpec;
15use nodedb_query::scan_filter::ScanFilter;
16use nodedb_types::result::QueryResult;
17use nodedb_types::value::Value;
18
19use crate::error::LiteError;
20
21use super::aggregate_numeric::Numeric;
22
23// ─── Type aliases ─────────────────────────────────────────────────────────────
24
25type GroupMap = HashMap<String, (Vec<Value>, Vec<HashMap<String, Value>>)>;
26
27// ─── Public entry points ─────────────────────────────────────────────────────
28
29/// Execute a full Aggregate: apply filters, group, aggregate, HAVING, sort.
30pub fn execute_aggregate(
31    rows: Vec<HashMap<String, Value>>,
32    group_by: &[GroupKeySpec],
33    aggregates: &[AggregateSpec],
34    filters: &[u8],
35    having: &[u8],
36    sort_keys: &[SortKeySpec],
37    grouping_sets: &[Vec<u32>],
38) -> Result<QueryResult, LiteError> {
39    let scan_filters = decode_filters(filters)?;
40    let having_filters = decode_filters(having)?;
41
42    let filtered = retain_matching(rows, &scan_filters)?;
43
44    if grouping_sets.is_empty() {
45        // Plain GROUP BY — single grouping set containing all group_by columns.
46        let grouped = group_rows(&filtered, group_by)?;
47        let mut result_rows = compute_aggregate_groups(grouped, group_by, aggregates)?;
48        apply_having(&mut result_rows, &having_filters, group_by, aggregates)?;
49        apply_sort(&mut result_rows, sort_keys, group_by, aggregates);
50        let columns = make_columns(group_by, aggregates);
51        Ok(QueryResult {
52            columns,
53            rows: result_rows,
54            rows_affected: 0,
55            command: None,
56        })
57    } else {
58        // GROUPING SETS: union results for each subset.
59        let mut all_rows: Vec<Vec<Value>> = Vec::new();
60        for set_indices in grouping_sets {
61            let subset: Vec<GroupKeySpec> = set_indices
62                .iter()
63                .filter_map(|&i| group_by.get(i as usize).cloned())
64                .collect();
65            let grouped = group_rows(&filtered, &subset)?;
66            let mut set_rows = compute_aggregate_groups(grouped, &subset, aggregates)?;
67            // Null-fill absent group columns.
68            let full_col_count = group_by.len();
69            for row in &mut set_rows {
70                while row.len() < full_col_count + aggregates.len() {
71                    row.insert(set_indices.len(), Value::Null);
72                }
73            }
74            apply_having(&mut set_rows, &having_filters, &subset, aggregates)?;
75            all_rows.extend(set_rows);
76        }
77        apply_sort(&mut all_rows, sort_keys, group_by, aggregates);
78        let columns = make_columns(group_by, aggregates);
79        Ok(QueryResult {
80            columns,
81            rows: all_rows,
82            rows_affected: 0,
83            command: None,
84        })
85    }
86}
87
88/// Execute a PartialAggregate.
89///
90/// On single-node Lite, data is already local so this produces the same
91/// final result as Aggregate (no HAVING/sort/grouping_sets) and encodes it
92/// in the partial-state format Origin expects: each output row is
93/// `[group_key_col_0, ..., group_key_col_N, count_i64, sum_f64, min, max, ...]`
94/// in the same column order as `make_columns`, with a leading `__partial`
95/// boolean column set to `true` so the Control-Plane merger can distinguish
96/// partial from final results.
97pub fn execute_partial_aggregate(
98    rows: Vec<HashMap<String, Value>>,
99    group_by: &[String],
100    aggregates: &[AggregateSpec],
101    filters: &[u8],
102) -> Result<QueryResult, LiteError> {
103    // `QueryOp::PartialAggregate` groups by bare column names only, so the
104    // column → spec lift is lossless here.
105    let group_by: Vec<GroupKeySpec> = group_by.iter().map(GroupKeySpec::column).collect();
106    let group_by = group_by.as_slice();
107    let scan_filters = decode_filters(filters)?;
108    let filtered = retain_matching(rows, &scan_filters)?;
109
110    let grouped = group_rows(&filtered, group_by)?;
111    let result_rows = compute_aggregate_groups(grouped, group_by, aggregates)?;
112
113    let mut columns = vec!["__partial".to_string()];
114    columns.extend(make_columns(group_by, aggregates));
115
116    let output = result_rows
117        .into_iter()
118        .map(|mut row| {
119            row.insert(0, Value::Bool(true));
120            row
121        })
122        .collect();
123
124    Ok(QueryResult {
125        columns,
126        rows: output,
127        rows_affected: 0,
128        command: None,
129    })
130}
131
132// ─── Internal helpers ─────────────────────────────────────────────────────────
133
134fn decode_filters(bytes: &[u8]) -> Result<Vec<ScanFilter>, LiteError> {
135    if bytes.is_empty() {
136        return Ok(Vec::new());
137    }
138    zerompk::from_msgpack(bytes).map_err(|e| LiteError::Serialization {
139        detail: format!("decode filters: {e}"),
140    })
141}
142
143/// Group rows by the specified key columns.
144///
145/// `Value` does not implement `Hash`/`Eq`, so we use a string serialisation of
146/// the key for the `HashMap` discriminant and carry the original `Vec<Value>`
147/// alongside for output.
148pub(crate) fn group_rows(
149    rows: &[HashMap<String, Value>],
150    group_by: &[GroupKeySpec],
151) -> Result<GroupMap, LiteError> {
152    let mut groups: GroupMap = HashMap::new();
153    for row in rows {
154        let key: Vec<Value> = group_by
155            .iter()
156            .map(|spec| group_key_value(row, spec))
157            .collect::<Result<Vec<_>, LiteError>>()?;
158        let key_str = value_key_str(&key);
159        groups
160            .entry(key_str)
161            .or_insert_with(|| (key, Vec::new()))
162            .1
163            .push(row.clone());
164    }
165    Ok(groups)
166}
167
168/// Keep only the rows matching every filter. Evaluation is fallible (a
169/// divide-by-zero predicate fails the statement rather than dropping the row),
170/// so this replaces the `filter(..).collect()` form it grew out of.
171fn retain_matching(
172    rows: Vec<HashMap<String, Value>>,
173    filters: &[ScanFilter],
174) -> Result<Vec<HashMap<String, Value>>, LiteError> {
175    if filters.is_empty() {
176        return Ok(rows);
177    }
178    let mut kept = Vec::with_capacity(rows.len());
179    for row in rows {
180        let doc = Value::Object(row.clone());
181        if ScanFilter::all_match_value(filters, &doc)? {
182            kept.push(row);
183        }
184    }
185    Ok(kept)
186}
187
188/// Resolve one group key for a row.
189///
190/// A spec carrying an expression (e.g. `GROUP BY date_trunc('day', ts)`) is
191/// evaluated against the row; otherwise the named field is extracted directly.
192/// A spec with neither groups everything under `Null`.
193fn group_key_value(row: &HashMap<String, Value>, spec: &GroupKeySpec) -> Result<Value, LiteError> {
194    match (&spec.expr, &spec.field) {
195        (Some(expr), _) => Ok(expr.eval(&Value::Object(row.clone()))?),
196        (None, Some(field)) => Ok(row.get(field).cloned().unwrap_or(Value::Null)),
197        (None, None) => Ok(Value::Null),
198    }
199}
200
201/// Stable string discriminant for a `Vec<Value>` group key.
202fn value_key_str(key: &[Value]) -> String {
203    key.iter()
204        .map(value_discriminant)
205        .collect::<Vec<_>>()
206        .join("|")
207}
208
209pub(crate) fn value_discriminant(v: &Value) -> String {
210    match v {
211        Value::Null => "null".to_string(),
212        Value::Bool(b) => format!("b:{b}"),
213        Value::Integer(n) => format!("i:{n}"),
214        Value::Float(f) => format!("f:{f}"),
215        Value::String(s) => format!("s:{s}"),
216        Value::Uuid(s) | Value::Ulid(s) | Value::Regex(s) => format!("str:{s}"),
217        Value::Bytes(b) => format!("by:{}", b.len()),
218        Value::Array(a) | Value::Set(a) => format!("a:{}", a.len()),
219        Value::Object(m) => format!("o:{}", m.len()),
220        _ => "other".to_string(),
221    }
222}
223
224fn compute_aggregate_groups(
225    groups: GroupMap,
226    _group_by: &[GroupKeySpec],
227    aggregates: &[AggregateSpec],
228) -> Result<Vec<Vec<Value>>, LiteError> {
229    let mut result = Vec::with_capacity(groups.len());
230    for (_key_str, (key_vals, group_rows)) in groups {
231        let mut row: Vec<Value> = key_vals;
232        for spec in aggregates {
233            let agg_val = compute_one_aggregate(&group_rows, spec)?;
234            row.push(agg_val);
235        }
236        result.push(row);
237    }
238    Ok(result)
239}
240
241fn compute_one_aggregate(
242    group: &[HashMap<String, Value>],
243    spec: &AggregateSpec,
244) -> Result<Value, LiteError> {
245    let func = spec.function.to_uppercase();
246    // Lite has no accumulator for a parameterized aggregate (`approx_topk`,
247    // `approx_percentile`, `percentile_cont`), so its parameter has nothing
248    // to feed.
249    if spec.param.is_some()
250        || PARAMETERIZED_AGGREGATE_FUNCTIONS.contains(&spec.function.to_lowercase().as_str())
251    {
252        return Err(unsupported_function(&spec.function));
253    }
254    if func == "COUNT" && spec.field == "*" && spec.expr.is_none() {
255        return Ok(Value::Integer(group.len() as i64));
256    }
257    // One input per row: the evaluated argument expression, or the field.
258    // A NULL input is skipped by every function below.
259    let inputs: Vec<Value> = group
260        .iter()
261        .map(|row| match &spec.expr {
262            Some(expr) => Ok(expr.eval(&Value::Object(row.clone()))?),
263            None => Ok(row.get(&spec.field).cloned().unwrap_or(Value::Null)),
264        })
265        .collect::<Result<_, LiteError>>()?;
266    let non_null = || inputs.iter().filter(|v| !matches!(v, Value::Null));
267    let distinct = || {
268        let mut seen: Vec<&Value> = Vec::new();
269        for v in non_null() {
270            if !seen.contains(&v) {
271                seen.push(v);
272            }
273        }
274        seen
275    };
276    match func.as_str() {
277        "COUNT" => Ok(Value::Integer(non_null().count() as i64)),
278        "COUNT_DISTINCT" => Ok(Value::Integer(distinct().len() as i64)),
279        "SUM" => Numeric::collect(non_null())
280            .sum()
281            .map_err(|e| total_error(spec, e)),
282        "SUM_DISTINCT" => Numeric::collect(distinct().into_iter())
283            .sum()
284            .map_err(|e| total_error(spec, e)),
285        "AVG" => Numeric::collect(non_null())
286            .avg()
287            .map_err(|e| total_error(spec, e)),
288        "AVG_DISTINCT" => Numeric::collect(distinct().into_iter())
289            .avg()
290            .map_err(|e| total_error(spec, e)),
291        "TS_STDDEV" => {
292            // Population standard deviation; NULL below two inputs.
293            let mut accum = nodedb_query::ts_functions::TsStddevAccum::new();
294            for v in Numeric::collect(non_null()).as_f64() {
295                accum.update(v);
296            }
297            Ok(accum
298                .evaluate_population()
299                .map_or(Value::Null, Value::Float))
300        }
301        "MIN" => {
302            let numeric = Numeric::collect(non_null());
303            if !numeric.is_empty() {
304                Ok(numeric.min())
305            } else {
306                // String MIN
307                let s = non_null()
308                    .filter_map(|v| v.as_str().map(|s| s.to_string()))
309                    .min();
310                Ok(s.map(Value::String).unwrap_or(Value::Null))
311            }
312        }
313        "MAX" => {
314            let numeric = Numeric::collect(non_null());
315            if !numeric.is_empty() {
316                Ok(numeric.max())
317            } else {
318                let s = non_null()
319                    .filter_map(|v| v.as_str().map(|s| s.to_string()))
320                    .max();
321                Ok(s.map(Value::String).unwrap_or(Value::Null))
322            }
323        }
324        _ => Err(unsupported_function(&spec.function)),
325    }
326}
327
328/// The error of a SUM or AVG total. An exact total past the `Decimal` range
329/// is `NumericValueOutOfRange`, SQLSTATE `22003` as on Origin, and names the
330/// aggregate and its column.
331fn total_error(spec: &AggregateSpec, e: EvalError) -> LiteError {
332    match e {
333        EvalError::NumericOverflow { .. } => LiteError::NumericValueOutOfRange {
334            detail: format!(
335                "{}({}): integer total out of range",
336                spec.function.to_lowercase(),
337                spec.field
338            ),
339        },
340        other @ (EvalError::DivisionByZero
341        | EvalError::UnknownFunction { .. }
342        | EvalError::VectorDimensionMismatch { .. }
343        | EvalError::ArgumentType { .. }
344        | EvalError::InvalidJsonPath { .. }) => LiteError::from(other),
345    }
346}
347
348fn unsupported_function(function: &str) -> LiteError {
349    LiteError::Unsupported {
350        detail: format!("aggregate function {function}() is not supported in Lite"),
351    }
352}
353
354fn apply_having(
355    rows: &mut Vec<Vec<Value>>,
356    having: &[ScanFilter],
357    group_by: &[GroupKeySpec],
358    aggregates: &[AggregateSpec],
359) -> Result<(), LiteError> {
360    if having.is_empty() {
361        return Ok(());
362    }
363    let columns = make_columns(group_by, aggregates);
364    // HAVING evaluation is fallible, so the keep-decision is computed before
365    // the retain rather than inside a closure that cannot propagate.
366    let mut keep = Vec::with_capacity(rows.len());
367    for row in rows.iter() {
368        let mut map = HashMap::new();
369        for (col, val) in columns.iter().zip(row.iter()) {
370            map.insert(col.clone(), val.clone());
371        }
372        let doc = Value::Object(map);
373        keep.push(ScanFilter::all_match_value(having, &doc)?);
374    }
375    let mut iter = keep.into_iter();
376    rows.retain(|_| iter.next().unwrap_or(true));
377    Ok(())
378}
379
380fn apply_sort(
381    rows: &mut [Vec<Value>],
382    sort_keys: &[SortKeySpec],
383    group_by: &[GroupKeySpec],
384    aggregates: &[AggregateSpec],
385) {
386    if sort_keys.is_empty() {
387        return;
388    }
389    let columns = make_columns(group_by, aggregates);
390    rows.sort_by(|a, b| {
391        for key in sort_keys {
392            // Aggregate output rows are addressed by output column name, so a
393            // computed key has nothing to match here — skip it, matching the
394            // Origin executor's `sort_aggregated_rows`.
395            let Some(col) = key.as_column() else {
396                continue;
397            };
398            let idx = columns.iter().position(|c| c == col);
399            if let Some(i) = idx {
400                let av = a.get(i).unwrap_or(&Value::Null);
401                let bv = b.get(i).unwrap_or(&Value::Null);
402                // NULL placement is absolute; the direction never flips it.
403                let ord = match key.order_nulls(av == &Value::Null, bv == &Value::Null) {
404                    Some(ord) => ord,
405                    None => key.direct(value_cmp(av, bv)),
406                };
407                if ord != std::cmp::Ordering::Equal {
408                    return ord;
409                }
410            }
411        }
412        std::cmp::Ordering::Equal
413    });
414}
415
416pub(crate) fn value_cmp(a: &Value, b: &Value) -> std::cmp::Ordering {
417    match (a, b) {
418        (Value::Integer(x), Value::Integer(y)) => x.cmp(y),
419        (Value::Float(x), Value::Float(y)) => x.partial_cmp(y).unwrap_or(std::cmp::Ordering::Equal),
420        (Value::Integer(x), Value::Float(y)) => (*x as f64)
421            .partial_cmp(y)
422            .unwrap_or(std::cmp::Ordering::Equal),
423        (Value::Float(x), Value::Integer(y)) => x
424            .partial_cmp(&(*y as f64))
425            .unwrap_or(std::cmp::Ordering::Equal),
426        (Value::String(x), Value::String(y)) => x.cmp(y),
427        (Value::Bool(x), Value::Bool(y)) => x.cmp(y),
428        (Value::Null, Value::Null) => std::cmp::Ordering::Equal,
429        (Value::Null, _) => std::cmp::Ordering::Less,
430        (_, Value::Null) => std::cmp::Ordering::Greater,
431        _ => std::cmp::Ordering::Equal,
432    }
433}
434
435pub(crate) fn make_columns(group_by: &[GroupKeySpec], aggregates: &[AggregateSpec]) -> Vec<String> {
436    let mut cols: Vec<String> = group_by.iter().map(|s| s.output_name.clone()).collect();
437    for spec in aggregates {
438        if let Some(alias) = &spec.user_alias {
439            cols.push(alias.clone());
440        } else {
441            cols.push(spec.alias.clone());
442        }
443    }
444    cols
445}
446
447// ─── Tests ────────────────────────────────────────────────────────────────────
448
449#[cfg(test)]
450mod tests {
451    use super::*;
452    use nodedb_physical::physical_plan::query::AggregateSpec;
453    use nodedb_query::expr::{BinaryOp, SqlExpr};
454
455    fn make_spec(function: &str, field: &str, alias: &str) -> AggregateSpec {
456        AggregateSpec {
457            function: function.to_string(),
458            alias: alias.to_string(),
459            user_alias: None,
460            field: field.to_string(),
461            expr: None,
462            param: None,
463        }
464    }
465
466    fn make_rows(data: &[(&str, i64, f64)]) -> Vec<HashMap<String, Value>> {
467        data.iter()
468            .map(|(cat, n, price)| {
469                let mut m = HashMap::new();
470                m.insert("category".into(), Value::String(cat.to_string()));
471                m.insert("count".into(), Value::Integer(*n));
472                m.insert("price".into(), Value::Float(*price));
473                m
474            })
475            .collect()
476    }
477
478    #[test]
479    fn group_by_with_having() {
480        let rows = make_rows(&[("a", 1, 10.0), ("a", 2, 20.0), ("b", 3, 5.0)]);
481        let aggregates = vec![
482            make_spec("COUNT", "*", "cnt"),
483            make_spec("SUM", "price", "total"),
484        ];
485        // HAVING total > 10
486        let having_filter = ScanFilter {
487            field: "total".into(),
488            op: nodedb_query::scan_filter::FilterOp::Gt,
489            value: Value::Float(10.0),
490            clauses: Vec::new(),
491            expr: None,
492        };
493        let having_bytes = zerompk::to_msgpack_vec(&vec![having_filter]).unwrap();
494
495        let result = execute_aggregate(
496            rows,
497            &[GroupKeySpec::column("category")],
498            &aggregates,
499            &[],
500            &having_bytes,
501            &[],
502            &[],
503        )
504        .unwrap();
505        // Only group "a" has total=30 > 10; group "b" has total=5.0
506        assert_eq!(result.rows.len(), 1);
507        let cat = &result.rows[0][0];
508        assert_eq!(cat, &Value::String("a".into()));
509    }
510
511    #[test]
512    fn integer_sum_past_i64_does_not_clamp() {
513        let rows = make_rows(&[("a", i64::MAX, 1.0), ("a", 1, 1.0)]);
514        let aggregates = vec![make_spec("SUM", "count", "total")];
515        let result = execute_aggregate(rows, &[], &aggregates, &[], &[], &[], &[]).unwrap();
516        assert_eq!(
517            result.rows[0][0],
518            Value::Decimal(rust_decimal::Decimal::from_i128_with_scale(
519                i128::from(i64::MAX) + 1,
520                0
521            ))
522        );
523    }
524
525    #[test]
526    fn integer_sum_overflow_names_aggregate_and_column() {
527        let row = |amount| {
528            let mut m = HashMap::new();
529            m.insert("amount".to_string(), Value::Decimal(amount));
530            m
531        };
532        let max = rust_decimal::Decimal::MAX;
533        let aggregates = vec![make_spec("SUM", "amount", "total")];
534        let err = execute_aggregate(
535            vec![row(max), row(max)],
536            &[],
537            &aggregates,
538            &[],
539            &[],
540            &[],
541            &[],
542        )
543        .unwrap_err();
544        match err {
545            LiteError::NumericValueOutOfRange { detail } => {
546                assert!(detail.contains("sum(amount)"), "{detail}");
547            }
548            other => panic!("expected NumericValueOutOfRange, got {other:?}"),
549        }
550    }
551
552    #[test]
553    fn partial_aggregate_prefix_column() {
554        let rows = make_rows(&[("x", 1, 1.0)]);
555        let aggregates = vec![make_spec("COUNT", "*", "cnt")];
556        let result =
557            execute_partial_aggregate(rows, &["category".to_string()], &aggregates, &[]).unwrap();
558        assert_eq!(result.columns[0], "__partial");
559        assert_eq!(result.rows[0][0], Value::Bool(true));
560    }
561
562    #[test]
563    fn group_by_computed_expression_key() {
564        // GROUP BY on a computed key (count + count) rather than a bare column:
565        // rows with counts 1 and 3 collapse into distinct groups 2 and 6, while
566        // the two rows with count 1 share a group.
567        let rows = make_rows(&[("a", 1, 10.0), ("b", 1, 20.0), ("c", 3, 5.0)]);
568        let aggregates = vec![make_spec("COUNT", "*", "cnt")];
569
570        let doubled = GroupKeySpec {
571            output_name: "doubled".to_string(),
572            field: None,
573            expr: Some(SqlExpr::BinaryOp {
574                left: Box::new(SqlExpr::Column("count".to_string())),
575                op: BinaryOp::Add,
576                right: Box::new(SqlExpr::Column("count".to_string())),
577            }),
578        };
579
580        let result = execute_aggregate(rows, &[doubled], &aggregates, &[], &[], &[], &[]).unwrap();
581
582        assert_eq!(result.columns[0], "doubled");
583        assert_eq!(result.rows.len(), 2);
584
585        let mut got: Vec<(i64, i64)> = result
586            .rows
587            .iter()
588            .filter_map(|r| match (&r[0], &r[1]) {
589                (Value::Integer(k), Value::Integer(n)) => Some((*k, *n)),
590                _ => None,
591            })
592            .collect();
593        got.sort_unstable();
594        assert_eq!(got, vec![(2, 2), (6, 1)]);
595    }
596}
nodedb_lite/query/query_ops/aggregate.rs

nodedb_lite/query/query_ops/
aggregate.rs
