Skip to main content

nodedb_query/window/
running.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Peer-aware running-aggregate for RANGE BETWEEN UNBOUNDED PRECEDING AND
4//! CURRENT ROW.
5//!
6//! This is the specialised fast path for the most common ordered-window
7//! pattern (the PostgreSQL default frame for windows with ORDER BY).
8//!
9//! Peer semantics: two rows with equal ORDER BY values are in the same peer
10//! group. All peers in a group see the same aggregate result: the value
11//! computed at the *last* peer in the group (i.e., they all include each
12//! other). This matches PostgreSQL behaviour for RANGE CURRENT ROW.
13
14use nodedb_types::Value;
15
16use super::arg::{ArgValues, arg_at};
17use super::extremum::{extremum_direction, value_replaces};
18use super::helpers::{order_keys_equal, set_window_col};
19use super::spec::WindowFuncSpec;
20use crate::numeric_sum::{ExactSum, sum_input};
21
22/// Apply a peer-aware running aggregate over a sorted partition.
23///
24/// `indices` is the sorted slice of row indices within the partition.
25/// `arg_values` holds the function argument already evaluated once per
26/// partition position; `None` is the no-argument form (`COUNT(*)`).
27pub(super) fn running_aggregate(
28    rows: &mut [(String, Value)],
29    indices: &[usize],
30    spec: &WindowFuncSpec,
31    arg_values: &ArgValues,
32) -> Result<(), crate::expr::EvalError> {
33    let len = indices.len();
34    if len == 0 {
35        return Ok(());
36    }
37
38    // Accumulate state incrementally row-by-row, but defer writing results
39    // until the end of each peer group (so all peers see the group's final
40    // value). We track where the current peer group started.
41    // SUM / AVG total exactly per `ExactSum`.
42    let mut running_sum = ExactSum::new();
43    let mut running_count = 0u64;
44    // MIN / MAX hold the original argument value, compared exactly.
45    let extremum_dir = extremum_direction(&spec.func_name);
46    let mut running_extreme: Option<Value> = None;
47
48    // Indices of rows belonging to the *current* peer group (deferred write).
49    let mut peer_start = 0usize;
50
51    for pos in 0..len {
52        let i = indices[pos];
53        let val = arg_at(arg_values, pos);
54        if sum_input(&val).is_some_and(|n| running_sum.add_value(&n)) {
55            running_count += 1;
56            if let Some(want_max) = extremum_dir
57                && value_replaces(&val, running_extreme.as_ref(), want_max)
58            {
59                running_extreme = Some(val);
60            }
61        } else if spec.func_name == "count" && (arg_values.is_none() || !val.is_null()) {
62            // `COUNT(*)` counts every row; `COUNT(expr)` counts rows whose
63            // argument is non-NULL, including non-numeric values.
64            running_count += 1;
65        }
66
67        // Check if the *next* row starts a new peer group (or we're at the end).
68        let is_last_in_group =
69            pos + 1 == len || !order_keys_equal(rows, i, indices[pos + 1], &spec.order_by)?;
70
71        if is_last_in_group {
72            // Compute the result at the end of this peer group.
73            let result = match spec.func_name.as_str() {
74                "sum" => running_sum.sum()?,
75                "count" => Value::from_u64(running_count),
76                "avg" => running_sum.avg()?,
77                "min" | "max" => running_extreme.clone().unwrap_or(Value::Null),
78                "first_value" => arg_at(arg_values, 0),
79                "last_value" => arg_at(arg_values, pos),
80                _ => Value::Null,
81            };
82
83            // Write the *same* result to every row in the peer group.
84            for &peer_idx in &indices[peer_start..=pos] {
85                set_window_col(&mut rows[peer_idx].1, &spec.alias, result.clone());
86            }
87
88            peer_start = pos + 1;
89        }
90    }
91    Ok(())
92}