use std::collections::{HashMap, HashSet};
use reifydb_core::{
common::JoinType,
value::column::{columns::Columns, headers::ColumnHeaders},
};
use reifydb_transaction::transaction::Transaction;
use reifydb_value::{
fragment::Fragment,
reifydb_assertions,
util::hash::Hash128,
value::{Value, row_number::RowNumber},
};
use tracing::instrument;
use super::common::{JoinContext, compute_join_hash, load_and_merge_all, resolve_column_names};
use crate::{
Result,
vm::volcano::query::{QueryContext, QueryNode},
};
pub struct NaturalJoinNode {
left: Box<dyn QueryNode>,
right: Box<dyn QueryNode>,
join_type: JoinType,
alias: Option<Fragment>,
headers: Option<ColumnHeaders>,
context: JoinContext,
}
impl NaturalJoinNode {
pub(crate) fn new(
left: Box<dyn QueryNode>,
right: Box<dyn QueryNode>,
join_type: JoinType,
alias: Option<Fragment>,
) -> Self {
Self {
left,
right,
join_type,
alias,
headers: None,
context: JoinContext::new(),
}
}
fn find_common_columns(left_columns: &Columns, right_columns: &Columns) -> Vec<(String, usize, usize)> {
let mut common_columns = Vec::new();
for (left_idx, left_col) in left_columns.iter().enumerate() {
for (right_idx, right_col) in right_columns.iter().enumerate() {
if left_col.name() == right_col.name() {
common_columns.push((left_col.name().text().to_string(), left_idx, right_idx));
}
}
}
common_columns
}
}
impl QueryNode for NaturalJoinNode {
#[instrument(name = "volcano::join::natural::initialize", level = "trace", skip_all)]
fn initialize<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &QueryContext) -> Result<()> {
self.context.set(ctx);
self.left.initialize(rx, ctx)?;
self.right.initialize(rx, ctx)?;
Ok(())
}
#[instrument(name = "volcano::join::natural::next", level = "trace", skip_all)]
fn next<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &mut QueryContext) -> Result<Option<Columns>> {
reifydb_assertions! {
assert!(self.context.is_initialized(), "NaturalJoinNode::next() called before initialize()");
}
if self.headers.is_some() {
return Ok(None);
}
let left_columns = load_and_merge_all(&mut self.left, rx, ctx)?;
let right_columns = load_and_merge_all(&mut self.right, rx, ctx)?;
let left_rows = left_columns.row_count();
let left_row_numbers = left_columns.row_numbers().to_vec();
let common_columns = Self::find_common_columns(&left_columns, &right_columns);
if common_columns.is_empty() {
return Ok(None);
}
let excluded_right_cols: HashSet<usize> =
common_columns.iter().map(|(_, _, right_idx)| *right_idx).collect();
let excluded_indices: Vec<usize> = excluded_right_cols.iter().copied().collect();
let resolved =
resolve_column_names(&left_columns, &right_columns, &self.alias, Some(&excluded_indices));
let right_col_indices: Vec<usize> = common_columns.iter().map(|(_, _, ri)| *ri).collect();
let mut hash_buf = Vec::with_capacity(256);
let hash_table = Self::build(&right_columns, &right_col_indices, &mut hash_buf);
let left_col_indices: Vec<usize> = common_columns.iter().map(|(_, li, _)| *li).collect();
let (result_rows, result_row_numbers) = self.probe(
&left_columns,
&right_columns,
&hash_table,
&common_columns,
&excluded_right_cols,
&left_col_indices,
&left_row_numbers,
left_rows,
&mut hash_buf,
);
let columns = Self::materialize(&resolved.qualified_names, result_rows, result_row_numbers);
self.headers = Some(ColumnHeaders::from_columns(&columns));
Ok(Some(columns))
}
fn headers(&self) -> Option<ColumnHeaders> {
self.headers.clone()
}
}
impl NaturalJoinNode {
#[instrument(level = "trace", skip_all, name = "volcano::join::natural::build")]
fn build(
right_columns: &Columns,
right_col_indices: &[usize],
hash_buf: &mut Vec<u8>,
) -> HashMap<Hash128, Vec<usize>> {
let mut hash_table: HashMap<Hash128, Vec<usize>> = HashMap::new();
let right_rows = right_columns.row_count();
for j in 0..right_rows {
if let Some(h) = compute_join_hash(right_columns, right_col_indices, j, hash_buf) {
hash_table.entry(h).or_default().push(j);
}
}
hash_table
}
#[allow(clippy::too_many_arguments)]
#[instrument(level = "trace", skip_all, name = "volcano::join::natural::probe")]
fn probe(
&self,
left_columns: &Columns,
right_columns: &Columns,
hash_table: &HashMap<Hash128, Vec<usize>>,
common_columns: &[(String, usize, usize)],
excluded_right_cols: &HashSet<usize>,
left_col_indices: &[usize],
left_row_numbers: &[RowNumber],
left_rows: usize,
hash_buf: &mut Vec<u8>,
) -> (Vec<Vec<Value>>, Vec<RowNumber>) {
let mut result_rows = Vec::new();
let mut result_row_numbers: Vec<RowNumber> = Vec::new();
for i in 0..left_rows {
let left_row = left_columns.get_row(i);
let mut matched = false;
let candidates = compute_join_hash(left_columns, left_col_indices, i, hash_buf)
.and_then(|h| hash_table.get(&h));
if let Some(indices) = candidates {
for &j in indices {
let right_row = right_columns.get_row(j);
let all_match = common_columns.iter().all(|(_, left_idx, right_idx)| {
left_row[*left_idx] == right_row[*right_idx]
});
if all_match {
let mut combined = left_row.clone();
for (idx, value) in right_row.iter().enumerate() {
if !excluded_right_cols.contains(&idx) {
combined.push(value.clone());
}
}
result_rows.push(combined);
matched = true;
if !left_row_numbers.is_empty() {
result_row_numbers.push(left_row_numbers[i]);
}
}
}
}
if !matched && matches!(self.join_type, JoinType::Left) {
let mut combined = left_row.clone();
let undefined_count = right_columns.len() - excluded_right_cols.len();
combined.extend(vec![Value::none(); undefined_count]);
result_rows.push(combined);
if !left_row_numbers.is_empty() {
result_row_numbers.push(left_row_numbers[i]);
}
}
}
(result_rows, result_row_numbers)
}
#[instrument(level = "trace", skip_all, name = "volcano::join::natural::materialize")]
fn materialize(
qualified_names: &[String],
result_rows: Vec<Vec<Value>>,
result_row_numbers: Vec<RowNumber>,
) -> Columns {
let names_refs: Vec<&str> = qualified_names.iter().map(|s| s.as_str()).collect();
if result_row_numbers.is_empty() {
Columns::from_rows(&names_refs, &result_rows)
} else {
Columns::from_rows(&names_refs, &result_rows).with_row_numbers(result_row_numbers)
}
}
}