spiral-sdk 0.1.8

//! Merge aggregate partial states returned by scan splits without the native runtime.

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

use arrow_array::cast::AsArray as _;
use arrow_array::types::UInt64Type;
use arrow_array::{
    ArrayRef as ArrowArrayRef, BooleanArray, RecordBatch, UInt64Array, new_null_array,
};
use arrow_row::{RowConverter, SortField};
use arrow_schema::SchemaRef;
use arrow_select::concat::concat;
use vortex_array::aggregate_fn::fns::count::Count;
use vortex_array::aggregate_fn::fns::max::Max;
use vortex_array::aggregate_fn::fns::min::Min;
use vortex_array::aggregate_fn::fns::sum::Sum;
use vortex_array::aggregate_fn::{
    AccumulatorRef, AggregateFnRef, AggregateFnVTableExt as _, NumericalAggregateOpts,
};
use vortex_array::builders::builder_with_capacity_in;
use vortex_array::dtype::{DType, DecimalDType, Nullability, PType};
use vortex_array::scalar::Scalar;
use vortex_array::{ExecutionCtx, RepeatedArrayProbe, VortexSessionExecute as _};
use vortex_arrow::ArrowSessionExt as _;
use vortex_session::VortexSession;

use super::super::error::ScandalError;
use super::{
    Aggregate, AggregateFunction, DEFAULT_MAX_GROUPS, MAX_GROUPS, MAX_HISTOGRAM_BINS, OTHER_COLUMN,
    ROWS_COLUMN, internal, invalid,
};

/// An aggregate checked against its partial results' schema.
#[derive(Clone, Debug)]
pub struct BoundAggregate {
    /// Each group's output column, in group order.
    groups: Vec<String>,
    functions: Vec<BoundFunction>,
    max_groups: usize,
    schema: SchemaRef,
    session: VortexSession,
}

/// An aggregate's partial type and the operation that combines it.
#[derive(Clone, Debug)]
struct BoundFunction {
    merge: PartialMerge,
    partial: DType,
}

#[derive(Clone, Debug)]
enum PartialMerge {
    Vortex {
        function: AggregateFnRef,
        input: DType,
    },
    Histogram {
        slots: usize,
    },
}

impl BoundFunction {
    fn bind(function: &AggregateFunction, partial: DType) -> Result<Self, ScandalError> {
        let numeric = NumericalAggregateOpts::include_nans();
        let aggregate_fn = match function {
            AggregateFunction::Count(_) => Count.bind(numeric),
            AggregateFunction::Sum(_) => Sum.bind(numeric),
            AggregateFunction::Min(_) => Min.bind(numeric),
            AggregateFunction::Max(_) => Max.bind(numeric),
            AggregateFunction::Histogram(histogram) => {
                if !histogram.start.is_finite()
                    || !histogram.width.is_finite()
                    || histogram.width <= 0.0
                {
                    return Err(invalid(
                        "a histogram needs a finite start and a positive, finite width",
                    ));
                }
                if !(1..=MAX_HISTOGRAM_BINS).contains(&histogram.bins) {
                    return Err(invalid(format!(
                        "a histogram counts between 1 and {MAX_HISTOGRAM_BINS} bins"
                    )));
                }
                let expected = DType::FixedSizeList(
                    Arc::new(DType::Primitive(PType::U64, Nullability::NonNullable)),
                    histogram.bins + 2,
                    Nullability::NonNullable,
                );
                if partial != expected {
                    return Err(invalid("histogram results are not its partial counts"));
                }
                return Ok(Self {
                    merge: PartialMerge::Histogram {
                        slots: histogram.bins as usize + 2,
                    },
                    partial,
                });
            }
            AggregateFunction::Unknown { tag, .. } => {
                return Err(invalid(format!("unknown aggregate `{tag}`")));
            }
        };
        // Infer an input type only to select the public Vortex partial-state merger.
        let input = match (function, &partial) {
            (AggregateFunction::Count(_), _) => DType::Bool(Nullability::NonNullable),
            (AggregateFunction::Sum(_), DType::Primitive(ptype, _)) => {
                DType::Primitive(*ptype, Nullability::NonNullable)
            }
            (AggregateFunction::Sum(_), DType::Decimal(decimal, _)) => {
                // A sum widens a decimal's precision by ten, to at most 38.
                let precision = decimal.precision().saturating_sub(10).max(1);
                DType::Decimal(
                    DecimalDType::new(precision, decimal.scale()),
                    Nullability::NonNullable,
                )
            }
            _ => partial.clone(),
        };
        if aggregate_fn.state_dtype(&input).as_ref() != Some(&partial)
            || aggregate_fn.return_dtype(&input).as_ref() != Some(&partial)
        {
            return Err(invalid("aggregate results are not its partial states"));
        }
        Ok(Self {
            merge: PartialMerge::Vortex {
                function: aggregate_fn,
                input,
            },
            partial,
        })
    }

