Skip to main content

uqa_execution/relational/
limit.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Streaming LIMIT/OFFSET.
8
9use super::{Batch, ExecResult, PhysicalOperator, RowSchema};
10
11pub struct Limit<'a> {
12    child: Box<dyn PhysicalOperator + 'a>,
13    offset: u64,
14    limit: Option<u64>,
15    skipped: u64,
16    emitted: u64,
17    schema: RowSchema,
18}
19
20impl<'a> Limit<'a> {
21    pub fn new(child: Box<dyn PhysicalOperator + 'a>, offset: u64, limit: Option<u64>) -> Self {
22        let schema = child.row_schema().clone();
23        Self {
24            child,
25            offset,
26            limit,
27            skipped: 0,
28            emitted: 0,
29            schema,
30        }
31    }
32}
33
34impl PhysicalOperator for Limit<'_> {
35    fn row_schema(&self) -> &RowSchema {
36        &self.schema
37    }
38
39    fn output_ordering(&self) -> &[crate::PhysicalOrder] {
40        self.child.output_ordering()
41    }
42
43    fn open(&mut self) -> ExecResult<()> {
44        self.skipped = 0;
45        self.emitted = 0;
46        self.child.open()
47    }
48
49    fn next(&mut self) -> ExecResult<Option<Batch>> {
50        if self.limit.is_some_and(|limit| self.emitted >= limit) {
51            return Ok(None);
52        }
53        loop {
54            let Some(batch) = self.child.next()? else {
55                return Ok(None);
56            };
57            let mut buf = Vec::new();
58            for row in batch.rows {
59                if self.skipped < self.offset {
60                    self.skipped += 1;
61                    continue;
62                }
63                if let Some(lim) = self.limit {
64                    if self.emitted >= lim {
65                        return if buf.is_empty() {
66                            Ok(None)
67                        } else {
68                            Ok(Some(Batch::from_physical_rows(self.schema.clone(), buf)))
69                        };
70                    }
71                }
72                buf.push(row);
73                self.emitted += 1;
74            }
75            if !buf.is_empty() {
76                return Ok(Some(Batch::from_physical_rows(self.schema.clone(), buf)));
77            }
78        }
79    }
80
81    fn close(&mut self) -> ExecResult<()> {
82        self.child.close()
83    }
84}