use rudb_common::{Error, Field, LogicalType, Memory, Reservation, Result, Value};
use rudb_kernels::{Accumulator, is_true};
use rudb_plan::{Expr, ExprRef, Plan, Slice};
use rudb_vector::{Chunk, VECTOR_SIZE, Vector};
use crate::expr::{evaluate, evaluate_all};
use crate::key::{Key, RowSet};
use crate::operator::Operator;
use crate::rows;
use crate::schema::Schema;
use crate::spill::{Reader, Spill};
use crate::table::{Probe, Table};
#[derive(Debug, Clone)]
struct Call {
name: String,
args: Vec<ExprRef>,
distinct: bool,
filter: Option<ExprRef>,
returns: LogicalType,
}
#[derive(Debug)]
pub(crate) struct Aggregate<'a> {
input: Box<dyn Operator + 'a>,
plan: &'a Plan,
input_schema: Schema,
groups: Vec<ExprRef>,
calls: Vec<Call>,
schema: Schema,
built: bool,
chunks: Vec<Chunk>,
at: usize,
memory: Memory,
held: Reservation,
}
impl<'a> Aggregate<'a> {
pub(crate) fn new(
plan: &'a Plan,
input: Box<dyn Operator + 'a>,
index: u32,
groups: Slice,
aggregates: Slice,
memory: &Memory,
) -> Result<Self> {
let input_schema = input.schema().clone();
let groups: Vec<ExprRef> = plan.expr_list(groups).to_vec();
let mut calls = Vec::new();
for &reference in plan.expr_list(aggregates) {
let Expr::Aggregate { name, args, distinct, filter } = *plan.expr(reference) else {
return Err(Error::internal(format!(
"expression {reference} is in the aggregate list of an Aggregate and is not an aggregate"
)));
};
calls.push(Call {
name: plan.string(name).to_string(),
args: plan.expr_list(args).to_vec(),
distinct,
filter,
returns: plan.expr_type(reference).clone(),
});
}
let mut fields = Vec::with_capacity(groups.len() + calls.len());
for (at, &group) in groups.iter().enumerate() {
fields.push(Field::new(
group_name(plan, group, &input_schema, at),
plan.expr_type(group).clone(),
));
}
for call in &calls {
fields.push(Field::new(call.name.clone(), call.returns.clone()));
}
let schema = Schema::numbered(fields, index);
Ok(Self {
input,
plan,
input_schema,
groups,
calls,
schema,
built: false,
chunks: Vec::new(),
at: 0,
memory: memory.clone(),
held: memory.reservation(),
})
}
fn build(&mut self) -> Result<()> {
let pass = Pass {
plan: self.plan,
input_schema: &self.input_schema,
schema: &self.schema,
groups: &self.groups,
calls: &self.calls,
memory: &self.memory,
};
let mut source = Source::Input(&mut *self.input);
let mut left = pass.once(&mut source, &mut self.chunks, &mut self.held)?;
while let Some(mut file) = left {
let mut source = Source::Spilled(Spilled::new(file.read()?, pass.spilled_types()));
left = pass.once(&mut source, &mut self.chunks, &mut self.held)?;
}
Ok(())
}
}
struct Pass<'p> {
plan: &'p Plan,
input_schema: &'p Schema,
schema: &'p Schema,
groups: &'p [ExprRef],
calls: &'p [Call],
memory: &'p Memory,
}
impl Pass<'_> {
fn once(
&self,
source: &mut Source<'_, '_>,
chunks: &mut Vec<Chunk>,
held: &mut Reservation,
) -> Result<Option<Spill>> {
let mut scratch = self.memory.reservation();
let mut containers = self.memory.reservation();
let mut charged = 0;
let mut charged_keys = 0;
let mut table = Table::new(self.groups.len());
let mut states: Vec<Accumulator> = Vec::new();
let mut seen: Vec<RowSet> = Vec::new();
let calls = self.calls.len();
let sets = self.calls.iter().any(|call| call.distinct);
let alone = self.groups.is_empty();
let mut groups = 0;
if alone {
groups = 1;
self.fresh(&mut states)?;
if sets {
seen.resize_with(calls, RowSet::default);
}
}
let by_vector: Vec<bool> = self
.calls
.iter()
.map(|call| alone && !call.distinct && call.filter.is_none())
.collect();
let every = by_vector.iter().all(|&yes| yes);
let mut given: Vec<Key> = vec![Key(Vec::new()); calls];
let mut hashes: Vec<u64> = Vec::new();
let mut over: Option<Spill> = None;
let mut away: Vec<Value> = Vec::new();
while let Some(seen_rows) = source.next(self)? {
let Rows { keys, arguments, filters, rows: length } = &seen_rows;
let mut aside = 0;
for at in 0..calls {
if by_vector[at] {
states[at].update_run(&arguments[at], *length)?;
}
}
if alone && every {
continue;
}
if !alone {
crate::table::hash(keys, *length, &mut hashes);
}
#[expect(
clippy::needless_range_loop,
reason = "hashes is empty when there is no key to hash, and every other thing in \
the body is indexed by the row as well"
)]
for row in 0..*length {
let slot = if alone {
0
} else {
match table.probe(hashes[row], keys, row) {
Probe::Found(slot) => slot,
Probe::Vacant(bucket) => {
if let Some(file) = over.as_mut() {
put_away(file, &seen_rows, row, &mut away)?;
continue;
}
let slot = table.insert(bucket, hashes[row], keys, row)?;
groups = table.len();
self.fresh(&mut states)?;
if sets {
seen.resize_with(seen.len() + calls, RowSet::default);
}
slot
}
}
};
for (at, call) in self.calls.iter().enumerate() {
if by_vector[at] {
continue;
}
if let Some(flags) = &filters[at] {
if !is_true(&flags.value_at(row)) {
continue;
}
}
let args = &mut given[at];
fill(args, &arguments[at], row);
if call.distinct {
let set = &mut seen[slot * calls + at];
if set.contains(args) {
continue;
}
let stored = args.clone();
aside += rows::footprint(&stored.0);
set.insert(stored);
}
states[slot * calls + at].update(&args.0)?;
}
}
rows::capacity(table.owned(), &mut charged_keys, &mut scratch)?;
containers.grow(aside)?;
let now = tables(&table, &states, &seen);
rows::capacity(now, &mut charged, &mut containers)?;
match over.as_ref() {
None if !alone && crowded(self.memory) => {
over = Some(Spill::new("aggregate", self.spilled_types())?);
}
Some(file) => hopeless(file, groups)?,
None => {}
}
}
drop(seen);
let alive = table.footprint() + width_of(states.capacity() * size_of::<Accumulator>());
containers.shrink(containers.bytes().saturating_sub(alive));
let types = self.schema.types();
let width = self.groups.len();
scratch.grow(width_of(VECTOR_SIZE.min(groups) * size_of::<Value>()))?;
let mut results: Vec<Value> = Vec::new();
for start in (0..groups).step_by(VECTOR_SIZE) {
let end = (start + VECTOR_SIZE).min(groups);
let mut columns = Vec::with_capacity(width + calls);
for (at, ty) in types.iter().take(width).enumerate() {
columns.push(Vector::from_values(ty.clone(), &table.column(at)[start..end])?);
}
for (at, ty) in types.iter().skip(width).enumerate() {
let mut taken = 0;
results.clear();
for slot in start..end {
let value = states[slot * calls + at].finish()?;
taken += rows::owned(&value);
results.push(value);
}
scratch.grow(taken)?;
columns.push(Vector::from_values(ty.clone(), &results)?);
scratch.shrink(taken);
}
let chunk = Chunk::with_rows(columns, end - start)?;
held.grow(width_of(chunk.footprint()))?;
chunks.push(chunk);
}
drop(results);
drop(states);
drop(table);
containers.release();
match over {
Some(file) if groups == 0 && file.rows() > 0 => Err(Error::out_of_memory(format!(
"the memory limit does not leave room for a single group of this aggregate, \
{} rows and {} bytes went to a spill file and none of them could be finished",
file.rows(),
file.bytes()
))),
Some(file) if file.rows() > 0 => Ok(Some(file)),
_ => Ok(None),
}
}
fn fresh(&self, states: &mut Vec<Accumulator>) -> Result<()> {
for call in self.calls {
states.push(Accumulator::new(&call.name, &call.returns)?);
}
Ok(())
}
fn spilled_types(&self) -> Vec<LogicalType> {
let mut types = Vec::new();
for &group in self.groups {
types.push(self.plan.expr_type(group).clone());
}
for call in self.calls {
for &argument in &call.args {
types.push(self.plan.expr_type(argument).clone());
}
}
for call in self.calls {
if let Some(filter) = call.filter {
types.push(self.plan.expr_type(filter).clone());
}
}
types
}
}
struct Rows {
keys: Vec<Vector>,
arguments: Vec<Vec<Vector>>,
filters: Vec<Option<Vector>>,
rows: usize,
}
impl Rows {
fn width(&self) -> usize {
self.keys.len()
+ self.arguments.iter().map(Vec::len).sum::<usize>()
+ self.filters.iter().flatten().count()
}
}
enum Source<'s, 'o> {
Input(&'s mut (dyn Operator + 'o)),
Spilled(Spilled<'s>),
}
impl Source<'_, '_> {
fn next(&mut self, pass: &Pass<'_>) -> Result<Option<Rows>> {
match self {
Source::Input(input) => {
let Some(chunk) = input.next()? else {
return Ok(None);
};
let rows = chunk.len();
let keys = evaluate_all(pass.plan, pass.groups, pass.input_schema, &chunk)?;
let mut arguments = Vec::with_capacity(pass.calls.len());
let mut filters = Vec::with_capacity(pass.calls.len());
for call in pass.calls {
arguments.push(evaluate_all(pass.plan, &call.args, pass.input_schema, &chunk)?);
filters.push(match call.filter {
Some(filter) => {
Some(evaluate(pass.plan, filter, pass.input_schema, &chunk)?)
}
None => None,
});
}
Ok(Some(Rows { keys, arguments, filters, rows }))
}
Source::Spilled(spilled) => spilled.next(pass),
}
}
}
struct Spilled<'s> {
reader: Reader<'s>,
types: Vec<LogicalType>,
row: Vec<Value>,
columns: Vec<Vec<Value>>,
}
impl<'s> Spilled<'s> {
fn new(reader: Reader<'s>, types: Vec<LogicalType>) -> Self {
let columns = vec![Vec::new(); types.len()];
Self { reader, types, row: Vec::new(), columns }
}
fn next(&mut self, pass: &Pass<'_>) -> Result<Option<Rows>> {
let Self { reader, types, row, columns } = self;
for column in columns.iter_mut() {
column.clear();
}
let mut rows = 0;
while rows < VECTOR_SIZE && reader.next_into(row)? {
for (at, value) in row.iter_mut().enumerate() {
columns[at].push(std::mem::replace(value, Value::Null));
}
rows += 1;
}
if rows == 0 {
return Ok(None);
}
let mut built = Vec::with_capacity(columns.len());
for (values, ty) in columns.iter().zip(&*types) {
built.push(Vector::from_values(ty.clone(), values)?);
}
let mut taking = built.into_iter();
let keys: Vec<Vector> = taking.by_ref().take(pass.groups.len()).collect();
let mut arguments = Vec::with_capacity(pass.calls.len());
for call in pass.calls {
arguments.push(taking.by_ref().take(call.args.len()).collect());
}
let mut filters = Vec::with_capacity(pass.calls.len());
for call in pass.calls {
filters.push(if call.filter.is_some() { taking.next() } else { None });
}
Ok(Some(Rows { keys, arguments, filters, rows }))
}
}
fn put_away(file: &mut Spill, seen: &Rows, row: usize, away: &mut Vec<Value>) -> Result<()> {
let columns = seen
.keys
.iter()
.chain(seen.arguments.iter().flatten())
.chain(seen.filters.iter().flatten());
away.truncate(seen.width());
for (at, column) in columns.enumerate() {
match away.get_mut(at) {
Some(slot) => set(slot, column, row),
None => away.push(column.value_at(row)),
}
}
file.write(away)
}
const PASSES: u64 = 64;
fn hopeless(file: &Spill, groups: usize) -> Result<()> {
let left = file.rows() / width_of(groups).max(1);
if left > PASSES {
return Err(Error::out_of_memory(format!(
"the memory limit leaves room for {groups} groups at a time and {} rows have already \
gone to a spill file, which is more passes over it than this will finish in",
file.rows()
)));
}
Ok(())
}
fn crowded(memory: &Memory) -> bool {
match memory.limit() {
Some(limit) => memory.used() >= limit / 2,
None => false,
}
}
fn fill(key: &mut Key, columns: &[Vector], row: usize) {
key.0.truncate(columns.len());
for (at, column) in columns.iter().enumerate() {
match key.0.get_mut(at) {
Some(slot) => set(slot, column, row),
None => key.0.push(column.value_at(row)),
}
}
}
fn set(slot: &mut Value, column: &Vector, row: usize) {
if let (Value::Varchar(buffer), Some(text)) = (&mut *slot, column.text_at(row)) {
buffer.clear();
buffer.push_str(text);
return;
}
*slot = column.value_at(row);
}
impl Operator for Aggregate<'_> {
fn schema(&self) -> &Schema {
&self.schema
}
fn next(&mut self) -> Result<Option<Chunk>> {
if !self.built {
self.build()?;
self.built = true;
}
if self.at >= self.chunks.len() {
return Ok(None);
}
let chunk = self.chunks[self.at].clone();
self.at += 1;
Ok(Some(chunk))
}
}
fn tables(table: &Table, states: &Vec<Accumulator>, seen: &Vec<RowSet>) -> u64 {
let width = |count: usize, size: usize| {
u64::try_from(count).unwrap_or(u64::MAX).saturating_mul(width_of(size))
};
table.footprint()
+ width(states.capacity(), size_of::<Accumulator>())
+ width(seen.capacity(), size_of::<RowSet>())
}
fn width_of(size: usize) -> u64 {
u64::try_from(size).unwrap_or(u64::MAX)
}
fn group_name(plan: &Plan, group: ExprRef, input: &Schema, at: usize) -> String {
if let Expr::Column(binding) = *plan.expr(group) {
if let Some(position) = input.position_of(binding) {
return input.fields()[position].name.clone();
}
}
format!("group{at}")
}
#[derive(Debug)]
pub(crate) struct Distinct<'a> {
input: Box<dyn Operator + 'a>,
plan: &'a Plan,
on: Vec<ExprRef>,
schema: Schema,
built: bool,
chunks: Vec<Chunk>,
at: usize,
memory: Memory,
held: Reservation,
}
impl<'a> Distinct<'a> {
pub(crate) fn new(
plan: &'a Plan,
input: Box<dyn Operator + 'a>,
on: Slice,
memory: &Memory,
) -> Self {
let schema = input.schema().clone();
Self {
input,
plan,
on: plan.expr_list(on).to_vec(),
schema,
built: false,
chunks: Vec::new(),
at: 0,
memory: memory.clone(),
held: memory.reservation(),
}
}
fn build(&mut self) -> Result<()> {
let mut scratch = self.memory.reservation();
let mut table = self.memory.reservation();
let mut charged = 0;
let mut charged_table = 0;
let mut seen: RowSet = RowSet::default();
let mut kept: Vec<Vec<Value>> = Vec::new();
let mut key = Key(Vec::new());
while let Some(chunk) = self.input.next()? {
let keys = if self.on.is_empty() {
Vec::new()
} else {
evaluate_all(self.plan, &self.on, &self.schema, &chunk)?
};
let mut taken = 0;
let mut aside = 0;
for row in 0..chunk.len() {
if self.on.is_empty() {
key.0.clear();
key.0.extend(chunk.row(row));
} else {
fill(&mut key, &keys, row);
}
if seen.contains(&key) {
continue;
}
let values: Vec<Value> =
if self.on.is_empty() { key.0.clone() } else { chunk.row(row).collect() };
let stored = key.clone();
taken += rows::heap(&values);
aside += rows::heap(&stored.0);
seen.insert(stored);
kept.push(values);
}
scratch.grow(taken)?;
table.grow(aside)?;
let rows = width_of(kept.capacity() * size_of::<Vec<Value>>());
rows::capacity(rows, &mut charged, &mut scratch)?;
let now = rows::buckets(seen.capacity()) * (width_of(size_of::<Key>()) + 1);
rows::capacity(now, &mut charged_table, &mut table)?;
}
drop(seen);
table.release();
self.chunks = rows::chunks(&self.schema.types(), &kept, &mut self.held)?;
Ok(())
}
}
impl Operator for Distinct<'_> {
fn schema(&self) -> &Schema {
&self.schema
}
fn next(&mut self) -> Result<Option<Chunk>> {
if !self.built {
self.build()?;
self.built = true;
}
if self.at >= self.chunks.len() {
return Ok(None);
}
let chunk = self.chunks[self.at].clone();
self.at += 1;
Ok(Some(chunk))
}
}