Skip to main content

nodedb_query/window/
helpers.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Shared helpers for window-function evaluation over document rows.
4//!
5//! A row is `(id, Value::Object)`. `Value` holds every result as computed,
6//! NaN and ±Infinity included, which a JSON number cannot.
7
8use std::collections::HashMap;
9
10use nodedb_types::Value;
11
12use crate::expr::types::SqlExpr;
13use crate::value_ops::{compare_sort_values, sort_peers, value_to_display_string};
14
15/// Group row indices by partition key, preserving first-seen partition order,
16/// then sort each partition's indices by the spec's ORDER BY.
17///
18/// The returned index lists are ordered by `order_by`; the `rows` array
19/// itself keeps its input order — only the per-partition index lists move.
20///
21/// A division/modulo-by-zero in a PARTITION BY or ORDER BY expression
22/// propagates as `Err(EvalError::DivisionByZero)` rather than being folded to
23/// NULL.
24pub(super) fn build_partitions(
25    rows: &[(String, Value)],
26    partition_by: &[SqlExpr],
27    order_by: &[(SqlExpr, bool)],
28) -> Result<Vec<Vec<usize>>, crate::expr::EvalError> {
29    let mut partitions = if partition_by.is_empty() {
30        vec![(0..rows.len()).collect::<Vec<usize>>()]
31    } else {
32        let mut groups: HashMap<String, Vec<usize>> = HashMap::new();
33        let mut order = Vec::new();
34
35        for (i, (_id, doc)) in rows.iter().enumerate() {
36            let key: String = partition_by
37                .iter()
38                .map(|expr| expr.eval(doc).map(|v| partition_key_part(&v)))
39                .collect::<Result<Vec<_>, _>>()?
40                .join("\x00");
41            let entry = groups.entry(key.clone()).or_default();
42            if entry.is_empty() {
43                order.push(key);
44            }
45            entry.push(i);
46        }
47
48        order.iter().filter_map(|k| groups.remove(k)).collect()
49    };
50
51    if !order_by.is_empty() {
52        let mut keys: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
53        for (_id, doc) in rows.iter() {
54            keys.push(
55                order_by
56                    .iter()
57                    .map(|(expr, _)| expr.eval(doc))
58                    .collect::<Result<Vec<_>, _>>()?,
59            );
60        }
61
62        for partition in &mut partitions {
63            partition.sort_by(|&a, &b| compare_order_keys(&keys[a], &keys[b], order_by));
64        }
65    }
66
67    Ok(partitions)
68}
69
70/// One PARTITION BY value as a key fragment. The type name keeps the text
71/// `"1"` apart from the integer `1`.
72fn partition_key_part(v: &Value) -> String {
73    format!("{}:{}", v.type_name(), value_to_display_string(v))
74}
75
76/// Decide NULL placement for one ORDER BY column, shared by every window
77/// evaluator's `compare_order_keys`.
78///
79/// NULL placement follows PostgreSQL's default: ASC places NULLs last, DESC
80/// places NULLs first. A window spec carries no explicit NULLS FIRST/LAST
81/// override, so this default is fixed by direction alone. Returns `None`
82/// when neither value is NULL, leaving the non-null comparison to the
83/// caller.
84pub(super) fn null_order(
85    a_null: bool,
86    b_null: bool,
87    ascending: bool,
88) -> Option<std::cmp::Ordering> {
89    use std::cmp::Ordering;
90    let nulls_first = !ascending;
91    match (a_null, b_null) {
92        (true, true) => Some(Ordering::Equal),
93        (true, false) => Some(if nulls_first {
94            Ordering::Less
95        } else {
96            Ordering::Greater
97        }),
98        (false, true) => Some(if nulls_first {
99            Ordering::Greater
100        } else {
101            Ordering::Less
102        }),
103        (false, false) => None,
104    }
105}
106
107/// Compare two rows' pre-evaluated ORDER BY keys.
108fn compare_order_keys(
109    a: &[Value],
110    b: &[Value],
111    order_by: &[(SqlExpr, bool)],
112) -> std::cmp::Ordering {
113    use std::cmp::Ordering;
114    for (idx, (_, ascending)) in order_by.iter().enumerate() {
115        let (Some(va), Some(vb)) = (a.get(idx), b.get(idx)) else {
116            continue;
117        };
118        let ord = null_order(va.is_null(), vb.is_null(), *ascending).unwrap_or_else(|| {
119            let c = compare_sort_values(va, vb);
120            if *ascending { c } else { c.reverse() }
121        });
122        if ord != Ordering::Equal {
123            return ord;
124        }
125    }
126    Ordering::Equal
127}
128
129pub(super) fn set_window_col(row: &mut Value, alias: &str, val: Value) {
130    if let Value::Object(map) = row {
131        map.insert(alias.to_string(), val);
132    }
133}
134
135/// Numeric view of a value for RANGE offsets and MIN / MAX filtering: a
136/// number or numeric text. `None` for any other value.
137pub(super) fn as_f64(v: &Value) -> Option<f64> {
138    crate::value_ops::value_to_f64(v, false)
139}
140
141/// Returns true when row at index `b` has the same ORDER BY key as row at
142/// index `a` (used by peer-aware ranking like RANK and PERCENT_RANK).
143pub(super) fn order_keys_equal(
144    rows: &[(String, Value)],
145    a: usize,
146    b: usize,
147    order_by: &[(SqlExpr, bool)],
148) -> Result<bool, crate::expr::EvalError> {
149    for (expr, _) in order_by {
150        let va = expr.eval(&rows[a].1)?;
151        let vb = expr.eval(&rows[b].1)?;
152        if !sort_peers(&va, &vb) {
153            return Ok(false);
154        }
155    }
156    Ok(true)
157}