Use approx footer size for caching instead of walking the layout tree

vortex-datafusion/src/persistent/cache.rs

@@ -1,32 +1,21 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use std::mem::size_of;
-
-use datafusion_common::ScalarValue;
 use datafusion_execution::cache::cache_manager::FileMetadata;
-use vortex::expr::stats::Precision;
-use vortex::expr::stats::Stat;
 use vortex::file::Footer;
-use vortex::file::SegmentSpec;
 use vortex::file::VortexFile;
-use vortex::layout::Layout;
-use vortex::layout::segments::SegmentId;
 
 /// Cached Vortex file metadata for use with DataFusion's [`FileMetadataCache`].
 pub struct CachedVortexMetadata {
     footer: Footer,
-    memory_size: usize,
 }
 
 impl CachedVortexMetadata {
     /// Create a new cached metadata entry from a VortexFile.
     pub fn new(vortex_file: &VortexFile) -> Self {
         let footer = vortex_file.footer();
-        let memory_size = estimate_footer_size(footer);
         Self {
             footer: footer.clone(),
-            memory_size,
         }
     }
 
@@ -42,42 +31,15 @@ impl FileMetadata for CachedVortexMetadata {
     }
 
     fn memory_size(&self) -> usize {
-        self.memory_size
+        self.footer
+            .approx_byte_size()
+            // 64KB is not an insane estimate...
+            // We just want to avoid returning zero and _never_ being evicted from the cache.
+            .unwrap_or(1024 * 64)
     }
 
     #[allow(clippy::disallowed_types)]
     fn extra_info(&self) -> std::collections::HashMap<String, String> {
         Default::default()
     }
 }
-
-/// Approximate the in-memory size of a footer.
-fn estimate_footer_size(footer: &Footer) -> usize {
-    let segments_size = footer.segment_map().len() * size_of::<SegmentSpec>();
-    let stats_size = footer
-        .statistics()
-        .map(|file_statistics| {
-            file_statistics
-                .stats_sets()
-                .iter()
-                .map(|s| {
-                    s.iter().count() * (size_of::<Stat>() + size_of::<Precision<ScalarValue>>())
-                })
-                .sum::<usize>()
-        })
-        .unwrap_or(0);
-
-    let layout_size = footer
-        .layout()
-        .depth_first_traversal()
-        .filter_map(|l| l.ok().map(|l| layout_size(l.as_ref())))
-        .sum::<usize>();
-
-    segments_size + stats_size + layout_size
-}
-
-fn layout_size(layout: &dyn Layout) -> usize {
-    size_of_val(layout.dtype())
-        + layout.metadata().len()
-        + layout.segment_ids().len() * size_of::<SegmentId>()
-}

vortex-file/public-api.lock

@@ -120,6 +120,8 @@ pub struct vortex_file::Footer
 
 impl vortex_file::Footer
 
+pub fn vortex_file::Footer::approx_byte_size(&self) -> core::option::Option<usize>
+
 pub fn vortex_file::Footer::deserializer(eof_buffer: vortex_buffer::ByteBuffer, session: vortex_session::VortexSession) -> vortex_file::FooterDeserializer
 
 pub fn vortex_file::Footer::dtype(&self) -> &vortex_array::dtype::DType

vortex-file/src/footer/mod.rs

@@ -48,6 +48,8 @@ pub struct Footer {
     statistics: Option<FileStatistics>,
     // The specific arrays used within the file, in the order they were registered.
     array_ctx: ArrayContext,
+    // The approximate size of the footer in bytes, used for caching and memory management.
+    approx_byte_size: Option<usize>,
 }
 
 impl Footer {
@@ -62,6 +64,7 @@ impl Footer {
             segments,
             statistics,
             array_ctx,
+            approx_byte_size: None,
         }
     }
 
@@ -73,6 +76,7 @@ impl Footer {
         statistics: Option<FileStatistics>,
         session: &VortexSession,
     ) -> VortexResult<Self> {
+        let approx_byte_size = footer_bytes.len() + layout_bytes.len();
         let fb_footer = root::<fb::Footer>(&footer_bytes)?;
 
         // Create a LayoutContext from the registry.
@@ -118,6 +122,7 @@ impl Footer {
             segments,
             statistics,
             array_ctx,
+            approx_byte_size: Some(approx_byte_size),
         })
     }
 
@@ -141,6 +146,11 @@ impl Footer {
         self.root_layout.dtype()
     }
 
+    /// Returns the approximate size of the footer in bytes, used for caching and memory management.
+    pub fn approx_byte_size(&self) -> Option<usize> {
+        self.approx_byte_size
+    }
+
     /// Returns the number of rows in the file.
     pub fn row_count(&self) -> u64 {
         self.root_layout.row_count()