spiral-sdk 0.1.11

//! 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::builtins::ArrayBuiltins as _;
use vortex_array::dtype::{DType, DecimalDType, Nullability, PType};
use vortex_array::{
    Canonical, ExecutionCtx, IntoArray as _, RepeatedArrayProbe, VortexSessionExecute as _,
};
use vortex_arrow::ArrowSessionExt as _;
use vortex_session::VortexSession;

use super::super::error::ScandalError;
use super::{
    Aggregate, AggregateFunction, DEFAULT_MAX_GROUPS, Histogram, HistogramOptions, 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 {
    function: AggregateFnRef,
    input: DType,
    partial: DType,
}

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"
                    )));
                }
                Histogram.bind(HistogramOptions {
                    start: histogram.start,
                    width: histogram.width,
                    bins: histogram.bins,
                })
            }
            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::Histogram(_), _) => {
                DType::Primitive(PType::F64, 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 {
            function: aggregate_fn,
            input,
            partial,
        })
    }

    fn accumulator(&self) -> Result<AccumulatorRef, ScandalError> {
        self.function.accumulator(&self.input).map_err(internal)
    }
}

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<AccumulatorRef>>,
}

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),
                )
                .and_then(|column| column.cast(function.partial.clone()))
                .map_err(internal)?
                .execute::<Canonical>(&mut self.ctx)
                .map_err(internal)?
                .into_array();
            // The Vortex merge API takes Scalars; prepare the column once for repeated reads.
            let mut partials = RepeatedArrayProbe::new(column);
            for (row, group) in indices.iter().enumerate() {
                let partial = partials
                    .execute_scalar(row, &mut self.ctx)
                    .map_err(internal)?;
                self.accumulators[*group][offset]
                    .combine_partials(&partial)
                    .map_err(internal)?;
            }
        }
        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.as_mut()).map_err(internal)?;
                }
                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 state = self.accumulators[*group][index]
                    .final_scalar()
                    .map_err(internal)?;
                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)
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use arrow_array::{Array as _, FixedSizeListArray, StringArray};
    use arrow_schema::{DataType, Field, Schema};

    use super::super::{CountAggregate, Emit, HistogramAggregate, NamedAggregate};
    use super::*;

    fn histogram(groups: Vec<String>) -> Emit {
        Emit::Aggregate(Aggregate {
            groups,
            max_groups: Some(1),
            aggregates: vec![
                NamedAggregate {
                    name: "histogram".into(),
                    function: AggregateFunction::Histogram(HistogramAggregate {
                        column: "value".into(),
                        start: 0.0,
                        width: 1.0,
                        bins: 2,
                    }),
                },
                NamedAggregate {
                    name: "count".into(),
                    function: AggregateFunction::Count(CountAggregate { column: None }),
                },
            ],
        })
    }

    fn partials(
        groups: &[Option<&str>],
        counts: &[[u64; 4]],
        rows: &[u64],
        other: Vec<bool>,
    ) -> RecordBatch {
        let item = Arc::new(
            Field::new("bin", DataType::UInt64, false)
                .with_metadata(HashMap::from([("unit".into(), "rows".into())])),
        );
        let counts = FixedSizeListArray::try_new(
            item,
            4,
            Arc::new(counts.iter().flatten().copied().collect::<UInt64Array>()),
            None,
        )
        .unwrap();
        RecordBatch::try_new(
            Arc::new(Schema::new(vec![
                Field::new("group", DataType::Utf8, true),
                Field::new("histogram", counts.data_type().clone(), false),
                Field::new("count", DataType::UInt64, false),
                Field::new(ROWS_COLUMN, DataType::UInt64, false),
                Field::new(OTHER_COLUMN, DataType::Boolean, false),
            ])),
            vec![
                Arc::new(StringArray::from(groups.to_vec())),
                Arc::new(counts),
                Arc::new(UInt64Array::from(rows.to_vec())),
                Arc::new(UInt64Array::from(rows.to_vec())),
                Arc::new(BooleanArray::from(other)),
            ],
        )
        .unwrap()
    }

    #[test]
    fn histogram_merge_preserves_slices_exact_counts_and_folded_groups() {
        let large = 1_u64 << 60;
        let first = partials(
            &[Some("ignored"), Some("a"), Some("b")],
            &[[99; 4], [1, 2, 3, 4], [4, 3, 2, 1]],
            &[396, 10, 10],
            vec![false; 3],
        )
        .slice(1, 2);
        let second = partials(
            &[Some("a"), Some("a"), Some("c"), None],
            &[[large, 0, 0, 0], [0, 1, 0, 0], [1, 0, 1, 0], [0, 1, 0, 1]],
            &[large, 1, 2, 2],
            vec![false, false, false, true],
        );
        let bound = histogram(vec!["group".into()])
            .bind_results(&first.schema())
            .unwrap();
        let merged = bound.merge(&[first.clone(), second]).unwrap();
        assert_eq!(merged.schema(), first.schema());
        assert_eq!(merged.num_rows(), 2);
        let counts = merged.column(1).as_fixed_size_list().values();
        assert_eq!(
            counts.as_primitive::<UInt64Type>().values().as_ref(),
            &[large + 1, 3, 3, 4, 5, 4, 3, 2]
        );
        assert_eq!(
            merged.column(0).as_string::<i32>(),
            &StringArray::from(vec![Some("a"), None])
        );
        for index in [2, 3] {
            assert_eq!(
                merged
                    .column(index)
                    .as_primitive::<UInt64Type>()
                    .values()
                    .as_ref(),
                &[large + 11, 14]
            );
        }
        assert_eq!(
            merged.column(4).as_boolean(),
            &BooleanArray::from(vec![false, true])
        );
        assert_eq!(bound.merge(&[first.slice(0, 0)]).unwrap().num_rows(), 0);

        let ungrouped = first.project(&[1, 2, 3, 4]).unwrap();
        let empty = histogram(Vec::new())
            .bind_results(&ungrouped.schema())
            .unwrap()
            .merge(&[])
            .unwrap();
        assert_eq!(empty.num_rows(), 1);
        assert_eq!(
            empty
                .column(0)
                .as_fixed_size_list()
                .values()
                .as_primitive::<UInt64Type>()
                .values()
                .as_ref(),
            &[0; 4]
        );
    }

    #[test]
    fn histogram_merge_refuses_malformed_partial_counts() {
        let valid = partials(&[Some("a")], &[[0; 4]], &[0], vec![false]);
        let bound = histogram(vec!["group".into()])
            .bind_results(&valid.schema())
            .unwrap();
        for counts in [
            FixedSizeListArray::from_iter_primitive::<UInt64Type, _, _>(
                [Some(vec![Some(0); 3])],
                3,
            ),
            FixedSizeListArray::from_iter_primitive::<UInt64Type, _, _>(
                [None::<Vec<Option<u64>>>],
                4,
            ),
            FixedSizeListArray::from_iter_primitive::<UInt64Type, _, _>(
                [Some(vec![Some(0), None, Some(0), Some(0)])],
                4,
            ),
        ] {
            let mut fields = valid.schema().fields().to_vec();
            fields[1] = Arc::new(Field::new("histogram", counts.data_type().clone(), true));
            let mut columns = valid.columns().to_vec();
            columns[1] = Arc::new(counts);
            let malformed = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap();
            assert!(bound.merge(&[malformed]).is_err());
        }
    }
}