nodedb_query/window/
eval.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Top-level dispatch for window-function evaluation.
4
5use super::aggregate::apply_aggregate_window;
6use super::helpers::build_partitions;
7use super::offset::{apply_lag, apply_lead, apply_nth_value};
8use super::ranking::{
9    apply_cume_dist, apply_dense_rank, apply_ntile, apply_percent_rank, apply_rank,
10    apply_row_number,
11};
12use super::spec::WindowFuncSpec;
13
14/// Evaluate window functions over sorted, partitioned rows.
15///
16/// `rows` is the result set. Each row is a `(doc_id, Value::Object)`. The
17/// same rows are mutated in place with window columns appended to each
18/// document. A window result keeps NaN and ±Infinity. The row array keeps
19/// its input order; each spec's partitions are ordered by that spec's own
20/// ORDER BY, independent of the row array order.
21///
22/// Unknown window function names must be rejected by the planner before
23/// reaching this dispatcher; an unrecognised name here is an internal bug
24/// and panics rather than silently dropping the projection.
25///
26/// A division/modulo-by-zero in any PARTITION BY / ORDER BY / argument
27/// expression propagates as `Err(EvalError::DivisionByZero)` rather than
28/// folding to NULL.
29pub fn evaluate_window_functions(
30    rows: &mut [(String, nodedb_types::Value)],
31    specs: &[WindowFuncSpec],
32) -> Result<(), crate::expr::EvalError> {
33    for spec in specs {
34        let partitions = build_partitions(rows, &spec.partition_by, &spec.order_by)?;
35
36        for partition_indices in &partitions {
37            match spec.func_name.as_str() {
38                "row_number" => apply_row_number(rows, partition_indices, &spec.alias),
39                "rank" => apply_rank(rows, partition_indices, &spec.alias, &spec.order_by)?,
40                "dense_rank" => {
41                    apply_dense_rank(rows, partition_indices, &spec.alias, &spec.order_by)?
42                }
43                "ntile" => apply_ntile(rows, partition_indices, spec),
44                "percent_rank" => {
45                    apply_percent_rank(rows, partition_indices, &spec.alias, &spec.order_by)?
46                }
47                "cume_dist" => {
48                    apply_cume_dist(rows, partition_indices, &spec.alias, &spec.order_by)?
49                }
50                "lag" => apply_lag(rows, partition_indices, spec)?,
51                "lead" => apply_lead(rows, partition_indices, spec)?,
52                "nth_value" => apply_nth_value(rows, partition_indices, spec)?,
53                "sum" | "count" | "avg" | "min" | "max" | "first_value" | "last_value" => {
54                    apply_aggregate_window(rows, partition_indices, spec)?
55                }
56                other => {
57                    unreachable!(
58                        "invariant: SQL planner validates window function names before dispatch; '{other}' is unrecognized and should have been rejected at planning time"
59                    )
60                }
61            }
62        }
63    }
64    Ok(())
65}
66
67#[cfg(test)]
68mod tests {
69    use super::super::spec::{WindowFrame, WindowFuncSpec};
70    use super::evaluate_window_functions;
71    use crate::expr::SqlExpr;
72    use nodedb_types::Value;
73    use serde_json::json;
74
75    type Row = (String, Value);
76
77    fn row(id: &str, doc: serde_json::Value) -> Row {
78        (id.to_string(), Value::from(doc))
79    }
80
81    /// Column `name` of `row` in its JSON form, NULL when absent.
82    fn col(row: &Row, name: &str) -> serde_json::Value {
83        serde_json::Value::from(row.1.get(name).cloned().unwrap_or(Value::Null))
84    }
85
86    fn make_rows() -> Vec<Row> {
87        vec![
88            row("1", json!({"dept": "eng", "salary": 100, "name": "Alice"})),
89            row("2", json!({"dept": "eng", "salary": 120, "name": "Bob"})),
90            row("3", json!({"dept": "eng", "salary": 90, "name": "Carol"})),
91            row("4", json!({"dept": "sales", "salary": 80, "name": "Dave"})),
92            row("5", json!({"dept": "sales", "salary": 110, "name": "Eve"})),
93        ]
94    }
95
96    fn numbered(n: usize) -> Vec<Row> {
97        (1..=n)
98            .map(|i| row(&i.to_string(), json!({ "n": i })))
99            .collect()
100    }
101
102    fn peers() -> Vec<Row> {
103        vec![
104            row("a", json!({"n": 1})),
105            row("b", json!({"n": 1})),
106            row("c", json!({"n": 2})),
107            row("d", json!({"n": 3})),
108        ]
109    }
110
111    fn ordered_by_n(alias: &str, func: &str) -> WindowFuncSpec {
112        WindowFuncSpec {
113            alias: alias.into(),
114            func_name: func.into(),
115            args: vec![],
116            partition_by: vec![],
117            order_by: vec![(SqlExpr::Column("n".into()), true)],
118            frame: WindowFrame::default(),
119        }
120    }
121
122    #[test]
123    fn row_number_single_partition() {
124        let mut rows = make_rows();
125        let spec = WindowFuncSpec {
126            alias: "rn".into(),
127            func_name: "row_number".into(),
128            args: vec![],
129            partition_by: vec![],
130            order_by: vec![],
131            frame: WindowFrame::default(),
132        };
133        evaluate_window_functions(&mut rows, &[spec]).unwrap();
134        assert_eq!(col(&rows[0], "rn"), json!(1));
135        assert_eq!(col(&rows[4], "rn"), json!(5));
136    }
137
138    #[test]
139    fn row_number_partitioned() {
140        let mut rows = make_rows();
141        let spec = WindowFuncSpec {
142            alias: "rn".into(),
143            func_name: "row_number".into(),
144            args: vec![],
145            partition_by: vec![SqlExpr::Column("dept".into())],
146            order_by: vec![],
147            frame: WindowFrame::default(),
148        };
149        evaluate_window_functions(&mut rows, &[spec]).unwrap();
150        assert_eq!(col(&rows[0], "rn"), json!(1));
151        assert_eq!(col(&rows[2], "rn"), json!(3));
152        assert_eq!(col(&rows[3], "rn"), json!(1));
153        assert_eq!(col(&rows[4], "rn"), json!(2));
154    }
155
156    #[test]
157    fn running_sum() {
158        let mut rows = make_rows();
159        let spec = WindowFuncSpec {
160            alias: "running_total".into(),
161            func_name: "sum".into(),
162            args: vec![SqlExpr::Column("salary".into())],
163            partition_by: vec![SqlExpr::Column("dept".into())],
164            order_by: vec![(SqlExpr::Column("salary".into()), true)],
165            frame: WindowFrame::default(),
166        };
167        evaluate_window_functions(&mut rows, &[spec]).unwrap();
168        // The frame runs in salary order within each dept, not in row
169        // arrival order: eng = Carol(90) → Alice(100) → Bob(120).
170        // Integer salaries total exactly as integers.
171        assert_eq!(col(&rows[0], "running_total"), json!(190));
172        assert_eq!(col(&rows[1], "running_total"), json!(310));
173        assert_eq!(col(&rows[2], "running_total"), json!(90));
174        assert_eq!(col(&rows[3], "running_total"), json!(80));
175        assert_eq!(col(&rows[4], "running_total"), json!(190));
176    }
177
178    /// A float overflow reaches the window column as `Infinity`, not NULL.
179    #[test]
180    fn running_sum_keeps_infinity() {
181        let mut rows = vec![
182            (
183                "a".to_string(),
184                Value::Object(
185                    [
186                        ("n".to_string(), Value::Integer(1)),
187                        ("x".to_string(), Value::Float(1e308)),
188                    ]
189                    .into(),
190                ),
191            ),
192            (
193                "b".to_string(),
194                Value::Object(
195                    [
196                        ("n".to_string(), Value::Integer(2)),
197                        ("x".to_string(), Value::Float(1e308)),
198                    ]
199                    .into(),
200                ),
201            ),
202        ];
203        let mut spec = ordered_by_n("total", "sum");
204        spec.args = vec![SqlExpr::Column("x".into())];
205        evaluate_window_functions(&mut rows, &[spec]).unwrap();
206        assert_eq!(rows[0].1.get("total"), Some(&Value::Float(1e308)));
207        assert_eq!(rows[1].1.get("total"), Some(&Value::Float(f64::INFINITY)));
208    }
209
210    /// NaN order keys are peers of each other and sort after every number.
211    #[test]
212    fn nan_order_keys_rank_last_as_peers() {
213        let doc = |n: f64| Value::Object([("n".to_string(), Value::Float(n))].into());
214        let mut rows = vec![
215            ("a".to_string(), doc(f64::NAN)),
216            ("b".to_string(), doc(2.0)),
217            ("c".to_string(), doc(f64::NAN)),
218            ("d".to_string(), doc(f64::INFINITY)),
219        ];
220        evaluate_window_functions(&mut rows, &[ordered_by_n("rnk", "rank")]).unwrap();
221        assert_eq!(col(&rows[1], "rnk"), json!(1));
222        assert_eq!(col(&rows[3], "rnk"), json!(2));
223        assert_eq!(col(&rows[0], "rnk"), json!(3));
224        assert_eq!(col(&rows[2], "rnk"), json!(3));
225    }
226
227    #[test]
228    fn percent_rank_distinct_keys() {
229        let mut rows = numbered(5);
230        evaluate_window_functions(&mut rows, &[ordered_by_n("pr", "percent_rank")]).unwrap();
231        assert_eq!(col(&rows[0], "pr"), json!(0.0));
232        assert_eq!(col(&rows[1], "pr"), json!(0.25));
233        assert_eq!(col(&rows[2], "pr"), json!(0.5));
234        assert_eq!(col(&rows[3], "pr"), json!(0.75));
235        assert_eq!(col(&rows[4], "pr"), json!(1.0));
236    }
237
238    #[test]
239    fn percent_rank_with_peers() {
240        // Peers share the leader's rank, so [1, 1, 2, 3] yields ranks
241        // 1, 1, 3, 4 → percent_rank = 0, 0, 2/3, 3/3.
242        let mut rows = peers();
243        evaluate_window_functions(&mut rows, &[ordered_by_n("pr", "percent_rank")]).unwrap();
244        assert_eq!(col(&rows[0], "pr"), json!(0.0));
245        assert_eq!(col(&rows[1], "pr"), json!(0.0));
246        assert_eq!(col(&rows[2], "pr"), json!(2.0 / 3.0));
247        assert_eq!(col(&rows[3], "pr"), json!(1.0));
248    }
249
250    #[test]
251    fn cume_dist_distinct_keys() {
252        let mut rows = numbered(5);
253        evaluate_window_functions(&mut rows, &[ordered_by_n("cd", "cume_dist")]).unwrap();
254        assert_eq!(col(&rows[0], "cd"), json!(0.2));
255        assert_eq!(col(&rows[1], "cd"), json!(0.4));
256        assert_eq!(col(&rows[2], "cd"), json!(0.6));
257        assert_eq!(col(&rows[3], "cd"), json!(0.8));
258        assert_eq!(col(&rows[4], "cd"), json!(1.0));
259    }
260
261    #[test]
262    fn cume_dist_with_peers() {
263        let mut rows = peers();
264        evaluate_window_functions(&mut rows, &[ordered_by_n("cd", "cume_dist")]).unwrap();
265        // Peers share value of last peer's position / N.
266        assert_eq!(col(&rows[0], "cd"), json!(0.5));
267        assert_eq!(col(&rows[1], "cd"), json!(0.5));
268        assert_eq!(col(&rows[2], "cd"), json!(0.75));
269        assert_eq!(col(&rows[3], "cd"), json!(1.0));
270    }
271
272    #[test]
273    fn nth_value_returns_nth_then_holds() {
274        let mut rows = numbered(5);
275        let mut spec = ordered_by_n("nv", "nth_value");
276        spec.args = vec![
277            SqlExpr::Column("n".into()),
278            SqlExpr::Literal(Value::Integer(2)),
279        ];
280        evaluate_window_functions(&mut rows, &[spec]).unwrap();
281        assert_eq!(col(&rows[0], "nv"), json!(null));
282        assert_eq!(col(&rows[1], "nv"), json!(2));
283        assert_eq!(col(&rows[2], "nv"), json!(2));
284        assert_eq!(col(&rows[3], "nv"), json!(2));
285        assert_eq!(col(&rows[4], "nv"), json!(2));
286    }
287
288    #[test]
289    fn rank_orders_by_spec_order_by_not_row_arrival_order() {
290        // Rows arrive as Alice(100), Bob(120), Carol(90) within dept "eng" —
291        // not sorted by salary. RANK() OVER (ORDER BY salary DESC) must rank
292        // by salary, and the row array order must stay unchanged.
293        let mut rows = make_rows();
294        let spec = WindowFuncSpec {
295            alias: "rnk".into(),
296            func_name: "rank".into(),
297            args: vec![],
298            partition_by: vec![SqlExpr::Column("dept".into())],
299            order_by: vec![(SqlExpr::Column("salary".into()), false)],
300            frame: WindowFrame::default(),
301        };
302        evaluate_window_functions(&mut rows, &[spec]).unwrap();
303        assert_eq!(col(&rows[0], "name"), json!("Alice"));
304        assert_eq!(col(&rows[1], "name"), json!("Bob"));
305        assert_eq!(col(&rows[2], "name"), json!("Carol"));
306        assert_eq!(col(&rows[3], "name"), json!("Dave"));
307        assert_eq!(col(&rows[4], "name"), json!("Eve"));
308        assert_eq!(col(&rows[0], "rnk"), json!(2)); // Alice, salary 100
309        assert_eq!(col(&rows[1], "rnk"), json!(1)); // Bob, salary 120
310        assert_eq!(col(&rows[2], "rnk"), json!(3)); // Carol, salary 90
311        assert_eq!(col(&rows[3], "rnk"), json!(2)); // Dave, salary 80
312        assert_eq!(col(&rows[4], "rnk"), json!(1)); // Eve, salary 110
313    }
314
315    #[test]
316    #[should_panic(expected = "should have been rejected at planning time")]
317    fn unknown_function_panics_at_evaluator() {
318        let mut rows = numbered(2);
319        let spec = WindowFuncSpec {
320            alias: "x".into(),
321            func_name: "frobnicate".into(),
322            args: vec![],
323            partition_by: vec![],
324            order_by: vec![],
325            frame: WindowFrame::default(),
326        };
327        evaluate_window_functions(&mut rows, &[spec]).unwrap();
328    }
329}
nodedb_query/window/eval.rs

nodedb_query/window/
eval.rs