    fn accumulator(&self) -> Result<PartialAccumulator, ScandalError> {
        match &self.merge {
            PartialMerge::Vortex { function, input } => function
                .accumulator(input)
                .map(PartialAccumulator::Vortex)
                .map_err(internal),
            PartialMerge::Histogram { slots } => Ok(PartialAccumulator::Histogram(vec![0; *slots])),
        }
    }
}

/// Merges partial results only; scan-row accumulation belongs to the native runtime.
enum PartialAccumulator {
    Vortex(AccumulatorRef),
    Histogram(Vec<u64>),
}

impl PartialAccumulator {
    fn combine_partials(&mut self, partial: &Scalar) -> Result<(), ScandalError> {
        match self {
            Self::Vortex(accumulator) => accumulator.combine_partials(partial).map_err(internal),
            Self::Histogram(counts) => {
                let partial = partial
                    .as_list()
                    .elements()
                    .ok_or_else(|| internal("histogram partial is null"))?;
                if partial.len() != counts.len() {
                    return Err(internal("histogram partial has the wrong number of counts"));
                }
                for (count, other) in counts.iter_mut().zip(partial) {
                    *count += other
                        .as_primitive()
                        .typed_value::<u64>()
                        .ok_or_else(|| internal("histogram count is null"))?;
                }
                Ok(())
            }
        }
    }

    fn merge_from(&mut self, other: &mut Self) -> Result<(), ScandalError> {
        match (self, other) {
            (Self::Vortex(into), Self::Vortex(from)) => {
                into.merge_from(from.as_mut()).map_err(internal)
            }
            (Self::Histogram(into), Self::Histogram(from)) => {
                for (count, other) in into.iter_mut().zip(from) {
                    *count += *other;
                }
                Ok(())
            }
            _ => Err(internal("cannot merge different aggregate states")),
        }
    }

    fn final_scalar(&self) -> Result<Scalar, ScandalError> {
        match self {
            Self::Vortex(accumulator) => accumulator.final_scalar().map_err(internal),
            Self::Histogram(counts) => Ok(Scalar::fixed_size_list(
                Arc::new(DType::Primitive(PType::U64, Nullability::NonNullable)),
                counts
                    .iter()
                    .map(|count| Scalar::primitive(*count, Nullability::NonNullable))
                    .collect(),
                Nullability::NonNullable,
            )),
        }
    }
}

impl BoundAggregate {
    /// Check `aggregate` against the schema of its results, which a client merges without the
    /// rows the results came from.
    pub(super) fn bind_results(
        aggregate: &Aggregate,
        results: &SchemaRef,
        session: &VortexSession,
    ) -> Result<Self, ScandalError> {
        let max_groups = Self::check(aggregate)?;
        let expected = aggregate
            .groups
            .iter()
            .map(String::as_str)
            .chain(aggregate.aggregates.iter().map(|named| named.name.as_str()))
            .chain([ROWS_COLUMN, OTHER_COLUMN]);
        if !results
            .fields()
            .iter()
            .map(|field| field.name().as_str())
            .eq(expected)
        {
            return Err(invalid("results do not match the aggregate"));
        }
        let functions = aggregate
            .aggregates
            .iter()
            .enumerate()
            .map(|(offset, named)| {
                let partial = session
                    .arrow()
                    .from_arrow_field(results.field(aggregate.groups.len() + offset))
                    .map_err(internal)?;
                BoundFunction::bind(&named.function, partial)
            })
            .collect::<Result<Vec<_>, ScandalError>>()?;
        Ok(Self {
            groups: aggregate.groups.clone(),
            functions,
            max_groups,
            schema: Arc::clone(results),
            session: session.clone(),
        })
    }

