vortex_array/aggregate_fn/fns/sum_v2/

mod.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

//! SQL-style sums with explicit overflow and empty-input state.
//!
//! [`SumV2`] returns null when there are no valid values, while [`Sum`] returns zero. Tracking empty
//! input separately from overflow requires a different partial representation. A separate aggregate
//! preserves compatibility with the scalar partials stored by older Vortex files.

mod grouped;
pub(crate) use grouped::PrimitiveGroupedSumV2EncodingKernel;
use vortex_error::VortexExpect;
use vortex_error::VortexResult;
use vortex_error::vortex_bail;
use vortex_error::vortex_ensure;
use vortex_error::vortex_err;
use vortex_session::VortexSession;
use vortex_session::registry::CachedId;

use crate::ArrayRef;
use crate::ArrayView;
use crate::Canonical;
use crate::Columnar;
use crate::ExecutionCtx;
use crate::aggregate_fn::Accumulator;
use crate::aggregate_fn::AggregateArgs;
use crate::aggregate_fn::AggregateFnId;
use crate::aggregate_fn::AggregateFnVTable;
use crate::aggregate_fn::DynAccumulator;
use crate::aggregate_fn::NumericalAggregateOpts;
use crate::aggregate_fn::fns::sum::Sum;
use crate::aggregate_fn::fns::sum::SumState;
use crate::aggregate_fn::fns::sum::accumulate_bool;
use crate::aggregate_fn::fns::sum::accumulate_decimal;
use crate::aggregate_fn::fns::sum::accumulate_primitive;
use crate::aggregate_fn::fns::sum::checked_add_sum_states;
use crate::aggregate_fn::fns::sum::make_zero_state;
use crate::aggregate_fn::fns::sum::multiply_constant;
use crate::arrays::Struct;
use crate::arrays::struct_::StructArrayExt;
use crate::builtins::ArrayBuiltins;
use crate::dtype::DType;
use crate::dtype::FieldName;
use crate::dtype::FieldNames;
use crate::dtype::Nullability;
use crate::dtype::PType;
use crate::dtype::StructFields;
use crate::expr::stats::Precision;
use crate::expr::stats::Stat;
use crate::expr::stats::StatsProviderExt;
use crate::scalar::Scalar;
use crate::scalar::ScalarValue;
use crate::scalar_fn::fns::operators::Operator;
use crate::validity::Validity;

const SUM_FIELD: &str = "sum";
const IS_OVERFLOW_FIELD: &str = "is_overflow";
const IS_EMPTY_FIELD: &str = "is_empty";

