use std::sync::Arc;
use reifydb_core::{
error::{CoreError, diagnostic::query},
value::column::{
ColumnWithName,
buffer::ColumnBuffer,
columns::Columns,
headers::ColumnHeaders,
view::group_by::{GroupId, GroupKeyDict},
},
};
use reifydb_routine_abi::{
Accumulator, FunctionKind, context::FunctionContext, error::RoutineError, registry::Routines,
};
use reifydb_rql::expression::{Expression, name::display_label};
use reifydb_transaction::transaction::Transaction;
use reifydb_value::{error, fragment::Fragment, reifydb_assertions, value::value_type::ValueType};
use tracing::instrument;
use crate::{
Result,
vm::volcano::query::{QueryContext, QueryNode},
};
enum Projection {
Aggregate {
column: String,
column_fragment: Fragment,
alias: Fragment,
accumulator: Box<dyn Accumulator>,
},
Group {
column: String,
alias: Fragment,
},
}
pub(crate) struct AggregateNode {
input: Box<dyn QueryNode>,
by: Vec<Expression>,
map: Vec<Expression>,
headers: Option<ColumnHeaders>,
context: Option<Arc<QueryContext>>,
}
impl AggregateNode {
pub fn new(
input: Box<dyn QueryNode>,
by: Vec<Expression>,
map: Vec<Expression>,
context: Arc<QueryContext>,
) -> Self {
Self {
input,
by,
map,
headers: None,
context: Some(context),
}
}
#[instrument(level = "trace", skip_all, name = "volcano::aggregate::accumulate")]
fn accumulate<'a>(
input: &mut Box<dyn QueryNode>,
rx: &mut Transaction<'a>,
ctx: &mut QueryContext,
keys: &[&str],
projections: &mut [Projection],
dict: &mut GroupKeyDict,
) -> Result<()> {
while let Some(columns) = input.next(rx, ctx)? {
let groups = columns.group_by_ids(keys, dict)?;
for projection in projections.iter_mut() {
if let Projection::Aggregate {
accumulator,
column,
column_fragment,
..
} = projection
{
let column_ref = columns.column(column).ok_or_else(|| {
error!(query::column_not_found(column_fragment.clone()))
})?;
let cwn = ColumnWithName::new(
column_ref.name().clone(),
column_ref.data().clone(),
);
accumulator.update(&Columns::new(vec![cwn]), &groups)?;
}
}
}
Ok(())
}
#[instrument(level = "trace", skip_all, name = "volcano::aggregate::finalize")]
fn finalize(projections: Vec<Projection>, keys: &[&str], dict: &GroupKeyDict) -> Vec<ColumnWithName> {
let mut result_columns = Vec::new();
for projection in projections {
match projection {
Projection::Group {
alias,
column,
..
} => {
let col_idx = keys.iter().position(|k| k == &column).unwrap();
let first_key_type = dict.values(GroupId(0)).map(|key| key[col_idx].get_type());
let mut c = ColumnWithName {
name: Fragment::internal(alias.fragment()),
data: ColumnBuffer::none_typed(
first_key_type.unwrap_or(ValueType::Boolean),
0,
),
};
for (_, key) in dict.iter() {
c.data_mut().push_value(key[col_idx].clone());
}
result_columns.push(c);
}
Projection::Aggregate {
alias,
mut accumulator,
..
} => {
let (keys_out, mut data) = accumulator.finalize().unwrap();
align_column_data(dict, &keys_out, &mut data).unwrap();
result_columns.push(ColumnWithName {
name: Fragment::internal(alias.fragment()),
data,
});
}
}
}
result_columns
}
}
impl QueryNode for AggregateNode {
#[instrument(level = "trace", skip_all, name = "volcano::aggregate::initialize")]
fn initialize<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &QueryContext) -> Result<()> {
self.input.initialize(rx, ctx)?;
Ok(())
}
#[instrument(level = "trace", skip_all, name = "volcano::aggregate::next")]
fn next<'a>(&mut self, rx: &mut Transaction<'a>, ctx: &mut QueryContext) -> Result<Option<Columns>> {
reifydb_assertions! {
assert!(self.context.is_some(), "AggregateNode::next() called before initialize()");
}
let stored_ctx = self.context.as_ref().unwrap();
if self.headers.is_some() {
return Ok(None);
}
let (keys, mut projections) =
parse_keys_and_aggregates(&self.by, &self.map, &stored_ctx.services.routines, stored_ctx)?;
let mut dict = GroupKeyDict::new();
Self::accumulate(&mut self.input, rx, ctx, &keys, &mut projections, &mut dict)?;
let result_columns = Self::finalize(projections, &keys, &dict);
let columns = Columns::new(result_columns);
self.headers = Some(ColumnHeaders::from_columns(&columns));
Ok(Some(columns))
}
fn headers(&self) -> Option<ColumnHeaders> {
self.headers.clone().or(self.input.headers())
}
}
fn parse_keys_and_aggregates<'a>(
by: &'a [Expression],
project: &'a [Expression],
routines: &'a Routines,
ctx: &QueryContext,
) -> Result<(Vec<&'a str>, Vec<Projection>)> {
let mut keys = Vec::new();
let mut projections = Vec::new();
for gb in by {
match gb {
Expression::Column(c) => {
keys.push(c.0.name.text());
projections.push(Projection::Group {
column: c.0.name.text().to_string(),
alias: c.0.name.clone(),
})
}
Expression::AccessSource(access) => {
keys.push(access.column.name.text());
projections.push(Projection::Group {
column: access.column.name.text().to_string(),
alias: access.column.name.clone(),
})
}
expr => panic!("Non-column group by not supported: {expr:#?}"),
}
}
for p in project {
let (actual_expr, alias) = match p {
Expression::Alias(alias_expr) => (alias_expr.expression.as_ref(), alias_expr.alias.0.clone()),
expr => (expr, display_label(expr)),
};
match actual_expr {
Expression::Call(call) => {
let func_name = call.func.0.text();
let function = routines.get_aggregate_function(func_name).ok_or_else(|| {
RoutineError::FunctionNotFound {
function: call.func.0.clone(),
}
})?;
let _ = FunctionKind::Aggregate;
let mut fn_ctx = FunctionContext {
fragment: call.func.0.clone(),
identity: ctx.identity,
row_count: 0,
runtime_context: &ctx.services.runtime_context,
};
let accumulator = function.accumulator(&mut fn_ctx).ok_or_else(|| {
RoutineError::FunctionExecutionFailed {
function: call.func.0.clone(),
reason: format!("Function {} is not an aggregate", func_name),
}
})?;
match call.args.first() {
Some(Expression::Column(c)) => {
projections.push(Projection::Aggregate {
column: c.0.name.text().to_string(),
column_fragment: c.0.name.clone(),
alias,
accumulator,
});
}
Some(Expression::AccessSource(access)) => {
projections.push(Projection::Aggregate {
column: access.column.name.text().to_string(),
column_fragment: access.column.name.clone(),
alias,
accumulator,
});
}
None => {
return Err(RoutineError::FunctionArityMismatch {
function: call.func.0.clone(),
expected: 1,
actual: 0,
}
.into());
}
Some(arg) => {
let actual_type = arg.infer_type().ok_or_else(|| {
RoutineError::FunctionExecutionFailed {
function: call.func.0.clone(),
reason: "aggregate function arguments must be column references".to_string(),
}
})?;
let expected = function.accepted_types().expected_at(0).to_vec();
return Err(RoutineError::FunctionInvalidArgumentType {
function: call.func.0.clone(),
argument_index: 0,
expected,
actual: actual_type,
}
.into());
}
}
}
_ => panic!("Expected aggregate call expression, got: {actual_expr:#?}"),
}
}
Ok((keys, projections))
}
fn align_column_data(dict: &GroupKeyDict, produced: &[GroupId], data: &mut ColumnBuffer) -> Result<()> {
let mut position_of: Vec<Option<usize>> = vec![None; dict.len()];
for (position, group) in produced.iter().enumerate() {
if let Some(slot) = position_of.get_mut(group.index()) {
*slot = Some(position);
}
}
let reorder_indices: Vec<usize> = (0..dict.len())
.map(|index| {
position_of[index].ok_or_else(|| {
CoreError::FrameError {
message: format!(
"Group key {:?} missing in aggregate output",
dict.values(GroupId(index as u32))
),
}
.into()
})
})
.collect::<Result<Vec<_>>>()?;
data.reorder(&reorder_indices);
Ok(())
}