uqa_execution/relational/
limit.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Streaming LIMIT/OFFSET, including ordered `FETCH ... WITH TIES`.
8
9use super::{
10    compare_sort_key_values, BackwardScanSupport, Batch, ExecError, ExecResult, PhysicalOperator,
11    PhysicalScanDirection, RowSchema, ScalarExpr, SharedExpressionEvaluator, SortKey, Value,
12};
13use crate::PhysicalRow;
14
15struct WithTies<'a> {
16    keys: Vec<SortKey>,
17    evaluator: SharedExpressionEvaluator<'a>,
18    boundary: Option<Vec<Value>>,
19    finished: bool,
20}
21
22#[derive(Clone, Copy)]
23enum DirectionalLimitState {
24    Initial,
25    Rescan,
26    Empty,
27    InWindow,
28    SubplanEof,
29    WindowEnd,
30    WindowEndTies,
31    WindowStart,
32}
33
34pub struct Limit<'a> {
35    child: Box<dyn PhysicalOperator + 'a>,
36    offset: u64,
37    limit: Option<u64>,
38    skipped: u64,
39    emitted: u64,
40    schema: RowSchema,
41    with_ties: Option<WithTies<'a>>,
42    directional_state: DirectionalLimitState,
43    directional_position: u64,
44    directional_current: Option<PhysicalRow>,
45}
46
47impl<'a> Limit<'a> {
48    pub fn new(child: Box<dyn PhysicalOperator + 'a>, offset: u64, limit: Option<u64>) -> Self {
49        let schema = child.row_schema().clone();
50        Self {
51            child,
52            offset,
53            limit,
54            skipped: 0,
55            emitted: 0,
56            schema,
57            with_ties: None,
58            directional_state: DirectionalLimitState::Initial,
59            directional_position: 0,
60            directional_current: None,
61        }
62    }
63
64    /// Build an ordered `FETCH ... WITH TIES` boundary. The caller must pass the complete effective `ORDER BY` key list and a non-null row count.
65    pub fn with_ties(
66        child: Box<dyn PhysicalOperator + 'a>,
67        offset: u64,
68        limit: u64,
69        mut keys: Vec<SortKey>,
70        evaluator: SharedExpressionEvaluator<'a>,
71    ) -> Self {
72        let schema = child.row_schema().clone();
73        for key in &mut keys {
74            let expression = std::mem::replace(&mut key.expr, ScalarExpr::Literal(Value::Null));
75            key.expr = evaluator.bind_type_introspection(expression, &schema);
76        }
77        Self {
78            child,
79            offset,
80            limit: Some(limit),
81            skipped: 0,
82            emitted: 0,
83            schema,
84            with_ties: Some(WithTies {
85                keys,
86                evaluator,
87                boundary: None,
88                finished: false,
89            }),
90            directional_state: DirectionalLimitState::Initial,
91            directional_position: 0,
92            directional_current: None,
93        }
94    }
95
96    fn reset_directional_state(&mut self) {
97        self.directional_state = DirectionalLimitState::Rescan;
98        self.directional_position = 0;
99        self.directional_current = None;
100        if let Some(with_ties) = self.with_ties.as_mut() {
101            with_ties.boundary = None;
102            with_ties.finished = false;
103        }
104    }
105
106    fn child_row(&mut self, direction: PhysicalScanDirection) -> ExecResult<Option<PhysicalRow>> {
107        let Some(batch) = self.child.next_direction(direction)? else {
108            return Ok(None);
109        };
110        if batch.schema != self.schema {
111            return Err(ExecError::Other(format!(
112                "directional LIMIT input schema mismatch: expected {:?}, got {:?}",
113                self.schema, batch.schema
114            )));
115        }
116        let mut rows = batch.rows.into_iter();
117        let row = rows.next();
118        if rows.next().is_some() {
119            return Err(ExecError::Other(
120                "directional LIMIT input returned more than one row".into(),
121            ));
122        }
123        Ok(row)
124    }
125
126    fn capture_tie_boundary(&mut self, row: &PhysicalRow) -> ExecResult<()> {
127        let Some(with_ties) = self.with_ties.as_mut() else {
128            return Ok(());
129        };
130        with_ties.boundary = Some(
131            with_ties
132                .keys
133                .iter()
134                .map(|key| {
135                    with_ties
136                        .evaluator
137                        .evaluate_physical(&key.expr, &self.schema, row)
138                })
139                .collect::<ExecResult<Vec<_>>>()?,
140        );
141        Ok(())
142    }
143
144    fn tie_matches(&self, row: &PhysicalRow) -> ExecResult<bool> {
145        let with_ties = self
146            .with_ties
147            .as_ref()
148            .ok_or_else(|| ExecError::Other("LIMIT tie state is absent".into()))?;
149        let values = with_ties
150            .keys
151            .iter()
152            .map(|key| {
153                with_ties
154                    .evaluator
155                    .evaluate_physical(&key.expr, &self.schema, row)
156            })
157            .collect::<ExecResult<Vec<_>>>()?;
158        let boundary = with_ties
159            .boundary
160            .as_ref()
161            .ok_or_else(|| ExecError::Other("LIMIT tie boundary is absent".into()))?;
162        Ok(
163            compare_sort_key_values(&with_ties.keys, boundary, &values)
164                == std::cmp::Ordering::Equal,
165        )
166    }
167
168    fn directional_batch(&self, row: PhysicalRow) -> Batch {
169        Batch::from_physical_rows(self.schema.clone(), vec![row])
170    }
171
172    #[expect(
173        clippy::too_many_lines,
174        reason = "mirrors PostgreSQL's directional LIMIT state machine"
175    )]
176    fn next_directional(&mut self, direction: PhysicalScanDirection) -> ExecResult<Option<Batch>> {
177        if matches!(self.directional_state, DirectionalLimitState::Initial) {
178            self.reset_directional_state();
179        }
180        loop {
181            match self.directional_state {
182                DirectionalLimitState::Initial => unreachable!(),
183                DirectionalLimitState::Rescan => {
184                    if direction == PhysicalScanDirection::Backward {
185                        return Ok(None);
186                    }
187                    if self.limit == Some(0) {
188                        self.directional_state = DirectionalLimitState::Empty;
189                        return Ok(None);
190                    }
191                    loop {
192                        let Some(row) = self.child_row(PhysicalScanDirection::Forward)? else {
193                            self.directional_state = DirectionalLimitState::Empty;
194                            return Ok(None);
195                        };
196                        if self.with_ties.is_some()
197                            && self.limit.is_some_and(|limit| {
198                                self.directional_position.saturating_sub(self.offset)
199                                    == limit.saturating_sub(1)
200                            })
201                        {
202                            self.capture_tie_boundary(&row)?;
203                        }
204                        self.directional_current = Some(row.clone());
205                        self.directional_position = self
206                            .directional_position
207                            .checked_add(1)
208                            .ok_or_else(|| ExecError::Other("LIMIT position overflow".into()))?;
209                        if self.directional_position > self.offset {
210                            self.directional_state = DirectionalLimitState::InWindow;
211                            return Ok(Some(self.directional_batch(row)));
212                        }
213                    }
214                }
215                DirectionalLimitState::Empty => return Ok(None),
216                DirectionalLimitState::InWindow => {
217                    if direction == PhysicalScanDirection::Forward {
218                        if self.limit.is_some_and(|limit| {
219                            self.directional_position.saturating_sub(self.offset) >= limit
220                        }) {
221                            if self.with_ties.is_none() {
222                                self.directional_state = DirectionalLimitState::WindowEnd;
223                                return Ok(None);
224                            }
225                            self.directional_state = DirectionalLimitState::WindowEndTies;
226                            continue;
227                        }
228                        let Some(row) = self.child_row(PhysicalScanDirection::Forward)? else {
229                            self.directional_state = DirectionalLimitState::SubplanEof;
230                            return Ok(None);
231                        };
232                        if self.with_ties.is_some()
233                            && self.limit.is_some_and(|limit| {
234                                self.directional_position.saturating_sub(self.offset)
235                                    == limit.saturating_sub(1)
236                            })
237                        {
238                            self.capture_tie_boundary(&row)?;
239                        }
240                        self.directional_current = Some(row.clone());
241                        self.directional_position = self
242                            .directional_position
243                            .checked_add(1)
244                            .ok_or_else(|| ExecError::Other("LIMIT position overflow".into()))?;
245                        return Ok(Some(self.directional_batch(row)));
246                    }
247                    if self.directional_position <= self.offset.saturating_add(1) {
248                        self.directional_state = DirectionalLimitState::WindowStart;
249                        return Ok(None);
250                    }
251                    let row = self
252                        .child_row(PhysicalScanDirection::Backward)?
253                        .ok_or_else(|| {
254                            ExecError::Other("LIMIT input failed to scan backwards".into())
255                        })?;
256                    self.directional_current = Some(row.clone());
257                    self.directional_position -= 1;
258                    return Ok(Some(self.directional_batch(row)));
259                }
260                DirectionalLimitState::WindowEndTies => {
261                    if direction == PhysicalScanDirection::Forward {
262                        let Some(row) = self.child_row(PhysicalScanDirection::Forward)? else {
263                            self.directional_state = DirectionalLimitState::SubplanEof;
264                            return Ok(None);
265                        };
266                        if self.tie_matches(&row)? {
267                            self.directional_current = Some(row.clone());
268                            self.directional_position =
269                                self.directional_position.checked_add(1).ok_or_else(|| {
270                                    ExecError::Other("LIMIT position overflow".into())
271                                })?;
272                            return Ok(Some(self.directional_batch(row)));
273                        }
274                        self.directional_state = DirectionalLimitState::WindowEnd;
275                        return Ok(None);
276                    }
277                    if self.directional_position <= self.offset.saturating_add(1) {
278                        self.directional_state = DirectionalLimitState::WindowStart;
279                        return Ok(None);
280                    }
281                    let row = self
282                        .child_row(PhysicalScanDirection::Backward)?
283                        .ok_or_else(|| {
284                            ExecError::Other("LIMIT input failed to scan backwards".into())
285                        })?;
286                    self.directional_current = Some(row.clone());
287                    self.directional_position -= 1;
288                    self.directional_state = DirectionalLimitState::InWindow;
289                    return Ok(Some(self.directional_batch(row)));
290                }
291                DirectionalLimitState::SubplanEof => {
292                    if direction == PhysicalScanDirection::Forward {
293                        return Ok(None);
294                    }
295                    let row = self
296                        .child_row(PhysicalScanDirection::Backward)?
297                        .ok_or_else(|| {
298                            ExecError::Other("LIMIT input failed to leave end position".into())
299                        })?;
300                    self.directional_current = Some(row.clone());
301                    self.directional_state = DirectionalLimitState::InWindow;
302                    return Ok(Some(self.directional_batch(row)));
303                }
304                DirectionalLimitState::WindowEnd => {
305                    if direction == PhysicalScanDirection::Forward {
306                        return Ok(None);
307                    }
308                    let row = if self.with_ties.is_some() {
309                        self.child_row(PhysicalScanDirection::Backward)?
310                            .ok_or_else(|| {
311                                ExecError::Other("LIMIT input failed to leave tie boundary".into())
312                            })?
313                    } else {
314                        self.directional_current.clone().ok_or_else(|| {
315                            ExecError::Other("LIMIT window has no current row".into())
316                        })?
317                    };
318                    self.directional_current = Some(row.clone());
319                    self.directional_state = DirectionalLimitState::InWindow;
320                    return Ok(Some(self.directional_batch(row)));
321                }
322                DirectionalLimitState::WindowStart => {
323                    if direction == PhysicalScanDirection::Backward {
324                        return Ok(None);
325                    }
326                    let row = self.directional_current.clone().ok_or_else(|| {
327                        ExecError::Other("LIMIT window has no current row".into())
328                    })?;
329                    self.directional_state = DirectionalLimitState::InWindow;
330                    return Ok(Some(self.directional_batch(row)));
331                }
332            }
333        }
334    }
335}
336
337impl PhysicalOperator for Limit<'_> {
338    fn row_schema(&self) -> &RowSchema {
339        &self.schema
340    }
341
342    fn output_ordering(&self) -> &[crate::PhysicalOrder] {
343        self.child.output_ordering()
344    }
345
346    fn backward_scan_support(&self) -> BackwardScanSupport {
347        if self.child.backward_scan_support() == BackwardScanSupport::Native {
348            BackwardScanSupport::Native
349        } else {
350            BackwardScanSupport::Unsupported
351        }
352    }
353
354    fn open(&mut self) -> ExecResult<()> {
355        self.skipped = 0;
356        self.emitted = 0;
357        if let Some(with_ties) = self.with_ties.as_mut() {
358            with_ties.boundary = None;
359            with_ties.finished = false;
360        }
361        self.directional_state = DirectionalLimitState::Initial;
362        self.directional_position = 0;
363        self.directional_current = None;
364        self.child.open()
365    }
366
367    fn next(&mut self) -> ExecResult<Option<Batch>> {
368        if self.limit == Some(0)
369            || self
370                .with_ties
371                .as_ref()
372                .is_some_and(|with_ties| with_ties.finished)
373            || self.with_ties.is_none() && self.limit.is_some_and(|limit| self.emitted >= limit)
374        {
375            return Ok(None);
376        }
377        loop {
378            let Some(batch) = self.child.next()? else {
379                return Ok(None);
380            };
381            let mut buf = Vec::new();
382            for row in batch.rows {
383                if self.skipped < self.offset {
384                    self.skipped += 1;
385                    continue;
386                }
387                if let Some(lim) = self.limit {
388                    if self.emitted >= lim {
389                        let Some(with_ties) = self.with_ties.as_mut() else {
390                            return if buf.is_empty() {
391                                Ok(None)
392                            } else {
393                                Ok(Some(Batch::from_physical_rows(self.schema.clone(), buf)))
394                            };
395                        };
396                        let values = with_ties
397                            .keys
398                            .iter()
399                            .map(|key| {
400                                with_ties
401                                    .evaluator
402                                    .evaluate_physical(&key.expr, &self.schema, &row)
403                            })
404                            .collect::<ExecResult<Vec<_>>>()?;
405                        let boundary = with_ties.boundary.as_ref().ok_or_else(|| {
406                            crate::ExecError::Other(
407                                "WITH TIES boundary was not captured".to_string(),
408                            )
409                        })?;
410                        if compare_sort_key_values(&with_ties.keys, boundary, &values)
411                            != std::cmp::Ordering::Equal
412                        {
413                            with_ties.finished = true;
414                            return if buf.is_empty() {
415                                Ok(None)
416                            } else {
417                                Ok(Some(Batch::from_physical_rows(self.schema.clone(), buf)))
418                            };
419                        }
420                    } else if self.with_ties.is_some() && self.emitted + 1 == lim {
421                        let with_ties = self.with_ties.as_mut().expect("checked above");
422                        with_ties.boundary = Some(
423                            with_ties
424                                .keys
425                                .iter()
426                                .map(|key| {
427                                    with_ties.evaluator.evaluate_physical(
428                                        &key.expr,
429                                        &self.schema,
430                                        &row,
431                                    )
432                                })
433                                .collect::<ExecResult<Vec<_>>>()?,
434                        );
435                    }
436                }
437                buf.push(row);
438                self.emitted += 1;
439            }
440            if !buf.is_empty() {
441                return Ok(Some(Batch::from_physical_rows(self.schema.clone(), buf)));
442            }
443        }
444    }
445
446    fn next_direction(&mut self, direction: PhysicalScanDirection) -> ExecResult<Option<Batch>> {
447        self.next_directional(direction)
448    }
449
450    fn rewind(&mut self) -> ExecResult<()> {
451        self.reset_directional_state();
452        self.child.rewind()
453    }
454
455    fn close(&mut self) -> ExecResult<()> {
456        self.child.close()
457    }
458}
uqa_execution/relational/limit.rs

uqa_execution/relational/
limit.rs
