nodedb_query/window/
value_partition.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Partition building and per-partition ORDER BY sort for the Value-native
4//! window evaluator.
5
6use std::collections::HashMap;
7
8use nodedb_types::Value;
9
10use super::helpers::null_order;
11use super::spec::WindowFuncSpec;
12use super::value_eval::{WindowError, cmp_values, eval_arg_for_row};
13use crate::expr::types::SqlExpr;
14
15/// Group row indices by partition key, preserving first-seen partition
16/// order, then sort each partition's indices by the spec's ORDER BY.
17///
18/// The returned index lists are ordered by `spec.order_by`; `rows` itself
19/// keeps its input order — only the per-partition index lists move.
20pub(super) fn build_value_partitions(
21    rows: &[Vec<Value>],
22    column_index: &HashMap<String, usize>,
23    spec: &WindowFuncSpec,
24) -> Result<Vec<Vec<usize>>, WindowError> {
25    let mut partitions = if spec.partition_by.is_empty() {
26        vec![(0..rows.len()).collect::<Vec<usize>>()]
27    } else {
28        let mut groups: HashMap<String, Vec<usize>> = HashMap::new();
29        let mut order: Vec<String> = Vec::new();
30
31        for (i, row) in rows.iter().enumerate() {
32            let key = partition_key(row, column_index, &spec.partition_by)?;
33            let entry = groups.entry(key.clone()).or_default();
34            if entry.is_empty() {
35                order.push(key);
36            }
37            entry.push(i);
38        }
39
40        order.iter().filter_map(|k| groups.remove(k)).collect()
41    };
42
43    if !spec.order_by.is_empty() {
44        let mut keys: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
45        for row in rows.iter() {
46            keys.push(
47                spec.order_by
48                    .iter()
49                    .map(|(expr, _)| eval_arg_for_row(expr, row, column_index))
50                    .collect::<Result<Vec<_>, _>>()?,
51            );
52        }
53
54        for partition in &mut partitions {
55            partition.sort_by(|&a, &b| compare_order_keys(&keys[a], &keys[b], &spec.order_by));
56        }
57    }
58
59    Ok(partitions)
60}
61
62fn partition_key(
63    row: &[Value],
64    column_index: &HashMap<String, usize>,
65    partition_by: &[SqlExpr],
66) -> Result<String, WindowError> {
67    Ok(partition_by
68        .iter()
69        .map(|expr| {
70            let v = eval_arg_for_row(expr, row, column_index)?;
71            Ok(format!("{v:?}"))
72        })
73        .collect::<Result<Vec<_>, WindowError>>()?
74        .join("\x00"))
75}
76
77/// Compare two rows' pre-evaluated ORDER BY keys.
78fn compare_order_keys(
79    a: &[Value],
80    b: &[Value],
81    order_by: &[(SqlExpr, bool)],
82) -> std::cmp::Ordering {
83    use std::cmp::Ordering;
84    for (idx, (_, ascending)) in order_by.iter().enumerate() {
85        let (Some(va), Some(vb)) = (a.get(idx), b.get(idx)) else {
86            continue;
87        };
88        let ord = null_order(
89            matches!(va, Value::Null),
90            matches!(vb, Value::Null),
91            *ascending,
92        )
93        .unwrap_or_else(|| {
94            let c = cmp_values(va, vb);
95            if *ascending { c } else { c.reverse() }
96        });
97        if ord != Ordering::Equal {
98            return ord;
99        }
100    }
101    Ordering::Equal
102}
nodedb_query/window/value_partition.rs

nodedb_query/window/
value_partition.rs
