feat: add stat for uncompressed size in bytes (#1315)

fixes #1237

Author: lwwmanning

.github/workflows/ci.yml

@@ -162,7 +162,7 @@ jobs:
         with:
           version: v1.0.0
       - name: Rust Bench as test
-        run: cargo bench --bench '*[!noci]' -- --test
+        run: cargo bench --bench '*[!noci]' --profile benchtest -- --test
 
   generated-files:
     name: "Check generated proto/fbs files are up to date"

Cargo.lock

@@ -222,3 +222,7 @@ lto = "thin"      # attempts to perform optimizations across all crates within t
 codegen-units = 16 # default for "release", which "bench" inherits
 lto = false        # default
 debug = true
+
+[profile.benchtest]
+inherits = "bench"
+debug-assertions = true

Cargo.toml

@@ -8,10 +8,14 @@ use vortex_error::VortexResult;
 use crate::array::BoolArray;
 use crate::stats::{ArrayStatisticsCompute, Stat, StatsSet};
 use crate::validity::{ArrayValidity, LogicalValidity};
-use crate::{ArrayDType, IntoArrayVariant};
+use crate::{ArrayDType, ArrayTrait as _, IntoArrayVariant};
 
 impl ArrayStatisticsCompute for BoolArray {
     fn compute_statistics(&self, stat: Stat) -> VortexResult<StatsSet> {
+        if stat == Stat::UncompressedSizeInBytes {
+            return Ok(StatsSet::of(stat, self.nbytes()));
+        }
+
         if self.is_empty() {
             return Ok(StatsSet::from_iter([
                 (Stat::TrueCount, 0.into()),

vortex-array/src/array/bool/stats.rs

@@ -17,7 +17,7 @@ use crate::compute::unary::{scalar_at, scalar_at_unchecked, subtract_scalar, Sub
 use crate::compute::{search_sorted, SearchSortedSide};
 use crate::encoding::ids;
 use crate::iter::{ArrayIterator, ArrayIteratorAdapter};
-use crate::stats::StatsSet;
+use crate::stats::ArrayStatistics;
 use crate::stream::{ArrayStream, ArrayStreamAdapter};
 use crate::validity::Validity::NonNullable;
 use crate::validity::{ArrayValidity, LogicalValidity, Validity};
@@ -61,9 +61,19 @@ impl ChunkedArray {
             .collect_vec();
 
         let nchunks = chunk_offsets.len() - 1;
-        let length = *chunk_offsets.last().unwrap_or_else(|| {
-            unreachable!("Chunk ends is guaranteed to have at least one element")
-        }) as usize;
+        let length = *chunk_offsets
+            .last()
+            .vortex_expect("Chunk ends is guaranteed to have at least one element")
+            as usize;
+
+        let stats = chunks
+            .iter()
+            .map(|chunk| chunk.statistics().to_set())
+            .reduce(|mut acc, stats| {
+                acc.merge_ordered(&stats);
+                acc
+            })
+            .unwrap_or_default();
 
         let mut children = Vec::with_capacity(chunks.len() + 1);
         children.push(PrimitiveArray::from_vec(chunk_offsets, NonNullable).into_array());
@@ -74,7 +84,7 @@ impl ChunkedArray {
             length,
             ChunkedMetadata { nchunks },
             children.into(),
-            StatsSet::default(),
+            stats,
         )
     }
 

vortex-array/src/array/chunked/mod.rs

@@ -5,6 +5,8 @@ use crate::stats::{ArrayStatistics, ArrayStatisticsCompute, Stat, StatsSet};
 
 impl ArrayStatisticsCompute for ChunkedArray {
     fn compute_statistics(&self, stat: Stat) -> VortexResult<StatsSet> {
+        // for UncompressedSizeInBytes, we end up with sum of chunk uncompressed sizes
+        // this ignores the `chunk_offsets` array child, so it won't exactly match self.nbytes()
         Ok(self
             .chunks()
             .map(|c| {

vortex-array/src/array/chunked/stats.rs

@@ -54,7 +54,11 @@ impl ArrayValidity for NullArray {
 }
 
 impl ArrayStatisticsCompute for NullArray {
-    fn compute_statistics(&self, _stat: Stat) -> VortexResult<StatsSet> {
+    fn compute_statistics(&self, stat: Stat) -> VortexResult<StatsSet> {
+        if stat == Stat::UncompressedSizeInBytes {
+            return Ok(StatsSet::of(stat, self.nbytes()));
+        }
+
         Ok(StatsSet::nulls(self.len(), &DType::Null))
     }
 }

vortex-array/src/array/null/mod.rs

@@ -14,14 +14,18 @@ use crate::array::primitive::PrimitiveArray;
 use crate::stats::{ArrayStatisticsCompute, Stat, StatsSet};
 use crate::validity::{ArrayValidity, LogicalValidity};
 use crate::variants::PrimitiveArrayTrait;
-use crate::{ArrayDType, IntoArrayVariant};
+use crate::{ArrayDType, ArrayTrait as _, IntoArrayVariant};
 
 trait PStatsType: NativePType + Into<Scalar> + BitWidth {}
 
 impl<T: NativePType + Into<Scalar> + BitWidth> PStatsType for T {}
 
 impl ArrayStatisticsCompute for PrimitiveArray {
     fn compute_statistics(&self, stat: Stat) -> VortexResult<StatsSet> {
+        if stat == Stat::UncompressedSizeInBytes {
+            return Ok(StatsSet::of(stat, self.nbytes()));
+        }
+
         let mut stats = match_each_native_ptype!(self.ptype(), |$P| {
             match self.logical_validity() {
                 LogicalValidity::AllValid(_) => self.maybe_null_slice::<$P>().compute_statistics(stat),
@@ -77,7 +81,7 @@ impl<T: PStatsType> ArrayStatisticsCompute for &[T] {
                 self.iter().skip(1).for_each(|next| stats.next(*next));
                 stats.finish()
             }
-            Stat::TrueCount => StatsSet::default(),
+            Stat::TrueCount | Stat::UncompressedSizeInBytes => StatsSet::default(),
         })
     }
 }
@@ -87,7 +91,7 @@ struct NullableValues<'a, T: PStatsType>(&'a [T], &'a BooleanBuffer);
 impl<T: PStatsType> ArrayStatisticsCompute for NullableValues<'_, T> {
     fn compute_statistics(&self, stat: Stat) -> VortexResult<StatsSet> {
         let values = self.0;
-        if values.is_empty() || stat == Stat::TrueCount {
+        if values.is_empty() || stat == Stat::TrueCount || stat == Stat::UncompressedSizeInBytes {
             return Ok(StatsSet::default());
         }
 

vortex-array/src/array/primitive/stats.rs

@@ -7,7 +7,7 @@ use vortex_error::{vortex_bail, vortex_err, vortex_panic, VortexExpect as _, Vor
 
 use crate::array::visitor::{AcceptArrayVisitor, ArrayVisitor};
 use crate::encoding::ids;
-use crate::stats::{ArrayStatisticsCompute, StatsSet};
+use crate::stats::{ArrayStatistics, ArrayStatisticsCompute, Stat, StatsSet};
 use crate::validity::{ArrayValidity, LogicalValidity, Validity, ValidityMetadata};
 use crate::variants::{ArrayVariants, StructArrayTrait};
 use crate::{
@@ -191,7 +191,21 @@ impl AcceptArrayVisitor for StructArray {
     }
 }
 
-impl ArrayStatisticsCompute for StructArray {}
+impl ArrayStatisticsCompute for StructArray {
+    fn compute_statistics(&self, stat: Stat) -> VortexResult<StatsSet> {
+        Ok(match stat {
+            Stat::UncompressedSizeInBytes => self
+                .children()
+                .map(|f| f.statistics().compute_uncompressed_size_in_bytes())
+                .reduce(|acc, field_size| acc.zip(field_size).map(|(a, b)| a + b))
+                .flatten()
+                .map(|size| StatsSet::of(stat, size))
+                .unwrap_or_default(),
+            Stat::NullCount => StatsSet::of(stat, self.validity().null_count(self.len())?),
+            _ => StatsSet::default(),
+        })
+    }
+}
 
 #[cfg(test)]
 mod test {

vortex-array/src/array/struct_/mod.rs

@@ -7,10 +7,14 @@ use vortex_error::VortexResult;
 use crate::accessor::ArrayAccessor;
 use crate::array::varbin::{varbin_scalar, VarBinArray};
 use crate::stats::{ArrayStatisticsCompute, Stat, StatsSet};
-use crate::ArrayDType;
+use crate::{ArrayDType, ArrayTrait as _};
 
 impl ArrayStatisticsCompute for VarBinArray {
-    fn compute_statistics(&self, _stat: Stat) -> VortexResult<StatsSet> {
+    fn compute_statistics(&self, stat: Stat) -> VortexResult<StatsSet> {
+        if stat == Stat::UncompressedSizeInBytes {
+            return Ok(StatsSet::of(stat, self.nbytes()));
+        }
+
         if self.is_empty() {
             return Ok(StatsSet::default());
         }