nodedb_query/window/
aggregate.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Aggregate functions used as windows: sum, count, avg, min, max,
4//! first_value, last_value.
5//!
6//! Dispatch logic:
7//! - `RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW` → fast running path
8//!   (preserves SIMD accumulation).
9//! - All other frame combinations → per-row frame evaluator that computes the
10//!   concrete `[start_idx, end_idx]` slice for every row and aggregates over
11//!   it.
12//!
13//! A float result keeps NaN and ±Infinity, as in PostgreSQL.
14
15use nodedb_types::Value;
16
17use super::arg::{ArgValues, arg_at, eval_arg_values};
18use super::extremum::value_replaces;
19use super::frame::{build_peer_groups, evaluate_frame_bounds};
20use super::helpers::{as_f64, set_window_col};
21use super::running::running_aggregate;
22use super::spec::{FrameBound, WindowFuncSpec};
23use crate::numeric_sum::{ExactSum, sum_input};
24
25pub(super) fn apply_aggregate_window(
26    rows: &mut [(String, Value)],
27    indices: &[usize],
28    spec: &WindowFuncSpec,
29) -> Result<(), crate::expr::EvalError> {
30    // The argument is an expression evaluated per row, not a column name: a
31    // zero divisor inside it fails the statement here rather than folding the
32    // aggregate to a silently-wrong total.
33    let arg_values = eval_arg_values(rows, indices, spec, 0)?;
34
35    // Fast path: RANGE UNBOUNDED PRECEDING TO CURRENT ROW is the most common
36    // pattern (the PostgreSQL default for ordered windows). Use the running
37    // accumulator rather than re-aggregating the slice from scratch each row.
38    let use_running = spec.frame.mode == "range"
39        && matches!(spec.frame.start, FrameBound::UnboundedPreceding)
40        && matches!(spec.frame.end, FrameBound::CurrentRow);
41
42    if use_running {
43        running_aggregate(rows, indices, spec, &arg_values)?;
44        return Ok(());
45    }
46
47    per_row_aggregate(rows, indices, spec, &arg_values)
48}
49
50/// Per-row frame evaluator.
51///
52/// For each row position `pos` in the partition:
53/// 1. Resolve the concrete `[start_idx, end_idx]` frame slice via
54///    `evaluate_frame_bounds`.
55/// 2. Aggregate the evaluated argument over `indices[start_idx..=end_idx]`.
56/// 3. Write the result back under `spec.alias`.
57fn per_row_aggregate(
58    rows: &mut [(String, Value)],
59    indices: &[usize],
60    spec: &WindowFuncSpec,
61    arg_values: &ArgValues,
62) -> Result<(), crate::expr::EvalError> {
63    let len = indices.len();
64    if len == 0 {
65        return Ok(());
66    }
67
68    // Extract order-by values for RANGE numeric offsets.
69    let order_expr = spec.order_by.first().map(|(expr, _)| expr);
70    let order_values: Vec<Value> = indices
71        .iter()
72        .map(|&i| match order_expr {
73            Some(expr) => expr.eval(&rows[i].1),
74            None => Ok(Value::Null),
75        })
76        .collect::<Result<Vec<_>, _>>()?;
77
78    // Peer groups needed for GROUPS mode (and for RANGE CurrentRow peer
79    // awareness — reused from the frame module which handles both).
80    let peer_groups: Vec<usize> = if spec.frame.mode == "groups" {
81        build_peer_groups(&order_values)
82    } else {
83        Vec::new()
84    };
85
86    // Numeric view of the evaluated argument, one slot per partition position.
87    let all_vals: Vec<Option<f64>> = (0..len)
88        .map(|pos| as_f64(&arg_at(arg_values, pos)))
89        .collect();
90
91    let results: Vec<Value> = (0..len)
92        .map(|pos| {
93            let (start_idx, end_idx) =
94                evaluate_frame_bounds(&spec.frame, pos, len, &order_values, &peer_groups);
95
96            aggregate_slice(&all_vals, arg_values, spec, start_idx, end_idx)
97        })
98        .collect::<Result<_, _>>()?;
99
100    for (pos, result) in results.into_iter().enumerate() {
101        let row_idx = indices[pos];
102        set_window_col(&mut rows[row_idx].1, &spec.alias, result);
103    }
104    Ok(())
105}
106
107/// The exact SUM / AVG state of the frame slice `[start_idx, end_idx]`.
108fn frame_sum(arg_values: &ArgValues, start_idx: usize, end_idx: usize) -> ExactSum {
109    let mut acc = ExactSum::new();
110    if let Some(values) = arg_values {
111        for v in values.get(start_idx..=end_idx).unwrap_or(&[]) {
112            if let Some(n) = sum_input(v) {
113                acc.add_value(&n);
114            }
115        }
116    }
117    acc
118}
119
120/// Aggregate the evaluated argument over the frame slice
121/// `[start_idx, end_idx]` (partition positions, not row indices).
122fn aggregate_slice(
123    all_vals: &[Option<f64>],
124    arg_values: &ArgValues,
125    spec: &WindowFuncSpec,
126    start_idx: usize,
127    end_idx: usize,
128) -> Result<Value, crate::expr::EvalError> {
129    Ok(match spec.func_name.as_str() {
130        "sum" => frame_sum(arg_values, start_idx, end_idx).sum()?,
131        // `COUNT(*)` counts frame rows; `COUNT(expr)` counts the rows whose
132        // argument is non-NULL, so a NULL argument is excluded rather than
133        // inflating the count.
134        "count" => match arg_values {
135            None => Value::from_u64((end_idx - start_idx + 1) as u64),
136            Some(values) => Value::from_u64(
137                values[start_idx..=end_idx]
138                    .iter()
139                    .filter(|v| !v.is_null())
140                    .count() as u64,
141            ),
142        },
143        "avg" => frame_sum(arg_values, start_idx, end_idx).avg()?,
144        "min" | "max" => {
145            let want_max = spec.func_name == "max";
146            let Some(values) = arg_values else {
147                return Ok(Value::Null);
148            };
149            let frame_values = values.get(start_idx..=end_idx).unwrap_or(&[]);
150            let mut best: Option<&Value> = None;
151            for (numeric, candidate) in all_vals[start_idx..=end_idx].iter().zip(frame_values) {
152                if numeric.is_none() {
153                    continue;
154                }
155                if value_replaces(candidate, best, want_max) {
156                    best = Some(candidate);
157                }
158            }
159            best.cloned().unwrap_or(Value::Null)
160        }
161        "first_value" => arg_at(arg_values, start_idx),
162        "last_value" => arg_at(arg_values, end_idx),
163        _ => Value::Null,
164    })
165}
166
167#[cfg(test)]
168mod tests {
169    use super::super::spec::{FrameBound, WindowFrame, WindowFuncSpec};
170    use super::apply_aggregate_window;
171    use crate::expr::SqlExpr;
172    use nodedb_types::Value;
173    use serde_json::json;
174
175    fn v(j: serde_json::Value) -> Value {
176        Value::from(j)
177    }
178
179    fn res(row: &(String, Value)) -> Value {
180        row.1.get("result").cloned().unwrap_or(Value::Null)
181    }
182
183    fn numbered(n: usize) -> Vec<(String, Value)> {
184        (1..=n)
185            .map(|i| (i.to_string(), v(json!({ "n": i as i64 }))))
186            .collect()
187    }
188
189    fn make_spec(func: &str, field: &str, frame: WindowFrame) -> WindowFuncSpec {
190        WindowFuncSpec {
191            alias: "result".into(),
192            func_name: func.into(),
193            args: if field == "*" {
194                vec![]
195            } else {
196                vec![SqlExpr::Column(field.into())]
197            },
198            partition_by: vec![],
199            order_by: vec![(SqlExpr::Column("n".into()), true)],
200            frame,
201        }
202    }
203
204    fn rows_frame(start: FrameBound, end: FrameBound) -> WindowFrame {
205        WindowFrame {
206            mode: "rows".into(),
207            start,
208            end,
209        }
210    }
211
212    fn range_frame(start: FrameBound, end: FrameBound) -> WindowFrame {
213        WindowFrame {
214            mode: "range".into(),
215            start,
216            end,
217        }
218    }
219
220    fn groups_frame(start: FrameBound, end: FrameBound) -> WindowFrame {
221        WindowFrame {
222            mode: "groups".into(),
223            start,
224            end,
225        }
226    }
227
228    // ── ROWS ──────────────────────────────────────────────────────────────────
229
230    #[test]
231    fn rows_1_preceding_1_following_sum() {
232        let mut rows = numbered(5);
233        let indices: Vec<usize> = (0..5).collect();
234        let spec = make_spec(
235            "sum",
236            "n",
237            rows_frame(FrameBound::Preceding(1), FrameBound::Following(1)),
238        );
239        apply_aggregate_window(&mut rows, &indices, &spec).unwrap();
240        // row 0 (n=1): sum of [1,2] = 3
241        // row 1 (n=2): sum of [1,2,3] = 6
242        // row 2 (n=3): sum of [2,3,4] = 9
243        // row 3 (n=4): sum of [3,4,5] = 12
244        // row 4 (n=5): sum of [4,5] = 9
245        assert_eq!(res(&rows[0]), Value::Integer(3));
246        assert_eq!(res(&rows[1]), Value::Integer(6));
247        assert_eq!(res(&rows[2]), Value::Integer(9));
248        assert_eq!(res(&rows[3]), Value::Integer(12));
249        assert_eq!(res(&rows[4]), Value::Integer(9));
250    }
251
252    #[test]
253    fn rows_unbounded_preceding_current_sum() {
254        let mut rows = numbered(5);
255        let indices: Vec<usize> = (0..5).collect();
256        let spec = make_spec(
257            "sum",
258            "n",
259            rows_frame(FrameBound::UnboundedPreceding, FrameBound::CurrentRow),
260        );
261        apply_aggregate_window(&mut rows, &indices, &spec).unwrap();
262        assert_eq!(res(&rows[0]), Value::Integer(1));
263        assert_eq!(res(&rows[1]), Value::Integer(3));
264        assert_eq!(res(&rows[2]), Value::Integer(6));
265        assert_eq!(res(&rows[3]), Value::Integer(10));
266        assert_eq!(res(&rows[4]), Value::Integer(15));
267    }
268
269    #[test]
270    fn rows_current_unbounded_following_sum() {
271        let mut rows = numbered(5);
272        let indices: Vec<usize> = (0..5).collect();
273        let spec = make_spec(
274            "sum",
275            "n",
276            rows_frame(FrameBound::CurrentRow, FrameBound::UnboundedFollowing),
277        );
278        apply_aggregate_window(&mut rows, &indices, &spec).unwrap();
279        // row 0: sum 1+2+3+4+5=15
280        // row 1: sum 2+3+4+5=14
281        // ...
282        assert_eq!(res(&rows[0]), Value::Integer(15));
283        assert_eq!(res(&rows[1]), Value::Integer(14));
284        assert_eq!(res(&rows[4]), Value::Integer(5));
285    }
286
287    // ── RANGE ─────────────────────────────────────────────────────────────────
288
289    #[test]
290    fn range_unbounded_preceding_current_row_with_ties() {
291        // Values: n in [1, 1, 2, 3] — two rows with n=1 share same frame.
292        let mut rows = vec![
293            ("a".into(), v(json!({"n": 1i64}))),
294            ("b".into(), v(json!({"n": 1i64}))),
295            ("c".into(), v(json!({"n": 2i64}))),
296            ("d".into(), v(json!({"n": 3i64}))),
297        ];
298        let indices: Vec<usize> = (0..4).collect();
299        // Use the fast running path (RANGE UNBOUNDED PRECEDING TO CURRENT ROW)
300        // This is the default frame; both rows a and b must see SUM=2 (both
301        // peers are included up to CURRENT ROW which expands to the last peer).
302        let spec = make_spec(
303            "sum",
304            "n",
305            range_frame(FrameBound::UnboundedPreceding, FrameBound::CurrentRow),
306        );
307        apply_aggregate_window(&mut rows, &indices, &spec).unwrap();
308        // Row a (n=1, pos=0): CURRENT ROW expands to last peer at pos=1, sum=1+1=2
309        assert_eq!(res(&rows[0]), Value::Integer(2));
310        // Row b (n=1, pos=1): same
311        assert_eq!(res(&rows[1]), Value::Integer(2));
312        // Row c (n=2): sum=1+1+2=4
313        assert_eq!(res(&rows[2]), Value::Integer(4));
314        // Row d (n=3): sum=1+1+2+3=7
315        assert_eq!(res(&rows[3]), Value::Integer(7));
316    }
317
318    // ── SUM / AVG exactness ───────────────────────────────────────────────────
319
320    #[test]
321    fn frame_and_running_sum_keep_integers_above_2_pow_53_exact() {
322        let vals = [
323            Value::Integer(9_007_199_254_740_993),
324            Value::Integer(9_007_199_254_740_992),
325        ];
326        let indices: Vec<usize> = (0..2).collect();
327
328        let mut rows = keyed(&vals);
329        apply_aggregate_window(&mut rows, &indices, &make_spec("sum", "v", whole())).unwrap();
330        assert_eq!(res(&rows[0]), Value::Integer(18_014_398_509_481_985));
331
332        let running = range_frame(FrameBound::UnboundedPreceding, FrameBound::CurrentRow);
333        let mut rows = keyed(&vals);
334        apply_aggregate_window(&mut rows, &indices, &make_spec("sum", "v", running)).unwrap();
335        assert_eq!(
336            results(&rows),
337            vec![
338                Value::Integer(9_007_199_254_740_993),
339                Value::Integer(18_014_398_509_481_985)
340            ]
341        );
342
343        let mut rows = keyed(&vals);
344        apply_aggregate_window(&mut rows, &indices, &make_spec("avg", "v", whole())).unwrap();
345        assert_eq!(res(&rows[0]), Value::Float(9_007_199_254_740_992.0));
346    }
347
348    #[test]
349    fn frame_sum_past_i64_and_u64_is_exact() {
350        let vals = [
351            Value::Integer(i64::MAX),
352            Value::from_u64(u64::MAX),
353            Value::from_u64(u64::MAX),
354        ];
355        let indices: Vec<usize> = (0..3).collect();
356        let mut rows = keyed(&vals);
357        apply_aggregate_window(&mut rows, &indices, &make_spec("sum", "v", whole())).unwrap();
358        let want = i128::from(i64::MAX) + 2 * i128::from(u64::MAX);
359        assert_eq!(
360            res(&rows[0]),
361            Value::Decimal(rust_decimal::Decimal::from_i128_with_scale(want, 0))
362        );
363    }
364
365    #[test]
366    fn frame_sum_mixed_int_float_is_float_and_empty_is_null() {
367        let vals = [Value::Integer(2), Value::Float(0.5), Value::Null];
368        let indices: Vec<usize> = (0..3).collect();
369        let mut rows = keyed(&vals);
370        apply_aggregate_window(&mut rows, &indices, &make_spec("sum", "v", whole())).unwrap();
371        assert_eq!(res(&rows[0]), Value::Float(2.5));
372
373        let one_row = rows_frame(FrameBound::CurrentRow, FrameBound::CurrentRow);
374        let mut rows = keyed(&vals);
375        apply_aggregate_window(&mut rows, &indices, &make_spec("sum", "v", one_row)).unwrap();
376        assert_eq!(res(&rows[2]), Value::Null);
377    }
378
379    /// Float overflow is `Infinity` and `Infinity` plus `-Infinity` is NaN;
380    /// both reach the window column as floats, as in PostgreSQL.
381    #[test]
382    fn frame_and_running_sum_keep_non_finite_floats() {
383        let indices: Vec<usize> = (0..2).collect();
384        let overflow = [Value::Float(1e308), Value::Float(1e308)];
385
386        let mut rows = keyed(&overflow);
387        apply_aggregate_window(&mut rows, &indices, &make_spec("sum", "v", whole())).unwrap();
388        assert_eq!(res(&rows[0]), Value::Float(f64::INFINITY));
389
390        let running = range_frame(FrameBound::UnboundedPreceding, FrameBound::CurrentRow);
391        let mut rows = keyed(&overflow);
392        apply_aggregate_window(&mut rows, &indices, &make_spec("sum", "v", running)).unwrap();
393        assert_eq!(
394            results(&rows),
395            vec![Value::Float(1e308), Value::Float(f64::INFINITY)]
396        );
397
398        let cancel = [Value::Float(f64::INFINITY), Value::Float(f64::NEG_INFINITY)];
399        let mut rows = keyed(&cancel);
400        apply_aggregate_window(&mut rows, &indices, &make_spec("avg", "v", whole())).unwrap();
401        assert!(matches!(res(&rows[0]), Value::Float(f) if f.is_nan()));
402    }
403
404    // ── MIN / MAX exactness ───────────────────────────────────────────────────
405
406    /// Rows with order key `n` = 1..; `v` carries the given values.
407    fn keyed(vals: &[Value]) -> Vec<(String, Value)> {
408        vals.iter()
409            .enumerate()
410            .map(|(i, val)| {
411                let doc = std::collections::HashMap::from([
412                    ("n".to_string(), Value::Integer(i as i64 + 1)),
413                    ("v".to_string(), val.clone()),
414                ]);
415                (i.to_string(), Value::Object(doc))
416            })
417            .collect()
418    }
419
420    fn results(rows: &[(String, Value)]) -> Vec<Value> {
421        rows.iter().map(res).collect()
422    }
423
424    fn whole() -> WindowFrame {
425        rows_frame(
426            FrameBound::UnboundedPreceding,
427            FrameBound::UnboundedFollowing,
428        )
429    }
430
431    #[test]
432    fn frame_min_max_keep_integers_above_2_pow_53_exact() {
433        let vals = [
434            Value::Integer(9_007_199_254_740_993),
435            Value::Integer(9_007_199_254_740_992),
436        ];
437        let indices: Vec<usize> = (0..2).collect();
438
439        let mut rows = keyed(&vals);
440        apply_aggregate_window(&mut rows, &indices, &make_spec("min", "v", whole())).unwrap();
441        assert_eq!(res(&rows[0]), Value::Integer(9_007_199_254_740_992));
442
443        let mut rows = keyed(&vals);
444        apply_aggregate_window(&mut rows, &indices, &make_spec("max", "v", whole())).unwrap();
445        assert_eq!(res(&rows[0]), Value::Integer(9_007_199_254_740_993));
446    }
447
448    #[test]
449    fn frame_min_max_over_nanosecond_timestamps() {
450        let vals = [
451            Value::Integer(1_700_000_000_000_000_002),
452            Value::Integer(1_700_000_000_000_000_001),
453            Value::Integer(1_700_000_000_000_000_003),
454        ];
455        let indices: Vec<usize> = (0..3).collect();
456
457        let mut rows = keyed(&vals);
458        apply_aggregate_window(&mut rows, &indices, &make_spec("min", "v", whole())).unwrap();
459        assert_eq!(res(&rows[2]), Value::Integer(1_700_000_000_000_000_001));
460
461        let mut rows = keyed(&vals);
462        apply_aggregate_window(&mut rows, &indices, &make_spec("max", "v", whole())).unwrap();
463        assert_eq!(res(&rows[2]), Value::Integer(1_700_000_000_000_000_003));
464    }
465
466    #[test]
467    fn frame_min_max_mixed_int_float_return_original_type() {
468        let vals = [
469            Value::Integer(9_007_199_254_740_993),
470            Value::Float(9_007_199_254_740_992.0),
471            Value::Float(1.5),
472            Value::Integer(2),
473        ];
474        let indices: Vec<usize> = (0..4).collect();
475
476        let mut rows = keyed(&vals);
477        apply_aggregate_window(&mut rows, &indices, &make_spec("min", "v", whole())).unwrap();
478        assert_eq!(res(&rows[0]), Value::Float(1.5));
479
480        let mut rows = keyed(&vals);
481        apply_aggregate_window(&mut rows, &indices, &make_spec("max", "v", whole())).unwrap();
482        assert_eq!(res(&rows[0]), Value::Integer(9_007_199_254_740_993));
483    }
484
485    /// NaN sorts above every number: MAX over a frame with NaN is NaN, and
486    /// MIN skips it.
487    #[test]
488    fn frame_min_max_place_nan_above_every_number() {
489        let vals = [Value::Float(f64::NAN), Value::Integer(3), Value::Float(1.5)];
490        let indices: Vec<usize> = (0..3).collect();
491
492        let mut rows = keyed(&vals);
493        apply_aggregate_window(&mut rows, &indices, &make_spec("max", "v", whole())).unwrap();
494        assert!(matches!(res(&rows[0]), Value::Float(f) if f.is_nan()));
495
496        let mut rows = keyed(&vals);
497        apply_aggregate_window(&mut rows, &indices, &make_spec("min", "v", whole())).unwrap();
498        assert_eq!(res(&rows[0]), Value::Float(1.5));
499    }
500
501    #[test]
502    fn running_min_max_keep_integers_exact() {
503        let vals = [
504            Value::Integer(9_007_199_254_740_993),
505            Value::Integer(9_007_199_254_740_992),
506            Value::Integer(9_007_199_254_740_995),
507        ];
508        let indices: Vec<usize> = (0..3).collect();
509        let running = || range_frame(FrameBound::UnboundedPreceding, FrameBound::CurrentRow);
510
511        let mut rows = keyed(&vals);
512        apply_aggregate_window(&mut rows, &indices, &make_spec("min", "v", running())).unwrap();
513        assert_eq!(
514            results(&rows),
515            vec![
516                Value::Integer(9_007_199_254_740_993),
517                Value::Integer(9_007_199_254_740_992),
518                Value::Integer(9_007_199_254_740_992),
519            ]
520        );
521
522        let mut rows = keyed(&vals);
523        apply_aggregate_window(&mut rows, &indices, &make_spec("max", "v", running())).unwrap();
524        assert_eq!(
525            results(&rows),
526            vec![
527                Value::Integer(9_007_199_254_740_993),
528                Value::Integer(9_007_199_254_740_993),
529                Value::Integer(9_007_199_254_740_995),
530            ]
531        );
532    }
533
534    // ── GROUPS ────────────────────────────────────────────────────────────────
535
536    #[test]
537    fn groups_1_preceding_1_following_sum() {
538        // Values: [1, 1, 2, 3, 3] — groups [0, 0, 1, 2, 2]
539        let mut rows = vec![
540            ("a".into(), v(json!({"n": 1i64}))),
541            ("b".into(), v(json!({"n": 1i64}))),
542            ("c".into(), v(json!({"n": 2i64}))),
543            ("d".into(), v(json!({"n": 3i64}))),
544            ("e".into(), v(json!({"n": 3i64}))),
545        ];
546        let indices: Vec<usize> = (0..5).collect();
547        let spec = make_spec(
548            "sum",
549            "n",
550            groups_frame(FrameBound::Preceding(1), FrameBound::Following(1)),
551        );
552        apply_aggregate_window(&mut rows, &indices, &spec).unwrap();
553        // pos=0 (group 0): frame → groups 0..=1 → rows 0..=2 → sum=1+1+2=4
554        assert_eq!(res(&rows[0]), Value::Integer(4));
555        // pos=1 (group 0): same frame
556        assert_eq!(res(&rows[1]), Value::Integer(4));
557        // pos=2 (group 1): frame → groups 0..=2 → rows 0..=4 → sum=1+1+2+3+3=10
558        assert_eq!(res(&rows[2]), Value::Integer(10));
559        // pos=3 (group 2): frame → groups 1..=2 → rows 2..=4 → sum=2+3+3=8
560        assert_eq!(res(&rows[3]), Value::Integer(8));
561        // pos=4 (group 2): same
562        assert_eq!(res(&rows[4]), Value::Integer(8));
563    }
564}
nodedb_query/window/aggregate.rs

nodedb_query/window/
aggregate.rs
