linkrs_query/executor/streaming/operators/blocking/
window.rs

1use crate::executor::expression::evaluator::ExpressionEvaluator;
2use crate::executor::streaming::executor::{SortDirection, ValueRowContext};
3use crate::executor::streaming::helpers::compare_values;
4use crate::executor::streaming::spill::{HashPartitionSpiller, SpilledRun};
5use linkrs_core::columnar::MaterializedBatch;
6use linkrs_core::types::expr::Expression;
7use linkrs_core::value::NullType;
8use linkrs_core::Value;
9
10/// Sort the rows of a single window partition by the ORDER BY expressions.
11///
12/// Order keys are precomputed once per row (previously re-evaluated on
13/// every comparison); the sort is stable so ties keep the input order.
14pub(crate) fn sort_partition_rows(
15    partition_rows: &mut [(usize, Vec<Value>)],
16    col_names: &[String],
17    order_by_exprs: &[Expression],
18    order_by_directions: &[SortDirection],
19) {
20    if order_by_exprs.is_empty() {
21        return;
22    }
23    let keys: Vec<Vec<Value>> = partition_rows
24        .iter()
25        .map(|(_, row)| {
26            order_by_exprs
27                .iter()
28                .map(|expr| {
29                    let mut ctx = ValueRowContext::from_names(row.clone(), col_names.to_vec());
30                    ExpressionEvaluator::evaluate(expr, &mut ctx)
31                        .unwrap_or(Value::Null(NullType::Null))
32                })
33                .collect()
34        })
35        .collect();
36    let mut order: Vec<usize> = (0..partition_rows.len()).collect();
37    order.sort_by(|&i, &j| {
38        for (idx, _) in order_by_exprs.iter().enumerate() {
39            let direction = order_by_directions
40                .get(idx)
41                .copied()
42                .unwrap_or(SortDirection::Ascending);
43            let cmp = compare_values(&keys[i][idx], &keys[j][idx]);
44            let final_cmp = match direction {
45                SortDirection::Ascending => cmp,
46                SortDirection::Descending => cmp.reverse(),
47            };
48            if final_cmp != std::cmp::Ordering::Equal {
49                return final_cmp;
50            }
51        }
52        std::cmp::Ordering::Equal
53    });
54    let mut sorted = Vec::with_capacity(partition_rows.len());
55    for &i in &order {
56        sorted.push(std::mem::take(&mut partition_rows[i]));
57    }
58    partition_rows.clone_from_slice(&sorted);
59}
60
61#[derive(Debug)]
62pub struct WindowFunctionState {
63    /// Column-oriented input buffer (was `Vec<Vec<Value>>`).
64    pub batch: MaterializedBatch,
65    pub col_names: Vec<String>,
66    /// Output buffer replacing the drained row iterator.
67    pub result_batch: MaterializedBatch,
68    /// Output cursor into `result_batch`.
69    pub emitted_offset: usize,
70    pub partition_spiller: Option<HashPartitionSpiller>,
71    pub spilled_runs: Vec<Option<SpilledRun>>,
72    pub current_partition: usize,
73    pub has_spilled: bool,
74    /// True once all partitions have been emitted.
75    pub output_complete: bool,
76}
77
78#[derive(Debug)]
79pub struct WindowState {
80    /// Column-oriented input buffer (was `Vec<Vec<Value>>`).
81    pub batch: MaterializedBatch,
82    pub col_names: Vec<String>,
83    /// Output buffer replacing the drained row iterator.
84    pub result_batch: MaterializedBatch,
85    /// Output cursor into `result_batch`.
86    pub emitted_offset: usize,
87    pub partition_spiller: Option<HashPartitionSpiller>,
88    pub spilled_runs: Vec<Option<SpilledRun>>,
89    pub current_partition: usize,
90    pub has_spilled: bool,
91    /// True once all partitions have been emitted.
92    pub output_complete: bool,
93}
94
95/// Compute the result rows for one window partition.
96///
97/// The partition rows must already be ordered as desired (partition-order
98/// sort happens at the caller). Positional window values are computed per
99/// `window_exprs` and appended to each source row.
100pub(crate) fn compute_window_partition_result(
101    partition_rows: &[(usize, Vec<Value>)],
102    col_names: &[String],
103    window_exprs: &[Expression],
104) -> Vec<Vec<Value>> {
105    let mut result_rows = Vec::with_capacity(partition_rows.len());
106    for (pos, (_, row)) in partition_rows.iter().enumerate() {
107        let mut result_row = row.clone();
108        for window_expr in window_exprs.iter() {
109            if let Expression::WindowFunction { name, args, .. } = window_expr {
110                let func_args: Vec<Value> = args
111                    .iter()
112                    .map(|arg| {
113                        let mut ctx = ValueRowContext::from_names(row.clone(), col_names.to_vec());
114                        ExpressionEvaluator::evaluate(arg, &mut ctx)
115                            .unwrap_or(Value::Null(NullType::Null))
116                    })
117                    .collect();
118                let window_result = compute_window_function(name, &func_args, partition_rows, pos);
119                result_row.push(window_result);
120            }
121        }
122        result_rows.push(result_row);
123    }
124    result_rows
125}
126
127pub(crate) fn compute_window_function(
128    name: &str,
129    args: &[Value],
130    partition_rows: &[(usize, Vec<Value>)],
131    current_pos: usize,
132) -> Value {
133    let value_to_i64 = |val: &Value| -> i64 {
134        match val {
135            Value::Int(i) => *i as i64,
136            Value::BigInt(i) => *i,
137            _ => 1,
138        }
139    };
140
141    match name {
142        "row_number" => Value::BigInt((current_pos + 1) as i64),
143        "rank" => Value::BigInt((current_pos + 1) as i64),
144        "dense_rank" => Value::BigInt((current_pos + 1) as i64),
145        "lead" => {
146            let offset = if !args.is_empty() {
147                value_to_i64(&args[0]) as usize
148            } else {
149                1
150            };
151            let default_val = if args.len() > 1 {
152                args[1].clone()
153            } else {
154                Value::Null(NullType::Null)
155            };
156            if current_pos + offset < partition_rows.len() {
157                partition_rows[current_pos + offset]
158                    .1
159                    .first()
160                    .cloned()
161                    .unwrap_or(default_val)
162            } else {
163                default_val
164            }
165        }
166        "lag" => {
167            let offset = if !args.is_empty() {
168                value_to_i64(&args[0]) as usize
169            } else {
170                1
171            };
172            let default_val = if args.len() > 1 {
173                args[1].clone()
174            } else {
175                Value::Null(NullType::Null)
176            };
177            if current_pos >= offset {
178                partition_rows[current_pos - offset]
179                    .1
180                    .first()
181                    .cloned()
182                    .unwrap_or(default_val)
183            } else {
184                default_val
185            }
186        }
187        "first_value" => partition_rows
188            .first()
189            .and_then(|(_, r)| r.first().cloned())
190            .unwrap_or(Value::Null(NullType::Null)),
191        "last_value" => partition_rows
192            .last()
193            .and_then(|(_, r)| r.first().cloned())
194            .unwrap_or(Value::Null(NullType::Null)),
195        "nth_value" => {
196            let n = if !args.is_empty() {
197                value_to_i64(&args[0]) as usize
198            } else {
199                1
200            };
201            if n > 0 && n <= partition_rows.len() {
202                partition_rows[n - 1]
203                    .1
204                    .first()
205                    .cloned()
206                    .unwrap_or(Value::Null(NullType::Null))
207            } else {
208                Value::Null(NullType::Null)
209            }
210        }
211        "ntile" => {
212            let n = if !args.is_empty() {
213                value_to_i64(&args[0])
214            } else {
215                1
216            };
217            if n > 0 {
218                let bucket_size = (partition_rows.len() as i64 + n - 1) / n;
219                Value::BigInt((current_pos as i64 / bucket_size) + 1)
220            } else {
221                Value::Null(NullType::Null)
222            }
223        }
224        _ => Value::Null(NullType::Null),
225    }
226}
linkrs_query/executor/streaming/operators/blocking/window.rs

linkrs_query/executor/streaming/operators/blocking/
window.rs
