use std::{mem, sync::Arc};
use arrow_array::{ArrayRef, RecordBatch};
use arrow_schema::FieldRef;
use reifydb_core::{
error::diagnostic::query::extend_duplicate_column,
expression::{Expression, name::display_label},
interface::{
evaluate::TargetColumn,
resolved::{ResolvedColumn, ResolvedObject},
},
value::{batch::batch, column::headers::ColumnHeaders},
};
use reifydb_evaluate::{
expression::{context::EvalContext, eval::cast_for_write},
lower::LoweredExpr,
};
use reifydb_extension::transform::{Transform, context::TransformContext};
use reifydb_transaction::transaction::Transaction;
use reifydb_value::{
fragment::Fragment,
reifydb_assertions, return_error,
value::{column_view::ColumnView, system_columns::user_columns},
};
use tracing::instrument;
use super::{NoopNode, user_header_names, user_pairs, with_system_headers, with_user_columns};
use crate::{
Result,
vm::volcano::{
inline::{expand_aliases, expand_sumtype_ctor, resolve_is_variants},
query::{QueryContext, QueryNode, eval_context_from_query, eval_context_from_transform},
udf::{UdfEvalNode, evaluate_udfs_no_input, strip_udf_columns},
},
};
pub(crate) struct ExtendNode {
input: Box<dyn QueryNode>,
expressions: Vec<Expression>,
written: Vec<Fragment>,
source: Option<ResolvedObject>,
udf_names: Vec<String>,
headers: Option<ColumnHeaders>,
context: Option<(Arc<QueryContext>, Vec<LoweredExpr>)>,
}
impl ExtendNode {
pub fn new(input: Box<dyn QueryNode>, expressions: Vec<Expression>, source: Option<ResolvedObject>) -> Self {
Self {
input,
expressions,
written: Vec::new(),
source,
udf_names: Vec::new(),
headers: None,
context: None,
}
}
#[instrument(level = "trace", skip_all, name = "volcano::extend::eval_context")]
fn eval_context<'e>(
session: &EvalContext<'e>,
new_columns: &[(FieldRef, ArrayRef)],
row_count: usize,
) -> Result<EvalContext<'e>> {
Ok(session.with_eval(batch(new_columns.to_vec())?, row_count))
}
#[instrument(level = "trace", skip_all, name = "volcano::extend::eval")]
fn eval_projection(lowered: &LoweredExpr, exec_ctx: &EvalContext) -> Result<(FieldRef, ArrayRef)> {
lowered.evaluate(exec_ctx)
}
}
impl QueryNode for ExtendNode {
#[instrument(name = "volcano::extend::initialize", level = "trace", skip_all)]
fn initialize<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &QueryContext) -> Result<()> {
(self.expressions, self.written) =
expand_aliases(mem::take(&mut self.expressions), |alias_expr, expanded| {
expand_sumtype_ctor(ctx, rx, ctx.source.as_ref(), alias_expr, expanded)
})?;
if let Some(source) = self.source.as_ref() {
for expr in &mut self.expressions {
resolve_is_variants(&ctx.services.catalog, rx, source, expr)?;
}
}
let (input, expressions, udf_names) = UdfEvalNode::wrap_if_needed(
mem::replace(&mut self.input, Box::new(NoopNode)),
&self.expressions,
&ctx.symbols,
);
self.input = input;
self.expressions = expressions;
self.udf_names = udf_names;
let lowered = self.expressions.iter().map(|e| LoweredExpr::new(e.clone(), "extend")).collect();
self.context = Some((Arc::new(ctx.clone()), lowered));
self.input.initialize(rx, ctx)?;
Ok(())
}
#[instrument(name = "volcano::extend::next", level = "trace", skip_all)]
fn next<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &mut QueryContext) -> Result<Option<RecordBatch>> {
reifydb_assertions! {
assert!(self.context.is_some(), "ExtendNode::next() called before initialize()");
}
if let Some(columns) = self.input.next(rx, ctx)? {
let stored_ctx = &self.context.as_ref().unwrap().0;
let transform_ctx = TransformContext {
routines: &ctx.services.routines,
runtime_context: &stored_ctx.services.runtime_context,
params: &stored_ctx.params,
};
let result = self.apply(&transform_ctx, columns)?;
if self.headers.is_none() {
let (mut all_headers, source_headers) = if let Some(input_headers) =
self.input.headers()
{
(user_header_names(&input_headers), input_headers)
} else {
let result_headers = ColumnHeaders::from_batch(&result);
let input_column_count = user_columns(&result).count() - self.expressions.len();
let mut names = user_header_names(&result_headers);
names.truncate(input_column_count);
(names, result_headers)
};
let new_names: Vec<Fragment> = self.expressions.iter().map(display_label).collect();
all_headers.extend(new_names);
self.headers = Some(with_system_headers(all_headers, &source_headers));
}
let result = strip_udf_columns(result, &self.udf_names)?;
return Ok(Some(result));
}
if self.headers.is_none()
&& let Some(input_headers) = self.input.headers()
{
let mut all_headers = user_header_names(&input_headers);
let new_names: Vec<Fragment> = self.expressions.iter().map(display_label).collect();
reject_duplicate_columns(&all_headers, &new_names, &self.written)?;
all_headers.extend(new_names);
self.headers = Some(with_system_headers(all_headers, &input_headers));
}
Ok(None)
}
fn headers(&self) -> Option<ColumnHeaders> {
self.headers.clone().or(self.input.headers())
}
}
impl Transform for ExtendNode {
fn apply(&self, ctx: &TransformContext, input: RecordBatch) -> Result<RecordBatch> {
let (stored_ctx, lowered) =
self.context.as_ref().expect("ExtendNode::apply() called before initialize()");
let row_count = input.num_rows();
let existing_names: Vec<Fragment> =
user_columns(&input).map(|(field, _)| Fragment::internal(field.name())).collect();
let session = eval_context_from_transform(ctx, stored_ctx);
let mut new_columns = user_pairs(&input);
let mut new_names = Vec::with_capacity(lowered.len());
for (expr, lowered_expr) in self.expressions.iter().zip(lowered.iter()) {
let mut exec_ctx = Self::eval_context(&session, &new_columns, row_count)?;
if let (Expression::Alias(alias_expr), Some(source)) = (expr, &stored_ctx.source) {
let alias_name = alias_expr.alias.name();
if let Some(table_column) = source.columns().iter().find(|col| col.name == alias_name) {
let column_ident = Fragment::internal(&table_column.name);
let resolved_column =
ResolvedColumn::new(column_ident, source.clone(), table_column.clone());
exec_ctx.target = Some(TargetColumn::Resolved(resolved_column));
}
}
let mut column = Self::eval_projection(lowered_expr, &exec_ctx)?;
if let Some(target_type) = exec_ctx.target.as_ref().map(|t| t.column_type()) {
let view = ColumnView::try_from(&column)?;
if view.get_type() != target_type {
column = cast_for_write(&exec_ctx, &view, target_type, &expr.lazy_fragment())?;
}
}
new_columns.push(column);
new_names.push(display_label(expr));
}
reject_duplicate_columns(&existing_names, &new_names, &self.written)?;
with_user_columns(new_columns, &input)
}
}
pub(crate) struct ExtendWithoutInputNode {
expressions: Vec<Expression>,
written: Vec<Fragment>,
headers: Option<ColumnHeaders>,
udf_columns: Option<RecordBatch>,
context: Option<(Arc<QueryContext>, Vec<LoweredExpr>)>,
}
impl ExtendWithoutInputNode {
pub fn new(expressions: Vec<Expression>) -> Self {
Self {
expressions,
written: Vec::new(),
headers: None,
udf_columns: None,
context: None,
}
}
}
impl QueryNode for ExtendWithoutInputNode {
#[instrument(name = "volcano::extend::noinput::initialize", level = "trace", skip_all)]
fn initialize<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &QueryContext) -> Result<()> {
(self.expressions, self.written) =
expand_aliases(mem::take(&mut self.expressions), |alias_expr, expanded| {
expand_sumtype_ctor(ctx, rx, ctx.source.as_ref(), alias_expr, expanded)
})?;
if let Some((rewritten, udf_cols)) = evaluate_udfs_no_input(&self.expressions, ctx, rx)? {
self.expressions = rewritten;
self.udf_columns = Some(udf_cols);
}
let lowered = self.expressions.iter().map(|e| LoweredExpr::new(e.clone(), "extend")).collect();
self.context = Some((Arc::new(ctx.clone()), lowered));
Ok(())
}
#[instrument(name = "volcano::extend::noinput::next", level = "trace", skip_all)]
fn next<'a>(&mut self, _rx: &mut Transaction<'a>, _ctx: &mut QueryContext) -> Result<Option<RecordBatch>> {
reifydb_assertions! {
assert!(self.context.is_some(), "ExtendWithoutInputNode::next() called before initialize()");
}
let (stored_ctx, lowered) = self.context.as_ref().unwrap();
if self.headers.is_some() {
return Ok(None);
}
let session = eval_context_from_query(stored_ctx);
let mut new_columns = Vec::with_capacity(self.expressions.len());
for lowered_expr in lowered {
let exec_ctx = match &self.udf_columns {
Some(udf_cols) => session.with_eval(udf_cols.clone(), 1),
None => session.with_eval_empty(),
};
let column = lowered_expr.evaluate(&exec_ctx)?;
new_columns.push(column);
}
let column_names: Vec<Fragment> = self.expressions.iter().map(display_label).collect();
reject_duplicate_columns(&[], &column_names, &self.written)?;
self.headers = Some(ColumnHeaders {
columns: column_names,
});
Ok(Some(batch(new_columns)?))
}
fn headers(&self) -> Option<ColumnHeaders> {
self.headers.clone()
}
}
fn reject_duplicate_columns(existing: &[Fragment], new_names: &[Fragment], written: &[Fragment]) -> Result<()> {
for (index, name) in new_names.iter().enumerate() {
if existing.iter().chain(&new_names[..index]).any(|other| other.text() == name.text()) {
return_error!(extend_duplicate_column(written[index].clone(), name.text()));
}
}
Ok(())
}