rudb-exec 0.2.17

//! Grouping and duplicate elimination.
//!
//! Both are hash tables over [`Key`], which is what makes them agree about what one row is. A
//! `GROUP BY x` that put two nulls in two groups and a `SELECT DISTINCT x` that collapsed them into
//! one would be two answers to the same question, and the only way to be sure that never happens is
//! for both to ask the same type.
//!
//! The table is a `HashMap` from key to slot alongside a `Vec` of keys in first arrival order, so
//! the output comes out in the order the groups were first seen. SQL does not promise that and
//! DuckDB does not either, but a deterministic order costs one push per group and makes a failing
//! test a diff instead of an investigation.

use std::collections::{HashMap, HashSet};

use rudb_common::{ALLOCATION, Error, Field, LogicalType, Memory, Reservation, Result, Value};
use rudb_kernels::{Accumulator, is_true};
use rudb_plan::{Expr, ExprRef, Plan, Slice};
use rudb_vector::Chunk;

use crate::expr::{evaluate, evaluate_all};
use crate::key::Key;
use crate::operator::Operator;
use crate::rows;
use crate::schema::Schema;

/// One aggregate call, taken apart once when the operator is built.
#[derive(Debug, Clone)]
struct Call {
    name: String,
    args: Vec<ExprRef>,
    distinct: bool,
    filter: Option<ExprRef>,
    returns: LogicalType,
}

/// A grouped or ungrouped aggregation.
///
/// The output is the group expressions followed by the aggregates, which is what a binding into
/// this operator's table index means and what the binder assumed when it made one.
///
/// An ungrouped aggregate produces exactly one row even over an empty input. That is done by
/// creating the single empty group when the operator is built rather than when the first row
/// arrives, which is the whole of the difference between `SELECT count(*) FROM empty` answering
/// zero and answering nothing.
#[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,
    /// What the finished chunks are charged, held for as long as this operator holds them.
    held: Reservation,
}

impl<'a> Aggregate<'a> {
    /// An aggregation over the plan's groups and aggregate calls.
    ///
    /// # Errors
    ///
    /// If an entry in the aggregate list is not an aggregate, which [`Plan::validate`] rejects and
    /// which is checked again here because this operator has no sensible behaviour if it is wrong.
    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(),
        })
    }

    /// Reads the whole input and builds the hash table.
    fn build(&mut self) -> Result<()> {
        // The hash table, charged separately from the chunks it produces, because it is given back
        // when the last group has been finished and they are not.
        let mut scratch = self.memory.reservation();
        let mut charged = 0;
        let mut order: Vec<Key> = Vec::new();
        let mut slots: HashMap<Key, usize> = HashMap::new();
        let mut states: Vec<Vec<Accumulator>> = Vec::new();
        let mut seen: Vec<Vec<HashSet<Key>>> = Vec::new();
        let alone = self.groups.is_empty();
        if alone {
            let key = Key(Vec::new());
            slots.insert(key.clone(), 0);
            order.push(key);
            states.push(self.fresh()?);
            seen.push(vec![HashSet::new(); self.calls.len()]);
        }
        // Which calls fold a vector at a time. An ungrouped aggregate has exactly one slot, so
        // there is no key to build, no hash to take and no lookup to do, and what is left of the
        // row loop is the fold itself. `DISTINCT` needs a value per row to put in a set and
        // `FILTER` needs the rows it kept, and neither has a vector form yet, so a call with either
        // stays on the row loop while the calls beside it do not.
        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);
        while let Some(chunk) = self.input.next()? {
            let keys = evaluate_all(self.plan, &self.groups, &self.input_schema, &chunk)?;
            let mut taken = 0;
            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,
                });
            }
            for at in 0..self.calls.len() {
                if by_vector[at] {
                    states[0][at].update_run(&arguments[at], chunk.len())?;
                }
            }
            if alone && every {
                continue;
            }
            // row at a time: the worst one in the tree, because grouping is where the rows are.
            // 2f (#60) gives it a table that hashes a column at a time and probes a vector at a
            // time, and 2g (#61) gives the aggregate an update that takes a vector and a run of
            // slots, at which point neither the key nor the argument is a `Value` any more.
            for row in 0..chunk.len() {
                let slot = if alone {
                    0
                } else {
                    let key = Key(keys.iter().map(|column| column.value_at(row)).collect());
                    match slots.get(&key) {
                        Some(&slot) => slot,
                        None => {
                            let slot = states.len();
                            // A group costs its key twice, once in the table and once in the list
                            // that keeps the arrival order, and its own state beside them. What the
                            // four containers took to hold all of that is charged separately, below
                            // and once per chunk, because it is a property of the containers rather
                            // than of this group.
                            taken += group_state(&key.0, self.calls.len());
                            slots.insert(key.clone(), slot);
                            order.push(key);
                            states.push(self.fresh()?);
                            seen.push(vec![HashSet::new(); self.calls.len()]);
                            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: Vec<Value> =
                        arguments[at].iter().map(|column| column.value_at(row)).collect();
                    if call.distinct {
                        let key = Key(args.clone());
                        let size = rows::footprint(&key.0);
                        if !seen[slot][at].insert(key) {
                            continue;
                        }
                        taken += size;
                    }
                    states[slot][at].update(&args)?;
                }
            }
            scratch.grow(taken)?;
            let now = tables(&slots, &order, &states, &seen);
            rows::capacity(now, &mut charged, &mut scratch)?;
        }
        let mut out = Vec::with_capacity(order.len());
        for (slot, key) in order.into_iter().enumerate() {
            let mut row = key.0;
            for accumulator in &states[slot] {
                row.push(accumulator.finish()?);
            }
            out.push(row);
        }
        self.chunks = rows::chunks(&self.schema.types(), &out, &mut self.held)?;
        Ok(())
    }

    /// A fresh accumulator per call, which is what one new group costs.
    fn fresh(&self) -> Result<Vec<Accumulator>> {
        self.calls.iter().map(|call| Accumulator::new(&call.name, &call.returns)).collect()
    }
}

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))
    }
}

