Global segment cache across DataFusion (#2958)

Share a single capped segment cache across DataFusion

Author: gatesn

vortex-datafusion/src/persistent/cache.rs

@@ -1,5 +1,6 @@
 use std::sync::Arc;
 
+use async_trait::async_trait;
 use chrono::{DateTime, Utc};
 use datafusion_common::ScalarValue;
 use moka::future::Cache;
@@ -8,77 +9,76 @@ use object_store::{ObjectMeta, ObjectStore};
 use vortex_array::ArrayRegistry;
 use vortex_array::aliases::DefaultHashBuilder;
 use vortex_array::stats::{Precision, Stat};
+use vortex_buffer::ByteBuffer;
 use vortex_dtype::DType;
 use vortex_error::{VortexError, VortexResult, vortex_err};
+use vortex_file::segments::SegmentCache;
 use vortex_file::{Footer, SegmentSpec, VortexFile, VortexOpenOptions};
 use vortex_layout::LayoutRegistry;
 use vortex_layout::segments::SegmentId;
 use vortex_metrics::VortexMetrics;
 
 #[derive(Clone)]
 pub(crate) struct VortexFileCache {
-    inner: Cache<Key, VortexFile, DefaultHashBuilder>,
+    file_cache: Cache<FileKey, VortexFile, DefaultHashBuilder>,
+    segment_cache: Cache<SegmentKey, ByteBuffer, DefaultHashBuilder>,
     array_registry: Arc<ArrayRegistry>,
     layout_registry: Arc<LayoutRegistry>,
     metrics: VortexMetrics,
 }
 
-#[derive(Hash, Eq, PartialEq, Debug)]
-pub(crate) struct Key {
-    location: Path,
+/// Cache key for a [`VortexFile`].
+#[derive(Hash, Eq, PartialEq, Debug, Clone)]
+struct FileKey {
+    location: Arc<Path>,
     m_time: DateTime<Utc>,
 }
 
-impl From<&ObjectMeta> for Key {
+impl From<&ObjectMeta> for FileKey {
     fn from(value: &ObjectMeta) -> Self {
         Self {
-            location: value.location.clone(),
+            location: Arc::new(value.location.clone()),
             m_time: value.last_modified,
         }
     }
 }
 
-/// Approximate the in-memory size of a layout
-fn estimate_layout_size(footer: &Footer) -> usize {
-    let segments_size = footer.segment_map().len() * size_of::<SegmentSpec>();
-    let stats_size = footer
-        .statistics()
-        .iter()
-        .map(|v| {
-            v.iter()
-                .map(|_| size_of::<Stat>() + size_of::<Precision<ScalarValue>>())
-                .sum::<usize>()
-        })
-        .sum::<usize>();
-
-    let root_layout = footer.layout();
-    let layout_size = size_of::<DType>()
-        + root_layout.metadata().map(|b| b.len()).unwrap_or_default()
-        + root_layout.nsegments() * size_of::<SegmentId>();
-
-    segments_size + stats_size + layout_size
+/// Global cache key for a segment.
+#[derive(Hash, Eq, PartialEq, Debug)]
+struct SegmentKey {
+    file: FileKey,
+    segment_id: SegmentId,
 }
 
 impl VortexFileCache {
     pub fn new(
         size_mb: usize,
+        segment_size_mb: usize,
         array_registry: Arc<ArrayRegistry>,
         layout_registry: Arc<LayoutRegistry>,
         metrics: VortexMetrics,
     ) -> Self {
-        let inner = Cache::builder()
+        let file_cache = Cache::builder()
             .max_capacity(size_mb as u64 * (2 << 20))
-            .eviction_listener(|k: Arc<Key>, _v: VortexFile, cause| {
-                log::trace!("Removed {} due to {:?}", k.location, cause);
+            .eviction_listener(|k: Arc<FileKey>, _v: VortexFile, cause| {
+                log::trace!("Removed {:?} due to {:?}", k, cause);
             })
             .weigher(|_k, vxf| {
-                let size = estimate_layout_size(vxf.footer());
-                u32::try_from(size).unwrap_or(u32::MAX)
+                u32::try_from(estimate_layout_size(vxf.footer())).unwrap_or(u32::MAX)
             })
             .build_with_hasher(DefaultHashBuilder::default());
 
+        let segment_cache = Cache::builder()
+            .max_capacity(segment_size_mb as u64 * (2 << 20))
+            .eviction_listener(|k: Arc<SegmentKey>, _v: ByteBuffer, cause| {
+                log::trace!("Removed {:?} due to {:?}", k, cause);
+            })
+            .weigher(|_k, v| u32::try_from(v.len()).unwrap_or(u32::MAX))
+            .build_with_hasher(DefaultHashBuilder::default());
+
         Self {
-            inner,
+            file_cache,
+            segment_cache,
             array_registry,
             layout_registry,
             metrics,
@@ -90,9 +90,10 @@ impl VortexFileCache {
         object: &ObjectMeta,
         object_store: Arc<dyn ObjectStore>,
     ) -> VortexResult<VortexFile> {
-        self.inner
+        let file_key = FileKey::from(object);
+        self.file_cache
             .try_get_with(
-                Key::from(object),
+                file_key.clone(),
                 VortexOpenOptions::file()
                     .with_array_registry(self.array_registry.clone())
                     .with_layout_registry(self.layout_registry.clone())
@@ -101,6 +102,10 @@ impl VortexFileCache {
                             .child_with_tags([("filename", object.location.to_string())]),
                     )
                     .with_file_size(object.size as u64)
+                    .with_segment_cache(Arc::new(VortexFileSegmentCache {
+                        file_key,
+                        segment_cache: self.segment_cache.clone(),
+                    }))
                     .open_object_store(&object_store, object.location.as_ref()),
             )
             .await
@@ -109,3 +114,56 @@ impl VortexFileCache {
             })
     }
 }
+
+/// A [`SegmentCache`] implementation that uses the shared global segment cache.
+struct VortexFileSegmentCache {
+    file_key: FileKey,
+    segment_cache: Cache<SegmentKey, ByteBuffer, DefaultHashBuilder>,
+}
+
+#[async_trait]
+impl SegmentCache for VortexFileSegmentCache {
+    async fn get(&self, segment_id: SegmentId) -> VortexResult<Option<ByteBuffer>> {
+        Ok(self
+            .segment_cache
+            .get(&SegmentKey {
+                file: self.file_key.clone(),
+                segment_id,
+            })
+            .await)
+    }
+
+    async fn put(&self, segment_id: SegmentId, buffer: ByteBuffer) -> VortexResult<()> {
+        self.segment_cache
+            .insert(
+                SegmentKey {
+                    file: self.file_key.clone(),
+                    segment_id,
+                },
+                buffer,
+            )
+            .await;
+        Ok(())
+    }
+}
+
+/// Approximate the in-memory size of a layout
+fn estimate_layout_size(footer: &Footer) -> usize {
+    let segments_size = footer.segment_map().len() * size_of::<SegmentSpec>();
+    let stats_size = footer
+        .statistics()
+        .iter()
+        .map(|v| {
+            v.iter()
+                .map(|_| size_of::<Stat>() + size_of::<Precision<ScalarValue>>())
+                .sum::<usize>()
+        })
+        .sum::<usize>();
+
+    let root_layout = footer.layout();
+    let layout_size = size_of::<DType>()
+        + root_layout.metadata().map(|b| b.len()).unwrap_or_default()
+        + root_layout.nsegments() * size_of::<SegmentId>();
+
+    segments_size + stats_size + layout_size
+}

vortex-datafusion/src/persistent/format.rs

@@ -57,12 +57,17 @@ impl Debug for VortexFormat {
 #[derive(Debug)]
 pub struct VortexFormatOptions {
     /// The size of the in-memory [`vortex_file::Footer`] cache.
-    pub cache_size_mb: usize,
+    pub footer_cache_size_mb: usize,
+    /// The size of the in-memory segment cache.
+    pub segment_cache_size_mb: usize,
 }
 
 impl Default for VortexFormatOptions {
     fn default() -> Self {
-        Self { cache_size_mb: 256 }
+        Self {
+            footer_cache_size_mb: 64,
+            segment_cache_size_mb: 1024,
+        }
     }
 }
 
@@ -141,7 +146,8 @@ impl VortexFormat {
         let opts = VortexFormatOptions::default();
         Self {
             file_cache: VortexFileCache::new(
-                opts.cache_size_mb,
+                opts.footer_cache_size_mb,
+                opts.segment_cache_size_mb,
                 array_registry,
                 layout_registry,
                 metrics.clone(),