Push decimal aggregates through Widen

vortex-array/src/arrays/widen/compute/aggregate.rs

@@ -1,12 +1,28 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+//! Aggregates over the narrow child of a [`Widen`] array.
+//!
+//! Results must match aggregation over the widened values. Widening preserves order and equality,
+//! but sums also depend on the logical dtype's precision and overflow rules.
+
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 
 use crate::ArrayRef;
+use crate::Columnar;
 use crate::ExecutionCtx;
+use crate::aggregate_fn::AggregateDTypes;
 use crate::aggregate_fn::AggregateFnRef;
+use crate::aggregate_fn::AggregateFnVTable;
+use crate::aggregate_fn::NumericalAggregateOpts;
+use crate::aggregate_fn::fns::is_constant::IsConstant;
+use crate::aggregate_fn::fns::is_sorted::IsSorted;
+use crate::aggregate_fn::fns::max::Max;
+use crate::aggregate_fn::fns::min::Min;
+use crate::aggregate_fn::fns::min_max::MinMax;
+use crate::aggregate_fn::fns::min_max::MinMaxResult;
+use crate::aggregate_fn::fns::min_max::min_max;
 use crate::aggregate_fn::fns::sum::Sum;
 use crate::aggregate_fn::kernels::DynAggregateKernel;
 use crate::arrays::Widen;
@@ -15,11 +31,7 @@ use crate::dtype::DType;
 use crate::expr::stats::Stat;
 use crate::scalar::Scalar;
 
-/// Runs an aggregate of a [`Widen`] on its values and widens the partial to the logical dtype.
-///
-/// The values run the same accumulator code as the canonical array, so edge cases such as an empty
-/// batch or a single null match by construction, and the values keep their own cached statistics
-/// and kernels. [`Sum`] is served only when [`is_exact_sum`] holds.
+/// Computes aggregate partials from a [`Widen`]'s child without materializing widened values.
 #[derive(Debug)]
 pub(crate) struct WidenArrayAggregate;
 
@@ -34,27 +46,123 @@ impl DynAggregateKernel for WidenArrayAggregate {
             return Ok(None);
         };
 
