vortex_array/aggregate_fn/fns/sum_v2/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod grouped;
5
6pub(crate) use grouped::PrimitiveGroupedSumV2EncodingKernel;
7use vortex_error::VortexResult;
8use vortex_error::vortex_bail;
9use vortex_error::vortex_ensure;
10use vortex_error::vortex_err;
11use vortex_session::VortexSession;
12use vortex_session::registry::CachedId;
13
14use crate::ArrayRef;
15use crate::ArrayView;
16use crate::Canonical;
17use crate::Columnar;
18use crate::ExecutionCtx;
19use crate::aggregate_fn::Accumulator;
20use crate::aggregate_fn::AggregateFnId;
21use crate::aggregate_fn::AggregateFnVTable;
22use crate::aggregate_fn::DynAccumulator;
23use crate::aggregate_fn::NumericalAggregateOpts;
24use crate::aggregate_fn::fns::sum::Sum;
25use crate::aggregate_fn::fns::sum::SumState;
26use crate::aggregate_fn::fns::sum::accumulate_bool;
27use crate::aggregate_fn::fns::sum::accumulate_decimal;
28use crate::aggregate_fn::fns::sum::accumulate_primitive;
29use crate::aggregate_fn::fns::sum::make_zero_state;
30use crate::aggregate_fn::fns::sum::multiply_constant;
31use crate::arrays::Struct;
32use crate::arrays::struct_::StructArrayExt;
33use crate::builtins::ArrayBuiltins;
34use crate::dtype::DType;
35use crate::dtype::FieldName;
36use crate::dtype::FieldNames;
37use crate::dtype::Nullability;
38use crate::dtype::StructFields;
39use crate::expr::stats::Precision;
40use crate::expr::stats::Stat;
41use crate::expr::stats::StatsProviderExt;
42use crate::scalar::DecimalValue;
43use crate::scalar::Scalar;
44use crate::scalar_fn::fns::operators::Operator;
45use crate::validity::Validity;
46
47const SUM_FIELD: &str = "sum";
48const IS_OVERFLOW_FIELD: &str = "is_overflow";
49const IS_EMPTY_FIELD: &str = "is_empty";
50
51/// Return the SQL-style sum of an array.
52///
53/// Unlike [`sum`](crate::aggregate_fn::fns::sum::sum), this returns null when the array has no
54/// valid values.
55pub fn sum_v2(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Scalar> {
56    let mut acc = Accumulator::try_new(
57        SumV2,
58        NumericalAggregateOpts::default(),
59        array.dtype().clone(),
60    )?;
61    acc.accumulate(array, ctx)?;
62    acc.finish()
63}
64
65/// Sum an array, returning null when it has no valid values or if the sum overflows.
66///
67/// This aggregate intentionally has a distinct ID and partial representation from the legacy
68/// [`Sum`]. Keeping `vortex.sum` unchanged preserves the scalar partials stored by older Vortex
69/// files, while `SumV2` can use an explicit `{ sum, is_overflow, is_empty }` state.
70///
71/// NaN handling for float inputs is controlled by [`NumericalAggregateOpts`]. With `skip_nans`
72/// (the default), NaN values contribute nothing but still make the input non-empty. Otherwise,
73/// any NaN value poisons the sum to NaN.
74#[derive(Clone, Copy, Debug)]
75pub struct SumV2;
76
77impl AggregateFnVTable for SumV2 {
78    type Options = NumericalAggregateOpts;
79    type Partial = SumV2Partial;
80
81    fn id(&self) -> AggregateFnId {
82        static ID: CachedId = CachedId::new("vortex.sum_v2");
83        *ID
84    }
85
86    fn serialize(&self, options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
87        Ok(Some(options.serialize()))
88    }
89
90    fn deserialize(
91        &self,
92        metadata: &[u8],
93        _session: &VortexSession,
94    ) -> VortexResult<Self::Options> {
95        NumericalAggregateOpts::deserialize(metadata)
96    }
97
98    fn return_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
99        Sum.return_dtype(options, input_dtype)
100    }
101
102    fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
103        self.return_dtype(options, input_dtype)
104            .map(sum_v2_partial_dtype)
105    }
106
107    fn empty_partial(
108        &self,
109        options: &Self::Options,
110        input_dtype: &DType,
111    ) -> VortexResult<Self::Partial> {
112        let return_dtype = self
113            .return_dtype(options, input_dtype)
114            .ok_or_else(|| vortex_err!("Unsupported sum_v2 dtype: {}", input_dtype))?;
115        let sum = make_zero_state(&return_dtype);
116        Ok(SumV2Partial {
117            return_dtype,
118            sum,
119            is_overflow: false,
120            is_empty: true,
121            skip_nans: options.skip_nans,
122        })
123    }
124
125    fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> {
126        let (other_sum, other_is_overflow, other_is_empty) = decode_partial_scalar(other)?;
127        validate_sum_field_dtype(&other_sum, &partial.return_dtype)?;
128
129        if partial.is_overflow {
130            return Ok(());
131        }
132        if other_is_overflow {
133            partial.is_overflow = true;
134            partial.is_empty = false;
135            return Ok(());
136        }
137        if other_is_empty {
138            return Ok(());
139        }
140
141        partial.is_overflow = checked_add_sum_state(&mut partial.sum, &other_sum)?;
142        partial.is_empty = false;
143        Ok(())
144    }
145
146    fn to_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
147        Ok(Scalar::struct_(
148            sum_v2_partial_dtype(partial.return_dtype.clone()),
149            vec![
150                sum_state_scalar(partial, Nullability::NonNullable),
151                Scalar::bool(partial.is_overflow, Nullability::NonNullable),
152                Scalar::bool(partial.is_empty, Nullability::NonNullable),
153            ],
154        ))
155    }
156
157    fn reset(&self, partial: &mut Self::Partial) {
158        partial.sum = make_zero_state(&partial.return_dtype);
159        partial.is_overflow = false;
160        partial.is_empty = true;
161    }
162
163    fn is_saturated(&self, partial: &Self::Partial) -> bool {
164        partial.is_overflow || matches!(&partial.sum, SumState::Float(value) if value.is_nan())
165    }
166
167    fn try_accumulate(
168        &self,
169        partial: &mut Self::Partial,
170        batch: &ArrayRef,
171        _ctx: &mut ExecutionCtx,
172    ) -> VortexResult<bool> {
173        if partial.skip_nans || !matches!(&partial.sum, SumState::Float(_)) {
174            return Ok(false);
175        }
176
177        match batch.statistics().get_as::<u64>(Stat::NaNCount) {
178            Precision::Exact(0) => Ok(false),
179            Precision::Exact(_) => {
180                let SumState::Float(sum) = &mut partial.sum else {
181                    unreachable!("checked float sum state")
182                };
183                *sum = f64::NAN;
184                partial.is_empty = false;
185                Ok(true)
186            }
187            _ => Ok(false),
188        }
189    }
190
191    fn accumulate(
192        &self,
193        partial: &mut Self::Partial,
194        batch: &Columnar,
195        ctx: &mut ExecutionCtx,
196    ) -> VortexResult<()> {
197        if partial.is_overflow {
198            return Ok(());
199        }
200
201        if let Columnar::Constant(constant) = batch {
202            if !constant.scalar().is_null() && !constant.is_empty() {
203                partial.is_empty = false;
204            }
205            if partial.skip_nans
206                && constant
207                    .scalar()
208                    .as_primitive_opt()
209                    .is_some_and(|primitive| primitive.is_nan())
210            {
211                return Ok(());
212            }
213            if let Some(product) =
214                multiply_constant(constant.scalar(), constant.len(), &partial.return_dtype)?
215            {
216                if product.is_null() {
217                    partial.is_overflow = true;
218                    partial.is_empty = false;
219                } else {
220                    partial.is_overflow = checked_add_sum_state(&mut partial.sum, &product)?;
221                    partial.is_empty = false;
222                }
223            }
224            return Ok(());
225        }
226
227        let any_valid = partial.is_empty && has_valid_value(batch, ctx)?;
228        let result = match batch {
229            Columnar::Canonical(canonical) => match canonical {
230                Canonical::Primitive(array) => {
231                    accumulate_primitive(&mut partial.sum, array, ctx, partial.skip_nans)
232                }
233                Canonical::Bool(array) => accumulate_bool(&mut partial.sum, array, ctx),
234                Canonical::Decimal(array) => accumulate_decimal(&mut partial.sum, array, ctx),
235                _ => vortex_bail!("Unsupported canonical type for sum_v2: {}", batch.dtype()),
236            },
237            Columnar::Constant(_) => unreachable!(),
238        };
239
240        if any_valid {
241            partial.is_empty = false;
242        }
243        if result? {
244            partial.is_overflow = true;
245            partial.is_empty = false;
246        }
247        Ok(())
248    }
249
250    fn finalize(&self, partials: ArrayRef) -> VortexResult<ArrayRef> {
251        if let Some(partials) = partials.as_opt::<Struct>() {
252            return finalize_struct(partials);
253        }
254
255        let sum = partials.get_item(SUM_FIELD)?;
256        let is_invalid = partials
257            .get_item(IS_OVERFLOW_FIELD)?
258            .binary(partials.get_item(IS_EMPTY_FIELD)?, Operator::Or)?
259            .fill_null(true)?;
260        sum.mask(is_invalid.not()?)
261    }
262
263    fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
264        if partial.is_overflow || partial.is_empty {
265            return Ok(Scalar::null(partial.return_dtype.as_nullable()));
266        }
267        Ok(sum_state_scalar(partial, Nullability::Nullable))
268    }
269}
270
271fn finalize_struct(partials: ArrayView<'_, Struct>) -> VortexResult<ArrayRef> {
272    let sum = partials.unmasked_field_by_name(SUM_FIELD)?.clone();
273    let is_overflow = partials.unmasked_field_by_name(IS_OVERFLOW_FIELD)?.clone();
274    let is_empty = partials.unmasked_field_by_name(IS_EMPTY_FIELD)?.clone();
275
276    let is_invalid = is_overflow
277        .binary(is_empty, Operator::Or)?
278        .fill_null(true)?;
279    let mut is_valid = is_invalid.not()?;
280    match partials.struct_validity() {
281        Validity::NonNullable | Validity::AllValid => {}
282        validity => {
283            is_valid = is_valid.binary(validity.to_array(partials.len()), Operator::And)?;
284        }
285    }
286
287    sum.mask(is_valid)
288}
289
290/// In-memory state for SumV2 accumulation.
291pub struct SumV2Partial {
292    return_dtype: DType,
293    sum: SumState,
294    is_overflow: bool,
295    is_empty: bool,
296    skip_nans: bool,
297}
298
299fn has_valid_value(batch: &Columnar, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
300    let (validity, len) = match batch {
301        Columnar::Canonical(Canonical::Primitive(array)) => {
302            (array.as_ref().validity()?, array.as_ref().len())
303        }
304        Columnar::Canonical(Canonical::Bool(array)) => {
305            (array.as_ref().validity()?, array.as_ref().len())
306        }
307        Columnar::Canonical(Canonical::Decimal(array)) => {
308            (array.as_ref().validity()?, array.as_ref().len())
309        }
310        Columnar::Canonical(_) => return Ok(false),
311        Columnar::Constant(constant) => {
312            return Ok(!constant.is_empty() && !constant.scalar().is_null());
313        }
314    };
315    Ok(validity.execute_mask(len, ctx)?.true_count() > 0)
316}
317
318fn decode_partial_scalar(scalar: Scalar) -> VortexResult<(Scalar, bool, bool)> {
319    vortex_ensure!(!scalar.is_null(), "SumV2 partial must not be null");
320
321    let Some(fields) = scalar.as_struct_opt() else {
322        vortex_bail!("SumV2 partial must be a struct, got {}", scalar.dtype());
323    };
324    let sum = fields
325        .field(SUM_FIELD)
326        .ok_or_else(|| vortex_err!("SumV2 partial is missing the sum field"))?;
327    let is_overflow = bool::try_from(
328        &fields
329            .field(IS_OVERFLOW_FIELD)
330            .ok_or_else(|| vortex_err!("SumV2 partial is missing the is_overflow field"))?,
331    )?;
332    let is_empty = bool::try_from(
333        &fields
334            .field(IS_EMPTY_FIELD)
335            .ok_or_else(|| vortex_err!("SumV2 partial is missing the is_empty field"))?,
336    )?;
337
338    Ok((sum, is_overflow, is_empty))
339}
340
341fn validate_sum_field_dtype(sum: &Scalar, return_dtype: &DType) -> VortexResult<()> {
342    vortex_ensure!(
343        sum.dtype().nullability() == Nullability::NonNullable
344            && sum.dtype().eq_ignore_nullability(return_dtype),
345        "SumV2 partial value has dtype {}, expected {}",
346        sum.dtype(),
347        return_dtype.as_nonnullable(),
348    );
349    Ok(())
350}
351
352fn checked_add_sum_state(state: &mut SumState, other: &Scalar) -> VortexResult<bool> {
353    Ok(match state {
354        SumState::Unsigned(sum) => checked_add_u64(sum, u64::try_from(other)?),
355        SumState::Signed(sum) => checked_add_i64(sum, i64::try_from(other)?),
356        SumState::Float(sum) => {
357            *sum += f64::try_from(other)?;
358            false
359        }
360        SumState::Decimal { value, dtype } => {
361            let other = DecimalValue::try_from(other)?;
362            match value.checked_add(&other) {
363                Some(result) if result.fits_in_precision(*dtype) => {
364                    *value = result;
365                    false
366                }
367                Some(_) | None => true,
368            }
369        }
370    })
371}
372
373fn sum_v2_partial_dtype(sum_dtype: DType) -> DType {
374    DType::Struct(sum_v2_partial_fields(sum_dtype), Nullability::Nullable)
375}
376
377fn sum_v2_partial_fields(sum_dtype: DType) -> StructFields {
378    StructFields::new(
379        FieldNames::from_iter([
380            FieldName::from(SUM_FIELD),
381            FieldName::from(IS_OVERFLOW_FIELD),
382            FieldName::from(IS_EMPTY_FIELD),
383        ]),
384        vec![
385            sum_dtype.as_nonnullable(),
386            DType::Bool(Nullability::NonNullable),
387            DType::Bool(Nullability::NonNullable),
388        ],
389    )
390}
391
392fn sum_state_scalar(partial: &SumV2Partial, nullability: Nullability) -> Scalar {
393    match &partial.sum {
394        SumState::Unsigned(value) => Scalar::primitive(*value, nullability),
395        SumState::Signed(value) => Scalar::primitive(*value, nullability),
396        SumState::Float(value) => Scalar::primitive(*value, nullability),
397        SumState::Decimal { value, dtype } => Scalar::decimal(*value, *dtype, nullability),
398    }
399}
400
401fn checked_add_u64(sum: &mut u64, value: u64) -> bool {
402    match sum.checked_add(value) {
403        Some(result) => {
404            *sum = result;
405            false
406        }
407        None => true,
408    }
409}
410
411fn checked_add_i64(sum: &mut i64, value: i64) -> bool {
412    match sum.checked_add(value) {
413        Some(result) => {
414            *sum = result;
415            false
416        }
417        None => true,
418    }
419}
420
421#[cfg(test)]
422mod tests;
vortex_array/aggregate_fn/fns/sum_v2/mod.rs

vortex_array/aggregate_fn/fns/sum_v2/
mod.rs
