use std::sync::Arc;
use arrow_array::{ArrayRef, RecordBatch, UInt64Array};
use arrow_schema::FieldRef;
use reifydb_core::{
error::diagnostic::operation,
value::{
batch::{batch, concat, data_size, take_rows, take_rows_or_none},
column::{
factory::{datetime, from_one, rename, typed_none},
view::group_by::common_key_type,
},
},
};
use reifydb_evaluate::lower::LoweredExpr;
use reifydb_transaction::transaction::Transaction;
use reifydb_value::{
error,
fragment::Fragment,
value::{
Value,
column_view::ColumnView,
datetime::DateTime,
system_columns::{
SystemColumn, is_system_field, stamp_system_columns, system_column, time, user_columns,
},
value_type::ValueType,
},
};
use crate::{
Result,
vm::volcano::query::{QueryContext, QueryNode, charge_query_memory_bytes, eval_context_from_query},
};
pub(crate) fn load_and_merge_all<'a>(
node: &mut Box<dyn QueryNode>,
rx: &mut Transaction<'a>,
ctx: &mut QueryContext,
) -> Result<RecordBatch> {
let mut batches = Vec::new();
let mut charged = 0usize;
let mut total = 0usize;
while let Some(columns) = node.next(rx, ctx)? {
total += data_size(&columns)?;
charge_query_memory_bytes(&ctx.memory, &mut charged, total)?;
batches.push(columns);
}
concat(&batches)
}
pub(crate) fn user_views(columns: &RecordBatch) -> Result<Vec<ColumnView<'_>>> {
user_columns(columns).map(|(field, array)| ColumnView::try_from((array, field.as_ref()))).collect()
}
pub(crate) fn user_row(views: &[ColumnView<'_>], index: usize) -> Vec<Value> {
views.iter().map(|view| view.get_value(index)).collect()
}
pub(crate) fn user_key_columns<'b>(columns: &'b RecordBatch, indices: &[usize]) -> Vec<(&'b FieldRef, &'b ArrayRef)> {
let user: Vec<(&FieldRef, &ArrayRef)> = user_columns(columns).collect();
indices.iter().map(|&index| user[index]).collect()
}
pub(crate) struct JoinSlot<'a> {
pub columns: &'a RecordBatch,
pub picks: &'a [usize],
}
pub(crate) fn materialize_join(
qualified_names: &[String],
left_slots: &[JoinSlot<'_>],
right_columns: &RecordBatch,
right_excluded: &[usize],
right_picks: &[Option<usize>],
emitted: u64,
) -> Result<RecordBatch> {
let parts = left_slots.iter().map(|slot| take_rows(slot.columns, slot.picks)).collect::<Result<Vec<_>>>()?;
let left = concat(&parts)?;
let right = take_rows_or_none(right_columns, right_picks)?;
let mut picked: Vec<(FieldRef, ArrayRef)> = Vec::with_capacity(left.num_columns() + right.num_columns());
for (field, array) in user_columns(&left) {
let name = &qualified_names[picked.len()];
picked.push(rename((field.clone(), array.clone()), name));
}
for (index, (field, array)) in user_columns(&right).enumerate() {
if right_excluded.contains(&index) {
continue;
}
let name = &qualified_names[picked.len()];
picked.push(rename((field.clone(), array.clone()), name));
}
for (field, array) in left.schema_ref().fields().iter().zip(left.columns()) {
if is_system_field(field) {
picked.push((field.clone(), array.clone()));
}
}
let mut stamps: Vec<(SystemColumn, ArrayRef)> = Vec::new();
if system_column(&left, SystemColumn::RowNumbers).is_some() {
let row_numbers = UInt64Array::from_iter_values((1..=right_picks.len() as u64).map(|i| emitted + i));
stamps.push((SystemColumn::RowNumbers, Arc::new(row_numbers)));
}
let left_time = time(&left)?;
let right_time = time(right_columns)?;
if !left_time.is_empty() && !right_time.is_empty() {
let merged: Vec<DateTime> = left_time
.iter()
.zip(right_picks)
.map(|(&time, &pick)| match pick {
None => time,
Some(index) => time.max(right_time[index]),
})
.collect();
let (_, array) = datetime(SystemColumn::Time.name(), merged);
stamps.push((SystemColumn::Time, array));
}
stamp_system_columns(batch(picked)?, stamps)
}
pub struct ResolvedColumnNames {
pub qualified_names: Vec<String>,
}
pub fn resolve_column_names(
left_columns: &RecordBatch,
right_columns: &RecordBatch,
alias: &Option<Fragment>,
excluded_right_indices: Option<&[usize]>,
) -> ResolvedColumnNames {
let mut qualified_names = Vec::new();
for (field, _) in user_columns(left_columns) {
qualified_names.push(field.name().to_string());
}
for (idx, (field, _)) in user_columns(right_columns).enumerate() {
if let Some(excluded) = excluded_right_indices
&& excluded.contains(&idx)
{
continue;
}
let col_name = field.name();
let alias_text = alias.as_ref().map(|a| a.text()).unwrap_or("other");
let prefixed_name = format!("{}_{}", alias_text, col_name);
let mut final_name = prefixed_name.clone();
if qualified_names.contains(&final_name) {
let mut counter = 2;
loop {
let candidate = format!("{}_{}", prefixed_name, counter);
if !qualified_names.contains(&candidate) {
final_name = candidate;
break;
}
counter += 1;
}
}
qualified_names.push(final_name);
}
ResolvedColumnNames {
qualified_names,
}
}
pub fn build_eval_columns(
left_columns: &[ColumnView<'_>],
right_columns: &[ColumnView<'_>],
left_row: &[Value],
right_row: &[Value],
alias: &Option<Fragment>,
) -> Vec<(FieldRef, ArrayRef)> {
let mut eval_columns = Vec::new();
for (idx, col) in left_columns.iter().enumerate() {
let name = col.field.name();
let data = match &left_row[idx] {
Value::None {
..
} => typed_none(name, &col.get_type()),
value => from_one(name, value.clone()),
};
eval_columns.push((Arc::new(col.field.clone()), data.1));
}
for (idx, col) in right_columns.iter().enumerate() {
let name = match alias {
Some(alias) => format!("{}.{}", alias.text(), col.field.name()),
None => col.field.name().clone(),
};
let data = match &right_row[idx] {
Value::None {
..
} => typed_none(&name, &col.get_type()),
value => from_one(&name, value.clone()),
};
eval_columns.push((Arc::new(col.field.clone().with_name(name)), data.1));
}
eval_columns
}
pub struct JoinContext {
pub context: Option<Arc<QueryContext>>,
pub lowered: Vec<LoweredExpr>,
}
impl Default for JoinContext {
fn default() -> Self {
Self::new()
}
}
impl JoinContext {
pub fn new() -> Self {
Self {
context: None,
lowered: vec![],
}
}
pub fn set(&mut self, ctx: &QueryContext) {
self.context = Some(Arc::new(ctx.clone()));
}
pub fn get(&self) -> &Arc<QueryContext> {
self.context.as_ref().expect("Join context not initialized")
}
pub fn is_initialized(&self) -> bool {
self.context.is_some()
}
}
pub(crate) fn ensure_join_keyable(
columns: &[ColumnView<'_>],
key_indices: &[usize],
fragment: impl Fn(usize) -> Fragment,
) -> Result<()> {
for (key, &idx) in key_indices.iter().enumerate() {
let ty = columns[idx].get_type();
if matches!(ty.inner_type(), ValueType::Digest { .. }) {
return Err(error!(operation::join_key_unkeyable(fragment(key), ty)));
}
}
Ok(())
}
pub(crate) fn join_key_types(
left: &[ColumnView<'_>],
left_indices: &[usize],
right: &[ColumnView<'_>],
right_indices: &[usize],
fragment: impl Fn(usize) -> Fragment,
) -> Result<Vec<ValueType>> {
let mut targets = Vec::with_capacity(left_indices.len());
for (key, (&li, &ri)) in left_indices.iter().zip(right_indices.iter()).enumerate() {
let left_ty = left[li].get_type();
let right_ty = right[ri].get_type();
match common_key_type(left_ty.inner_type(), right_ty.inner_type()) {
Some(target) => targets.push(target),
None => {
return Err(error!(operation::join_key_type_mismatch(
fragment(key),
left_ty.inner_type().clone(),
right_ty.inner_type().clone()
)));
}
}
}
Ok(targets)
}
pub(crate) fn eval_join_condition(
lowered: &[LoweredExpr],
left_columns: &[ColumnView<'_>],
right_columns: &[ColumnView<'_>],
left_row: &[Value],
right_row: &[Value],
alias: &Option<Fragment>,
ctx: &QueryContext,
) -> Result<bool> {
if lowered.is_empty() {
return Ok(true);
}
let eval_columns = build_eval_columns(left_columns, right_columns, left_row, right_row, alias);
let session = eval_context_from_query(ctx);
let exec_ctx = session.with_eval_join(batch(eval_columns)?);
for lowered_expr in lowered {
let col = lowered_expr.evaluate(&exec_ctx)?;
if !matches!(ColumnView::try_from(&col)?.get_value(0), Value::Boolean(true)) {
return Ok(false);
}
}
Ok(true)
}