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
288fn emit_builtin_window_partition(
289    partition: &mut crate::spill::IndexedSpill,
290    spec: &WindowSpec,
291    functions: &[(String, WindowKind)],
292    params: &[SQLParam],
293    schema: &RowSchema,
294    output: &mut crate::spill::SpillBuffer,
295) -> ExecResult<()> {
296    let aliases = functions
297        .iter()
298        .map(|(alias, _)| alias.clone())
299        .collect::<Vec<_>>();
300    let partition_schema = partition.row_schema().clone();
301    let appended_schema = RowSchema::append(&partition_schema, &aliases);
302    let (physical_output_schema, output_slots) = appended_schema.canonical_projection();
303    if &physical_output_schema != schema {
304        return Err(ExecError::Other(format!(
305            "window output schema mismatch: expected {:?}, got {:?}",
306            schema.columns(),
307            physical_output_schema.columns()
308        )));
309    }
310    let aggregate_values = functions
311        .iter()
312        .map(|(_, kind)| builtin_window_partition_value(kind, partition, params))
313        .collect::<ExecResult<Vec<_>>>()?;
314    let mut previous_order_key = None;
315    let mut rank = 0_i64;
316    let mut dense_rank = 0_i64;
317    let mut pending = Vec::with_capacity(crate::batch::DEFAULT_BATCH_SIZE);
318    for index in 0..partition.len() {
319        let row = partition.get(index)?;
320        let row_view = partition_schema.view(&row);
321        let order_key = builtin_window_order_key(&row_view, spec, params)?;
322        if previous_order_key.as_ref() != Some(&order_key) {
323            rank = i64::try_from(
324                index
325                    .checked_add(1)
326                    .ok_or_else(|| ExecError::Other("window rank overflow".into()))?,
327            )
328            .map_err(|_| ExecError::Other("window rank exceeds SQL integer range".into()))?;
329            dense_rank = dense_rank
330                .checked_add(1)
331                .ok_or_else(|| ExecError::Other("window dense rank overflow".into()))?;
332        }
333        let mut window_values = Vec::with_capacity(functions.len());
334        for ((_, kind), aggregate_value) in functions.iter().zip(&aggregate_values) {
335            let value = match kind {
336                WindowKind::RowNumber => {
337                    Value::Int(
338                        i64::try_from(index.checked_add(1).ok_or_else(|| {
339                            ExecError::Other("window row number overflow".into())
340                        })?)
341                        .map_err(|_| {
342                            ExecError::Other("window row number exceeds SQL integer range".into())
343                        })?,
344                    )
345                }
346                WindowKind::Rank => Value::Int(rank),
347                WindowKind::DenseRank => Value::Int(dense_rank),
348                WindowKind::Lag(expression, offset) | WindowKind::Lead(expression, offset) => {
349                    let direction = if matches!(kind, WindowKind::Lag(..)) {
350                        -1_i128
351                    } else {
352                        1_i128
353                    };
354                    let target = i128::from(index) + direction * i128::from(*offset);
355                    if target < 0 || target >= i128::from(partition.len()) {
356                        Value::Null
357                    } else {
358                        let target_row = partition.get(u64::try_from(target).map_err(|_| {
359                            ExecError::Other("window offset target is out of range".into())
360                        })?)?;
361                        let target_view = partition_schema.view(&target_row);
362                        eval_scalar(
363                            expression,
364                            &ScalarEvalContext::from_row_lookup(&target_view, params),
365                        )?
366                    }
367                }
368                WindowKind::Ntile(buckets) => builtin_ntile(index, partition.len(), *buckets)?,
369                WindowKind::AggSum(_)
370                | WindowKind::AggCount(_)
371                | WindowKind::AggAvg(_)
372                | WindowKind::AggMin(_)
373                | WindowKind::AggMax(_) => aggregate_value.clone().ok_or_else(|| {
374                    ExecError::Other("aggregate window value was not precomputed".into())
375                })?,
376            };
377            window_values.push(value);
378        }
379        previous_order_key = Some(order_key);
380        pending.push(
381            row.append_values(window_values)
382                .project_slots(&output_slots)
383                .without_lock_origins(),
384        );
385        if pending.len() == crate::batch::DEFAULT_BATCH_SIZE {
386            output.push(Batch::from_physical_rows(
387                schema.clone(),
388                std::mem::take(&mut pending),
389            ))?;
390            pending = Vec::with_capacity(crate::batch::DEFAULT_BATCH_SIZE);
391        }
392    }
393    if !pending.is_empty() {
394        output.push(Batch::from_physical_rows(schema.clone(), pending))?;
395    }
396    Ok(())
397}
398
399impl PhysicalOperator for Window<'_> {
400    fn row_schema(&self) -> &RowSchema {
401        &self.schema
402    }
403
404    fn open(&mut self) -> ExecResult<()> {
405        self.child.open()?;
406        self.output_spilled = false;
407        if let Some(executor) = self.executor.as_mut() {
408            while let Some(batch) = self.child.next()? {
409                executor.consume(batch)?;
410            }
411            let mut output = executor.finish()?;
412            self.output_spilled = output.has_spilled();
413            self.output = Some(output.drain()?);
414            return Ok(());
415        }
416
417        let phase_budget = (self.work_mem_bytes / 3).max(1);
418        let mut input = crate::spill::SpillBuffer::new(phase_budget);
419        while let Some(batch) = self.child.next()? {
420            input.push(batch)?;
421        }
422        let scan: Box<dyn PhysicalOperator> = Box::new(crate::spill_scan::SpillScan::new(
423            self.child.schema().to_vec(),
424            input,
425        ));
426        let mut keys = self
427            .spec
428            .partition_by
429            .iter()
430            .cloned()
431            .map(|expr| SortKey {
432                expr,
433                descending: false,
434                nulls_first: None,
435            })
436            .collect::<Vec<_>>();
437        keys.extend(self.spec.order_by.iter().cloned());
438        let evaluator = DefaultExpressionEvaluator::shared(self.params.clone());
439        let mut sorted =
440            crate::external_sort::ExternalSort::new(scan, keys, evaluator, None, phase_budget);
441        sorted.open()?;
442
443        let partition_schema = sorted.row_schema().clone();
444        let mut current_partition_key: Option<Vec<Value>> = None;
445        let mut partition = crate::spill::IndexedSpill::new(partition_schema.clone())?;
446        let mut output = crate::spill::SpillBuffer::new(phase_budget);
447        let execution = (|| -> ExecResult<()> {
448            while let Some(batch) = sorted.next()? {
449                for row in batch.rows {
450                    let view = batch.schema.view(&row);
451                    let context = ScalarEvalContext::from_row_lookup(&view, &self.params);
452                    let key = self
453                        .spec
454                        .partition_by
455                        .iter()
456                        .map(|expression| eval_scalar(expression, &context))
457                        .collect::<Result<Vec<_>, _>>()?;
458                    if current_partition_key
459                        .as_ref()
460                        .is_some_and(|current| current != &key)
461                    {
462                        emit_builtin_window_partition(
463                            &mut partition,
464                            &self.spec,
465                            &self.functions,
466                            &self.params,
467                            &self.schema,
468                            &mut output,
469                        )?;
470                        partition = crate::spill::IndexedSpill::new(partition_schema.clone())?;
471                    }
472                    current_partition_key = Some(key);
473                    partition.push(&row)?;
474                }
475            }
476            if !partition.is_empty() {
477                emit_builtin_window_partition(
478                    &mut partition,
479                    &self.spec,
480                    &self.functions,
481                    &self.params,
482                    &self.schema,
483                    &mut output,
484                )?;
485            }
486            Ok(())
487        })();
488        let close = sorted.close();
489        crate::physical::with_cleanup(execution, close, "close window sort after failure")?;
490        self.output_spilled = output.has_spilled();
491        self.output = Some(output.drain()?);
492        Ok(())
493    }
494
495    fn next(&mut self) -> ExecResult<Option<Batch>> {
496        let Some(output) = self.output.as_mut() else {
497            return Ok(None);
498        };
499        output.next().transpose()
500    }
501
502    fn close(&mut self) -> ExecResult<()> {
503        self.output = None;
504        self.child.close()
505    }
506}
uqa_execution/relational/window.rs

uqa_execution/relational/
window.rs