    /// Check the cap, and that output names are distinct; return the cap.
    fn check(aggregate: &Aggregate) -> Result<usize, ScandalError> {
        let max_groups = aggregate.max_groups.unwrap_or(DEFAULT_MAX_GROUPS);
        if !(1..=MAX_GROUPS).contains(&max_groups) {
            return Err(invalid(format!(
                "an aggregate keeps between 1 and {MAX_GROUPS} groups; group more finely in your \
                 own engine"
            )));
        }
        let mut names = HashSet::from([ROWS_COLUMN, OTHER_COLUMN]);
        for name in aggregate
            .groups
            .iter()
            .chain(aggregate.aggregates.iter().map(|named| &named.name))
        {
            if !names.insert(name) {
                return Err(invalid(format!(
                    "aggregate output `{name}` names another output column"
                )));
            }
        }
        Ok(max_groups as usize)
    }

    /// The schema of a split's results and of their merge.
    pub fn schema(&self) -> &SchemaRef {
        &self.schema
    }
}

/// The groups and partial states collected from scan splits.
pub(super) struct AggregateMerger {
    bound: BoundAggregate,
    ctx: ExecutionCtx,
    /// Encodes group values; absent for an aggregate without groups.
    converter: Option<RowConverter>,
    /// Each group by its encoded values.
    index: HashMap<Arc<[u8]>, usize>,
    /// Each group's encoded values; `None` for the row folding groups past the cap.
    keys: Vec<Option<Arc<[u8]>>>,
    other: Option<usize>,
    rows: Vec<u64>,
    /// Each group's accumulator for each aggregate.
    accumulators: Vec<Vec<PartialAccumulator>>,
}

impl AggregateMerger {
    /// An empty table. An aggregate without groups has its one group from the start, so it
    /// returns a row even when no rows were read.
    pub fn new(bound: &BoundAggregate) -> Result<Self, ScandalError> {
        let groups = bound.groups.len();
        let converter = if groups == 0 {
            None
        } else {
            Some(
                RowConverter::new(
                    bound.schema.fields()[..groups]
                        .iter()
                        .map(|field| SortField::new(field.data_type().clone()))
                        .collect(),
                )
                .map_err(internal)?,
            )
        };
        let mut table = Self {
            ctx: bound.session.create_execution_ctx(),
            bound: bound.clone(),
            converter,
            index: HashMap::new(),
            keys: Vec::new(),
            other: None,
            rows: Vec::new(),
            accumulators: Vec::new(),
        };
        if table.converter.is_none() {
            table.insert(Some(Arc::from([])))?;
        }
        Ok(table)
    }

    /// Merge one split's results with each aggregate's `merge_partials`. Groups are capped only
    /// when the merge finishes.
    pub fn merge(&mut self, results: &RecordBatch) -> Result<(), ScandalError> {
        if results.num_columns() != self.bound.schema.fields().len() {
            return Err(internal("aggregate results do not match the aggregate"));
        }
        let groups = self.bound.groups.len();
        let other = results
            .column(results.num_columns() - 1)
            .as_boolean_opt()
            .ok_or_else(|| internal("aggregate results lack their other flags"))?
            .clone();
        let indices = match &self.converter {
            None => vec![0; results.num_rows()],
            Some(converter) => {
                let keys = converter
                    .convert_columns(&results.columns()[..groups])
                    .map_err(internal)?;
                let mut indices = Vec::with_capacity(keys.num_rows());
                for (row, key) in keys.iter().enumerate() {
                    let group = if other.value(row) {
                        self.other_group()?
                    } else {
                        match self.index.get(key.as_ref()) {
                            Some(group) => *group,
                            None => self.insert(Some(Arc::from(key.as_ref())))?,
                        }
                    };
                    indices.push(group);
                }
                indices
            }
        };
        for (offset, function) in self.bound.functions.iter().enumerate() {
            let column = self
                .bound
                .session
                .arrow()
                .from_arrow_array(
                    Arc::clone(results.column(groups + offset)),
                    self.bound.schema.field(groups + offset),
                )
                .map_err(internal)?;
            // `combine_partials` takes one partial state at a time, so read the column's states
            // through a prepared probe.
            let mut partials = RepeatedArrayProbe::new(column);
            for (row, group) in indices.iter().enumerate() {
                let partial = partials
                    .execute_scalar(row, &mut self.ctx)
                    .and_then(|partial| partial.cast(&function.partial))
                    .map_err(internal)?;
                self.accumulators[*group][offset].combine_partials(&partial)?;
            }
        }
        let rows = results
            .column(results.num_columns() - 2)
            .as_primitive_opt::<UInt64Type>()
            .ok_or_else(|| internal("aggregate results lack their row counts"))?;
        for (row, group) in indices.into_iter().enumerate() {
            self.rows[group] += rows.value(row);
        }
        Ok(())
    }

