uqa_execution/relational/
limit.rs1use 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}