uqa_execution/join/
nested_loop.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Nested-loop physical join execution.
8
9use uqa_sql::ast::JoinKind;
10use uqa_sql::expr::truthy;
11use uqa_sql::ResultRow;
12
13use super::{
14    output_schema, push_output_row, HybridRowStore, MatchFlags, DEFAULT_JOIN_WORK_MEM_BYTES,
15};
16use crate::{
17    Batch, ExecError, ExecResult, PhysicalOperator, PhysicalRow, RowSchema, ScalarExpr,
18    SharedExpressionEvaluator, SpillBuffer,
19};
20
21/// Nested-loop implementation for arbitrary join predicates and every SQL
22/// outer-join shape. Predicate evaluation happens against the merged row, so
23/// qualified columns and engine-provided scalar/subquery semantics remain
24/// available through the shared expression evaluator.
25pub struct NestedLoopJoin<'a> {
26    left: Box<dyn PhysicalOperator + 'a>,
27    right: Box<dyn PhysicalOperator + 'a>,
28    kind: JoinKind,
29    predicate: Option<ScalarExpr>,
30    evaluator: SharedExpressionEvaluator<'a>,
31    left_nulls: PhysicalRow,
32    right_nulls: PhysicalRow,
33    schema: RowSchema,
34    work_mem_bytes: usize,
35    output: Option<crate::spill::SpillDrain>,
36    output_spilled: bool,
37    right_input_spilled: bool,
38}
39
40impl<'a> NestedLoopJoin<'a> {
41    pub fn new(
42        left: Box<dyn PhysicalOperator + 'a>,
43        right: Box<dyn PhysicalOperator + 'a>,
44        kind: JoinKind,
45        predicate: Option<ScalarExpr>,
46        evaluator: SharedExpressionEvaluator<'a>,
47        left_nulls: ResultRow,
48        right_nulls: ResultRow,
49    ) -> Self {
50        Self::new_with_work_mem(
51            left,
52            right,
53            kind,
54            predicate,
55            evaluator,
56            left_nulls,
57            right_nulls,
58            DEFAULT_JOIN_WORK_MEM_BYTES,
59        )
60    }
61
62    #[expect(
63        clippy::too_many_arguments,
64        reason = "keeps join keys and schema aligned"
65    )]
66    pub fn new_with_work_mem(
67        left: Box<dyn PhysicalOperator + 'a>,
68        right: Box<dyn PhysicalOperator + 'a>,
69        kind: JoinKind,
70        predicate: Option<ScalarExpr>,
71        evaluator: SharedExpressionEvaluator<'a>,
72        left_nulls: ResultRow,
73        right_nulls: ResultRow,
74        work_mem_bytes: usize,
75    ) -> Self {
76        let schema = output_schema(
77            left.row_schema(),
78            right.row_schema(),
79            &left_nulls,
80            &right_nulls,
81        );
82        let left_nulls = PhysicalRow::nulls(left.row_schema().physical_width());
83        let right_nulls = PhysicalRow::nulls(right.row_schema().physical_width());
84        Self {
85            left,
86            right,
87            kind,
88            predicate,
89            evaluator,
90            left_nulls,
91            right_nulls,
92            schema,
93            work_mem_bytes,
94            output: None,
95            output_spilled: false,
96            right_input_spilled: false,
97        }
98    }
99
100    pub fn output_has_spilled(&self) -> bool {
101        self.output_spilled
102    }
103
104    /// Whether the repeatable nested-loop build input exceeded its memory
105    /// budget and migrated to an indexed temporary row store.
106    pub fn right_input_has_spilled(&self) -> bool {
107        self.right_input_spilled
108    }
109
110    fn matches(&self, row: &PhysicalRow) -> ExecResult<bool> {
111        match self.predicate.as_ref() {
112            None => Ok(true),
113            Some(predicate) => Ok(truthy(&self.evaluator.evaluate_physical(
114                predicate,
115                &self.schema,
116                row,
117            )?)),
118        }
119    }
120}
121
122impl PhysicalOperator for NestedLoopJoin<'_> {
123    fn row_schema(&self) -> &RowSchema {
124        &self.schema
125    }
126
127    fn open(&mut self) -> ExecResult<()> {
128        self.output = None;
129        self.output_spilled = false;
130        self.right_input_spilled = false;
131
132        let right_budget = self.work_mem_bytes / 2;
133        let output_budget = self.work_mem_bytes.saturating_sub(right_budget);
134        let right_schema = self.right.row_schema().clone();
135        let mut right = HybridRowStore::new(right_schema, right_budget);
136        self.right.open()?;
137        while let Some(batch) = self.right.next()? {
138            for row in batch.rows {
139                right.push(row)?;
140            }
141        }
142        self.right_input_spilled = right.has_spilled();
143        let mut matched_right = matches!(self.kind, JoinKind::Right | JoinKind::Full)
144            .then(|| MatchFlags::new(right.len()))
145            .transpose()?;
146        let mut output = SpillBuffer::new(output_budget);
147        let mut pending = Vec::with_capacity(crate::batch::DEFAULT_BATCH_SIZE);
148
149        self.left.open()?;
150        while let Some(batch) = self.left.next()? {
151            for left_row in batch.rows {
152                let mut matched_left = false;
153                for right_index in 0..right.len() {
154                    right.with_row(right_index, |right_row| {
155                        let merged = PhysicalRow::concat(&left_row, right_row);
156                        if self.matches(&merged)? {
157                            push_output_row(&mut output, &mut pending, &self.schema, merged)?;
158                            matched_left = true;
159                            if let Some(flags) = matched_right.as_mut() {
160                                flags.mark(right_index)?;
161                            }
162                        }
163                        Ok(())
164                    })?;
165                }
166                if !matched_left && matches!(self.kind, JoinKind::Left | JoinKind::Full) {
167                    push_output_row(
168                        &mut output,
169                        &mut pending,
170                        &self.schema,
171                        PhysicalRow::concat_left_owned(left_row, &self.right_nulls),
172                    )?;
173                }
174            }
175        }
176
177        if matches!(self.kind, JoinKind::Right | JoinKind::Full) {
178            let matched_right = matched_right.as_mut().ok_or_else(|| {
179                ExecError::Other("right/full nested-loop join has no match flags".into())
180            })?;
181            for right_index in 0..right.len() {
182                if !matched_right.is_marked(right_index)? {
183                    right.with_row(right_index, |right_row| {
184                        push_output_row(
185                            &mut output,
186                            &mut pending,
187                            &self.schema,
188                            PhysicalRow::concat(&self.left_nulls, right_row),
189                        )
190                    })?;
191                }
192            }
193        }
194        if !pending.is_empty() {
195            output.push(Batch::from_physical_rows(self.schema.clone(), pending))?;
196        }
197        self.output_spilled = output.has_spilled();
198        self.output = Some(output.drain()?);
199        Ok(())
200    }
201
202    fn next(&mut self) -> ExecResult<Option<Batch>> {
203        self.output
204            .as_mut()
205            .map_or(Ok(None), |output| output.next().transpose())
206    }
207
208    fn close(&mut self) -> ExecResult<()> {
209        self.output = None;
210        let left = self.left.close();
211        let right = self.right.close();
212        crate::physical::with_cleanup(left, right, "close right nested-loop join input")
213    }
214}
uqa_execution/join/nested_loop.rs

uqa_execution/join/
nested_loop.rs
