Move ArrayStats to store Scalars instead.

interns the AggregateRef id to reduce the number of strings stored for
the aggregates.

Author: thorfour

vortex-array/src/expr/stats/mod.rs

@@ -25,7 +25,6 @@ pub use stat_bound::*;
 
 use crate::aggregate_fn;
 use crate::aggregate_fn::AggregateFnRef;
-use crate::aggregate_fn::AggregateFnSatisfaction;
 use crate::aggregate_fn::AggregateFnVTable;
 use crate::aggregate_fn::AggregateFnVTableExt;
 use crate::aggregate_fn::EmptyOptions;
@@ -240,18 +239,6 @@ impl Stat {
         None
     }
 
-    pub fn from_aggregate_fn_partial(
-        aggregate_fn: &AggregateFnRef,
-    ) -> Option<(Self, AggregateFnSatisfaction)> {
-        Self::all().find_map(|stat| {
-            let requested = stat.aggregate_fn()?;
-            match aggregate_fn.can_satisfy(&requested) {
-                AggregateFnSatisfaction::No => None,
-                satisfaction => Some((stat, satisfaction)),
-            }
-        })
-    }
-
     pub fn name(&self) -> &str {
         match self {
             Self::IsConstant => "is_constant",

vortex-datafusion/src/persistent/format.rs

@@ -45,24 +45,27 @@ use futures::stream;
 use object_store::ObjectMeta;
 use object_store::ObjectStore;
 use vortex::VortexSessionDefault;
+use vortex::aggregate_fn::AggregateFnVTableExt;
+use vortex::aggregate_fn::EmptyOptions;
+use vortex::aggregate_fn::NumericalAggregateOpts;
+use vortex::aggregate_fn::fns::is_constant::IsConstant;
+use vortex::aggregate_fn::fns::max::Max;
+use vortex::aggregate_fn::fns::min::Min;
+use vortex::aggregate_fn::fns::null_count::NullCount;
+use vortex::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes;
 use vortex::array::memory::MemorySessionExt;
-use vortex::dtype::DType;
 use vortex::dtype::FieldPath;
-use vortex::dtype::Nullability;
-use vortex::dtype::PType;
 use vortex::error::VortexExpect;
 use vortex::error::VortexResult;
 use vortex::error::vortex_err;
 use vortex::expr::stats::Precision;
-use vortex::expr::stats::Stat;
 use vortex::file::EOF_SIZE;
 use vortex::file::MAX_POSTSCRIPT_SIZE;
 use vortex::file::OpenOptionsSessionExt;
 use vortex::file::VORTEX_FILE_EXTENSION;
 use vortex::io::object_store::ObjectStoreReadAt;
 use vortex::io::session::RuntimeSessionExt;
 use vortex::scalar::Scalar;
-use vortex::scalar::ScalarValue as VortexScalarValue;
 use vortex::session::VortexSession;
 use vortex_arrow::ArrowSessionExt;
 
@@ -658,7 +661,7 @@ impl FileFormat for VortexFormat {
             for field in table_schema.fields().iter() {
                 // If the column does not exist, continue. This can happen if the schema has evolved
                 // but we have not yet updated the Vortex file.
-                let Some((stats_set, stats_dtype)) =
+                let Some((aggregates, _)) =
                     file_stats.get_by_path(&FieldPath::from_name(field.name().as_str()))
                 else {
                     // The default sets all statistics to `Precision<Absent>`.
@@ -667,35 +670,33 @@ impl FileFormat for VortexFormat {
                 };
 
                 // Update the total size in bytes.
-                let column_size =
-                    stats_set.get_as::<usize>(Stat::UncompressedSizeInBytes, &PType::U64.into());
+                let column_size = aggregates
+                    .get(&UncompressedSizeInBytes.bind(EmptyOptions))
+                    .and_then(|size| usize::try_from(&size).ok());
 
                 let min = scalar_stat_to_df(
-                    Stat::Min,
-                    stats_set.get(Stat::Min),
-                    stats_dtype,
+                    aggregates.get(&Min.bind(NumericalAggregateOpts::skip_nans())),
                     field.data_type(),
                 );
 
                 let max = scalar_stat_to_df(
-                    Stat::Max,
-                    stats_set.get(Stat::Max),
-                    stats_dtype,
+                    aggregates.get(&Max.bind(NumericalAggregateOpts::skip_nans())),
                     field.data_type(),
                 );
 
-                let null_count = stats_set.get_as::<usize>(Stat::NullCount, &PType::U64.into());
+                let null_count = aggregates
+                    .get(&NullCount.bind(EmptyOptions))
+                    .and_then(|count| usize::try_from(&count).ok());
 
                 column_statistics.push(ColumnStatistics {
                     null_count: null_count.to_df(),
                     min_value: min,
                     max_value: max,
                     sum_value: DFPrecision::Absent,
                     distinct_count: is_constant_to_distinct_count(
-                        stats_set.get_as::<bool>(
-                            Stat::IsConstant,
-                            &DType::Bool(Nullability::NonNullable),
-                        ),
+                        aggregates
+                            .get(&IsConstant.bind(EmptyOptions))
+                            .and_then(|is_constant| bool::try_from(&is_constant).ok()),
                     ),
                     byte_size: column_size.to_df(),
                 })
@@ -768,23 +769,11 @@ impl FileFormat for VortexFormat {
 }
 
 fn scalar_stat_to_df(
-    stat: Stat,
-    value: Precision<VortexScalarValue>,
-    stats_dtype: &DType,
+    value: Precision<Scalar>,
     target_dtype: &DataType,
 ) -> DFPrecision<DFScalarValue> {
-    let Some(stat_dtype) = stat.dtype(stats_dtype) else {
-        return DFPrecision::Absent;
-    };
-
     value
-        .and_then(|stat_value| {
-            let scalar = Scalar::try_new(stat_dtype, Some(stat_value))
-                .ok()?
-                .try_to_df()
-                .ok()?;
-            scalar.cast_to(target_dtype).ok()
-        })
+        .and_then(|scalar| scalar.try_to_df().ok()?.cast_to(target_dtype).ok())
         .to_df()
 }
 
@@ -849,26 +838,18 @@ mod tests {
     fn test_scalar_stat_to_df_cast(
         #[case] value: ScalarValue,
         #[case] expected: ScalarValue,
-        #[values(Stat::Min, Stat::Max)] stat: Stat,
         #[values(true, false)] exact: bool,
     ) -> VortexResult<()> {
         let session = VortexSession::default();
         let scalar = scalar_from_df(&value, &session)?;
-        let value = scalar
-            .value()
-            .cloned()
-            .ok_or_else(|| vortex_err!("expected non-null scalar"))?;
         let target_dtype = expected.data_type();
         let (value, expected) = if exact {
-            (Precision::Exact(value), DFPrecision::Exact(expected))
+            (Precision::Exact(scalar), DFPrecision::Exact(expected))
         } else {
-            (Precision::Inexact(value), DFPrecision::Inexact(expected))
+            (Precision::Inexact(scalar), DFPrecision::Inexact(expected))
         };
 
-        assert_eq!(
-            scalar_stat_to_df(stat, value, scalar.dtype(), &target_dtype),
-            expected
-        );
+        assert_eq!(scalar_stat_to_df(value, &target_dtype), expected);
         Ok(())
     }
 
@@ -886,16 +867,11 @@ mod tests {
     fn test_scalar_stat_to_df_failed_cast(
         #[case] value: ScalarValue,
         #[case] target_dtype: DataType,
-        #[values(Stat::Min, Stat::Max)] stat: Stat,
     ) -> VortexResult<()> {
         let session = VortexSession::default();
         let scalar = scalar_from_df(&value, &session)?;
-        let value = scalar
-            .value()
-            .cloned()
-            .ok_or_else(|| vortex_err!("expected non-null scalar"))?;
         assert_eq!(
-            scalar_stat_to_df(stat, Precision::Exact(value), scalar.dtype(), &target_dtype),
+            scalar_stat_to_df(Precision::Exact(scalar), &target_dtype),
             DFPrecision::Absent
         );
         Ok(())

vortex-duckdb/src/column_statistics.rs

@@ -1,12 +1,19 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex::array::stats::StatsSet;
+use vortex::aggregate_fn::AggregateFnRef;
+use vortex::aggregate_fn::AggregateFnVTableExt;
+use vortex::aggregate_fn::EmptyOptions;
+use vortex::aggregate_fn::NumericalAggregateOpts;
+use vortex::aggregate_fn::fns::max::Max;
+use vortex::aggregate_fn::fns::min::Min;
+use vortex::aggregate_fn::fns::null_count::NullCount;
+use vortex::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes;
 use vortex::dtype::DType;
 use vortex::error::VortexExpect as _;
 use vortex::error::VortexResult;
 use vortex::expr::stats::Precision;
-use vortex::expr::stats::Stat;
+use vortex::layout::layouts::file_stats::AggregateStats;
 use vortex::scalar::Scalar;
 use vortex::scalar::ScalarValue;
 
@@ -62,29 +69,22 @@ pub struct ColumnStatisticsAggregate {
 }
 
 impl ColumnStatisticsAggregate {
-    pub fn new(stats: &StatsSet) -> Self {
-        let min = match stats.get(Stat::Min) {
-            Precision::Exact(min) => Some(min),
-            _ => None,
-        };
-        let max = match stats.get(Stat::Max) {
-            Precision::Exact(max) => Some(max),
+    pub fn new(aggregates: &AggregateStats) -> Self {
+        let exact = |aggregate_fn: AggregateFnRef| match aggregates.get(&aggregate_fn) {
+            Precision::Exact(value) => value.into_value(),
             _ => None,
         };
 
-        let max_string_length =
-            if let Precision::Exact(value) = stats.get(Stat::UncompressedSizeInBytes) {
-                // DuckDB's string length is u32
-                #[allow(clippy::cast_possible_truncation)]
-                Some(value.as_primitive().as_u64().vortex_expect("not a u64") as u32)
-            } else {
-                None
-            };
+        let min = exact(Min.bind(NumericalAggregateOpts::skip_nans()));
+        let max = exact(Max.bind(NumericalAggregateOpts::skip_nans()));
 
-        let has_null = match stats.get(Stat::NullCount) {
-            Precision::Exact(cnt) => cnt.as_primitive().as_u64().vortex_expect("not a u64") > 0,
-            _ => true,
-        };
+        // DuckDB's string length is u32
+        #[allow(clippy::cast_possible_truncation)]
+        let max_string_length = exact(UncompressedSizeInBytes.bind(EmptyOptions))
+            .map(|value| value.as_primitive().as_u64().vortex_expect("not a u64") as u32);
+
+        let has_null = exact(NullCount.bind(EmptyOptions))
+            .is_none_or(|count| count.as_primitive().as_u64().vortex_expect("not a u64") > 0);
 
         Self {
             min,

vortex-duckdb/src/copy.rs

@@ -12,6 +12,14 @@ use object_store::ObjectStore;
 use object_store::registry::ObjectStoreRegistry;
 use parking_lot::Mutex;
 use static_assertions::assert_impl_all;
+use vortex::aggregate_fn::AggregateFnRef;
+use vortex::aggregate_fn::AggregateFnVTableExt;
+use vortex::aggregate_fn::EmptyOptions;
+use vortex::aggregate_fn::NumericalAggregateOpts;
+use vortex::aggregate_fn::fns::max::Max;
+use vortex::aggregate_fn::fns::min::Min;
+use vortex::aggregate_fn::fns::nan_count::NanCount;
+use vortex::aggregate_fn::fns::null_count::NullCount;
 use vortex::array::ArrayRef;
 use vortex::array::stream::ArrayStreamAdapter;
 use vortex::dtype::DType;
@@ -25,7 +33,6 @@ use vortex::error::VortexResult;
 use vortex::error::vortex_bail;
 use vortex::error::vortex_err;
 use vortex::expr::stats::Precision;
-use vortex::expr::stats::Stat;
 use vortex::file::CompressedFieldSizes;
 use vortex::file::FileStatistics;
 use vortex::file::WriteOptionsSessionExt;
@@ -289,16 +296,18 @@ fn column_stats_from_summary(
             stats_columns.len()
         )
     })?;
-    let (stats, dtype) = file_stats.get(stats_index);
+    let (aggregates, dtype) = file_stats.get(stats_index);
     let path = &file_stats.paths()[stats_index];
+    let value =
+        |aggregate_fn: AggregateFnRef| aggregates.get(&aggregate_fn).and_then(Scalar::into_value);
 
     Ok(WrittenColumnStats {
         column_key: ducklake_column_key(path),
-        min: exact_scalar_to_duckdb(stats.get(Stat::Min), dtype)?,
-        max: exact_scalar_to_duckdb(stats.get(Stat::Max), dtype)?,
-        null_count: exact_u64(stats.get(Stat::NullCount)),
+        min: exact_scalar_to_duckdb(value(Min.bind(NumericalAggregateOpts::skip_nans())), dtype)?,
+        max: exact_scalar_to_duckdb(value(Max.bind(NumericalAggregateOpts::skip_nans())), dtype)?,
+        null_count: exact_u64(value(NullCount.bind(EmptyOptions))),
         // NaNCount is exact only for float columns, so this is emitted just for them (as in parquet).
-        has_nan: exact_u64(stats.get(Stat::NaNCount)).map(|count| count > 0),
+        has_nan: exact_u64(value(NanCount.bind(EmptyOptions))).map(|count| count > 0),
         num_values: summary.row_count(),
         // On-disk compressed size; excludes bytes not attributable to a column (e.g. struct validity).
         column_size_bytes: column_sizes.and_then(|sizes| sizes.get(path)),
@@ -377,25 +386,19 @@ pub fn copy_to_initialize_global(
 #[cfg(test)]
 mod tests {
     use vortex::array::IntoArray;
-    use vortex::array::aggregate_fn::AggregateFnRef;
     use vortex::array::arrays::StructArray;
     use vortex::array::validity::Validity;
     use vortex::buffer::ByteBufferMut;
     use vortex::buffer::buffer;
+    use vortex::expr::stats::Stat;
 
     use super::*;
 
     fn pruning_aggregate_fns() -> Vec<AggregateFnRef> {
-        [
-            Stat::Min,
-            Stat::Max,
-            Stat::Sum,
-            Stat::NullCount,
-            Stat::NaNCount,
-        ]
-        .into_iter()
-        .filter_map(|stat| stat.aggregate_fn())
-        .collect()
+        [Stat::Min, Stat::Max, Stat::NullCount, Stat::NaNCount]
+            .into_iter()
+            .filter_map(|stat| stat.aggregate_fn())
+            .collect()
     }
 
     /// Writes a one-column file and returns its summary, with `file_statistics` controlling which

vortex-file/flatbuffers/vortex-file/footer.fbs

@@ -54,24 +54,50 @@ table PostscriptSegment {
 // [postscript]
 
 // [footer]
+/// An `AggregateSpec` describes an aggregate function recorded in the file statistics.
+table AggregateSpec {
+    id: string (required);
+}
+
+/// The accumulated state of one aggregate function over one file statistics entry.
+table AggregateState {
+    /// Index of the aggregate function's ID in `FileStatistics.aggregate_specs`.
+    aggregate_spec: uint16;
+    /// The aggregate function's serialized options.
+    options: [ubyte];
+    /// Protobuf serialized `vortex.scalar.ScalarValue` of the aggregate's partial state, typed by
+    /// the aggregate's partial dtype for the entry's dtype. A null value means no value was
+    /// accumulated.
+    partial: [ubyte] (required);
+}
+
+/// The aggregate states of one file statistics entry.
+table FieldAggregates {
+    aggregates: [AggregateState];
+}
+
 /// The `FileStatistics` object contains file-level statistics for the Vortex file.
 table FileStatistics {
     /// Statistics for each top-level field in the root schema, in schema order. If the root schema
     /// is not a struct, there will be a single entry in this array. Does not recurse into nested
-    /// struct fields; see `nested_field_stats` for that.
+    /// struct fields; see `nested_field_aggregates` for that.
     ///
     /// Kept in this legacy shape for backward compatibility with readers that predate
-    /// `nested_field_stats`, which strictly validate this array's length against the number of
-    /// top-level struct fields.
+    /// `nested_field_aggregates`, which strictly validate this array's length against the number
+    /// of top-level struct fields.
     field_stats: [ArrayStats];
-    /// Statistics following a post-order walk of the root `DType` tree: each leaf field gets an
-    /// entry, and each nullable struct additionally gets a trailing null-count entry inserted
-    /// after its children's entries.
+    /// Dictionary-encoded aggregate function IDs referenced by `nested_field_aggregates`, up to
+    /// u16::MAX.
+    aggregate_specs: [AggregateSpec];
+    /// Aggregate states following a post-order walk of the root `DType` tree: each leaf field gets
+    /// an entry, and each nullable struct additionally gets a trailing entry for its own null
+    /// count, inserted after its children's entries.
     ///
-    /// Readers that know about this field should prefer it over `field_stats` when present. It is
-    /// absent for files written before nested file stats were supported, and readers that don't
-    /// know about it fall back to `field_stats`'s legacy top-level-fields-only layout.
-    nested_field_stats: [ArrayStats];
+    /// Readers that know about this field should prefer it over `field_stats` when present,
+    /// skipping aggregates they don't recognize. It is absent for files written before nested
+    /// file stats were supported, and readers that don't know about it fall back to
+    /// `field_stats`'s legacy top-level-fields-only layout.
+    nested_field_aggregates: [FieldAggregates];
 }
 
 /// The `Registry` object stores dictionary-encoded configuration for segments,