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    #[allow(clippy::too_many_arguments)]
42    pub fn new(
43        left: Box<dyn PhysicalOperator + 'a>,
44        right: Box<dyn PhysicalOperator + 'a>,
45        kind: JoinKind,
46        predicate: Option<ScalarExpr>,
47        evaluator: SharedExpressionEvaluator<'a>,
48        left_nulls: ResultRow,
49        right_nulls: ResultRow,
50    ) -> Self {
51        Self::new_with_work_mem(
52            left,
53            right,
54            kind,
55            predicate,
56            evaluator,
57            left_nulls,
58            right_nulls,
59            DEFAULT_JOIN_WORK_MEM_BYTES,
60        )
61    }
62
63    #[allow(clippy::too_many_arguments)]
64    pub fn new_with_work_mem(
65        left: Box<dyn PhysicalOperator + 'a>,
66        right: Box<dyn PhysicalOperator + 'a>,
67        kind: JoinKind,
68        predicate: Option<ScalarExpr>,
69        evaluator: SharedExpressionEvaluator<'a>,
70        left_nulls: ResultRow,
71        right_nulls: ResultRow,
72        work_mem_bytes: usize,
73    ) -> Self {
74        let schema = output_schema(
75            left.row_schema(),
76            right.row_schema(),
77            &left_nulls,
78            &right_nulls,
79        );
80        let left_nulls = PhysicalRow::nulls(left.row_schema().physical_width());
81        let right_nulls = PhysicalRow::nulls(right.row_schema().physical_width());
82        Self {
83            left,
84            right,
85            kind,
86            predicate,
87            evaluator,
88            left_nulls,
89            right_nulls,
90            schema,
91            work_mem_bytes,
92            output: None,
93            output_spilled: false,
94            right_input_spilled: false,
95        }
96    }
97
98    pub fn output_has_spilled(&self) -> bool {
99        self.output_spilled
100    }
101
102    /// Whether the repeatable nested-loop build input exceeded its memory
103    /// budget and migrated to an indexed temporary row store.
104    pub fn right_input_has_spilled(&self) -> bool {
105        self.right_input_spilled
106    }
107
108    fn matches(&self, row: &PhysicalRow) -> ExecResult<bool> {
109        match self.predicate.as_ref() {
110            None => Ok(true),
111            Some(predicate) => Ok(truthy(&self.evaluator.evaluate_physical(
112                predicate,
113                &self.schema,
114                row,
115            )?)),
116        }
117    }
118}
119
120impl PhysicalOperator for NestedLoopJoin<'_> {
121    fn row_schema(&self) -> &RowSchema {
122        &self.schema
123    }
124
125    fn open(&mut self) -> ExecResult<()> {
126        self.output = None;
127        self.output_spilled = false;
128        self.right_input_spilled = false;
129
130        let right_budget = self.work_mem_bytes / 2;
131        let output_budget = self.work_mem_bytes.saturating_sub(right_budget);
132        let right_schema = self.right.row_schema().clone();
133        let mut right = HybridRowStore::new(right_schema, right_budget);
134        self.right.open()?;
135        while let Some(batch) = self.right.next()? {
136            for row in batch.rows {
137                right.push(row)?;
138            }
139        }
140        self.right_input_spilled = right.has_spilled();
141        let mut matched_right = matches!(self.kind, JoinKind::Right | JoinKind::Full)
142            .then(|| MatchFlags::new(right.len()))
143            .transpose()?;
144        let mut output = SpillBuffer::new(output_budget);
145        let mut pending = Vec::with_capacity(crate::batch::DEFAULT_BATCH_SIZE);
146
147        self.left.open()?;
148        while let Some(batch) = self.left.next()? {
149            for left_row in batch.rows {
150                let mut matched_left = false;
151                for right_index in 0..right.len() {
152                    right.with_row(right_index, |right_row| {
153                        let merged = PhysicalRow::concat(&left_row, right_row);
154                        if self.matches(&merged)? {
155                            push_output_row(&mut output, &mut pending, &self.schema, merged)?;
156                            matched_left = true;
157                            if let Some(flags) = matched_right.as_mut() {
158                                flags.mark(right_index)?;
159                            }
160                        }
161                        Ok(())
162                    })?;
163                }
164                if !matched_left && matches!(self.kind, JoinKind::Left | JoinKind::Full) {
165                    push_output_row(
166                        &mut output,
167                        &mut pending,
168                        &self.schema,
169                        PhysicalRow::concat_left_owned(left_row, &self.right_nulls),
170                    )?;
171                }
172            }
173        }
174
175        if matches!(self.kind, JoinKind::Right | JoinKind::Full) {
176            let matched_right = matched_right.as_mut().ok_or_else(|| {
177                ExecError::Other("right/full nested-loop join has no match flags".into())
178            })?;
179            for right_index in 0..right.len() {
180                if !matched_right.is_marked(right_index)? {
181                    right.with_row(right_index, |right_row| {
182                        push_output_row(
183                            &mut output,
184                            &mut pending,
185                            &self.schema,
186                            PhysicalRow::concat(&self.left_nulls, right_row),
187                        )
188                    })?;
189                }
190            }
191        }
192        if !pending.is_empty() {
193            output.push(Batch::from_physical_rows(self.schema.clone(), pending))?;
194        }
195        self.output_spilled = output.has_spilled();
196        self.output = Some(output.drain()?);
197        Ok(())
198    }
199
200    fn next(&mut self) -> ExecResult<Option<Batch>> {
201        self.output
202            .as_mut()
203            .map_or(Ok(None), |output| output.next().transpose())
204    }
205
206    fn close(&mut self) -> ExecResult<()> {
207        self.output = None;
208        let left = self.left.close();
209        let right = self.right.close();
210        crate::physical::with_cleanup(left, right, "close right nested-loop join input")
211    }
212}
uqa_execution/join/nested_loop.rs

uqa_execution/join/
nested_loop.rs
