use std::sync::Mutex;
use rudb_common::{Error, Field, LogicalType, Memory, Reservation, Result, Value};
use rudb_kernels::{Accumulator, NOWHERE, is_true, update_scattered};
use rudb_pipeline::{Progress, Sink};
use rudb_plan::{Expr, ExprRef, Plan, Slice};
use rudb_vector::{Chunk, VECTOR_SIZE, Vector};
use crate::buffer::Buffered;
use crate::expr::{evaluate, evaluate_all};
use crate::key::{Key, RowSet};
use crate::prepared::{Prepared, Scratch};
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> {
plan: &'a Plan,
input_schema: Schema,
groups: Vec<ExprRef>,
calls: Vec<Call>,
schema: Schema,
alone: bool,
sets: bool,
by_vector: Vec<bool>,
every: bool,
memory: Memory,
built: Mutex<Built>,
out: Buffered,
}
#[derive(Debug)]
struct Built {
chunks: Vec<Chunk>,
held: Reservation,
instances: usize,
}
impl<'a> Aggregate<'a> {
pub(crate) fn new(
plan: &'a Plan,
input: &Schema,
index: u32,
groups: Slice,
aggregates: Slice,
memory: &Memory,
) -> Result<(Self, Buffered)> {
let input_schema = input.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);
let alone = groups.is_empty();
let by_vector: Vec<bool> =
calls.iter().map(|call| alone && !call.distinct && call.filter.is_none()).collect();
let out = Buffered::new();
let aggregate = Self {
plan,
input_schema,
alone,
sets: calls.iter().any(|call| call.distinct),
every: by_vector.iter().all(|&yes| yes),
by_vector,
groups,
calls,
schema,
memory: memory.clone(),
built: Mutex::new(Built {
chunks: Vec::new(),
held: memory.reservation(),
instances: 0,
}),
out: out.clone(),
};
Ok((aggregate, out))
}
pub(crate) fn schema(&self) -> &Schema {
&self.schema
}
fn read(&self, chunk: &Chunk) -> Result<Rows> {
let rows = chunk.len();
let keys = evaluate_all(self.plan, &self.groups, &self.input_schema, chunk)?;
let mut arguments = Vec::with_capacity(self.calls.len());
let mut filters = Vec::with_capacity(self.calls.len());
for call in &self.calls {
arguments.push(evaluate_all(self.plan, &call.args, &self.input_schema, chunk)?);
filters.push(match call.filter {
Some(filter) => Some(evaluate(self.plan, filter, &self.input_schema, chunk)?),
None => None,
});
}
Ok(Rows { keys, arguments, filters, rows })
}
fn again(
&self,
file: &mut Spill,
chunks: &mut Vec<Chunk>,
held: &mut Reservation,
) -> Result<Option<Spill>> {
let mut spilled = Spilled::new(file.read()?, self.spilled_types());
let mut local = self.start();
if let Some(error) = local.failure.take() {
return Err(error);
}
while let Some(rows) = spilled.next(self)? {
self.fold(&rows, &mut local)?;
}
self.finish(local, chunks, held)
}
fn start(&self) -> Building {
let calls = self.calls.len();
let mut local = Building {
scratch: self.memory.reservation(),
containers: self.memory.reservation(),
charged: 0,
charged_keys: 0,
table: Table::new(self.groups.len()),
states: Vec::new(),
seen: Vec::new(),
groups: 0,
given: vec![Key(Vec::new()); calls],
hashes: Vec::new(),
slots: Vec::new(),
kept: Vec::new(),
over: None,
away: Vec::new(),
failure: None,
};
if self.alone {
local.groups = 1;
if let Err(error) = self.fresh(&mut local.states) {
local.failure = Some(error);
}
if self.sets {
local.seen.resize_with(calls, RowSet::default);
}
}
local
}
fn fold(&self, seen_rows: &Rows, local: &mut Building) -> Result<()> {
let Building {
scratch,
containers,
charged,
charged_keys,
table,
states,
seen,
groups,
given,
hashes,
slots,
kept,
over,
away,
failure: _,
} = local;
let calls = self.calls.len();
let alone = self.alone;
let Rows { keys, arguments, filters, rows: length } = seen_rows;
let mut aside = 0;
for at in 0..calls {
if self.by_vector[at] {
states[at].update_run(&arguments[at], *length)?;
}
}
if alone && self.every {
return Ok(());
}
if !alone {
crate::table::hash(keys, *length, hashes);
}
slots.clear();
slots.resize(*length, NOWHERE);
for row in 0..*length {
slots[row] = 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, away)?;
continue;
}
let slot = table.insert(bucket, hashes[row], keys, row)?;
*groups = table.len();
self.fresh(states)?;
if self.sets {
seen.resize_with(seen.len() + calls, RowSet::default);
}
slot
}
}
};
}
for (at, call) in self.calls.iter().enumerate() {
if self.by_vector[at] {
continue;
}
if call.distinct {
aside += self.distinct(states, seen, seen_rows, slots, at, given)?;
continue;
}
let picked = match &filters[at] {
None => &*slots,
Some(flags) => {
kept.clear();
kept.extend(slots.iter().enumerate().map(|(row, &slot)| {
if slot != NOWHERE && is_true(&flags.value_at(row)) {
slot
} else {
NOWHERE
}
}));
&*kept
}
};
update_scattered(states, picked, calls, at, arguments[at].first(), *length)?;
}
rows::capacity(table.owned(), charged_keys, scratch)?;
containers.grow(aside)?;
let now = tables(table, states, seen);
rows::capacity(now, charged, 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 => {}
}
Ok(())
}
fn finish(
&self,
local: Building,
chunks: &mut Vec<Chunk>,
held: &mut Reservation,
) -> Result<Option<Spill>> {
let Building { mut scratch, mut containers, table, states, seen, groups, over, .. } = local;
let calls = self.calls.len();
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 distinct(
&self,
states: &mut [Accumulator],
seen: &mut [RowSet],
rows: &Rows,
slots: &[usize],
at: usize,
given: &mut [Key],
) -> Result<u64> {
let calls = self.calls.len();
let mut aside = 0;
for (row, &slot) in slots.iter().enumerate() {
if slot == NOWHERE {
continue;
}
if let Some(flags) = &rows.filters[at] {
if !is_true(&flags.value_at(row)) {
continue;
}
}
let args = &mut given[at];
fill(args, &rows.arguments[at], row);
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)?;
}
Ok(aside)
}
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()
}
}
#[derive(Debug)]
pub(crate) struct Building {
scratch: Reservation,
containers: Reservation,
charged: u64,
charged_keys: u64,
table: Table,
states: Vec<Accumulator>,
seen: Vec<RowSet>,
groups: usize,
given: Vec<Key>,
hashes: Vec<u64>,
slots: Vec<usize>,
kept: Vec<usize>,
over: Option<Spill>,
away: Vec<Value>,
failure: Option<Error>,
}
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: &Aggregate<'_>) -> 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 Sink for Aggregate<'_> {
type Local = Building;
fn local(&self) -> Building {
self.start()
}
fn sink(&self, chunk: &Chunk, local: &mut Building) -> Result<Progress> {
if let Some(error) = local.failure.take() {
return Err(error);
}
let rows = self.read(chunk)?;
self.fold(&rows, local)?;
Ok(Progress::More)
}
fn combine(&self, local: Building) -> Result<()> {
if let Some(error) = local.failure {
return Err(error);
}
let mut built = self.built.lock().map_err(poisoned)?;
built.instances += 1;
if built.instances > 1 {
return Err(Error::not_implemented(
"two instances of one aggregate are two hash tables, and merging them needs a \
serialize and a combine per aggregate that nothing implements yet",
));
}
let Built { chunks, held, .. } = &mut *built;
let mut left = self.finish(local, chunks, held)?;
while let Some(mut file) = left {
left = self.again(&mut file, chunks, held)?;
}
Ok(())
}
fn finalize(&self) -> Result<()> {
let chunks = std::mem::take(&mut self.built.lock().map_err(poisoned)?.chunks);
self.out.fill(chunks)
}
}
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 {
on: Prepared,
whole: bool,
types: Vec<LogicalType>,
memory: Memory,
global: Mutex<Held>,
charged: Mutex<Vec<Reservation>>,
held: Mutex<Reservation>,
out: Buffered,
}
#[derive(Debug)]
struct Held {
seen: RowSet,
kept: Vec<Vec<Value>>,
counted: u64,
table: Reservation,
counted_rows: u64,
slots: Reservation,
}
#[derive(Debug)]
pub(crate) struct Keeping {
seen: RowSet,
kept: Vec<(Key, Vec<Value>)>,
scratch: Scratch,
key: Key,
rows: Reservation,
table: Reservation,
counted: u64,
counted_table: u64,
}
impl Distinct {
pub(crate) fn new(
plan: &Plan,
input: &Schema,
on: Slice,
memory: &Memory,
) -> Result<(Self, Buffered)> {
let on = plan.expr_list(on).to_vec();
let out = Buffered::new();
let distinct = Self {
whole: on.is_empty(),
on: Prepared::new(plan, &on, input)?,
types: input.types(),
memory: memory.clone(),
global: Mutex::new(Held {
seen: RowSet::default(),
kept: Vec::new(),
counted: 0,
table: memory.reservation(),
counted_rows: 0,
slots: memory.reservation(),
}),
charged: Mutex::new(Vec::new()),
held: Mutex::new(memory.reservation()),
out: out.clone(),
};
Ok((distinct, out))
}
}
impl Sink for Distinct {
type Local = Keeping;
fn local(&self) -> Keeping {
Keeping {
seen: RowSet::default(),
kept: Vec::new(),
scratch: self.on.scratch(),
key: Key(Vec::new()),
rows: self.memory.reservation(),
table: self.memory.reservation(),
counted: 0,
counted_table: 0,
}
}
fn sink(&self, chunk: &Chunk, local: &mut Keeping) -> Result<Progress> {
let mut keys = Vec::with_capacity(self.on.len());
self.on.evaluate(chunk, &mut local.scratch, &mut keys)?;
let mut taken = 0;
let mut aside = 0;
for row in 0..chunk.len() {
if self.whole {
local.key.0.clear();
local.key.0.extend(chunk.row(row));
} else {
fill(&mut local.key, &keys, row);
}
if local.seen.contains(&local.key) {
continue;
}
let values: Vec<Value> =
if self.whole { local.key.0.clone() } else { chunk.row(row).collect() };
let stored = local.key.clone();
taken += rows::heap(&values) + rows::heap(&stored.0);
aside += rows::heap(&stored.0);
local.seen.insert(stored.clone());
local.kept.push((stored, values));
}
local.rows.grow(taken)?;
local.table.grow(aside)?;
let held = width_of(local.kept.capacity() * size_of::<(Key, Vec<Value>)>());
rows::capacity(held, &mut local.counted, &mut local.rows)?;
let now = rows::buckets(local.seen.capacity()) * (width_of(size_of::<Key>()) + 1);
rows::capacity(now, &mut local.counted_table, &mut local.table)?;
Ok(Progress::More)
}
fn combine(&self, local: Keeping) -> Result<()> {
let Keeping { seen, kept, rows, mut table, .. } = local;
drop(seen);
table.release();
let mut global = self.global.lock().map_err(poisoned)?;
let global = &mut *global;
let mut aside = 0;
for (key, values) in kept {
let cost = rows::heap(&key.0);
if global.seen.insert(key) {
aside += cost;
global.kept.push(values);
}
}
global.table.grow(aside)?;
let now = rows::buckets(global.seen.capacity()) * (width_of(size_of::<Key>()) + 1);
rows::capacity(now, &mut global.counted, &mut global.table)?;
let held = width_of(global.kept.capacity() * size_of::<Vec<Value>>());
rows::capacity(held, &mut global.counted_rows, &mut global.slots)?;
self.charged.lock().map_err(poisoned)?.push(rows);
Ok(())
}
fn finalize(&self) -> Result<()> {
let mut global = self.global.lock().map_err(poisoned)?;
let kept = std::mem::take(&mut global.kept);
global.seen = RowSet::default();
global.counted = 0;
global.table.release();
let mut held = self.held.lock().map_err(poisoned)?;
let chunks = rows::chunks(&self.types, &kept, &mut held)?;
self.out.fill(chunks)?;
global.counted_rows = 0;
global.slots.release();
self.charged.lock().map_err(poisoned)?.clear();
Ok(())
}
}
fn poisoned<T>(_: T) -> Error {
Error::internal("a thread panicked while holding the rows a distinct is keeping")
}
#[cfg(test)]
mod tests {
use rudb_common::{Field, LogicalType, Memory, Value};
use rudb_pipeline::Sink;
use rudb_plan::{Plan, Slice};
use rudb_vector::{Chunk, Data, Vector};
use super::{Aggregate, Distinct};
use crate::buffer::Buffered;
use crate::schema::Schema;
fn chunk(values: &[i32]) -> Chunk {
let column = Vector::flat(LogicalType::Integer, Data::Int32(values.to_vec().into()))
.expect("integers are an i32 layout");
Chunk::new(vec![column]).expect("one column is one length")
}
fn distinct() -> (Distinct, Buffered) {
let schema = Schema::numbered(vec![Field::new("a", LogicalType::Integer)], 0);
Distinct::new(&Plan::new(), &schema, Slice::EMPTY, &Memory::unlimited())
.expect("there are no expressions to resolve")
}
fn column(out: &Buffered) -> Vec<Value> {
let chunk = out.at(0).expect("readable").expect("one chunk");
(0..chunk.len()).map(|row| chunk.value_at(row, 0)).collect()
}
#[test]
fn one_instance_keeps_the_first_of_each_row() {
let (distinct, out) = distinct();
let mut local = distinct.local();
distinct.sink(&chunk(&[1, 2, 1, 3, 2]), &mut local).expect("five rows");
distinct.combine(local).expect("the one instance");
distinct.finalize().expect("the answer");
assert_eq!(column(&out), [Value::Integer(1), Value::Integer(2), Value::Integer(3)]);
}
#[test]
fn two_instances_that_both_kept_a_row_keep_one_of_it_between_them() {
let (distinct, out) = distinct();
let mut left = distinct.local();
let mut right = distinct.local();
distinct.sink(&chunk(&[1, 2]), &mut left).expect("two rows");
distinct.sink(&chunk(&[2, 3]), &mut right).expect("two rows");
distinct.combine(left).expect("the first instance");
distinct.combine(right).expect("the second instance");
distinct.finalize().expect("the answer");
assert_eq!(column(&out), [Value::Integer(1), Value::Integer(2), Value::Integer(3)]);
}
#[test]
fn an_ungrouped_aggregate_answers_one_row_from_one_instance() {
let plan = Plan::new();
let schema = Schema::numbered(vec![Field::new("a", LogicalType::Integer)], 0);
let (aggregate, out) =
Aggregate::new(&plan, &schema, 1, Slice::EMPTY, Slice::EMPTY, &Memory::unlimited())
.expect("no aggregates to take apart");
aggregate.combine(aggregate.local()).expect("the one instance");
aggregate.finalize().expect("the answer");
assert_eq!(out.at(0).expect("readable").expect("one chunk").len(), 1);
}
#[test]
fn a_second_instance_of_an_aggregate_is_refused() {
let plan = Plan::new();
let schema = Schema::numbered(vec![Field::new("a", LogicalType::Integer)], 0);
let (aggregate, _) =
Aggregate::new(&plan, &schema, 1, Slice::EMPTY, Slice::EMPTY, &Memory::unlimited())
.expect("no aggregates to take apart");
aggregate.combine(aggregate.local()).expect("the first instance");
let why = aggregate.combine(aggregate.local()).expect_err("and not the second");
assert!(why.to_string().contains("two hash tables"), "{why}");
}
#[test]
fn a_distinct_over_nothing_produces_nothing() {
let (distinct, out) = distinct();
distinct.combine(distinct.local()).expect("an instance that saw no chunks");
distinct.finalize().expect("the answer");
assert_eq!(out.len().expect("readable"), 0);
}
}