/// Return the SQL-style sum of an array.
///
/// Unlike [`sum`](crate::aggregate_fn::fns::sum::sum), this returns null when the array has no
/// valid values.
pub fn sum_v2(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Scalar> {
    let mut acc = Accumulator::try_new(
        SumV2,
        NumericalAggregateOpts::default(),
        array.dtype().clone(),
    )?;
    acc.accumulate(array, ctx)?;
    acc.finish()
}

/// Sum an array, returning null when it has no valid values or if the sum overflows.
///
/// This aggregate intentionally has a distinct ID and partial representation from the legacy
/// [`Sum`]. Keeping `vortex.sum` unchanged preserves the scalar partials stored by older Vortex
/// files, while `SumV2` can use an explicit `{ sum, is_overflow, is_empty }` state.
///
/// NaN handling for float inputs is controlled by [`NumericalAggregateOpts`]. With `skip_nans`
/// (the default), NaN values contribute nothing but still make the input non-empty. Otherwise,
/// any NaN value poisons the sum to NaN.
#[derive(Clone, Copy, Debug)]
pub struct SumV2;

impl SumV2 {
    /// Build an encoding kernel's partial from a widened primitive sum.
    ///
    /// `sum` must have dtype `u64`, `i64`, or `f64`. A null sum records overflow and overrides
    /// `is_empty`. Set `is_empty` only when there were no valid inputs. Valid NaNs make the input
    /// non-empty even when skipped. The returned struct and its fields are non-null.
    pub fn partial_from_sum(sum: Scalar, is_empty: bool) -> VortexResult<Scalar> {
        vortex_ensure!(
            matches!(
                sum.dtype(),
                DType::Primitive(PType::U64 | PType::I64 | PType::F64, _)
            ),
            "Expected a widened primitive sum, got {}",
            sum.dtype(),
        );

        let sum_dtype = sum.dtype().as_nonnullable();
        let is_overflow = sum.is_null();
        let sum = if is_overflow {
            Scalar::zero_value(&sum_dtype)
        } else {
            sum.cast(&sum_dtype)?
        };

        Ok(Scalar::struct_(
            sum_v2_partial_dtype(sum_dtype),
            vec![
                sum,
                Scalar::bool(is_overflow, Nullability::NonNullable),
                Scalar::bool(is_empty && !is_overflow, Nullability::NonNullable),
            ],
        ))
    }
}

impl AggregateFnVTable for SumV2 {
    type Options = NumericalAggregateOpts;
    type Partial = SumV2Partial;

    fn id(&self) -> AggregateFnId {
        static ID: CachedId = CachedId::new("vortex.sum_v2");
        *ID
    }

    fn serialize(&self, options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
        Ok(Some(options.serialize()))
    }

    fn deserialize(
        &self,
        metadata: &[u8],
        _session: &VortexSession,
    ) -> VortexResult<Self::Options> {
        NumericalAggregateOpts::deserialize(metadata)
    }

    fn return_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
        Sum.return_dtype(options, input_dtype)
    }

    fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
        self.return_dtype(options, input_dtype)
            .map(sum_v2_partial_dtype)
    }

    fn empty_partial(&self, args: AggregateArgs<'_, Self::Options>) -> VortexResult<Self::Partial> {
        Ok(SumV2Partial::empty(args.return_dtype))
    }

    fn partial_from_scalar(
        &self,
        args: AggregateArgs<'_, Self::Options>,
        scalar: &Scalar,
    ) -> VortexResult<Self::Partial> {
        let mut partial = SumV2Partial::empty(args.return_dtype);
        let (sum, is_overflow, is_empty) = decode_partial_scalar(scalar, args.return_dtype)?;

        // Adding the parsed value to the zero state cannot overflow; treat a decimal value that
        // no longer fits its precision as an already-overflowed partial.
        let overflowed = checked_add_sum_value(&mut partial.sum, args.return_dtype, sum)?;
        partial.is_overflow = is_overflow || overflowed;
        partial.is_empty = is_empty && !partial.is_overflow;
        Ok(partial)
    }

    fn merge_partials(
        &self,
        args: AggregateArgs<'_, Self::Options>,
        mut acc: Self::Partial,
        partial: Self::Partial,
    ) -> VortexResult<Self::Partial> {
        if acc.is_overflow {
            return Ok(acc);
        }
        if partial.is_overflow {
            // An overflowed state keeps its last valid sum, so an empty accumulator adopts the
            // incoming sum rather than reporting zero.
            if acc.is_empty {
                acc.sum = partial.sum;
            }
            acc.is_overflow = true;
            acc.is_empty = false;
            return Ok(acc);
        }
        if partial.is_empty {
            return Ok(acc);
        }
        acc.is_overflow = checked_add_sum_states(&mut acc.sum, args.return_dtype, &partial.sum)?;
        acc.is_empty = false;
        Ok(acc)
    }

    fn to_scalar(
        &self,
        args: AggregateArgs<'_, Self::Options>,
        partial: &Self::Partial,
    ) -> VortexResult<Scalar> {
        Ok(Scalar::struct_(
            args.partial_dtype.clone(),
            vec![
                sum_state_scalar(partial, args.return_dtype, Nullability::NonNullable),
                Scalar::bool(partial.is_overflow, Nullability::NonNullable),
                Scalar::bool(partial.is_empty, Nullability::NonNullable),
            ],
        ))
    }

    fn is_saturated(
        &self,
        _args: AggregateArgs<'_, Self::Options>,
        partial: &Self::Partial,
    ) -> bool {
        partial.is_overflow || matches!(&partial.sum, SumState::Float(value) if value.is_nan())
    }

    fn try_accumulate(
        &self,
        args: AggregateArgs<'_, Self::Options>,
        partial: &mut Self::Partial,
        batch: &ArrayRef,
        _ctx: &mut ExecutionCtx,
    ) -> VortexResult<bool> {
        if args.options.skip_nans || !matches!(&partial.sum, SumState::Float(_)) {
            return Ok(false);
        }

        match batch.statistics().get_as::<u64>(Stat::NaNCount) {
            Precision::Exact(0) => Ok(false),
            Precision::Exact(_) => {
                let SumState::Float(sum) = &mut partial.sum else {
                    unreachable!("checked float sum state")
                };
                *sum = f64::NAN;
                partial.is_empty = false;
                Ok(true)
            }
            _ => Ok(false),
        }
    }

    fn accumulate(
        &self,
        args: AggregateArgs<'_, Self::Options>,
        partial: &mut Self::Partial,
        batch: &Columnar,
        ctx: &mut ExecutionCtx,
    ) -> VortexResult<()> {
        if partial.is_overflow {
            return Ok(());
        }

        if let Columnar::Constant(constant) = batch {
            if !constant.scalar().is_null() && !constant.is_empty() {
                partial.is_empty = false;
            }
            if args.options.skip_nans
                && constant
                    .scalar()
                    .as_primitive_opt()
                    .is_some_and(|primitive| primitive.is_nan())
            {
                return Ok(());
            }
            if let Some(product) =
                multiply_constant(constant.scalar(), constant.len(), args.return_dtype)?
            {
                if product.is_null() {
                    partial.is_overflow = true;
                    partial.is_empty = false;
                } else {
                    partial.is_overflow =
                        checked_add_sum_state(&mut partial.sum, args.return_dtype, &product)?;
                    partial.is_empty = false;
                }
            }
            return Ok(());
        }

        let any_valid = partial.is_empty && has_valid_value(batch, ctx)?;
        let result = match batch {
            Columnar::Canonical(canonical) => match canonical {
                Canonical::Primitive(array) => {
                    accumulate_primitive(&mut partial.sum, array, ctx, args.options.skip_nans)
                }
                Canonical::Bool(array) => accumulate_bool(&mut partial.sum, array, ctx),
                Canonical::Decimal(array) => {
                    accumulate_decimal(&mut partial.sum, args.return_dtype, array, ctx)
                }
                _ => vortex_bail!("Unsupported canonical type for sum_v2: {}", batch.dtype()),
            },
            Columnar::Constant(_) => unreachable!(),
        };

        if any_valid {
            partial.is_empty = false;
        }
        if result? {
            partial.is_overflow = true;
            partial.is_empty = false;
        }
        Ok(())
    }

    fn finalize(
        &self,
        _args: AggregateArgs<'_, Self::Options>,
        partials: ArrayRef,
    ) -> VortexResult<ArrayRef> {
        if let Some(partials) = partials.as_opt::<Struct>() {
            return finalize_struct(partials);
        }

        let sum = partials.get_item(SUM_FIELD)?;
        let is_invalid = partials
            .get_item(IS_OVERFLOW_FIELD)?
            .binary(partials.get_item(IS_EMPTY_FIELD)?, Operator::Or)?
            .fill_null(true)?;
        sum.mask(is_invalid.not()?)
    }

    fn finalize_scalar(
        &self,
        args: AggregateArgs<'_, Self::Options>,
        partial: &Self::Partial,
    ) -> VortexResult<Scalar> {
        if partial.is_overflow || partial.is_empty {
            return Ok(Scalar::null(args.return_dtype.as_nullable()));
        }
        Ok(sum_state_scalar(
            partial,
            args.return_dtype,
            Nullability::Nullable,
        ))
    }
}