-        // NB: Exactness rests on integer decoding. A decimal `Widen` (CD4 in #10733) needs its own
-        // argument, because `Sum` checks a decimal for overflow against the precision of its input.
-        if !batch.dtype().is_int() {
+        let values = widen.values();
+        let logical_dtype = batch.dtype();
+
+        if logical_dtype.is_decimal() {
+            return aggregate_decimal(aggregate_fn, values, logical_dtype, ctx);
+        }
+
+        if !logical_dtype.is_int() {
             return Ok(None);
         }
 
-        let values = widen.values();
         if aggregate_fn.is::<Sum>() && !is_exact_sum(aggregate_fn, batch, values) {
             return Ok(None);
         }
 
-        let mut accumulator = aggregate_fn.accumulator(values.dtype())?;
-        accumulator.accumulate(values, ctx)?;
+        let partial = aggregate_values(aggregate_fn, values, logical_dtype, ctx)?;
+
+        Ok(Some(partial))
+    }
+}
+
+/// Selects decimal aggregates whose child values produce the canonical partial.
+fn aggregate_decimal(
+    aggregate_fn: &AggregateFnRef,
+    values: &ArrayRef,
+    logical_dtype: &DType,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<Option<Scalar>> {
+    if let Some(options) = aggregate_fn.as_opt::<Sum>() {
+        // This kernel returns a batch total. Only the Stat::Sum path sums each batch from zero.
+        // Other variants add to the running sum and can overflow at different points.
+        if Stat::from_aggregate_fn(aggregate_fn) != Some(Stat::Sum) {
+            return Ok(None);
+        }
+
+        let partial = decimal_sum_partial(values, options, logical_dtype, ctx)?;
+
+        return Ok(Some(partial));
+    }
+
+    if !(aggregate_fn.is::<MinMax>()
+        || aggregate_fn.is::<Min>()
+        || aggregate_fn.is::<Max>()
+        || aggregate_fn.is::<IsConstant>()
+        || aggregate_fn.is::<IsSorted>())
+    {
+        return Ok(None);
+    }
+
+    let partial = aggregate_values(aggregate_fn, values, logical_dtype, ctx)?;
 
-        let partial_dtype = aggregate_fn.state_dtype(batch.dtype()).vortex_expect(
-            "an aggregate that accepts the values accepts the logical dtype, which is a wider \
-             integer of the same signedness",
-        );
+    Ok(Some(partial))
+}
 
-        Ok(Some(accumulator.flush()?.cast(&partial_dtype)?))
+/// Computes a child partial and casts it to the logical aggregate state dtype.
+fn aggregate_values(
+    aggregate_fn: &AggregateFnRef,
+    values: &ArrayRef,
+    logical_dtype: &DType,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<Scalar> {
+    let partial_dtype = aggregate_fn.state_dtype(logical_dtype).vortex_expect(
+        "the served aggregate accepts the logical dtype when it accepts the narrower values",
+    );
+
+    if let Some(options) = aggregate_fn.as_opt::<MinMax>() {
+        return min_max_partial(values, *options, logical_dtype, partial_dtype, ctx);
     }
+
+    let mut accumulator = aggregate_fn.accumulator(values.dtype())?;
+    accumulator.accumulate(values, ctx)?;
+
+    accumulator.flush()?.cast(&partial_dtype)
+}
+
+/// Computes a decimal [`Sum`] batch partial without materializing the widened array.
+///
+/// Checked arithmetic overflow or a batch total outside the result precision produces null.
+/// Null rows are skipped, and empty or all-null batches produce zero.
+fn decimal_sum_partial(
+    values: &ArrayRef,
+    options: &NumericalAggregateOpts,
+    logical_dtype: &DType,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<Scalar> {
+    // Sum converts each stored value to the accumulator width selected by these logical dtypes
+    // before adding it. The child width never limits the sum. Widen preserves values, scale, and
+    // validity, so the additions, precision checks, and output dtype match the widened array.
+    let dtypes = AggregateDTypes::try_new(&Sum, options, logical_dtype.clone())?;
+    let args = dtypes.args(options);
+    let values = values.clone().execute::<Columnar>(ctx)?;
+
+    let mut partial = Sum.empty_partial(args)?;
+    Sum.accumulate(args, &mut partial, &values, ctx)?;
+
+    Sum.to_scalar(args, &partial)
+}
+
+/// Builds a logical [`MinMax`] partial from the child's cached or computed extrema.
+///
+/// [`min_max`] uses exact cached extrema before scanning the child. Extrema have the non-nullable
+/// logical dtype, even for nullable inputs. A batch with no valid values has a null partial.
+fn min_max_partial(
+    values: &ArrayRef,
+    options: NumericalAggregateOpts,
+    logical_dtype: &DType,
+    partial_dtype: DType,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<Scalar> {
+    let Some(MinMaxResult { min, max }) = min_max(values, ctx, options)? else {
+        return Ok(Scalar::null(partial_dtype));
+    };
+
+    let element_dtype = logical_dtype.as_nonnullable();
+    let extrema = vec![min.cast(&element_dtype)?, max.cast(&element_dtype)?];
+
+    Ok(Scalar::struct_(partial_dtype, extrema))
 }
 
 /// Returns whether the [`Sum`] of the values of `batch` equals the sum of the canonical batch.

vortex-array/src/arrays/widen/tests/aggregates.rs

@@ -1,9 +1,15 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+//! Tests for aggregate results and child-statistics reuse.
+//!
+//! Shared batch assertions compare integer and decimal [`Widen`] arrays with fresh canonical inputs.
+
+use std::fmt::Debug;
 use std::num::NonZeroUsize;
 
 use rstest::rstest;
+use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 
 use super::SESSION;
@@ -32,21 +38,27 @@ use crate::aggregate_fn::fns::last::Last;
 use crate::aggregate_fn::fns::max::Max;
 use crate::aggregate_fn::fns::min::Min;
 use crate::aggregate_fn::fns::min_max::MinMax;
+use crate::aggregate_fn::fns::min_max::make_minmax_dtype;
 use crate::aggregate_fn::fns::null_count::NullCount;
 use crate::aggregate_fn::fns::sum::Sum;
 use crate::aggregate_fn::fns::sum_v2::SumV2;
 use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes;
 use crate::arrays::ChunkedArray;
 use crate::arrays::ConstantArray;
+use crate::arrays::DecimalArray;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::Widen;
 use crate::arrays::WidenArray;
 use crate::builtins::ArrayBuiltins;
 use crate::dtype::DType;
+use crate::dtype::DecimalDType;
 use crate::dtype::Nullability;
 use crate::dtype::Nullability::NonNullable;
 use crate::dtype::Nullability::Nullable;
 use crate::dtype::PType;
+use crate::expr::stats::Precision;
+use crate::expr::stats::Stat;
+use crate::scalar::DecimalValue;
 use crate::scalar::Scalar;
 
 /// Every aggregate that applies to integers, and whether the `Widen` aggregate kernel serves it.
@@ -89,9 +101,8 @@ fn canonical_batch(
         .execute::<PrimitiveArray>(ctx)
 }
 
-/// Runs the kernel that the session registers for `aggregate` on `Widen` over `array`, and returns
-/// its partial, or `None` when no kernel is registered or the kernel declines.
-fn kernel_partial(
+/// Returns the registered [`Widen`] kernel's partial, or `None` if it is absent or declines.
+pub(super) fn kernel_partial(
     aggregate: &AggregateFnRef,
     array: &ArrayRef,
     ctx: &mut ExecutionCtx,
@@ -105,40 +116,40 @@ fn kernel_partial(
     }
 }
 
-/// Accumulates `batches` of `dtype` with `aggregate` once as [`WidenArray::encode`] encodes them
-/// and once as canonical arrays, and asserts that the running partials and the results are equal.
+/// Checks encoded aggregates against canonical inputs, including kernel selection and partials.
 ///
-/// On each batch that encodes to a `Widen`, the kernel must return a partial exactly when
-/// `served`, and the partial must equal the partial of a fresh accumulator on the canonical batch.
-/// Encoding and each accumulator get their own canonical batch, because both cache statistics on
-/// the array they read, and an accumulator reads a cached statistic instead of the data.
+/// `served` says whether the [`Widen`] kernel must return a partial. `build_batch` **must** return a
+/// fresh array with its own statistics on every call. Reusing arrays could let aggregation read a
+/// cached answer instead of testing the kernel.
 #[track_caller]
-fn assert_aggregate_batches(
+pub(super) fn assert_aggregate_batches<T: Debug, C: IntoArray>(
     aggregate: &AggregateFnRef,
     served: bool,
-    batches: &[Vec<Option<i64>>],
+    batches: &[T],
     dtype: &DType,
+    mut build_batch: impl FnMut(&T, &DType, &mut ExecutionCtx) -> VortexResult<C>,
+    mut encode: impl FnMut(C, &mut ExecutionCtx) -> VortexResult<ArrayRef>,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<()> {
     let mut actual = aggregate.accumulator(dtype)?;
     let mut expected = aggregate.accumulator(dtype)?;
 
     for values in batches {
-        let encoded = WidenArray::encode(canonical_batch(values, dtype, ctx)?, ctx)?;
+        let encoded = encode(build_batch(values, dtype, ctx)?, ctx)?;
 
         if encoded.is::<Widen>() {
             let partial = kernel_partial(aggregate, &encoded, ctx)?;
             assert_eq!(partial.is_some(), served, "{aggregate} on {values:?}");
 
             if let Some(partial) = partial {
                 let mut reference = aggregate.accumulator(dtype)?;
-                reference.accumulate(&canonical_batch(values, dtype, ctx)?.into_array(), ctx)?;
+                reference.accumulate(&build_batch(values, dtype, ctx)?.into_array(), ctx)?;
                 assert_eq!(partial, reference.flush()?, "{aggregate} on {values:?}");
             }
         }
 
         actual.accumulate(&encoded, ctx)?;
-        expected.accumulate(&canonical_batch(values, dtype, ctx)?.into_array(), ctx)?;
+        expected.accumulate(&build_batch(values, dtype, ctx)?.into_array(), ctx)?;
         assert_eq!(
             actual.partial_scalar()?,
             expected.partial_scalar()?,
@@ -230,7 +241,15 @@ fn test_aggregate_partials(
     let dtype = DType::Primitive(logical, nullability);
 
     for (aggregate, served) in aggregates() {
-        assert_aggregate_batches(&aggregate, served, &batches, &dtype, &mut ctx)?;
+        assert_aggregate_batches(
+            &aggregate,
+            served,
+            &batches,
+            &dtype,
+            |values, dtype, ctx| canonical_batch(values, dtype, ctx),
+            WidenArray::encode,
+            &mut ctx,
+        )?;
     }
 
     Ok(())
@@ -304,3 +323,46 @@ fn test_chunked_matches_canonical() -> VortexResult<()> {
 
     Ok(())
 }
+
+/// Cached extrema deliberately differ from the data, so the result proves the cache was read.
+#[rstest]
+#[case::integer(
+    PrimitiveArray::from_iter([1i8, 2, 3]).into_array(),
+    DType::Primitive(PType::I64, NonNullable),
+    Scalar::primitive(-100i8, NonNullable),
+    Scalar::primitive(100i8, NonNullable),
+)]
+#[case::decimal(
+    DecimalArray::from_option_iter([Some(1i16), None], DecimalDType::new(4, 2)).into_array(),
+    DType::Decimal(DecimalDType::new(18, 2), Nullable),
+    Scalar::decimal(DecimalValue::I16(-9_999), DecimalDType::new(4, 2), NonNullable),
+    Scalar::decimal(DecimalValue::I16(9_999), DecimalDType::new(4, 2), NonNullable),
+)]
+fn test_min_max_reads_values_extrema(
+    #[case] values: ArrayRef,
+    #[case] logical_dtype: DType,
+    #[case] min: Scalar,
+    #[case] max: Scalar,
+) -> VortexResult<()> {
+    let mut ctx = SESSION.create_execution_ctx();
+
+    for (stat, extremum) in [(Stat::Min, &min), (Stat::Max, &max)] {
+        let value = extremum.value().cloned();
+        let value = value.vortex_expect("the planted extremum is not null");
+        values.statistics().set(stat, Precision::Exact(value));
+    }
+
+    let array = WidenArray::try_new(values, logical_dtype.clone())?.into_array();
+    let aggregate = MinMax.bind(NumericalAggregateOpts::default());
+
+    let partial = kernel_partial(&aggregate, &array, &mut ctx)?;
+
+    let element_dtype = logical_dtype.as_nonnullable();
+    let expected = Scalar::struct_(
+        make_minmax_dtype(&logical_dtype),
+        vec![min.cast(&element_dtype)?, max.cast(&element_dtype)?],
+    );
+    assert_eq!(partial, Some(expected));
+
+    Ok(())
+}