uqa_execution/relational/
window.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Byte-bounded window execution.
8
9use super::{
10    compare_values, eval_scalar, value_to_f64, Batch, DefaultExpressionEvaluator, ExecError,
11    ExecResult, PhysicalOperator, RowSchema, SQLParam, ScalarEvalContext, ScalarExpr, SortKey,
12    Value,
13};
14use uqa_sql::expr::RowLookup;
15
16#[derive(Debug, Clone)]
17pub enum WindowKind {
18    RowNumber,
19    Rank,
20    DenseRank,
21    Lag(ScalarExpr, i64),
22    Lead(ScalarExpr, i64),
23    Ntile(i64),
24    AggSum(ScalarExpr),
25    AggCount(Option<ScalarExpr>),
26    AggAvg(ScalarExpr),
27    AggMin(ScalarExpr),
28    AggMax(ScalarExpr),
29}
30
31#[derive(Debug, Clone)]
32pub struct WindowSpec {
33    pub partition_by: Vec<ScalarExpr>,
34    pub order_by: Vec<SortKey>,
35}
36
37/// Byte-bounded window operator. Input sorting, random-access partitions, and
38/// output rows use disk-backed buffers; only fixed-size batches are decoded at
39/// the Volcano boundary.
40pub trait WindowExecutor: Send {
41    /// Consume one child batch without materializing the complete input in the
42    /// physical operator.
43    fn consume(&mut self, batch: Batch) -> ExecResult<()>;
44
45    /// Finalize window columns into a byte-bounded, disk-backed output stream.
46    fn finish(&mut self) -> ExecResult<crate::spill::SpillBuffer>;
47}
48
49pub struct Window<'a> {
50    child: Box<dyn PhysicalOperator + 'a>,
51    spec: WindowSpec,
52    functions: Vec<(String, WindowKind)>,
53    params: Vec<SQLParam>,
54    schema: RowSchema,
55    executor: Option<Box<dyn WindowExecutor + 'a>>,
56    work_mem_bytes: usize,
57    output: Option<crate::spill::SpillDrain>,
58    output_spilled: bool,
59}
60
61impl Window<'static> {
62    const DEFAULT_WORK_MEM_BYTES: usize = 64 * 1024 * 1024;
63
64    pub fn new(
65        child: Box<dyn PhysicalOperator>,
66        spec: WindowSpec,
67        functions: Vec<(String, WindowKind)>,
68        params: Vec<SQLParam>,
69    ) -> Self {
70        Self::new_with_work_mem(child, spec, functions, params, Self::DEFAULT_WORK_MEM_BYTES)
71    }
72
73    pub fn new_with_work_mem(
74        child: Box<dyn PhysicalOperator>,
75        spec: WindowSpec,
76        functions: Vec<(String, WindowKind)>,
77        params: Vec<SQLParam>,
78        work_mem_bytes: usize,
79    ) -> Self {
80        let names = functions
81            .iter()
82            .map(|(name, _)| name.clone())
83            .collect::<Vec<_>>();
84        let (schema, _) = RowSchema::append(child.row_schema(), &names).canonical_projection();
85        Self {
86            child,
87            spec,
88            functions,
89            params,
90            schema,
91            executor: None,
92            work_mem_bytes,
93            output: None,
94            output_spilled: false,
95        }
96    }
97}
98
99impl<'a> Window<'a> {
100    /// Construct a physical window operator backed by the engine's complete
101    /// frame and function implementation.
102    pub fn with_executor(
103        child: Box<dyn PhysicalOperator + 'a>,
104        output_schema: Vec<String>,
105        executor: Box<dyn WindowExecutor + 'a>,
106    ) -> Self {
107        let types = vec![None; output_schema.len()];
108        Self::with_typed_executor(child, output_schema, types, executor)
109    }
110
111    pub fn with_typed_executor(
112        child: Box<dyn PhysicalOperator + 'a>,
113        output_schema: Vec<String>,
114        output_types: Vec<Option<uqa_sql::ast::ColumnType>>,
115        executor: Box<dyn WindowExecutor + 'a>,
116    ) -> Self {
117        Self::with_row_schema_executor(
118            child,
119            RowSchema::with_types(output_schema, output_types),
120            executor,
121        )
122    }
123
124    pub fn with_row_schema_executor(
125        child: Box<dyn PhysicalOperator + 'a>,
126        schema: RowSchema,
127        executor: Box<dyn WindowExecutor + 'a>,
128    ) -> Self {
129        Self {
130            child,
131            spec: WindowSpec {
132                partition_by: Vec::new(),
133                order_by: Vec::new(),
134            },
135            functions: Vec::new(),
136            params: Vec::new(),
137            schema,
138            executor: Some(executor),
139            work_mem_bytes: 0,
140            output: None,
141            output_spilled: false,
142        }
143    }
144
145    /// Whether final window rows exceeded their output budget and were written
146    /// to disk during the current/most recent invocation.
147    pub fn output_has_spilled(&self) -> bool {
148        self.output_spilled
149    }
150}
151
152fn builtin_window_order_key(
153    row: &dyn RowLookup,
154    spec: &WindowSpec,
155    params: &[SQLParam],
156) -> ExecResult<Vec<Value>> {
157    let context = ScalarEvalContext::from_row_lookup(row, params);
158    spec.order_by
159        .iter()
160        .map(|key| Ok(eval_scalar(&key.expr, &context)?))
161        .collect()
162}
163
164fn builtin_window_partition_value(
165    kind: &WindowKind,
166    partition: &mut crate::spill::IndexedSpill,
167    params: &[SQLParam],
168) -> ExecResult<Option<Value>> {
169    let mut count = 0_i64;
170    let mut sum = 0.0_f64;
171    let mut min = None;
172    let mut max = None;
173    let expression = match kind {
174        WindowKind::AggSum(expression)
175        | WindowKind::AggAvg(expression)
176        | WindowKind::AggMin(expression)
177        | WindowKind::AggMax(expression) => Some(expression),
178        WindowKind::AggCount(expression) => expression.as_ref(),
179        _ => return Ok(None),
180    };
181    let partition_schema = partition.row_schema().clone();
182    for index in 0..partition.len() {
183        let row = partition.get(index)?;
184        let view = partition_schema.view(&row);
185        let value = match expression {
186            Some(expression) => eval_scalar(
187                expression,
188                &ScalarEvalContext::from_row_lookup(&view, params),
189            )?,
190            None => Value::Int(1),
191        };
192        if matches!(value, Value::Null) {
193            continue;
194        }
195        count = count
196            .checked_add(1)
197            .ok_or_else(|| ExecError::Other("window aggregate row count overflow".into()))?;
198        match kind {
199            WindowKind::AggSum(_) | WindowKind::AggAvg(_) => {
200                let number = value_to_f64(&value).ok_or_else(|| {
201                    ExecError::Other(format!("non-numeric window aggregate input: {value:?}"))
202                })?;
203                sum += number;
204            }
205            WindowKind::AggMin(_) => {
206                min = Some(match min.take() {
207                    Some(previous) if compare_values(&previous, &value).is_le() => previous,
208                    _ => value,
209                });
210            }
211            WindowKind::AggMax(_) => {
212                max = Some(match max.take() {
213                    Some(previous) if compare_values(&previous, &value).is_ge() => previous,
214                    _ => value,
215                });
216            }
217            WindowKind::AggCount(_) => {}
218            _ => {
219                return Err(ExecError::Other(
220                    "non-aggregate window kind reached aggregate evaluation".into(),
221                ))
222            }
223        }
224    }
225    Ok(Some(match kind {
226        WindowKind::AggSum(_) => {
227            if count == 0 {
228                Value::Null
229            } else {
230                Value::Float(sum)
231            }
232        }
233        WindowKind::AggCount(_) => Value::Int(count),
234        WindowKind::AggAvg(_) => {
235            if count == 0 {
236                Value::Null
237            } else {
238                Value::Float(sum / count as f64)
239            }
240        }
241        WindowKind::AggMin(_) => min.unwrap_or(Value::Null),
242        WindowKind::AggMax(_) => max.unwrap_or(Value::Null),
243        _ => {
244            return Err(ExecError::Other(
245                "non-aggregate window kind reached aggregate result construction".into(),
246            ))
247        }
248    }))
249}
250
251fn builtin_ntile(index: u64, rows: u64, buckets: i64) -> ExecResult<Value> {
252    let buckets = u64::try_from(buckets.max(1))
253        .map_err(|_| ExecError::Other("NTILE bucket count is out of range".into()))?;
254    let base = rows / buckets;
255    let extra = rows % buckets;
256    let larger_rows = if extra == 0 {
257        0
258    } else {
259        base.checked_add(1)
260            .and_then(|value| value.checked_mul(extra))
261            .ok_or_else(|| ExecError::Other("NTILE partition size overflow".into()))?
262    };
263    let bucket = if index < larger_rows {
264        index
265            .checked_div(
266                base.checked_add(1)
267                    .ok_or_else(|| ExecError::Other("NTILE bucket width overflow".into()))?,
268            )
269            .and_then(|value| value.checked_add(1))
270            .ok_or_else(|| ExecError::Other("NTILE bucket number overflow".into()))?
271    } else if base == 0 {
272        extra.max(1)
273    } else {
274        extra
275            .checked_add(
276                (index - larger_rows)
277                    .checked_div(base)
278                    .ok_or_else(|| ExecError::Other("invalid NTILE bucket width".into()))?,
279            )
280            .and_then(|value| value.checked_add(1))
281            .ok_or_else(|| ExecError::Other("NTILE bucket number overflow".into()))?
282    };
283    Ok(Value::Int(i64::try_from(bucket).map_err(|_| {
284        ExecError::Other("NTILE bucket number exceeds SQL integer range".into())
285    })?))
286}
287
288#[expect(
289    clippy::too_many_lines,
290    reason = "window execution preserves frame and peer ordering in one pass"
291)]
292fn emit_builtin_window_partition(
293    partition: &mut crate::spill::IndexedSpill,
294    spec: &WindowSpec,
295    functions: &[(String, WindowKind)],
296    params: &[SQLParam],
297    schema: &RowSchema,
298    output: &mut crate::spill::SpillBuffer,
299) -> ExecResult<()> {
300    let aliases = functions
301        .iter()
302        .map(|(alias, _)| alias.clone())
303        .collect::<Vec<_>>();
304    let partition_schema = partition.row_schema().clone();
305    let appended_schema = RowSchema::append(&partition_schema, &aliases);
306    let (physical_output_schema, output_slots) = appended_schema.canonical_projection();
307    if &physical_output_schema != schema {
308        return Err(ExecError::Other(format!(
309            "window output schema mismatch: expected {:?}, got {:?}",
310            schema.columns(),
311            physical_output_schema.columns()
312        )));
313    }
314    let aggregate_values = functions
315        .iter()
316        .map(|(_, kind)| builtin_window_partition_value(kind, partition, params))
317        .collect::<ExecResult<Vec<_>>>()?;
318    let mut previous_order_key = None;
319    let mut rank = 0_i64;
320    let mut dense_rank = 0_i64;
321    let mut pending = Vec::with_capacity(crate::batch::DEFAULT_BATCH_SIZE);
322    for index in 0..partition.len() {
323        let row = partition.get(index)?;
324        let row_view = partition_schema.view(&row);
325        let order_key = builtin_window_order_key(&row_view, spec, params)?;
326        if previous_order_key.as_ref() != Some(&order_key) {
327            rank = i64::try_from(
328                index
329                    .checked_add(1)
330                    .ok_or_else(|| ExecError::Other("window rank overflow".into()))?,
331            )
332            .map_err(|_| ExecError::Other("window rank exceeds SQL integer range".into()))?;
333            dense_rank = dense_rank
334                .checked_add(1)
335                .ok_or_else(|| ExecError::Other("window dense rank overflow".into()))?;
336        }
337        let mut window_values = Vec::with_capacity(functions.len());
338        for ((_, kind), aggregate_value) in functions.iter().zip(&aggregate_values) {
339            let value = match kind {
340                WindowKind::RowNumber => {
341                    Value::Int(
342                        i64::try_from(index.checked_add(1).ok_or_else(|| {
343                            ExecError::Other("window row number overflow".into())
344                        })?)
345                        .map_err(|_| {
346                            ExecError::Other("window row number exceeds SQL integer range".into())
347                        })?,
348                    )
349                }
350                WindowKind::Rank => Value::Int(rank),
351                WindowKind::DenseRank => Value::Int(dense_rank),
352                WindowKind::Lag(expression, offset) | WindowKind::Lead(expression, offset) => {
353                    let direction = if matches!(kind, WindowKind::Lag(..)) {
354                        -1_i128
355                    } else {
356                        1_i128
357                    };
358                    let target = i128::from(index) + direction * i128::from(*offset);
359                    if target < 0 || target >= i128::from(partition.len()) {
360                        Value::Null
361                    } else {
362                        let target_row = partition.get(u64::try_from(target).map_err(|_| {
363                            ExecError::Other("window offset target is out of range".into())
364                        })?)?;
365                        let target_view = partition_schema.view(&target_row);
366                        eval_scalar(
367                            expression,
368                            &ScalarEvalContext::from_row_lookup(&target_view, params),
369                        )?
370                    }
371                }
372                WindowKind::Ntile(buckets) => builtin_ntile(index, partition.len(), *buckets)?,
373                WindowKind::AggSum(_)
374                | WindowKind::AggCount(_)
375                | WindowKind::AggAvg(_)
376                | WindowKind::AggMin(_)
377                | WindowKind::AggMax(_) => aggregate_value.clone().ok_or_else(|| {
378                    ExecError::Other("aggregate window value was not precomputed".into())
379                })?,
380            };
381            window_values.push(value);
382        }
383        previous_order_key = Some(order_key);
384        pending.push(
385            row.append_values(window_values)
386                .project_slots(&output_slots)
387                .without_lock_origins(),
388        );
389        if pending.len() == crate::batch::DEFAULT_BATCH_SIZE {
390            output.push(Batch::from_physical_rows(
391                schema.clone(),
392                std::mem::take(&mut pending),
393            ))?;
394            pending = Vec::with_capacity(crate::batch::DEFAULT_BATCH_SIZE);
395        }
396    }
397    if !pending.is_empty() {
398        output.push(Batch::from_physical_rows(schema.clone(), pending))?;
399    }
400    Ok(())
401}
402
403impl PhysicalOperator for Window<'_> {
404    fn row_schema(&self) -> &RowSchema {
405        &self.schema
406    }
407
408    fn open(&mut self) -> ExecResult<()> {
409        self.child.open()?;
410        self.output_spilled = false;
411        if let Some(executor) = self.executor.as_mut() {
412            while let Some(batch) = self.child.next()? {
413                executor.consume(batch)?;
414            }
415            let mut output = executor.finish()?;
416            self.output_spilled = output.has_spilled();
417            self.output = Some(output.drain()?);
418            return Ok(());
419        }
420
421        let phase_budget = (self.work_mem_bytes / 3).max(1);
422        let mut input = crate::spill::SpillBuffer::new(phase_budget);
423        while let Some(batch) = self.child.next()? {
424            input.push(batch)?;
425        }
426        let scan: Box<dyn PhysicalOperator> = Box::new(crate::spill_scan::SpillScan::new(
427            self.child.schema().to_vec(),
428            input,
429        ));
430        let mut keys = self
431            .spec
432            .partition_by
433            .iter()
434            .cloned()
435            .map(|expr| SortKey {
436                expr,
437                descending: false,
438                nulls_first: None,
439            })
440            .collect::<Vec<_>>();
441        keys.extend(self.spec.order_by.iter().cloned());
442        let evaluator = DefaultExpressionEvaluator::shared(self.params.clone());
443        let mut sorted =
444            crate::external_sort::ExternalSort::new(scan, keys, evaluator, None, phase_budget);
445        sorted.open()?;
446
447        let partition_schema = sorted.row_schema().clone();
448        let mut current_partition_key: Option<Vec<Value>> = None;
449        let mut partition = crate::spill::IndexedSpill::new(partition_schema.clone())?;
450        let mut output = crate::spill::SpillBuffer::new(phase_budget);
451        let execution = (|| -> ExecResult<()> {
452            while let Some(batch) = sorted.next()? {
453                for row in batch.rows {
454                    let view = batch.schema.view(&row);
455                    let context = ScalarEvalContext::from_row_lookup(&view, &self.params);
456                    let key = self
457                        .spec
458                        .partition_by
459                        .iter()
460                        .map(|expression| eval_scalar(expression, &context))
461                        .collect::<Result<Vec<_>, _>>()?;
462                    if current_partition_key
463                        .as_ref()
464                        .is_some_and(|current| current != &key)
465                    {
466                        emit_builtin_window_partition(
467                            &mut partition,
468                            &self.spec,
469                            &self.functions,
470                            &self.params,
471                            &self.schema,
472                            &mut output,
473                        )?;
474                        partition = crate::spill::IndexedSpill::new(partition_schema.clone())?;
475                    }
476                    current_partition_key = Some(key);
477                    partition.push(&row)?;
478                }
479            }
480            if !partition.is_empty() {
481                emit_builtin_window_partition(
482                    &mut partition,
483                    &self.spec,
484                    &self.functions,
485                    &self.params,
486                    &self.schema,
487                    &mut output,
488                )?;
489            }
490            Ok(())
491        })();
492        let close = sorted.close();
493        crate::physical::with_cleanup(execution, close, "close window sort after failure")?;
494        self.output_spilled = output.has_spilled();
495        self.output = Some(output.drain()?);
496        Ok(())
497    }
498
499    fn next(&mut self) -> ExecResult<Option<Batch>> {
500        let Some(output) = self.output.as_mut() else {
501            return Ok(None);
502        };
503        output.next().transpose()
504    }
505
506    fn close(&mut self) -> ExecResult<()> {
507        self.output = None;
508        self.child.close()
509    }
510}
uqa_execution/relational/window.rs

uqa_execution/relational/
window.rs
