uqa_execution/
scan.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Scan operators.
8//!
9//! [`TableScan`] pulls rows from any [`RowSource`] in fixed-size
10//! batches. The trait is the integration seam: the engine implements
11//! it over its in-memory and SQLite-backed table state, the FDW layer
12//! implements it over `MemoryHandler` / `DuckDBHandler` /
13//! `ArrowHandler`, and tests can implement it directly over an
14//! in-memory `Vec<ResultRow>`.
15
16use crate::batch::{Batch, PhysicalRow, RowSchema, DEFAULT_BATCH_SIZE};
17use crate::physical::{BackwardScanSupport, ExecResult, PhysicalOperator, PhysicalOrder};
18use uqa_sql::ast::ColumnType;
19use uqa_sql::ResultRow;
20
21/// Source of rows feeding a [`TableScan`]. Implementors typically own
22/// a snapshot of the underlying table or external relation; the scan
23/// operator holds the source as a boxed trait object so callers can
24/// mix and match implementations across one query.
25pub trait RowSource: Send {
26    /// Stable column order for the rows produced by [`Self::next_row`].
27    fn schema(&self) -> &[String];
28
29    /// Optional non-identity schema used by positional sources. This carries
30    /// hidden lookup aliases and slot remaps that cannot be represented by the
31    /// legacy column-name slice.
32    fn physical_schema(&self) -> Option<&RowSchema> {
33        None
34    }
35
36    /// Estimated total rows available from this source.
37    fn estimated_cardinality(&self) -> Option<u64> {
38        None
39    }
40
41    /// Leading row order guaranteed by the source.
42    fn output_ordering(&self) -> &[PhysicalOrder] {
43        &[]
44    }
45
46    /// Pull the next row. Returns `None` when the source is exhausted.
47    fn next_row(&mut self) -> ExecResult<Option<ResultRow>>;
48
49    /// Pull up to `max_rows` without forcing batch-capable sources through a
50    /// row-at-a-time lock or backend call. The default preserves compatibility
51    /// for iterator-like sources.
52    fn next_batch(&mut self, max_rows: usize) -> ExecResult<Vec<ResultRow>> {
53        let mut rows = Vec::with_capacity(max_rows);
54        while rows.len() < max_rows {
55            match self.next_row()? {
56                Some(row) => rows.push(row),
57                None => break,
58            }
59        }
60        Ok(rows)
61    }
62
63    /// Pull a positional batch directly. Backend-native sources override this
64    /// to avoid constructing named maps at the scan boundary; compatibility
65    /// sources are converted exactly once here.
66    fn next_physical_batch(&mut self, max_rows: usize) -> ExecResult<Vec<PhysicalRow>> {
67        let schema = RowSchema::new(self.schema().to_vec());
68        self.next_batch(max_rows).map(|rows| {
69            rows.into_iter()
70                .map(|row| PhysicalRow::from_result_row(&schema, row))
71                .collect()
72        })
73    }
74
75    /// Feed backend-native projected rows directly to an aggregate. Sources
76    /// that cannot preserve their normal filter and virtual-column semantics
77    /// return `false` without advancing their cursor.
78    fn consume_into_aggregate(
79        &mut self,
80        _executor: &mut dyn crate::relational::AggregateExecutor,
81    ) -> ExecResult<bool> {
82        Ok(false)
83    }
84}
85
86/// In-memory source from a precomputed `Vec<ResultRow>`. Useful for
87/// tests and for materialising CTE bodies.
88pub struct VecSource {
89    schema: Vec<String>,
90    physical_schema: RowSchema,
91    rows: std::vec::IntoIter<ResultRow>,
92}
93
94/// In-memory positional source that preserves a structured [`RowSchema`] and shared physical row fragments without round-tripping through named maps.
95pub struct PhysicalVecSource {
96    schema: RowSchema,
97    rows: std::vec::IntoIter<PhysicalRow>,
98}
99
100/// Physical scan over a fallible row iterator. Unlike [`VecSource`], this
101/// adapter preserves producer backpressure and late errors without requiring a
102/// cardinality-sized staging vector.
103pub struct RowIteratorScan<'a> {
104    schema: RowSchema,
105    rows: Box<dyn Iterator<Item = ExecResult<ResultRow>> + Send + 'a>,
106    exhausted: bool,
107}
108
109/// Physical scan over a fallible positional-row iterator. This is the native
110/// adapter for producers whose output can contain duplicate or unnamed SQL
111/// columns and therefore cannot be represented by [`ResultRow`].
112pub struct PhysicalRowIteratorScan<'a> {
113    schema: RowSchema,
114    rows: Box<dyn Iterator<Item = ExecResult<PhysicalRow>> + Send + 'a>,
115    exhausted: bool,
116}
117
118impl<'a> RowIteratorScan<'a> {
119    pub fn new(
120        schema: Vec<String>,
121        rows: Box<dyn Iterator<Item = ExecResult<ResultRow>> + Send + 'a>,
122    ) -> Self {
123        Self {
124            schema: RowSchema::new(schema),
125            rows,
126            exhausted: false,
127        }
128    }
129
130    pub fn with_types(
131        schema: Vec<String>,
132        types: Vec<Option<ColumnType>>,
133        rows: Box<dyn Iterator<Item = ExecResult<ResultRow>> + Send + 'a>,
134    ) -> Self {
135        Self {
136            schema: RowSchema::with_types(schema, types),
137            rows,
138            exhausted: false,
139        }
140    }
141
142    pub fn with_row_schema(
143        schema: RowSchema,
144        rows: Box<dyn Iterator<Item = ExecResult<ResultRow>> + Send + 'a>,
145    ) -> Self {
146        Self {
147            schema,
148            rows,
149            exhausted: false,
150        }
151    }
152}
153
154impl PhysicalOperator for RowIteratorScan<'_> {
155    fn row_schema(&self) -> &RowSchema {
156        &self.schema
157    }
158
159    fn open(&mut self) -> ExecResult<()> {
160        self.exhausted = false;
161        Ok(())
162    }
163
164    fn next(&mut self) -> ExecResult<Option<Batch>> {
165        if self.exhausted {
166            return Ok(None);
167        }
168        let mut batch = Vec::with_capacity(DEFAULT_BATCH_SIZE);
169        while batch.len() < DEFAULT_BATCH_SIZE {
170            match self.rows.next() {
171                Some(Ok(row)) => batch.push(row),
172                Some(Err(error)) => return Err(error),
173                None => {
174                    self.exhausted = true;
175                    break;
176                }
177            }
178        }
179        if batch.is_empty() {
180            Ok(None)
181        } else {
182            Ok(Some(Batch::new(self.schema.clone(), batch)))
183        }
184    }
185
186    fn close(&mut self) -> ExecResult<()> {
187        self.exhausted = true;
188        Ok(())
189    }
190}
191
192impl<'a> PhysicalRowIteratorScan<'a> {
193    pub fn new(
194        schema: RowSchema,
195        rows: Box<dyn Iterator<Item = ExecResult<PhysicalRow>> + Send + 'a>,
196    ) -> Self {
197        Self {
198            schema,
199            rows,
200            exhausted: false,
201        }
202    }
203}
204
205impl PhysicalOperator for PhysicalRowIteratorScan<'_> {
206    fn row_schema(&self) -> &RowSchema {
207        &self.schema
208    }
209
210    fn backward_scan_support(&self) -> BackwardScanSupport {
211        BackwardScanSupport::Materialize
212    }
213
214    fn open(&mut self) -> ExecResult<()> {
215        self.exhausted = false;
216        Ok(())
217    }
218
219    fn next(&mut self) -> ExecResult<Option<Batch>> {
220        if self.exhausted {
221            return Ok(None);
222        }
223        let mut batch = Vec::with_capacity(DEFAULT_BATCH_SIZE);
224        while batch.len() < DEFAULT_BATCH_SIZE {
225            match self.rows.next() {
226                Some(Ok(row)) => batch.push(row),
227                Some(Err(error)) => return Err(error),
228                None => {
229                    self.exhausted = true;
230                    break;
231                }
232            }
233        }
234        if batch.is_empty() {
235            Ok(None)
236        } else {
237            Ok(Some(Batch::from_physical_rows(self.schema.clone(), batch)))
238        }
239    }
240
241    fn close(&mut self) -> ExecResult<()> {
242        self.exhausted = true;
243        Ok(())
244    }
245}
246
247impl VecSource {
248    pub fn new(schema: Vec<String>, rows: Vec<ResultRow>) -> Self {
249        let physical_schema = RowSchema::from_named_columns(schema.clone());
250        Self {
251            schema,
252            physical_schema,
253            rows: rows.into_iter(),
254        }
255    }
256
257    pub fn with_row_schema(physical_schema: RowSchema, rows: Vec<ResultRow>) -> Self {
258        Self {
259            schema: physical_schema.columns().to_vec(),
260            physical_schema,
261            rows: rows.into_iter(),
262        }
263    }
264
265    pub fn with_types(
266        schema: Vec<String>,
267        types: Vec<Option<ColumnType>>,
268        rows: Vec<ResultRow>,
269    ) -> Self {
270        let physical_schema = RowSchema::with_types(schema.clone(), types);
271        Self {
272            schema,
273            physical_schema,
274            rows: rows.into_iter(),
275        }
276    }
277}
278
279impl RowSource for VecSource {
280    fn schema(&self) -> &[String] {
281        &self.schema
282    }
283
284    fn physical_schema(&self) -> Option<&RowSchema> {
285        Some(&self.physical_schema)
286    }
287
288    fn estimated_cardinality(&self) -> Option<u64> {
289        u64::try_from(self.rows.len()).ok()
290    }
291
292    fn next_row(&mut self) -> ExecResult<Option<ResultRow>> {
293        Ok(self.rows.next())
294    }
295}
296
297impl PhysicalVecSource {
298    pub fn new(schema: RowSchema, rows: Vec<PhysicalRow>) -> Self {
299        Self {
300            schema,
301            rows: rows.into_iter(),
302        }
303    }
304}
305
306impl RowSource for PhysicalVecSource {
307    fn schema(&self) -> &[String] {
308        self.schema.columns()
309    }
310
311    fn physical_schema(&self) -> Option<&RowSchema> {
312        Some(&self.schema)
313    }
314
315    fn estimated_cardinality(&self) -> Option<u64> {
316        u64::try_from(self.rows.len()).ok()
317    }
318
319    fn next_row(&mut self) -> ExecResult<Option<ResultRow>> {
320        Ok(self
321            .rows
322            .next()
323            .map(|row| self.schema.view(&row).to_result_row()))
324    }
325
326    fn next_physical_batch(&mut self, max_rows: usize) -> ExecResult<Vec<PhysicalRow>> {
327        Ok(self.rows.by_ref().take(max_rows).collect())
328    }
329}
330
331/// `TableScan`: a leaf operator that drains its [`RowSource`].
332///
333/// Emits batches of at most [`DEFAULT_BATCH_SIZE`] rows. Idempotent
334/// `open` / `close` so the operator can be re-opened in tests.
335pub struct TableScan {
336    source: Option<Box<dyn RowSource>>,
337    schema: RowSchema,
338    ordering: Vec<PhysicalOrder>,
339    estimated_cardinality: Option<u64>,
340    exhausted: bool,
341}
342
343impl TableScan {
344    pub fn new(source: Box<dyn RowSource>) -> Self {
345        let schema = source
346            .physical_schema()
347            .cloned()
348            .unwrap_or_else(|| RowSchema::new(source.schema().to_vec()));
349        let ordering = source.output_ordering().to_vec();
350        let estimated_cardinality = source.estimated_cardinality();
351        Self {
352            source: Some(source),
353            schema,
354            ordering,
355            estimated_cardinality,
356            exhausted: false,
357        }
358    }
359
360    pub fn from_rows(schema: Vec<String>, rows: Vec<ResultRow>) -> Self {
361        Self::new(Box::new(VecSource::new(schema, rows)))
362    }
363
364    pub fn from_typed_rows(
365        schema: Vec<String>,
366        types: Vec<Option<ColumnType>>,
367        rows: Vec<ResultRow>,
368    ) -> Self {
369        Self::new(Box::new(VecSource::with_types(schema, types, rows)))
370    }
371
372    pub fn from_rows_with_schema(schema: RowSchema, rows: Vec<ResultRow>) -> Self {
373        Self::new(Box::new(VecSource::with_row_schema(schema, rows)))
374    }
375
376    pub fn from_physical_rows(schema: RowSchema, rows: Vec<PhysicalRow>) -> Self {
377        Self::new(Box::new(PhysicalVecSource::new(schema, rows)))
378    }
379}
380
381impl PhysicalOperator for TableScan {
382    fn row_schema(&self) -> &RowSchema {
383        &self.schema
384    }
385
386    fn estimated_cardinality(&self) -> Option<u64> {
387        self.estimated_cardinality
388    }
389
390    fn output_ordering(&self) -> &[PhysicalOrder] {
391        &self.ordering
392    }
393
394    fn backward_scan_support(&self) -> BackwardScanSupport {
395        BackwardScanSupport::Materialize
396    }
397
398    fn consume_into_aggregate(
399        &mut self,
400        executor: &mut dyn crate::relational::AggregateExecutor,
401    ) -> ExecResult<bool> {
402        let Some(source) = self.source.as_mut() else {
403            return Ok(false);
404        };
405        source.consume_into_aggregate(executor)
406    }
407
408    fn open(&mut self) -> ExecResult<()> {
409        self.exhausted = false;
410        Ok(())
411    }
412
413    fn next(&mut self) -> ExecResult<Option<Batch>> {
414        if self.exhausted {
415            return Ok(None);
416        }
417        let Some(src) = self.source.as_mut() else {
418            return Ok(None);
419        };
420        let buf = src.next_physical_batch(DEFAULT_BATCH_SIZE)?;
421        if buf.is_empty() {
422            self.exhausted = true;
423            return Ok(None);
424        }
425        Ok(Some(Batch::from_physical_rows(self.schema.clone(), buf)))
426    }
427
428    fn close(&mut self) -> ExecResult<()> {
429        self.source = None;
430        self.exhausted = true;
431        Ok(())
432    }
433}
434
435#[cfg(test)]
436mod tests {
437    use super::*;
438    use crate::physical::run_to_rows;
439    use uqa_core::Value;
440
441    fn row<const N: usize>(pairs: [(&str, Value); N]) -> ResultRow {
442        pairs.into_iter().map(|(k, v)| (k.to_string(), v)).collect()
443    }
444
445    #[test]
446    fn table_scan_drains_source() {
447        let rows = vec![
448            row([("id", Value::Int(1)), ("name", Value::Str("a".into()))]),
449            row([("id", Value::Int(2)), ("name", Value::Str("b".into()))]),
450        ];
451        let mut scan = TableScan::from_rows(vec!["id".into(), "name".into()], rows);
452        let (cols, rows) = run_to_rows(&mut scan).unwrap();
453        assert_eq!(cols, vec!["id", "name"]);
454        assert_eq!(rows.len(), 2);
455    }
456
457    #[test]
458    fn table_scan_empty_source_returns_no_batch() {
459        let mut scan = TableScan::from_rows(vec!["id".into()], Vec::new());
460        let (_cols, rows) = run_to_rows(&mut scan).unwrap();
461        assert!(rows.is_empty());
462    }
463
464    #[test]
465    fn iterator_scan_propagates_a_late_producer_error() {
466        let rows = vec![
467            Ok(row([("id", Value::Int(1))])),
468            Err(crate::ExecError::Other("late producer failure".into())),
469        ];
470        let mut scan = RowIteratorScan::new(vec!["id".into()], Box::new(rows.into_iter()));
471        let error = run_to_rows(&mut scan).unwrap_err();
472        assert!(error.to_string().contains("late producer failure"));
473    }
474}
uqa_execution/scan.rs

uqa_execution/
scan.rs
