uqa-engine 0.2.3

Engine: schema-aware table store, catalog restore, transactions
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

/// Preserve row-at-a-time demand across the scalar projection immediately below `LockRows`. `PostgreSQL` evaluates that projection for the candidate it is about to lock, but it does not evaluate a whole vectorized batch after an enclosing LIMIT has already obtained enough locked rows.
pub(super) struct RowAtATime<'a> {
    input: Box<dyn uqa_execution::PhysicalOperator + 'a>,
    schema: uqa_execution::RowSchema,
    ordering: Vec<uqa_execution::PhysicalOrder>,
    pending: std::vec::IntoIter<uqa_execution::PhysicalRow>,
}

impl<'a> RowAtATime<'a> {
    pub(super) fn new(input: Box<dyn uqa_execution::PhysicalOperator + 'a>) -> Self {
        let schema = input.row_schema().clone();
        let ordering = input.output_ordering().to_vec();
        Self {
            input,
            schema,
            ordering,
            pending: Vec::new().into_iter(),
        }
    }
}

impl uqa_execution::PhysicalOperator for RowAtATime<'_> {
    fn row_schema(&self) -> &uqa_execution::RowSchema {
        &self.schema
    }

    fn estimated_cardinality(&self) -> Option<u64> {
        self.input.estimated_cardinality()
    }

    fn output_ordering(&self) -> &[uqa_execution::PhysicalOrder] {
        &self.ordering
    }

    fn backward_scan_support(&self) -> uqa_execution::BackwardScanSupport {
        self.input.backward_scan_support()
    }

    fn open(&mut self) -> uqa_execution::ExecResult<()> {
        self.pending = Vec::new().into_iter();
        self.input.open()
    }

    fn next(&mut self) -> uqa_execution::ExecResult<Option<uqa_execution::Batch>> {
        loop {
            if let Some(row) = self.pending.next() {
                return Ok(Some(uqa_execution::Batch::from_physical_rows(
                    self.schema.clone(),
                    vec![row],
                )));
            }
            let Some(batch) = self.input.next()? else {
                return Ok(None);
            };
            if batch.schema != self.schema {
                return Err(uqa_execution::ExecError::Other(format!(
                    "row-at-a-time input schema mismatch: expected {:?}, got {:?}",
                    self.schema, batch.schema
                )));
            }
            self.pending = batch.rows.into_iter();
        }
    }

    fn next_direction(
        &mut self,
        direction: uqa_execution::PhysicalScanDirection,
    ) -> uqa_execution::ExecResult<Option<uqa_execution::Batch>> {
        if self.pending.len() != 0 {
            return Err(uqa_execution::ExecError::Other(
                "row-at-a-time operator cannot mix batched and directional pulls".into(),
            ));
        }
        let Some(batch) = self.input.next_direction(direction)? else {
            return Ok(None);
        };
        if batch.schema != self.schema {
            return Err(uqa_execution::ExecError::Other(format!(
                "row-at-a-time input schema mismatch: expected {:?}, got {:?}",
                self.schema, batch.schema
            )));
        }
        if batch.rows.len() != 1 {
            return Err(uqa_execution::ExecError::Other(format!(
                "directional row-at-a-time input returned {} rows",
                batch.rows.len()
            )));
        }
        Ok(Some(batch))
    }

    fn rewind(&mut self) -> uqa_execution::ExecResult<()> {
        self.pending = Vec::new().into_iter();
        self.input.rewind()
    }

    fn close(&mut self) -> uqa_execution::ExecResult<()> {
        self.pending = Vec::new().into_iter();
        self.input.close()
    }
}