reifydb_engine/vm/volcano/join/
nested_loop.rs1use 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}