Make pruning stats pluggable

Default selects based on the dtypes. Modeled in the same way that zone
maps configure aggregate partials.

Author: thorfour

vortex-array/src/stats/mod.rs

@@ -31,27 +31,8 @@ pub use array::*;
 pub use session::*;
 use vortex_error::VortexExpect;
 
-use crate::aggregate_fn::AggregateFnRef;
 use crate::expr::stats::Stat;
 
-/// Statistics that are used for pruning files (i.e., we want to ensure they are computed when compressing/writing).
-/// Sum is included for boolean arrays.
-pub const PRUNING_STATS: &[Stat] = &[
-    Stat::Min,
-    Stat::Max,
-    Stat::Sum,
-    Stat::NullCount,
-    Stat::NaNCount,
-];
-
-/// [`PRUNING_STATS`], converted to their aggregate function form.
-pub fn pruning_aggregate_fns() -> Vec<AggregateFnRef> {
-    PRUNING_STATS
-        .iter()
-        .filter_map(|stat| stat.aggregate_fn())
-        .collect()
-}
-
 pub fn as_stat_bitset_bytes(stats: &[Stat]) -> Vec<u8> {
     let max_stat = u8::from(last::<Stat>().vortex_expect("last stat")) as usize + 1;
     // TODO(ngates): use vortex-buffer::BitBuffer

vortex-duckdb/src/copy.rs

@@ -340,13 +340,25 @@ mod tests {
     use vortex::array::IntoArray;
     use vortex::array::aggregate_fn::AggregateFnRef;
     use vortex::array::arrays::StructArray;
-    use vortex::array::stats::pruning_aggregate_fns;
     use vortex::array::validity::Validity;
     use vortex::buffer::ByteBufferMut;
     use vortex::buffer::buffer;
 
     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()
+    }
+
     /// Writes a one-column file and returns its summary, with `file_statistics` controlling which
     /// statistics the footer carries (empty means none at all).
     fn write_summary(file_statistics: Vec<AggregateFnRef>) -> WriteSummary {

vortex-duckdb/src/file_reader.rs

@@ -338,17 +338,31 @@ pub fn footer_get_statistics(footer: &Footer, index: usize) -> Option<ColumnStat
 #[cfg(test)]
 mod tests {
     use vortex::array::IntoArray;
+    use vortex::array::aggregate_fn::AggregateFnRef;
     use vortex::array::arrays::StructArray;
-    use vortex::array::stats::pruning_aggregate_fns;
     use vortex::array::validity::Validity;
     use vortex::buffer::ByteBufferMut;
     use vortex::buffer::buffer;
+    use vortex::expr::stats::Stat;
     use vortex::file::WriteOptionsSessionExt;
 
     use super::*;
     use crate::RUNTIME;
     use crate::SESSION;
 
+    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()
+    }
+
     #[test]
     fn footer_get_statistics_resolves_by_name_past_a_nullable_struct_field() {
         // Regression test: a nullable top-level struct field ("s") contributes both a `s.b` entry

vortex-file/src/pruning.rs

@@ -80,7 +80,7 @@ impl FileStatsBinder<'_> {
     fn stat_ref(&self, field_path: &FieldPath, stat: Stat) -> Option<BoundExpression> {
         let (field_stats, field_dtype) = self.file_stats.get_by_path(field_path)?;
 
-        let stat_value = field_stats.get(stat).as_exact()?;
+        let stat_value = field_stats.get(stat).into_inner()?;
         let stat_dtype = stat.dtype(field_dtype)?;
         let stat_scalar = Scalar::try_new(stat_dtype, Some(stat_value)).ok()?;
 

vortex-file/src/tests.rs

@@ -15,6 +15,7 @@ use rstest::rstest;
 use vortex_array::ArrayRef;
 use vortex_array::IntoArray;
 use vortex_array::VortexSessionExecute;
+use vortex_array::aggregate_fn::AggregateFnRef;
 use vortex_array::array_session;
 use vortex_array::arrays::BoolArray;
 use vortex_array::arrays::ChunkedArray;
@@ -64,7 +65,6 @@ use vortex_array::scalar::ScalarValue;
 use vortex_array::scalar_fn::ScalarFnVTableExt;
 use vortex_array::scalar_fn::fns::pack::Pack;
 use vortex_array::scalar_fn::fns::pack::PackOptions;
-use vortex_array::stats::pruning_aggregate_fns;
 use vortex_array::stream::ArrayStreamAdapter;
 use vortex_array::stream::ArrayStreamExt;
 use vortex_array::validity::Validity;
@@ -115,6 +115,19 @@ static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
     session
 });
 
+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()
+}
+
 fn strict_sorted(indices: Buffer<u64>) -> StrictSortedBuffer<u64> {
     StrictSortedBuffer::try_new(indices).expect("test indices should be strictly increasing")
 }

vortex-file/src/writer.rs

@@ -25,7 +25,6 @@ use vortex_array::dtype::FieldPath;
 use vortex_array::iter::ArrayIterator;
 use vortex_array::iter::ArrayIteratorExt;
 use vortex_array::session::ArraySessionExt;
-use vortex_array::stats::pruning_aggregate_fns;
 use vortex_array::stream::ArrayStream;
 use vortex_array::stream::ArrayStreamAdapter;
 use vortex_array::stream::ArrayStreamExt;
@@ -84,7 +83,10 @@ pub struct VortexWriteOptions {
     buffered_bytes: BufferedBytesTracker,
     exclude_dtype: bool,
     max_variable_length_statistics_size: usize,
-    file_statistics: Vec<AggregateFnRef>,
+    /// The aggregates to compute for file-level statistics.
+    ///
+    /// If unset, the writer chooses pruning aggregates from the file dtype.
+    file_statistics: Option<Vec<AggregateFnRef>>,
     write_legacy_statistics: bool,
     metadata: HashMap<String, ByteBuffer>,
 }
@@ -107,7 +109,7 @@ impl VortexWriteOptions {
             buffered_bytes: BufferedBytesTracker::new(),
             session,
             exclude_dtype: false,
-            file_statistics: pruning_aggregate_fns(),
+            file_statistics: None,
             write_legacy_statistics: true,
             max_variable_length_statistics_size: 64,
             metadata: HashMap::default(),
@@ -158,9 +160,10 @@ impl VortexWriteOptions {
 
     /// Configure which statistics to compute at the file level.
     ///
-    /// Pass an empty vector to omit file-level statistics.
+    /// Pass an empty vector to omit file-level statistics. If left unset, the writer chooses
+    /// pruning aggregates from the file dtype.
     pub fn with_file_statistics(mut self, file_statistics: Vec<AggregateFnRef>) -> Self {
-        self.file_statistics = file_statistics;
+        self.file_statistics = Some(file_statistics);
         self
     }
 
@@ -286,9 +289,12 @@ impl VortexWriteOptions {
                 .map(move |result| result.map(|chunk| (ptr.advance(), chunk))),
         )
         .sendable();
+        // When unset, `accumulate_stats` resolves each leaf's own default from its dtype, the way
+        // `default_zoned_aggregate_fns` does for zoned layouts (e.g. bounded min/max for
+        // variable-length columns), rather than a single dtype-blind default applied uniformly.
         let (file_stats, stream) = accumulate_stats(
             stream,
-            self.file_statistics.clone().into(),
+            self.file_statistics.clone().map(Arc::from),
             self.max_variable_length_statistics_size,
             &self.session,
             self.write_legacy_statistics,
@@ -334,7 +340,7 @@ impl VortexWriteOptions {
         let (layout, segment_specs) = layout_fut.await?;
 
         // Assemble the Footer object now that we have all the segments.
-        let statistics = if self.file_statistics.is_empty() {
+        let statistics = if matches!(&self.file_statistics, Some(stats) if stats.is_empty()) {
             None
         } else {
             Some(FileStatistics::new_with_dtype(