fn finalize_struct(partials: ArrayView<'_, Struct>) -> VortexResult<ArrayRef> {
    let sum = partials.unmasked_field_by_name(SUM_FIELD)?.clone();
    let is_overflow = partials.unmasked_field_by_name(IS_OVERFLOW_FIELD)?.clone();
    let is_empty = partials.unmasked_field_by_name(IS_EMPTY_FIELD)?.clone();

    let is_invalid = is_overflow
        .binary(is_empty, Operator::Or)?
        .fill_null(true)?;
    let mut is_valid = is_invalid.not()?;
    match partials.struct_validity() {
        Validity::NonNullable | Validity::AllValid => {}
        validity => {
            is_valid = is_valid.binary(validity.to_array(partials.len()), Operator::And)?;
        }
    }

    sum.mask(is_valid)
}

/// In-memory state for SumV2 accumulation.
pub struct SumV2Partial {
    sum: SumState,
    is_overflow: bool,
    is_empty: bool,
}

impl SumV2Partial {
    /// The state of a group with no accumulated values.
    fn empty(return_dtype: &DType) -> Self {
        Self {
            sum: make_zero_state(return_dtype),
            is_overflow: false,
            is_empty: true,
        }
    }
}

fn has_valid_value(batch: &Columnar, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
    let (validity, len) = match batch {
        Columnar::Canonical(Canonical::Primitive(array)) => {
            (array.as_ref().validity()?, array.as_ref().len())
        }
        Columnar::Canonical(Canonical::Bool(array)) => {
            (array.as_ref().validity()?, array.as_ref().len())
        }
        Columnar::Canonical(Canonical::Decimal(array)) => {
            (array.as_ref().validity()?, array.as_ref().len())
        }
        Columnar::Canonical(_) => return Ok(false),
        Columnar::Constant(constant) => {
            return Ok(!constant.is_empty() && !constant.scalar().is_null());
        }
    };
    Ok(validity.execute_mask(len, ctx)?.true_count() > 0)
}

