Skip to main content

reifydb_engine/vm/volcano/join/
nested_loop.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use reifydb_core::value::column::{columns::Columns, headers::ColumnHeaders};
5use reifydb_rql::expression::Expression;
6use reifydb_transaction::transaction::Transaction;
7use reifydb_value::{
8	fragment::Fragment,
9	reifydb_assertions,
10	value::{Value, row_number::RowNumber},
11};
12use tracing::instrument;
13
14use super::common::{JoinContext, build_eval_columns, load_and_merge_all, resolve_column_names};
15use crate::{
16	Result,
17	expression::{
18		compile::compile_expression,
19		context::{CompileContext, EvalContext},
20	},
21	vm::volcano::query::{QueryContext, QueryNode},
22};
23
24#[derive(Clone, Copy, PartialEq)]
25enum NestedLoopMode {
26	Inner,
27	Left,
28}
29
30pub struct NestedLoopJoinNode {
31	left: Box<dyn QueryNode>,
32	right: Box<dyn QueryNode>,
33	on: Vec<Expression>,
34	alias: Option<Fragment>,
35	mode: NestedLoopMode,
36	headers: Option<ColumnHeaders>,
37	context: JoinContext,
38}
39
40impl NestedLoopJoinNode {
41	pub(crate) fn new_inner(
42		left: Box<dyn QueryNode>,
43		right: Box<dyn QueryNode>,
44		on: Vec<Expression>,
45		alias: Option<Fragment>,
46	) -> Self {
47		Self {
48			left,
49			right,
50			on,
51			alias,
52			mode: NestedLoopMode::Inner,
53			headers: None,
54			context: JoinContext::new(),
55		}
56	}
57
58	pub(crate) fn new_left(
59		left: Box<dyn QueryNode>,
60		right: Box<dyn QueryNode>,
61		on: Vec<Expression>,
62		alias: Option<Fragment>,
63	) -> Self {
64		Self {
65			left,
66			right,
67			on,
68			alias,
69			mode: NestedLoopMode::Left,
70			headers: None,
71			context: JoinContext::new(),
72		}
73	}
74}
75
76impl QueryNode for NestedLoopJoinNode {
77	#[instrument(level = "trace", skip_all, name = "volcano::join::nested_loop::initialize")]
78	fn initialize<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &QueryContext) -> Result<()> {
79		let compile_ctx = CompileContext {
80			symbols: &ctx.symbols,
81		};
82		self.context.compiled =
83			self.on.iter().map(|e| compile_expression(&compile_ctx, e).expect("compile")).collect();
84		self.context.set(ctx);
85		self.left.initialize(rx, ctx)?;
86		self.right.initialize(rx, ctx)?;
87		Ok(())
88	}
89
90	#[instrument(level = "trace", skip_all, name = "volcano::join::nested_loop::next")]
91	fn next<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &mut QueryContext) -> Result<Option<Columns>> {
92		reifydb_assertions! {
93			assert!(self.context.is_initialized(), "NestedLoopJoinNode::next() called before initialize()");
94		}
95		let _stored_ctx = self.context.get();
96
97		if self.headers.is_some() {
98			return Ok(None);
99		}
100
101		let left_columns = load_and_merge_all(&mut self.left, rx, ctx)?;
102		let right_columns = load_and_merge_all(&mut self.right, rx, ctx)?;
103
104		let left_rows = left_columns.row_count();
105		let right_rows = right_columns.row_count();
106		let right_width = right_columns.len();
107		let left_row_numbers = left_columns.row_numbers.to_vec();
108
109		let resolved = resolve_column_names(&left_columns, &right_columns, &self.alias, None);
110
111		let session = EvalContext::from_query(ctx);
112		let mut result_rows = Vec::new();
113		let mut result_row_numbers: Vec<RowNumber> = Vec::new();
114
115		for i in 0..left_rows {
116			let left_row = left_columns.get_row(i);
117
118			let mut matched = false;
119			for j in 0..right_rows {
120				let right_row = right_columns.get_row(j);
121
122				let eval_columns = build_eval_columns(
123					&left_columns,
124					&right_columns,
125					&left_row,
126					&right_row,
127					&self.alias,
128				);
129
130				let exec_ctx = session.with_eval_join(Columns::new(eval_columns));
131
132				let all_true = self.context.compiled.iter().fold(true, |acc, compiled_expr| {
133					let col = compiled_expr.execute(&exec_ctx).unwrap();
134					matches!(col.data().get_value(0), Value::Boolean(true)) && acc
135				});
136
137				if all_true {
138					let mut combined = left_row.clone();
139					combined.extend(right_row.clone());
140					result_rows.push(combined);
141					matched = true;
142					if !left_row_numbers.is_empty() {
143						result_row_numbers.push(left_row_numbers[i]);
144					}
145				}
146			}
147
148			if self.mode == NestedLoopMode::Left && !matched {
149				let mut combined = left_row.clone();
150				combined.extend(vec![Value::none(); right_width]);
151				result_rows.push(combined);
152				if !left_row_numbers.is_empty() {
153					result_row_numbers.push(left_row_numbers[i]);
154				}
155			}
156		}
157
158		let names_refs: Vec<&str> = resolved.qualified_names.iter().map(|s| s.as_str()).collect();
159		let columns = if result_row_numbers.is_empty() {
160			Columns::from_rows(&names_refs, &result_rows)
161		} else {
162			Columns::from_rows(&names_refs, &result_rows).with_row_numbers(result_row_numbers)
163		};
164
165		self.headers = Some(ColumnHeaders::from_columns(&columns));
166		Ok(Some(columns))
167	}
168
169	fn headers(&self) -> Option<ColumnHeaders> {
170		self.headers.clone()
171	}
172}