    /// The merged result: groups by their rows, most first, then by their values; at most
    /// `maxGroups` of them, with the rest folded into the other row, which comes last.
    pub fn finish(mut self) -> Result<RecordBatch, ScandalError> {
        let mut order = (0..self.keys.len())
            .filter(|group| Some(*group) != self.other)
            .collect::<Vec<_>>();
        order.sort_by(|left, right| {
            self.rows[*right]
                .cmp(&self.rows[*left])
                .then_with(|| self.keys[*left].cmp(&self.keys[*right]))
        });
        if order.len() > self.bound.max_groups {
            let other = self.other_group()?;
            for group in order.split_off(self.bound.max_groups) {
                let mut folded = std::mem::take(&mut self.accumulators[group]);
                for (into, from) in self.accumulators[other].iter_mut().zip(&mut folded) {
                    into.merge_from(from)?;
                }
                self.rows[other] += self.rows[group];
            }
        }
        self.output(order)
    }

    fn insert(&mut self, key: Option<Arc<[u8]>>) -> Result<usize, ScandalError> {
        let group = self.keys.len();
        let accumulators = self
            .bound
            .functions
            .iter()
            .map(BoundFunction::accumulator)
            .collect::<Result<Vec<_>, _>>()?;
        if let Some(key) = &key {
            self.index.insert(Arc::clone(key), group);
        }
        self.keys.push(key);
        self.rows.push(0);
        self.accumulators.push(accumulators);
        Ok(group)
    }

    fn other_group(&mut self) -> Result<usize, ScandalError> {
        match self.other {
            Some(other) => Ok(other),
            None => {
                let other = self.insert(None)?;
                self.other = Some(other);
                Ok(other)
            }
        }
    }

    /// The rows of `order`'s groups, then the other row when there is one, with each
    /// aggregate's merged result.
    fn output(mut self, order: Vec<usize>) -> Result<RecordBatch, ScandalError> {
        let mut columns = Vec::with_capacity(self.bound.schema.fields().len());
        if let Some(converter) = &self.converter {
            let parser = converter.parser();
            let values =
                converter
                    .convert_rows(order.iter().map(|group| {
                        parser.parse(self.keys[*group].as_deref().unwrap_or_default())
                    }))
                    .map_err(internal)?;
            for value in values {
                columns.push(match self.other {
                    Some(_) => concat(&[&value, &new_null_array(value.data_type(), 1)])
                        .map_err(internal)?,
                    None => value,
                });
            }
        }
        let groups = order.iter().copied().chain(self.other).collect::<Vec<_>>();
        let offset = self.bound.groups.len();
        for (index, function) in self.bound.functions.iter().enumerate() {
            let mut builder =
                builder_with_capacity_in(&function.partial, groups.len(), self.ctx.allocator());
            for group in &groups {
                let accumulator = &self.accumulators[*group][index];
                let state = accumulator.final_scalar()?;
                builder.append_scalar(&state).map_err(internal)?;
            }
            let column: ArrowArrayRef = self
                .bound
                .session
                .arrow()
                .execute_arrow(
                    builder.finish(),
                    Some(self.bound.schema.field(offset + index)),
                    &mut self.ctx,
                )
                .map_err(internal)?;
            columns.push(column);
        }
        columns.push(Arc::new(
            groups
                .iter()
                .map(|group| self.rows[*group])
                .collect::<UInt64Array>(),
        ));
        columns.push(Arc::new(
            groups
                .iter()
                .map(|group| Some(Some(*group) == self.other))
                .collect::<BooleanArray>(),
        ));
        RecordBatch::try_new(Arc::clone(&self.bound.schema), columns).map_err(internal)
    }
}