Centralize aggregate stat bridge

Signed-off-by: Nicholas Gates <nick@nickgates.com>

Author: gatesn

vortex-array/public-api.lock

@@ -15,6 +15,8 @@ use crate::aggregate_fn::session::AggregateFnSessionExt;
 use crate::columnar::AnyColumnar;
 use crate::dtype::DType;
 use crate::executor::max_iterations;
+use crate::expr::stats::Precision;
+use crate::expr::stats::StatsProvider;
 use crate::scalar::Scalar;
 
 /// Reference-counted type-erased accumulator.
@@ -116,17 +118,20 @@ impl<V: AggregateFnVTable> DynAccumulator for Accumulator<V> {
             batch.dtype()
         );
 
-        // 0. Stats-driven shortcut: if the aggregate can be derived directly from the batch's
-        //    cached statistics, use that and skip both kernel dispatch and decode. This is the
-        //    only layer that consults `batch.statistics()`; encoding kernels must not.
-        if let Some(result) = self.vtable.try_partial_from_stats(batch)? {
+        // 0. Legacy stats bridge: if this aggregate is still cached under a legacy Stat slot,
+        //    consume that exact stat before kernel dispatch or decode.
+        if let Some(stat) = self.vtable.maybe_stat()
+            && let Some(Precision::Exact(partial)) = batch.statistics().get(stat)
+        {
             vortex_ensure!(
-                result.dtype() == &self.partial_dtype,
-                "Aggregate try_partial_from_stats returned {}, expected {}",
-                result.dtype(),
+                partial.dtype() == &self.partial_dtype,
+                "Aggregate {} read legacy stat {} with dtype {}, expected {}",
+                self.aggregate_fn,
+                stat,
+                partial.dtype(),
                 self.partial_dtype,
             );
-            self.vtable.combine_partials(&mut self.partial, result)?;
+            self.vtable.combine_partials(&mut self.partial, partial)?;
             return Ok(());
         }
 

vortex-array/src/aggregate_fn/accumulator.rs

@@ -6,7 +6,6 @@ mod primitive;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
-use vortex_error::vortex_err;
 
 use self::primitive::accumulate_primitive;
 use crate::ArrayRef;
@@ -23,20 +22,13 @@ use crate::dtype::Nullability::NonNullable;
 use crate::dtype::PType;
 use crate::expr::stats::Precision;
 use crate::expr::stats::Stat;
-use crate::expr::stats::StatsProvider;
 use crate::scalar::Scalar;
 use crate::scalar::ScalarValue;
 
 /// Return the number of NaN values in an array.
 ///
 /// See [`NanCount`] for details.
 pub fn nan_count(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<usize> {
-    // Short-circuit using cached array statistics.
-    if let Some(Precision::Exact(nan_count_scalar)) = array.statistics().get(Stat::NaNCount) {
-        return usize::try_from(&nan_count_scalar)
-            .map_err(|e| vortex_err!("Failed to convert NaN count stat to usize: {e}"));
-    }
-
     // Short-circuit for non-float types.
     if NanCount
         .return_dtype(&EmptyOptions, array.dtype())
@@ -132,6 +124,10 @@ impl AggregateFnVTable for NanCount {
         false
     }
 
+    fn maybe_stat(&self) -> Option<Stat> {
+        Some(Stat::NaNCount)
+    }
+
     fn accumulate(
         &self,
         partial: &mut Self::Partial,

vortex-array/src/aggregate_fn/fns/nan_count/mod.rs

@@ -32,19 +32,13 @@ use crate::dtype::Nullability;
 use crate::dtype::PType;
 use crate::expr::stats::Precision;
 use crate::expr::stats::Stat;
-use crate::expr::stats::StatsProvider;
 use crate::scalar::DecimalValue;
 use crate::scalar::Scalar;
 
 /// Return the sum of an array.
 ///
 /// See [`Sum`] for details.
 pub fn sum(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Scalar> {
-    // Short-circuit using cached array statistics.
-    if let Some(Precision::Exact(sum_scalar)) = array.statistics().get(Stat::Sum) {
-        return Ok(sum_scalar);
-    }
-
     // Compute using Accumulator<Sum>.
     // TODO(ngates): we may want to wrap this three-step dance up into an extension crate maybe.
     let mut acc = Accumulator::try_new(Sum, EmptyOptions, array.dtype().clone())?;
@@ -213,6 +207,10 @@ impl AggregateFnVTable for Sum {
         }
     }
 
+    fn maybe_stat(&self) -> Option<Stat> {
+        Some(Stat::Sum)
+    }
+
     fn accumulate(
         &self,
         partial: &mut Self::Partial,

vortex-array/src/aggregate_fn/fns/sum/mod.rs

@@ -48,7 +48,6 @@ use crate::dtype::Nullability::NonNullable;
 use crate::dtype::PType;
 use crate::expr::stats::Precision;
 use crate::expr::stats::Stat;
-use crate::expr::stats::StatsProvider;
 use crate::scalar::Scalar;
 use crate::scalar::ScalarValue;
 
@@ -63,13 +62,6 @@ pub fn uncompressed_size_in_bytes(array: &ArrayRef, ctx: &mut ExecutionCtx) -> V
 }
 
 fn uncompressed_size_in_bytes_u64(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<u64> {
-    if let Some(Precision::Exact(size_scalar)) =
-        array.statistics().get(Stat::UncompressedSizeInBytes)
-    {
-        return u64::try_from(&size_scalar)
-            .map_err(|e| vortex_err!("Failed to convert uncompressed size stat to u64: {e}"));
-    }
-
     let mut acc =
         Accumulator::try_new(UncompressedSizeInBytes, EmptyOptions, array.dtype().clone())?;
     acc.accumulate(array, ctx)?;
@@ -150,15 +142,8 @@ impl AggregateFnVTable for UncompressedSizeInBytes {
         false
     }
 
-    fn try_partial_from_stats(&self, batch: &ArrayRef) -> VortexResult<Option<Scalar>> {
-        let Some(Precision::Exact(size_scalar)) =
-            batch.statistics().get(Stat::UncompressedSizeInBytes)
-        else {
-            return Ok(None);
-        };
-        let size = u64::try_from(&size_scalar)
-            .map_err(|e| vortex_err!("Failed to convert uncompressed size stat to u64: {e}"))?;
-        Ok(Some(Scalar::primitive(size, NonNullable)))
+    fn maybe_stat(&self) -> Option<Stat> {
+        Some(Stat::UncompressedSizeInBytes)
     }
 
     fn accumulate(

vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs

@@ -18,6 +18,7 @@ use crate::aggregate_fn::AggregateFn;
 use crate::aggregate_fn::AggregateFnId;
 use crate::aggregate_fn::AggregateFnRef;
 use crate::dtype::DType;
+use crate::expr::stats::Stat;
 use crate::scalar::Scalar;
 
 /// Defines the interface for aggregate function vtables.
@@ -102,23 +103,12 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync {
     /// final result is fully determined.
     fn is_saturated(&self, state: &Self::Partial) -> bool;
 
-    /// Try to derive a partial scalar from the batch's cached statistics, before any
-    /// kernel dispatch or canonicalization.
+    /// Return the legacy [`Stat`] slot that stores this aggregate, if one exists.
     ///
-    /// Returns `Some(partial_scalar)` if the answer can be read directly from `batch.statistics()`,
-    /// otherwise `Ok(None)` to fall through to the rest of dispatch. The returned scalar must
-    /// have the dtype reported by `partial_dtype`.
-    ///
-    /// This is the single place stats-based shortcuts live; encoding kernels must not consult
-    /// stats themselves. Runs first so that an upstream producer who pre-populates the relevant
-    /// stat (e.g. a layout reader hydrating `Stat::UncompressedSizeInBytes` from file metadata)
-    /// can skip both kernel dispatch and decode.
-    ///
-    /// TODO: this hook may be removed once `ArrayStats` stores aggregate partials internally —
-    /// at that point stat-driven shortcuts can be resolved automatically by the dispatch layer
-    /// without each aggregate vtable opting in.
-    fn try_partial_from_stats(&self, _batch: &ArrayRef) -> VortexResult<Option<Scalar>> {
-        Ok(None)
+    /// This is a temporary bridge while some aggregate partials are still cached under the
+    /// legacy [`Stat`] enum rather than by aggregate function identity.
+    fn maybe_stat(&self) -> Option<Stat> {
+        None
     }
 
     /// Try to accumulate the raw array before decompression.