feat: don't write leading/trailing zero histograms into file (#1372)

fixes #1362

Author: lwwmanning

vortex-array/src/data/mod.rs

@@ -51,7 +51,7 @@ impl ArrayData {
             metadata,
             buffer,
             children,
-            stats_map: Arc::new(RwLock::new(statistics)),
+            stats_set: Arc::new(RwLock::new(statistics)),
         };
 
         let array = ArrayData(InnerArrayData::Owned(data));

vortex-array/src/data/owned.rs

@@ -19,7 +19,7 @@ pub(crate) struct OwnedArrayData {
     pub(crate) metadata: Arc<dyn ArrayMetadata>,
     pub(crate) buffer: Option<Buffer>,
     pub(crate) children: Arc<[ArrayData]>,
-    pub(crate) stats_map: Arc<RwLock<StatsSet>>,
+    pub(crate) stats_set: Arc<RwLock<StatsSet>>,
 }
 
 impl OwnedArrayData {
@@ -92,7 +92,7 @@ impl OwnedArrayData {
 
 impl Statistics for OwnedArrayData {
     fn get(&self, stat: Stat) -> Option<Scalar> {
-        self.stats_map
+        self.stats_set
             .read()
             .unwrap_or_else(|_| {
                 vortex_panic!(
@@ -105,14 +105,14 @@ impl Statistics for OwnedArrayData {
     }
 
     fn to_set(&self) -> StatsSet {
-        self.stats_map
+        self.stats_set
             .read()
             .unwrap_or_else(|_| vortex_panic!("Failed to acquire read lock on stats map"))
             .clone()
     }
 
     fn set(&self, stat: Stat, value: Scalar) {
-        self.stats_map
+        self.stats_set
             .write()
             .unwrap_or_else(|_| {
                 vortex_panic!(
@@ -124,6 +124,13 @@ impl Statistics for OwnedArrayData {
             .set(stat, value);
     }
 
+    fn clear(&self, stat: Stat) {
+        self.stats_set
+            .write()
+            .unwrap_or_else(|_| vortex_panic!("Failed to acquire write lock on stats map"))
+            .clear(stat);
+    }
+
     fn compute(&self, stat: Stat) -> Option<Scalar> {
         if let Some(s) = self.get(stat) {
             return Some(s);
@@ -133,14 +140,21 @@ impl Statistics for OwnedArrayData {
             .with_dyn(|a| a.compute_statistics(stat))
             .ok()?;
 
-        self.stats_map
+        self.stats_set
             .write()
             .unwrap_or_else(|_| {
                 vortex_panic!("Failed to write to stats map while computing {}", stat)
             })
             .extend(computed);
         self.get(stat)
     }
+
+    fn retain_only(&self, stats: &[Stat]) {
+        self.stats_set
+            .write()
+            .unwrap_or_else(|_| vortex_panic!("Failed to acquire write lock on stats map"))
+            .retain_only(stats);
+    }
 }
 
 impl From<OwnedArrayData> for ArrayData {

vortex-array/src/data/viewed.rs

@@ -213,7 +213,15 @@ impl Statistics for ViewedArrayData {
     /// We want to avoid any sort of allocation on instantiation of the ArrayView, so we
     /// do not allocate a stats_set to cache values.
     fn set(&self, _stat: Stat, _value: Scalar) {
-        // We cannot set stats on a view
+        // We cannot modify stats on a view
+    }
+
+    fn clear(&self, _stat: Stat) {
+        // We cannot modify stats on a view
+    }
+
+    fn retain_only(&self, _stats: &[Stat]) {
+        // We cannot modify stats on a view
     }
 
     fn compute(&self, stat: Stat) -> Option<Scalar> {

vortex-array/src/stats/mod.rs

@@ -18,7 +18,7 @@ pub mod flatbuffers;
 mod statsset;
 
 /// Statistics that are used for pruning files (i.e., we want to ensure they are computed when compressing/writing).
-pub(crate) const PRUNING_STATS: &[Stat] = &[Stat::Min, Stat::Max, Stat::TrueCount, Stat::NullCount];
+pub const PRUNING_STATS: &[Stat] = &[Stat::Min, Stat::Max, Stat::TrueCount, Stat::NullCount];
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Sequence, Enum)]
 #[non_exhaustive]
@@ -44,6 +44,7 @@ pub enum Stat {
     TrueCount,
     /// The number of null values in the array
     NullCount,
+    /// The uncompressed size of the array in bytes
     UncompressedSizeInBytes,
 }
 
@@ -95,8 +96,12 @@ pub trait Statistics {
     /// Get all existing statistics
     fn to_set(&self) -> StatsSet;
 
+    /// Set the value of the statistic
     fn set(&self, stat: Stat, value: Scalar);
 
+    /// Clear the value of the statistic
+    fn clear(&self, stat: Stat);
+
     /// Computes the value of the stat if it's not present
     fn compute(&self, stat: Stat) -> Option<Scalar>;
 
@@ -111,6 +116,8 @@ pub trait Statistics {
         }
         Ok(stats_set)
     }
+
+    fn retain_only(&self, stats: &[Stat]);
 }
 
 pub trait ArrayStatistics {

vortex-array/src/stats/statsset.rs

@@ -1,3 +1,5 @@
+use core::mem;
+
 use enum_iterator::all;
 use enum_map::EnumMap;
 use itertools::{EitherOrBoth, Itertools};
@@ -130,6 +132,13 @@ impl StatsSet {
         self.values[stat] = None;
     }
 
+    pub fn retain_only(&mut self, stats: &[Stat]) {
+        let mut old_map = mem::take(&mut self.values);
+        for stat in stats {
+            self.values[*stat] = old_map[*stat].take();
+        }
+    }
+
     /// Merge stats set `other` into `self`, with the semantic assumption that `other`
     /// contains stats from an array that is *appended* to the array represented by `self`.
     pub fn merge_ordered(&mut self, other: &Self) -> &Self {

vortex-file/src/write/writer.rs

@@ -3,6 +3,7 @@ use std::{io, mem};
 use flatbuffers::FlatBufferBuilder;
 use futures::TryStreamExt;
 use vortex_array::array::{ChunkedArray, StructArray};
+use vortex_array::stats::{ArrayStatistics, Stat};
 use vortex_array::stream::ArrayStream;
 use vortex_array::{ArrayDType as _, ArrayData};
 use vortex_buffer::io_buf::IoBuf;
@@ -19,6 +20,18 @@ use crate::write::metadata_accumulators::{new_metadata_accumulator, MetadataAccu
 use crate::write::postscript::Postscript;
 use crate::{EOF_SIZE, MAGIC_BYTES, MAX_FOOTER_SIZE, VERSION};
 
+const STATS_TO_WRITE: &[Stat] = &[
+    Stat::Min,
+    Stat::Max,
+    Stat::TrueCount,
+    Stat::NullCount,
+    Stat::RunCount,
+    Stat::IsConstant,
+    Stat::IsSorted,
+    Stat::IsStrictSorted,
+    Stat::UncompressedSizeInBytes,
+];
+
 pub struct VortexFileWriter<W> {
     msgs: MessageWriter<W>,
 
@@ -217,7 +230,13 @@ impl ColumnWriter {
 
         while let Some(chunk) = stream.try_next().await? {
             rows_written += chunk.len() as u64;
+
+            // accumulate the stats for the stats table
             self.metadata.push_chunk(&chunk);
+
+            // clear the stats that we don't want to serialize into the file
+            chunk.statistics().retain_only(STATS_TO_WRITE);
+
             msgs.write_batch(chunk).await?;
             offsets.push(msgs.tell());
             row_offsets.push(rows_written);
@@ -292,11 +311,13 @@ mod tests {
     use flatbuffers::FlatBufferBuilder;
     use futures_executor::block_on;
     use vortex_array::array::{PrimitiveArray, StructArray, VarBinArray};
+    use vortex_array::stats::PRUNING_STATS;
     use vortex_array::validity::Validity;
     use vortex_array::IntoArrayData;
     use vortex_flatbuffers::WriteFlatBuffer;
 
     use crate::write::postscript::Postscript;
+    use crate::write::writer::STATS_TO_WRITE;
     use crate::{VortexFileWriter, V1_FOOTER_FBS_SIZE};
 
     #[test]
@@ -328,4 +349,11 @@ mod tests {
 
         assert_eq!(buffer[buffer_begin..buffer_end].len(), V1_FOOTER_FBS_SIZE);
     }
+
+    #[test]
+    fn stats_to_write() {
+        for stat in PRUNING_STATS {
+            assert!(STATS_TO_WRITE.contains(stat));
+        }
+    }
 }