reifydb_engine/vm/volcano/join/
nested_loop.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use arrow_array::RecordBatch;
5use reifydb_core::{
6	expression::Expression,
7	value::{batch::batch, column::headers::ColumnHeaders},
8};
9use reifydb_evaluate::{expression::context::EvalContext, lower::LoweredExpr};
10use reifydb_transaction::transaction::Transaction;
11use reifydb_value::{
12	fragment::Fragment,
13	reifydb_assertions,
14	value::{Value, column_view::ColumnView},
15};
16use tracing::instrument;
17
18use super::common::{
19	JoinContext, JoinSlot, build_eval_columns, load_and_merge_all, materialize_join, resolve_column_names,
20	user_row, user_views,
21};
22use crate::{
23	Result,
24	vm::volcano::query::{QueryContext, QueryNode, eval_context_from_query},
25};
26
27#[derive(Clone, Copy, PartialEq)]
28enum NestedLoopMode {
29	Inner,
30	Left,
31}
32
33pub struct NestedLoopJoinNode {
34	left: Box<dyn QueryNode>,
35	right: Box<dyn QueryNode>,
36	on: Vec<Expression>,
37	alias: Option<Fragment>,
38	mode: NestedLoopMode,
39	headers: Option<ColumnHeaders>,
40	context: JoinContext,
41}
42
43impl NestedLoopJoinNode {
44	pub(crate) fn new_inner(
45		left: Box<dyn QueryNode>,
46		right: Box<dyn QueryNode>,
47		on: Vec<Expression>,
48		alias: Option<Fragment>,
49	) -> Self {
50		Self {
51			left,
52			right,
53			on,
54			alias,
55			mode: NestedLoopMode::Inner,
56			headers: None,
57			context: JoinContext::new(),
58		}
59	}
60
61	pub(crate) fn new_left(
62		left: Box<dyn QueryNode>,
63		right: Box<dyn QueryNode>,
64		on: Vec<Expression>,
65		alias: Option<Fragment>,
66	) -> Self {
67		Self {
68			left,
69			right,
70			on,
71			alias,
72			mode: NestedLoopMode::Left,
73			headers: None,
74			context: JoinContext::new(),
75		}
76	}
77}
78
79impl QueryNode for NestedLoopJoinNode {
80	#[instrument(level = "trace", skip_all, name = "volcano::join::nested_loop::initialize")]
81	fn initialize<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &QueryContext) -> Result<()> {
82		self.context.lowered = self.on.iter().map(|e| LoweredExpr::new(e.clone(), "join_on")).collect();
83		self.context.set(ctx);
84		self.left.initialize(rx, ctx)?;
85		self.right.initialize(rx, ctx)?;
86		Ok(())
87	}
88
89	#[instrument(level = "trace", skip_all, name = "volcano::join::nested_loop::next")]
90	fn next<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &mut QueryContext) -> Result<Option<RecordBatch>> {
91		reifydb_assertions! {
92			assert!(self.context.is_initialized(), "NestedLoopJoinNode::next() called before initialize()");
93		}
94		let _stored_ctx = self.context.get();
95
96		if self.headers.is_some() {
97			return Ok(None);
98		}
99
100		let left_columns = load_and_merge_all(&mut self.left, rx, ctx)?;
101		let right_columns = load_and_merge_all(&mut self.right, rx, ctx)?;
102
103		let left_rows = left_columns.num_rows();
104		let right_rows = right_columns.num_rows();
105
106		let resolved = resolve_column_names(&left_columns, &right_columns, &self.alias, None);
107
108		let session = eval_context_from_query(ctx);
109		let (left_picks, right_picks) =
110			self.probe(&session, &left_columns, &right_columns, left_rows, right_rows)?;
111
112		let columns = materialize_join(
113			&resolved.qualified_names,
114			&[JoinSlot {
115				columns: &left_columns,
116				picks: &left_picks,
117			}],
118			&right_columns,
119			&[],
120			&right_picks,
121			0,
122		)?;
123
124		self.headers = Some(ColumnHeaders::from_batch(&columns));
125		Ok(Some(columns))
126	}
127
128	fn headers(&self) -> Option<ColumnHeaders> {
129		self.headers.clone()
130	}
131}
132
133impl NestedLoopJoinNode {
134	#[instrument(level = "trace", skip_all, name = "volcano::join::nested_loop::probe")]
135	fn probe(
136		&self,
137		session: &EvalContext,
138		left_columns: &RecordBatch,
139		right_columns: &RecordBatch,
140		left_rows: usize,
141		right_rows: usize,
142	) -> Result<(Vec<usize>, Vec<Option<usize>>)> {
143		let mut left_picks: Vec<usize> = Vec::new();
144		let mut right_picks: Vec<Option<usize>> = Vec::new();
145		let left_views = user_views(left_columns)?;
146		let right_views = user_views(right_columns)?;
147
148		for i in 0..left_rows {
149			let left_row = user_row(&left_views, i);
150
151			let mut matched = false;
152			for j in 0..right_rows {
153				let right_row = user_row(&right_views, j);
154
155				let eval_columns = build_eval_columns(
156					&left_views,
157					&right_views,
158					&left_row,
159					&right_row,
160					&self.alias,
161				);
162
163				let exec_ctx = session.with_eval_join(batch(eval_columns)?);
164
165				let mut all_true = true;
166				for lowered_expr in &self.context.lowered {
167					let col = lowered_expr.evaluate(&exec_ctx)?;
168					all_true &= matches!(
169						ColumnView::try_from(&col)?.get_value(0),
170						Value::Boolean(true)
171					);
172				}
173
174				if all_true {
175					left_picks.push(i);
176					right_picks.push(Some(j));
177					matched = true;
178				}
179			}
180
181			if self.mode == NestedLoopMode::Left && !matched {
182				left_picks.push(i);
183				right_picks.push(None);
184			}
185		}
186
187		Ok((left_picks, right_picks))
188	}
189}
reifydb_engine/vm/volcano/join/nested_loop.rs

reifydb_engine/vm/volcano/join/
nested_loop.rs
