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

@@ -31,12 +31,6 @@ use crate::scalar::ScalarValue;
 ///
 /// 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())
@@ -45,6 +39,12 @@ pub fn nan_count(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<usize
         return Ok(0);
     }
 
+    // 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 empty arrays or all-null arrays.
     if array.is_empty() || array.valid_count(ctx)? == 0 {
         return Ok(0);
@@ -132,6 +132,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

@@ -213,6 +213,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

@@ -150,15 +150,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.