/// Reads the `sum`, `is_overflow` and `is_empty` fields of a partial scalar.
///
/// This reads the struct values directly, because building a [`Scalar`] for each field costs more
/// than the merge itself.
fn decode_partial_scalar<'a>(
    scalar: &'a Scalar,
    return_dtype: &DType,
) -> VortexResult<(&'a ScalarValue, bool, bool)> {
    vortex_ensure!(!scalar.is_null(), "SumV2 partial must not be null");

    let (DType::Struct(fields, _), Some(ScalarValue::Tuple(values))) =
        (scalar.dtype(), scalar.value())
    else {
        vortex_bail!("SumV2 partial must be a struct, got {}", scalar.dtype());
    };

    let field_index = |name: &str| {
        fields
            .find(name)
            .ok_or_else(|| vortex_err!("SumV2 partial is missing the {name} field"))
    };
    let bool_field = |name: &str| -> VortexResult<bool> {
        match &values[field_index(name)?] {
            Some(ScalarValue::Bool(value)) => Ok(*value),
            _ => vortex_bail!("SumV2 partial field {name} must be a non-null bool"),
        }
    };

    let sum_index = field_index(SUM_FIELD)?;
    let sum_dtype = fields
        .field_dtypes()
        .nth(sum_index)
        .vortex_expect("field index from find is in bounds");
    validate_sum_field_dtype(sum_dtype, return_dtype)?;

    let sum = values[sum_index]
        .as_ref()
        .ok_or_else(|| vortex_err!("SumV2 partial sum must not be null"))?;

    Ok((
        sum,
        bool_field(IS_OVERFLOW_FIELD)?,
        bool_field(IS_EMPTY_FIELD)?,
    ))
}

fn validate_sum_field_dtype(sum_dtype: &DType, return_dtype: &DType) -> VortexResult<()> {
    vortex_ensure!(
        sum_dtype.nullability() == Nullability::NonNullable
            && sum_dtype.eq_ignore_nullability(return_dtype),
        "SumV2 partial value has dtype {}, expected {}",
        sum_dtype,
        return_dtype.as_nonnullable(),
    );
    Ok(())
}

fn checked_add_sum_state(
    state: &mut SumState,
    return_dtype: &DType,
    other: &Scalar,
) -> VortexResult<bool> {
    let other = other
        .value()
        .ok_or_else(|| vortex_err!("Can't extract present value from null scalar"))?;
    checked_add_sum_value(state, return_dtype, other)
}

fn checked_add_sum_value(
    state: &mut SumState,
    return_dtype: &DType,
    other: &ScalarValue,
) -> VortexResult<bool> {
    Ok(match (state, other) {
        (SumState::Unsigned(sum), ScalarValue::Primitive(other)) => {
            checked_add_u64(sum, other.cast::<u64>()?)
        }
        (SumState::Signed(sum), ScalarValue::Primitive(other)) => {
            checked_add_i64(sum, other.cast::<i64>()?)
        }
        (SumState::Float(sum), ScalarValue::Primitive(other)) => {
            *sum += other.cast::<f64>()?;
            false
        }
        (SumState::Decimal(value), ScalarValue::Decimal(other)) => {
            let dtype = return_dtype
                .as_decimal_opt()
                .vortex_expect("decimal sum result dtype");
            match value.checked_add(other) {
                Some(result) if result.fits_in_precision(*dtype) => {
                    *value = result;
                    false
                }
                Some(_) | None => true,
            }
        }
        (_, other) => vortex_bail!("SumV2 cannot add {other} to a {return_dtype} sum"),
    })
}

fn sum_v2_partial_dtype(sum_dtype: DType) -> DType {
    DType::Struct(sum_v2_partial_fields(sum_dtype), Nullability::Nullable)
}

fn sum_v2_partial_fields(sum_dtype: DType) -> StructFields {
    StructFields::new(
        FieldNames::from_iter([
            FieldName::from(SUM_FIELD),
            FieldName::from(IS_OVERFLOW_FIELD),
            FieldName::from(IS_EMPTY_FIELD),
        ]),
        vec![
            sum_dtype.as_nonnullable(),
            DType::Bool(Nullability::NonNullable),
            DType::Bool(Nullability::NonNullable),
        ],
    )
}

fn sum_state_scalar(
    partial: &SumV2Partial,
    return_dtype: &DType,
    nullability: Nullability,
) -> Scalar {
    match &partial.sum {
        SumState::Unsigned(value) => Scalar::primitive(*value, nullability),
        SumState::Signed(value) => Scalar::primitive(*value, nullability),
        SumState::Float(value) => Scalar::primitive(*value, nullability),
        SumState::Decimal(value) => Scalar::decimal(
            *value,
            *return_dtype
                .as_decimal_opt()
                .vortex_expect("decimal sum result dtype"),
            nullability,
        ),
    }
}

fn checked_add_u64(sum: &mut u64, value: u64) -> bool {
    match sum.checked_add(value) {
        Some(result) => {
            *sum = result;
            false
        }
        None => true,
    }
}

fn checked_add_i64(sum: &mut i64, value: i64) -> bool {
    match sum.checked_add(value) {
        Some(result) => {
            *sum = result;
            false
        }
        None => true,
    }
}

#[cfg(test)]
mod tests;