reifydb_engine/vm/volcano/join/
nested_loop.rs1use reifydb_core::value::column::{columns::Columns, headers::ColumnHeaders};
5use reifydb_evaluate::expression::{
6 compile::compile_expression,
7 context::{CompileContext, EvalContext},
8};
9use reifydb_rql::expression::Expression;
10use reifydb_transaction::transaction::Transaction;
11use reifydb_value::{
12 fragment::Fragment,
13 reifydb_assertions,
14 value::{Value, row_number::RowNumber},
15};
16use tracing::instrument;
17
18use super::common::{JoinContext, build_eval_columns, load_and_merge_all, resolve_column_names};
19use crate::{
20 Result,
21 vm::volcano::query::{QueryContext, QueryNode, eval_context_from_query},
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 = eval_context_from_query(ctx);
112 let (result_rows, result_row_numbers) = self.probe(
113 &session,
114 &left_columns,
115 &right_columns,
116 &left_row_numbers,
117 left_rows,
118 right_rows,
119 right_width,
120 );
121
122 let columns = Self::materialize(&resolved.qualified_names, result_rows, result_row_numbers);
123
124 self.headers = Some(ColumnHeaders::from_columns(&columns));
125 Ok(Some(columns))
126 }
127
128 fn headers(&self) -> Option<ColumnHeaders> {
129 self.headers.clone()
130 }
131}
132
133impl NestedLoopJoinNode {
134 #[allow(clippy::too_many_arguments)]
135 #[instrument(level = "trace", skip_all, name = "volcano::join::nested_loop::probe")]
136 fn probe(
137 &self,
138 session: &EvalContext,
139 left_columns: &Columns,
140 right_columns: &Columns,
141 left_row_numbers: &[RowNumber],
142 left_rows: usize,
143 right_rows: usize,
144 right_width: usize,
145 ) -> (Vec<Vec<Value>>, Vec<RowNumber>) {
146 let mut result_rows = Vec::new();
147 let mut result_row_numbers: Vec<RowNumber> = Vec::new();
148
149 for i in 0..left_rows {
150 let left_row = left_columns.get_row(i);
151
152 let mut matched = false;
153 for j in 0..right_rows {
154 let right_row = right_columns.get_row(j);
155
156 let eval_columns = build_eval_columns(
157 left_columns,
158 right_columns,
159 &left_row,
160 &right_row,
161 &self.alias,
162 );
163
164 let exec_ctx = session.with_eval_join(Columns::new(eval_columns));
165
166 let all_true = self.context.compiled.iter().fold(true, |acc, compiled_expr| {
167 let col = compiled_expr.execute(&exec_ctx).unwrap();
168 matches!(col.data().get_value(0), Value::Boolean(true)) && acc
169 });
170
171 if all_true {
172 let mut combined = left_row.clone();
173 combined.extend(right_row.clone());
174 result_rows.push(combined);
175 matched = true;
176 if !left_row_numbers.is_empty() {
177 result_row_numbers.push(left_row_numbers[i]);
178 }
179 }
180 }
181
182 if self.mode == NestedLoopMode::Left && !matched {
183 let mut combined = left_row.clone();
184 combined.extend(vec![Value::none(); right_width]);
185 result_rows.push(combined);
186 if !left_row_numbers.is_empty() {
187 result_row_numbers.push(left_row_numbers[i]);
188 }
189 }
190 }
191
192 (result_rows, result_row_numbers)
193 }
194
195 #[instrument(level = "trace", skip_all, name = "volcano::join::nested_loop::materialize")]
196 fn materialize(
197 qualified_names: &[String],
198 result_rows: Vec<Vec<Value>>,
199 result_row_numbers: Vec<RowNumber>,
200 ) -> Columns {
201 let names_refs: Vec<&str> = qualified_names.iter().map(|s| s.as_str()).collect();
202 if result_row_numbers.is_empty() {
203 Columns::from_rows(&names_refs, &result_rows)
204 } else {
205 Columns::from_rows(&names_refs, &result_rows).with_row_numbers(result_row_numbers)
206 }
207 }
208}