feat: Add uncompressed size stat for Vortex/Datafusion (#1512)

AdamGS · web-flow · commit 4855ff2c5ab5 · 2024-12-02T23:01:18.000Z
As a followup to #1455 as @robert3005 noted.
diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs
@@ -17,17 +17,18 @@ use datafusion_physical_expr::{LexRequirement, PhysicalExpr};
 use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
 use datafusion_physical_plan::ExecutionPlan;
 use object_store::{ObjectMeta, ObjectStore};
+use vortex_array::array::StructArray;
 use vortex_array::arrow::infer_schema;
 use vortex_array::Context;
 use vortex_file::metadata::MetadataFetcher;
 use vortex_file::{
-    read_initial_bytes, LayoutContext, LayoutDeserializer, LayoutMessageCache, RelativeLayoutCache,
-    Scan, VORTEX_FILE_EXTENSION,
+    read_initial_bytes, read_layout_from_initial, LayoutContext, LayoutDeserializer,
+    LayoutMessageCache, RelativeLayoutCache, Scan, VORTEX_FILE_EXTENSION,
 };
 use vortex_io::{IoDispatcher, ObjectStoreReadAt};
 
 use super::execution::VortexExec;
-use super::statistics::array_to_col_statistics;
+use super::statistics::{array_to_col_statistics, uncompressed_col_size};
 use crate::can_be_pushed_down;
 
 #[derive(Debug, Default)]
@@ -108,7 +109,7 @@ impl FileFormat for VortexFormat {
         let relative_message_cache =
             RelativeLayoutCache::new(layout_message_cache.clone(), dtype.into());
 
-        let root_layout = vortex_file::read_layout_from_initial(
+        let root_layout = read_layout_from_initial(
             &initial_read,
             &layout_deserializer,
             Scan::empty(),
@@ -125,17 +126,26 @@ impl FileFormat for VortexFormat {
 
         if let Some(metadata) = metadata_table {
             let mut column_statistics = Vec::with_capacity(table_schema.fields().len());
+            let mut total_size = 0_u64;
 
             for col_stats in metadata.into_iter() {
                 let col_stats = match col_stats {
-                    Some(array) => array_to_col_statistics(array.try_into()?)?,
+                    Some(array) => {
+                        let col_metadata_array = StructArray::try_from(array)?;
+                        let col_stats = array_to_col_statistics(&col_metadata_array)?;
+
+                        total_size +=
+                            uncompressed_col_size(&col_metadata_array)?.unwrap_or_default();
+                        col_stats
+                    }
                     None => ColumnStatistics::new_unknown(),
                 };
 
                 column_statistics.push(col_stats);
             }
 
             stats.column_statistics = column_statistics;
+            stats.total_byte_size = Precision::Inexact(total_size as usize);
         }
 
         Ok(stats)
diff --git a/vortex-datafusion/src/persistent/statistics.rs b/vortex-datafusion/src/persistent/statistics.rs
@@ -9,7 +9,7 @@ use vortex_array::variants::StructArrayTrait as _;
 use vortex_array::IntoCanonical;
 use vortex_error::VortexResult;
 
-pub fn array_to_col_statistics(array: StructArray) -> VortexResult<ColumnStatistics> {
+pub fn array_to_col_statistics(array: &StructArray) -> VortexResult<ColumnStatistics> {
     let mut stats = ColumnStatistics::new_unknown();
 
     if let Some(null_count_array) = array.field_by_name("null_count") {
@@ -40,3 +40,17 @@ pub fn array_to_col_statistics(array: StructArray) -> VortexResult<ColumnStatist
 
     Ok(stats)
 }
+
+pub fn uncompressed_col_size(array: &StructArray) -> VortexResult<Option<u64>> {
+    match array.field_by_name("uncompressed_size") {
+        None => Ok(None),
+        Some(array) => {
+            let array = array.into_canonical()?.into_arrow()?;
+            let array = array.as_primitive::<UInt64Type>();
+
+            let uncompressed_size = array.iter().map(|v| v.unwrap_or_default()).sum::<u64>();
+
+            Ok(Some(uncompressed_size))
+        }
+    }
+}
diff --git a/vortex-file/src/write/metadata_accumulators.rs b/vortex-file/src/write/metadata_accumulators.rs
@@ -41,6 +41,7 @@ struct BoolAccumulator {
     minima: UnwrappedStatAccumulator<bool>,
     true_count: UnwrappedStatAccumulator<u64>,
     null_count: UnwrappedStatAccumulator<u64>,
+    uncompressed_size: UnwrappedStatAccumulator<u64>,
 }
 
 impl BoolAccumulator {
@@ -50,6 +51,10 @@ impl BoolAccumulator {
             minima: UnwrappedStatAccumulator::new(Stat::Min, "min".into()),
             true_count: UnwrappedStatAccumulator::new(Stat::TrueCount, "true_count".into()),
             null_count: UnwrappedStatAccumulator::new(Stat::NullCount, "null_count".into()),
+            uncompressed_size: UnwrappedStatAccumulator::new(
+                Stat::UncompressedSizeInBytes,
+                "uncompressed_size".into(),
+            ),
         }
     }
 }
@@ -60,6 +65,7 @@ impl MetadataAccumulator for BoolAccumulator {
         self.minima.push_chunk(array);
         self.true_count.push_chunk(array);
         self.null_count.push_chunk(array);
+        self.uncompressed_size.push_chunk(array);
     }
 
     fn into_array(self: Box<Self>) -> VortexResult<Option<ArrayData>> {
@@ -68,6 +74,7 @@ impl MetadataAccumulator for BoolAccumulator {
             self.minima.into_column(),
             self.true_count.into_column(),
             self.null_count.into_column(),
+            self.uncompressed_size.into_column(),
         ]
         .into_iter()
         .flatten()
@@ -90,6 +97,7 @@ struct StandardAccumulator<T> {
     maxima: UnwrappedStatAccumulator<T>,
     minima: UnwrappedStatAccumulator<T>,
     null_count: UnwrappedStatAccumulator<u64>,
+    uncompressed_size: UnwrappedStatAccumulator<u64>,
 }
 
 impl<T> StandardAccumulator<T> {
@@ -98,6 +106,10 @@ impl<T> StandardAccumulator<T> {
             maxima: UnwrappedStatAccumulator::new(Stat::Max, "max".into()),
             minima: UnwrappedStatAccumulator::new(Stat::Min, "min".into()),
             null_count: UnwrappedStatAccumulator::new(Stat::NullCount, "null_count".into()),
+            uncompressed_size: UnwrappedStatAccumulator::new(
+                Stat::UncompressedSizeInBytes,
+                "uncompressed_size".into(),
+            ),
         }
     }
 }
@@ -111,13 +123,15 @@ where
         self.maxima.push_chunk(array);
         self.minima.push_chunk(array);
         self.null_count.push_chunk(array);
+        self.uncompressed_size.push_chunk(array);
     }
 
     fn into_array(self: Box<Self>) -> VortexResult<Option<ArrayData>> {
         let (names, fields): (Vec<FieldName>, Vec<ArrayData>) = [
             self.maxima.into_column(),
             self.minima.into_column(),
             self.null_count.into_column(),
+            self.uncompressed_size.into_column(),
         ]
         .into_iter()
         .flatten()
@@ -134,29 +148,38 @@ where
     }
 }
 
-/// A minimal accumulator which only tracks null counts.
+/// A minimal accumulator which only tracks null counts and total uncompressed size.
 struct BasicAccumulator {
     null_count: UnwrappedStatAccumulator<u64>,
+    uncompressed_size: UnwrappedStatAccumulator<u64>,
 }
 
 impl BasicAccumulator {
     fn new() -> Self {
         Self {
             null_count: UnwrappedStatAccumulator::new(Stat::NullCount, "null_count".into()),
+            uncompressed_size: UnwrappedStatAccumulator::new(
+                Stat::UncompressedSizeInBytes,
+                "uncompressed_size".into(),
+            ),
         }
     }
 }
 
 impl MetadataAccumulator for BasicAccumulator {
     fn push_chunk(&mut self, array: &ArrayData) {
-        self.null_count.push_chunk(array)
+        self.null_count.push_chunk(array);
+        self.uncompressed_size.push_chunk(array);
     }
 
     fn into_array(self: Box<Self>) -> VortexResult<Option<ArrayData>> {
-        let (names, fields): (Vec<FieldName>, Vec<ArrayData>) = [self.null_count.into_column()]
-            .into_iter()
-            .flatten()
-            .unzip();
+        let (names, fields): (Vec<FieldName>, Vec<ArrayData>) = [
+            self.null_count.into_column(),
+            self.uncompressed_size.into_column(),
+        ]
+        .into_iter()
+        .flatten()
+        .unzip();
         if fields.is_empty() {
             Ok(None)
         } else {
@@ -246,7 +269,16 @@ mod tests {
             StructArray::try_from(Box::new(bool_accumulator).into_array().unwrap().unwrap())
                 .unwrap();
         assert_eq!(struct_array.len(), 1);
-        assert_field_names(&struct_array, &["max", "min", "true_count", "null_count"]);
+        assert_field_names(
+            &struct_array,
+            &[
+                "max",
+                "min",
+                "true_count",
+                "null_count",
+                "uncompressed_size",
+            ],
+        );
     }
 
     #[test]
@@ -263,7 +295,10 @@ mod tests {
         )
         .unwrap();
         assert_eq!(struct_array.len(), 1);
-        assert_field_names(&struct_array, &["max", "min", "null_count"]);
+        assert_field_names(
+            &struct_array,
+            &["max", "min", "null_count", "uncompressed_size"],
+        );
     }
 
     #[test]
@@ -328,6 +363,6 @@ mod tests {
             StructArray::try_from(Box::new(basic_accumulator).into_array().unwrap().unwrap())
                 .unwrap();
         assert_eq!(struct_array.len(), 1);
-        assert_field_names(&struct_array, &["null_count"]);
+        assert_field_names(&struct_array, &["null_count", "uncompressed_size"]);
     }
 }