/// What one new group costs, not counting the room the four containers made for it.
///
/// The key twice, because the table and the arrival order each own a copy, and each copy is a
/// separate block from the allocator. Then the two vectors this group's own state lives in, one of
/// accumulators and one of distinct sets, one per aggregate call, and each of those is a block too.
/// A query with no aggregate calls, which is what `DISTINCT` binds to, allocates neither, because an
/// empty `Vec` does not go to the allocator at all.
///
/// An accumulator is charged as its own width and not as what it holds. That is a knowing undercount
/// and it is the one left: what a `list()` or a `string_agg()` holds grows with the input and there
/// is no way to ask one how large it has become.
fn group_state(key: &[Value], calls: usize) -> u64 {
    let mut bytes = 2 * rows::heap(key);
    if calls > 0 {
        let calls = u64::try_from(calls).unwrap_or(u64::MAX);
        let width = |size: usize| u64::try_from(size).unwrap_or(u64::MAX);
        bytes += calls * width(size_of::<Accumulator>()) + ALLOCATION;
        bytes += calls * width(size_of::<HashSet<Key>>()) + ALLOCATION;
    }
    bytes
}

/// What the four containers have taken from the allocator between them.
///
/// Capacity rather than length in all four, which is the point of #227. A `Vec` doubles and so sits
/// between half empty and full, and a `HashMap` fills to seven eighths before doubling as well, so
/// a table of seventeen million groups has paid for somewhere between seventeen and thirty four
/// million slots and the old charge counted seventeen.
///
/// The hash table also has a control byte per bucket beside the buckets themselves, which is how it
/// answers a lookup without touching the keys, and there are more buckets than the capacity it
/// reports. [`rows::buckets`] has that arithmetic.
///
/// The keys and the states these have room for are not counted here. They are charged as each group
/// arrives, by [`group_state`], and the two have to divide the group between them without
/// overlapping.
fn tables(
    slots: &HashMap<Key, usize>,
    order: &Vec<Key>,
    states: &Vec<Vec<Accumulator>>,
    seen: &Vec<Vec<HashSet<Key>>>,
) -> u64 {
    let width = |count: usize, size: usize| {
        u64::try_from(count).unwrap_or(u64::MAX).saturating_mul(width_of(size))
    };
    rows::buckets(slots.capacity()) * (width_of(size_of::<(Key, usize)>()) + 1)
        + width(order.capacity(), size_of::<Key>())
        + width(states.capacity(), size_of::<Vec<Accumulator>>())
        + width(seen.capacity(), size_of::<Vec<HashSet<Key>>>())
}

/// A `size_of` in the width the budget is counted in.
fn width_of(size: usize) -> u64 {
    u64::try_from(size).unwrap_or(u64::MAX)
}

/// What a group column is called in this operator's schema.
///
/// A group over a plain column keeps that column's name, because the thing somebody reading a plan
/// dump or a mid-pipeline schema wants to know is which column it is. Anything else gets a
/// positional name, since the projection above an aggregate is what names the query's output and
/// these names never reach a result set.
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}")
}

/// Duplicate elimination over the whole row or over named expressions.
///
/// `DISTINCT ON (a) b` keeps the first row of each `a`, whole, which is why the kept rows are the
/// input's columns and not the key's. Plain `DISTINCT` is the same operator with the key being
/// every column, and writing it that way rather than as a separate path is what keeps the two from
/// disagreeing about nulls.
#[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,
    /// What the kept chunks are charged, held for as long as this operator holds them.
    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 charged = 0;
        let mut seen: HashSet<Key> = HashSet::new();
        let mut kept: Vec<Vec<Value>> = 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;
            // row at a time: `DISTINCT` is a grouping that keeps no aggregate, so it gets its
            // answer from the same table 2f (#60) builds and stops building a key here then.
            for row in 0..chunk.len() {
                let values: Vec<Value> = chunk.row(row).collect();
                let key = if self.on.is_empty() {
                    Key(values.clone())
                } else {
                    Key(keys.iter().map(|column| column.value_at(row)).collect())
                };
                // The row is kept twice, once as the key in the table and once in the output, and
                // each copy is its own block. What the table and the output took to have room for
                // them is charged below, once per chunk.
                let size = rows::heap(&key.0) + rows::heap(&values);
                if seen.insert(key) {
                    kept.push(values);
                    taken += size;
                }
            }
            scratch.grow(taken)?;
            let now = rows::buckets(seen.capacity()) * (width_of(size_of::<Key>()) + 1)
                + width_of(kept.capacity() * size_of::<Vec<Value>>());
            rows::capacity(now, &mut charged, &mut scratch)?;
        }
        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))
    }
}