uqa_execution/
lateral_join.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Correlated physical join operator.
8
9use uqa_sql::ast::JoinKind;
10use uqa_sql::ResultRow;
11
12use crate::batch::DEFAULT_BATCH_SIZE;
13use crate::{
14    Batch, ExecError, ExecResult, OwnedPhysicalRow, PhysicalOperator, PhysicalRow, RowSchema,
15};
16
17/// Engine seam for a correlated right-hand relation.
18///
19/// The physical operator owns join iteration, `ON` filtering, and outer-row
20/// preservation. The engine callback only evaluates the right relation and
21/// scalar predicate in the current correlated scope.
22pub type LateralRows = Box<dyn Iterator<Item = ExecResult<OwnedPhysicalRow>> + Send>;
23
24pub trait LateralSource: Send {
25    fn rows_for(&mut self, left: &OwnedPhysicalRow) -> ExecResult<LateralRows>;
26
27    fn matches(&mut self, joined: &OwnedPhysicalRow) -> ExecResult<bool>;
28}
29
30fn output_schema(
31    left: &RowSchema,
32    right: &RowSchema,
33    left_nulls: &ResultRow,
34    right_nulls: &ResultRow,
35) -> RowSchema {
36    RowSchema::join(
37        left,
38        right,
39        left_nulls.keys().chain(right_nulls.keys()).cloned(),
40    )
41}
42
43/// Streaming physical implementation of a SQL `LATERAL` join.
44///
45/// The left child is pulled in batches. The correlated right relation is
46/// evaluated once per left row, and produced rows are drained in bounded
47/// output batches instead of materialising the complete join result.
48pub struct LateralJoin<'a> {
49    left: Box<dyn PhysicalOperator + 'a>,
50    source: Box<dyn LateralSource + 'a>,
51    kind: JoinKind,
52    schema: RowSchema,
53    left_schema: RowSchema,
54    right_schema: RowSchema,
55    left_rows: std::vec::IntoIter<OwnedPhysicalRow>,
56    current_left: Option<OwnedPhysicalRow>,
57    right_rows: Option<LateralRows>,
58    matched_left: bool,
59    exhausted: bool,
60}
61
62impl<'a> LateralJoin<'a> {
63    pub fn new(
64        left: Box<dyn PhysicalOperator + 'a>,
65        source: Box<dyn LateralSource + 'a>,
66        kind: JoinKind,
67        left_nulls: ResultRow,
68        right_nulls: ResultRow,
69    ) -> Self {
70        let right_schema = RowSchema::new(right_nulls.keys().cloned().collect());
71        Self::new_with_right_schema(left, source, kind, left_nulls, right_nulls, right_schema)
72    }
73
74    pub fn new_with_right_schema(
75        left: Box<dyn PhysicalOperator + 'a>,
76        source: Box<dyn LateralSource + 'a>,
77        kind: JoinKind,
78        left_nulls: ResultRow,
79        right_nulls: ResultRow,
80        right_schema: RowSchema,
81    ) -> Self {
82        let schema = output_schema(left.row_schema(), &right_schema, &left_nulls, &right_nulls);
83        let left_schema = left.row_schema().clone();
84        Self {
85            left,
86            source,
87            kind,
88            schema,
89            left_schema,
90            right_schema,
91            left_rows: Vec::new().into_iter(),
92            current_left: None,
93            right_rows: None,
94            matched_left: false,
95            exhausted: false,
96        }
97    }
98
99    fn next_left(&mut self) -> ExecResult<Option<OwnedPhysicalRow>> {
100        loop {
101            if let Some(row) = self.left_rows.next() {
102                return Ok(Some(row));
103            }
104            let Some(batch) = self.left.next()? else {
105                return Ok(None);
106            };
107            self.left_rows = batch.into_owned_rows().into_iter();
108        }
109    }
110
111    fn begin_left_row(&mut self, left: OwnedPhysicalRow) -> ExecResult<()> {
112        self.right_rows = Some(self.source.rows_for(&left)?);
113        self.current_left = Some(left);
114        self.matched_left = false;
115        Ok(())
116    }
117}
118
119impl PhysicalOperator for LateralJoin<'_> {
120    fn row_schema(&self) -> &RowSchema {
121        &self.schema
122    }
123
124    fn open(&mut self) -> ExecResult<()> {
125        self.left_rows = Vec::new().into_iter();
126        self.current_left = None;
127        self.right_rows = None;
128        self.matched_left = false;
129        self.exhausted = false;
130        self.left.open()
131    }
132
133    fn next(&mut self) -> ExecResult<Option<Batch>> {
134        if self.exhausted && self.current_left.is_none() {
135            return Ok(None);
136        }
137
138        let mut output = Vec::with_capacity(DEFAULT_BATCH_SIZE);
139        while output.len() < DEFAULT_BATCH_SIZE {
140            if self.current_left.is_none() {
141                match self.next_left()? {
142                    Some(left) => self.begin_left_row(left)?,
143                    None => {
144                        self.exhausted = true;
145                        break;
146                    }
147                }
148            }
149
150            let next_right = self
151                .right_rows
152                .as_mut()
153                .and_then(Iterator::next)
154                .transpose()?;
155            if let Some(right) = next_right {
156                let left = self.current_left.as_ref().ok_or_else(|| {
157                    ExecError::Other(
158                        "lateral join produced a right row without a current left row".into(),
159                    )
160                })?;
161                let joined = OwnedPhysicalRow::new(
162                    self.schema.clone(),
163                    PhysicalRow::concat(&left.row, &right.row),
164                );
165                let matched =
166                    matches!(self.kind, JoinKind::Cross) || self.source.matches(&joined)?;
167                if matched {
168                    self.matched_left = true;
169                    output.push(joined.row);
170                } else if matches!(self.kind, JoinKind::Right | JoinKind::Full) {
171                    output.push(PhysicalRow::concat(
172                        &PhysicalRow::nulls(self.left_schema.physical_width()),
173                        &right.row,
174                    ));
175                }
176                continue;
177            }
178
179            self.right_rows = None;
180            if !self.matched_left && matches!(self.kind, JoinKind::Left | JoinKind::Full) {
181                let left = self.current_left.take().ok_or_else(|| {
182                    ExecError::Other(
183                        "lateral join completed a right stream without a current left row".into(),
184                    )
185                })?;
186                output.push(PhysicalRow::concat(
187                    &left.row,
188                    &PhysicalRow::nulls(self.right_schema.physical_width()),
189                ));
190            } else {
191                self.current_left = None;
192            }
193        }
194
195        if output.is_empty() {
196            return Ok(None);
197        }
198        Ok(Some(Batch::from_physical_rows(self.schema.clone(), output)))
199    }
200
201    fn close(&mut self) -> ExecResult<()> {
202        self.left_rows = Vec::new().into_iter();
203        self.current_left = None;
204        self.right_rows = None;
205        self.matched_left = false;
206        self.exhausted = true;
207        self.left.close()
208    }
209}
210
211#[cfg(test)]
212mod tests {
213    use super::*;
214    use crate::physical::run_to_rows;
215    use crate::scan::TableScan;
216    use uqa_core::Value;
217
218    fn row(values: &[(&str, Value)]) -> ResultRow {
219        values
220            .iter()
221            .map(|(column, value)| ((*column).to_string(), value.clone()))
222            .collect()
223    }
224
225    fn right_row(value: i64) -> OwnedPhysicalRow {
226        let schema = RowSchema::new(vec!["r.n".into()]);
227        let row = PhysicalRow::from_values(vec![Value::Int(value)]);
228        OwnedPhysicalRow::new(schema, row)
229    }
230
231    struct RangeSource;
232
233    impl LateralSource for RangeSource {
234        fn rows_for(&mut self, left: &OwnedPhysicalRow) -> ExecResult<LateralRows> {
235            let Value::Int(end) = left.get("l.n").cloned().unwrap_or(Value::Null) else {
236                return Ok(Box::new(std::iter::empty()));
237            };
238            Ok(Box::new((1..=end).map(|value| Ok(right_row(value)))))
239        }
240
241        fn matches(&mut self, joined: &OwnedPhysicalRow) -> ExecResult<bool> {
242            Ok(joined.get("r.n") == Some(&Value::Int(2)))
243        }
244    }
245
246    #[test]
247    fn left_lateral_preserves_a_left_row_without_an_on_match() {
248        let left = TableScan::from_rows(
249            vec!["l.n".into()],
250            vec![
251                row(&[("l.n", Value::Int(1))]),
252                row(&[("l.n", Value::Int(2))]),
253            ],
254        );
255        let mut join = LateralJoin::new(
256            Box::new(left),
257            Box::new(RangeSource),
258            JoinKind::Left,
259            row(&[("l.n", Value::Null)]),
260            row(&[("r.n", Value::Null)]),
261        );
262        let (_, rows) = run_to_rows(&mut join).unwrap();
263        assert_eq!(rows.len(), 2);
264        assert_eq!(rows[0].get("l.n"), Some(&Value::Int(1)));
265        assert_eq!(rows[0].get("r.n"), Some(&Value::Null));
266        assert_eq!(rows[1].get("l.n"), Some(&Value::Int(2)));
267        assert_eq!(rows[1].get("r.n"), Some(&Value::Int(2)));
268    }
269
270    struct LargeRangeSource;
271
272    impl LateralSource for LargeRangeSource {
273        fn rows_for(&mut self, _left: &OwnedPhysicalRow) -> ExecResult<LateralRows> {
274            Ok(Box::new((0..10_000).map(|value| Ok(right_row(value)))))
275        }
276
277        fn matches(&mut self, _joined: &OwnedPhysicalRow) -> ExecResult<bool> {
278            Ok(true)
279        }
280    }
281
282    #[test]
283    fn large_correlated_relation_is_pulled_one_output_batch_at_a_time() {
284        let left = TableScan::from_rows(vec!["l.n".into()], vec![row(&[("l.n", Value::Int(1))])]);
285        let mut join = LateralJoin::new(
286            Box::new(left),
287            Box::new(LargeRangeSource),
288            JoinKind::Cross,
289            row(&[("l.n", Value::Null)]),
290            row(&[("r.n", Value::Null)]),
291        );
292        join.open().unwrap();
293        assert_eq!(join.next().unwrap().unwrap().len(), DEFAULT_BATCH_SIZE);
294        assert_eq!(join.next().unwrap().unwrap().len(), DEFAULT_BATCH_SIZE);
295        join.close().unwrap();
296    }
297
298    struct FailingSource;
299
300    impl LateralSource for FailingSource {
301        fn rows_for(&mut self, _left: &OwnedPhysicalRow) -> ExecResult<LateralRows> {
302            Ok(Box::new(
303                vec![
304                    Ok(right_row(1)),
305                    Err(crate::ExecError::Other("injected lateral failure".into())),
306                ]
307                .into_iter(),
308            ))
309        }
310
311        fn matches(&mut self, _joined: &OwnedPhysicalRow) -> ExecResult<bool> {
312            Ok(true)
313        }
314    }
315
316    #[test]
317    fn late_correlated_source_error_is_propagated() {
318        let left = TableScan::from_rows(vec!["l.n".into()], vec![row(&[("l.n", Value::Int(1))])]);
319        let mut join = LateralJoin::new(
320            Box::new(left),
321            Box::new(FailingSource),
322            JoinKind::Cross,
323            ResultRow::new(),
324            ResultRow::new(),
325        );
326        join.open().unwrap();
327        let error = join.next().unwrap_err();
328        assert!(error.to_string().contains("injected lateral failure"));
329    }
330}
uqa_execution/lateral_join.rs

uqa_execution/
lateral_join.rs
