Use aggregate descriptors for zoned stats

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

Author: gatesn

vortex-array/src/scalar_fn/fns/stat.rs

@@ -56,7 +56,7 @@ impl Display for StatOptions {
     }
 }
 
-/// Scalar function that broadcasts a stored aggregate statistic over the input rows.
+/// Scalar function that broadcasts a stored aggregate partial over the input rows.
 ///
 /// The only current consumer is **row-wise pruning**: substituting `stat(col, agg)` into a
 /// predicate produces a cheap, row-aligned approximation whose constant runs let downstream
@@ -124,7 +124,7 @@ impl ScalarFnVTable for StatFn {
 }
 
 fn stat_dtype(aggregate_fn: &AggregateFnRef, input_dtype: &DType) -> VortexResult<DType> {
-    let Some(dtype) = aggregate_fn.return_dtype(input_dtype) else {
+    let Some(dtype) = aggregate_fn.state_dtype(input_dtype) else {
         vortex_bail!(
             "Aggregate function {} does not support input dtype {}",
             aggregate_fn,

vortex-layout/public-api.lock

@@ -766,13 +766,9 @@ pub mod vortex_layout::layouts::zoned::writer
 
 pub struct vortex_layout::layouts::zoned::writer::ZonedLayoutOptions
 
-pub vortex_layout::layouts::zoned::writer::ZonedLayoutOptions::block_size: usize
-
-pub vortex_layout::layouts::zoned::writer::ZonedLayoutOptions::concurrency: usize
+pub vortex_layout::layouts::zoned::writer::ZonedLayoutOptions::aggregate_fns: alloc::sync::Arc<[vortex_array::aggregate_fn::erased::AggregateFnRef]>
 
-pub vortex_layout::layouts::zoned::writer::ZonedLayoutOptions::max_variable_length_statistics_size: usize
-
-pub vortex_layout::layouts::zoned::writer::ZonedLayoutOptions::stats: alloc::sync::Arc<[vortex_array::expr::stats::Stat]>
+pub vortex_layout::layouts::zoned::writer::ZonedLayoutOptions::block_size: usize
 
 impl core::default::Default for vortex_layout::layouts::zoned::writer::ZonedLayoutOptions
 
@@ -800,7 +796,7 @@ pub fn vortex_layout::layouts::zoned::zone_map::ZoneMap::dtype_for_stats_table(&
 
 pub fn vortex_layout::layouts::zoned::zone_map::ZoneMap::prune(&self, &vortex_array::expr::expression::Expression, &vortex_session::VortexSession) -> vortex_error::VortexResult<vortex_mask::Mask>
 
-pub fn vortex_layout::layouts::zoned::zone_map::ZoneMap::try_new(vortex_array::dtype::DType, vortex_array::arrays::struct_::vtable::StructArray, alloc::sync::Arc<[vortex_array::expr::stats::Stat]>, u64, u64) -> vortex_error::VortexResult<Self>
+pub fn vortex_layout::layouts::zoned::zone_map::ZoneMap::try_new(vortex_array::dtype::DType, vortex_array::arrays::struct_::vtable::StructArray, alloc::sync::Arc<[vortex_array::aggregate_fn::erased::AggregateFnRef]>, u64, u64) -> vortex_error::VortexResult<Self>
 
 impl core::clone::Clone for vortex_layout::layouts::zoned::zone_map::ZoneMap
 
@@ -848,11 +844,11 @@ pub struct vortex_layout::layouts::zoned::ZonedLayout
 
 impl vortex_layout::layouts::zoned::ZonedLayout
 
-pub fn vortex_layout::layouts::zoned::ZonedLayout::new(vortex_layout::LayoutRef, vortex_layout::LayoutRef, usize, alloc::sync::Arc<[vortex_array::expr::stats::Stat]>) -> Self
+pub fn vortex_layout::layouts::zoned::ZonedLayout::new(vortex_layout::LayoutRef, vortex_layout::LayoutRef, usize, alloc::sync::Arc<[vortex_array::aggregate_fn::erased::AggregateFnRef]>) -> Self
 
 pub fn vortex_layout::layouts::zoned::ZonedLayout::nzones(&self) -> usize
 
-pub fn vortex_layout::layouts::zoned::ZonedLayout::present_stats(&self) -> &alloc::sync::Arc<[vortex_array::expr::stats::Stat]>
+pub fn vortex_layout::layouts::zoned::ZonedLayout::present_aggregates(&self) -> &alloc::sync::Arc<[alloc::string::String]>
 
 pub fn vortex_layout::layouts::zoned::ZonedLayout::zone_len(&self) -> usize
 

vortex-layout/src/layouts/zoned/builder.rs

@@ -12,6 +12,7 @@ use vortex_array::ExecutionCtx;
 use vortex_array::IntoArray;
 use vortex_array::LEGACY_SESSION;
 use vortex_array::VortexSessionExecute;
+use vortex_array::aggregate_fn::AggregateFnRef;
 use vortex_array::aggregate_fn::fns::sum::sum;
 use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::StructArray;
@@ -25,7 +26,6 @@ use vortex_array::dtype::Nullability;
 use vortex_array::dtype::PType;
 use vortex_array::expr::stats::Precision;
 use vortex_array::expr::stats::Stat;
-use vortex_array::expr::stats::StatsProvider;
 use vortex_array::scalar::Scalar;
 use vortex_array::scalar::ScalarTruncation;
 use vortex_array::scalar::lower_bound;
@@ -35,9 +35,12 @@ use vortex_array::validity::Validity;
 use vortex_buffer::BufferString;
 use vortex_buffer::ByteBuffer;
 use vortex_error::VortexResult;
+use vortex_error::vortex_ensure_eq;
 
 use crate::layouts::zoned::schema::MAX_IS_TRUNCATED;
 use crate::layouts::zoned::schema::MIN_IS_TRUNCATED;
+use crate::layouts::zoned::schema::aggregate_descriptor;
+use crate::layouts::zoned::schema::aggregate_state_dtype;
 
 /// Accumulates write-time statistics for each logical zone.
 pub struct StatsAccumulator {
@@ -67,18 +70,6 @@ impl StatsAccumulator {
         }
     }
 
-    pub fn push_chunk_without_compute(&mut self, array: &ArrayRef) -> VortexResult<()> {
-        for builder in &mut self.builders {
-            if let Some(Precision::Exact(value)) = array.statistics().get(builder.stat()) {
-                builder.append_scalar(value.cast(&value.dtype().as_nullable())?)?;
-            } else {
-                builder.append_null();
-            }
-        }
-        self.length += 1;
-        Ok(())
-    }
-
     pub fn push_chunk(&mut self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> {
         for builder in &mut self.builders {
             if let Some(value) = array.statistics().compute_stat(builder.stat(), ctx)? {
@@ -165,6 +156,102 @@ impl StatsAccumulator {
     }
 }
 
+/// Accumulates aggregate-function partials for each logical zone.
+pub(crate) struct AggregateStatsAccumulator {
+    builders: Vec<AggregateStatsArrayBuilder>,
+    length: usize,
+}
+
+impl AggregateStatsAccumulator {
+    pub(crate) fn new(dtype: &DType, aggregate_fns: &[AggregateFnRef]) -> Self {
+        let builders = aggregate_fns
+            .iter()
+            .filter_map(|aggregate_fn| {
+                aggregate_state_dtype(dtype, aggregate_fn).map(|partial_dtype| {
+                    AggregateStatsArrayBuilder::new(
+                        aggregate_fn.clone(),
+                        &partial_dtype.as_nullable(),
+                        1024,
+                    )
+                })
+            })
+            .collect::<Vec<_>>();
+
+        Self {
+            builders,
+            length: 0,
+        }
+    }
+
+    pub(crate) fn aggregate_fns(&self) -> Arc<[AggregateFnRef]> {
+        self.builders
+            .iter()
+            .map(|builder| builder.aggregate_fn.clone())
+            .collect::<Vec<_>>()
+            .into()
+    }
+
+    pub(crate) fn push_partials(&mut self, partials: Vec<Scalar>) -> VortexResult<()> {
+        vortex_ensure_eq!(
+            partials.len(),
+            self.builders.len(),
+            "aggregate partial count must match zone stats builder count"
+        );
+
+        for (builder, value) in self.builders.iter_mut().zip_eq(partials) {
+            builder.append_scalar(value)?;
+        }
+        self.length += 1;
+        Ok(())
+    }
+
+    pub(crate) fn as_array(
+        &mut self,
+    ) -> VortexResult<Option<(StructArray, Arc<[AggregateFnRef]>)>> {
+        let mut names = Vec::new();
+        let mut fields = Vec::new();
+        let mut aggregate_fns = Vec::new();
+
+        for builder in self
+            .builders
+            .iter_mut()
+            .sorted_unstable_by(|lhs, rhs| lhs.descriptor.cmp(&rhs.descriptor))
+        {
+            let values = builder.finish();
+
+            if values.all_invalid()? {
+                continue;
+            }
+
+            aggregate_fns.push(builder.aggregate_fn.clone());
+            names.extend(values.names);
+            fields.extend(values.arrays);
+        }
+
+        if names.is_empty() {
+            return Ok(None);
+        }
+
+        let array = StructArray::try_new(names.into(), fields, self.length, Validity::NonNullable)?;
+        Ok(Some((array, aggregate_fns.into())))
+    }
+}
+
+pub(crate) fn aggregate_partials(
+    array: &ArrayRef,
+    aggregate_fns: &[AggregateFnRef],
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<Vec<Scalar>> {
+    aggregate_fns
+        .iter()
+        .map(|aggregate_fn| {
+            let mut accumulator = aggregate_fn.accumulator(array.dtype())?;
+            accumulator.accumulate(array, ctx)?;
+            accumulator.partial_scalar()
+        })
+        .collect()
+}
+
 fn stats_builder_with_capacity(
     stat: Stat,
     dtype: &DType,
@@ -203,6 +290,35 @@ fn stats_builder_with_capacity(
     }
 }
 
+struct AggregateStatsArrayBuilder {
+    aggregate_fn: AggregateFnRef,
+    descriptor: String,
+    dtype: DType,
+    builder: Box<dyn ArrayBuilder>,
+}
+
+impl AggregateStatsArrayBuilder {
+    fn new(aggregate_fn: AggregateFnRef, dtype: &DType, capacity: usize) -> Self {
+        Self {
+            descriptor: aggregate_descriptor(&aggregate_fn),
+            aggregate_fn,
+            dtype: dtype.clone(),
+            builder: builder_with_capacity(dtype, capacity),
+        }
+    }
+
+    fn append_scalar(&mut self, value: Scalar) -> VortexResult<()> {
+        self.builder.append_scalar(&value.cast(&self.dtype)?)
+    }
+
+    fn finish(&mut self) -> NamedArrays {
+        NamedArrays {
+            names: vec![self.descriptor.clone().into()],
+            arrays: vec![self.builder.finish()],
+        }
+    }
+}
+
 /// Arrays with their associated names, reduced version of a `StructArray`.
 struct NamedArrays {
     names: Vec<FieldName>,