Segment caching (#2934)

gatesn · web-flow · commit 0ef8c0347597 · 2025-04-09T06:52:06.000-07:00
Initial read segments are now explicitly cached separately from the
segment cache, and we default to _always_ store segments in the segment
cache as they are read.
diff --git a/pyvortex/src/file.rs b/pyvortex/src/file.rs
@@ -13,6 +13,7 @@ use vortex::dtype::{DType, PType};
 use vortex::error::{VortexExpect, vortex_err};
 use vortex::expr::{ExprRef, ident, select};
 use vortex::file::scan::SplitBy;
+use vortex::file::segments::MokaSegmentCache;
 use vortex::file::{VortexFile, VortexOpenOptions};
 use vortex::io::TokioFile;
 use vortex::stream::{ArrayStream, ArrayStreamAdapter, ArrayStreamExt};
@@ -38,7 +39,11 @@ pub(crate) fn init(py: Python, parent: &Bound<PyModule>) -> PyResult<()> {
 
 #[pyfunction]
 pub fn open(path: &str) -> PyResult<PyVortexFile> {
-    let vxf = TOKIO_RUNTIME.block_on(VortexOpenOptions::file().open(TokioFile::open(path)?))?;
+    let vxf = TOKIO_RUNTIME.block_on(
+        VortexOpenOptions::file()
+            .with_segment_cache(Arc::new(MokaSegmentCache::new(256 << 20)))
+            .open(TokioFile::open(path)?),
+    )?;
     Ok(PyVortexFile { vxf })
 }
 
diff --git a/vortex-file/src/generic.rs b/vortex-file/src/generic.rs
@@ -1,14 +1,16 @@
 use std::sync::{Arc, Mutex};
 
 use futures::{StreamExt, pin_mut};
-use moka::future::CacheBuilder;
 use vortex_error::{VortexExpect, VortexResult};
 use vortex_io::{Dispatch, InstrumentedReadAt, IoDispatcher, VortexReadAt};
 use vortex_layout::segments::{SegmentEvents, SegmentSource};
 use vortex_metrics::VortexMetrics;
 
 use crate::driver::CoalescedDriver;
-use crate::segments::{CachedSegmentSource, InMemorySegmentCache, SegmentCache};
+use crate::segments::{
+    InitialReadSegmentCache, MokaSegmentCache, SegmentCache, SegmentCacheMetrics,
+    SegmentCacheSourceAdapter,
+};
 use crate::{FileType, SegmentSourceFactory, SegmentSpec, VortexFile, VortexOpenOptions};
 
 /// A type of Vortex file that supports any [`VortexReadAt`] implementation.
@@ -27,10 +29,9 @@ impl VortexOpenOptions<GenericVortexFile> {
     /// Open a file using the provided [`VortexReadAt`] implementation.
     pub fn file() -> Self {
         Self::new(Default::default())
-            .with_segment_cache(Arc::new(InMemorySegmentCache::new(
-                // For now, use a fixed 1GB overhead.
-                CacheBuilder::new(1 << 30),
-            )))
+            // Start with an initial in-memory cache of 256MB.
+            // TODO(ngates): would it be better to default to a home directory disk cache?
+            .with_segment_cache(Arc::new(MokaSegmentCache::new(256 << 20)))
             .with_initial_read_size(Self::INITIAL_READ_SIZE)
     }
 
@@ -42,10 +43,18 @@ impl VortexOpenOptions<GenericVortexFile> {
     pub async fn open<R: VortexReadAt + Send>(self, read: R) -> VortexResult<VortexFile> {
         let footer = self.read_footer(&read).await?;
 
+        let segment_cache = Arc::new(SegmentCacheMetrics::new(
+            InitialReadSegmentCache {
+                initial: self.initial_read_segments,
+                fallback: self.segment_cache,
+            },
+            self.metrics.clone(),
+        ));
+
         let segment_source_factory = Arc::new(GenericVortexFileIo {
             read: Mutex::new(read),
             segment_map: footer.segment_map().clone(),
-            segment_cache: self.segment_cache,
+            segment_cache,
             options: self.options,
         });
 
@@ -70,7 +79,7 @@ impl<R: VortexReadAt + Send> SegmentSourceFactory for GenericVortexFileIo<R> {
         let (segment_source, events) = SegmentEvents::create();
 
         // Wrap the source to resolve segments from the initial read cache.
-        let segment_source = Arc::new(CachedSegmentSource::new(
+        let segment_source = Arc::new(SegmentCacheSourceAdapter::new(
             self.segment_cache.clone(),
             segment_source,
         ));
diff --git a/vortex-file/src/open.rs b/vortex-file/src/open.rs
@@ -1,9 +1,8 @@
-use std::sync::Arc;
+use std::sync::{Arc, RwLock};
 
 use flatbuffers::root;
-use futures::stream::FuturesUnordered;
-use futures::{StreamExt, TryStreamExt, stream};
 use vortex_array::ArrayRegistry;
+use vortex_array::aliases::hash_map::HashMap;
 use vortex_buffer::{Alignment, ByteBuffer, ByteBufferMut};
 use vortex_dtype::DType;
 use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
@@ -18,7 +17,7 @@ use crate::segments::{NoOpSegmentCache, SegmentCache};
 use crate::{DEFAULT_REGISTRY, EOF_SIZE, MAGIC_BYTES, MAX_FOOTER_SIZE, VERSION};
 
 pub trait FileType: Sized {
-    type Options: Clone;
+    type Options;
 }
 
 /// Open options for a Vortex file reader.
@@ -36,9 +35,9 @@ pub struct VortexOpenOptions<F: FileType> {
     /// An optional, externally provided, file layout.
     // TODO(ngates): add an optional DType so we only read the layout segment.
     footer: Option<Footer>,
-    /// TODO(ngates): rename to initial_segments: HashMap<SegmentId, ByteBuffer>
     pub(crate) segment_cache: Arc<dyn SegmentCache>,
     pub(crate) initial_read_size: u64,
+    pub(crate) initial_read_segments: RwLock<HashMap<SegmentId, ByteBuffer>>,
     pub(crate) metrics: VortexMetrics,
 }
 
@@ -53,6 +52,7 @@ impl<F: FileType> VortexOpenOptions<F> {
             footer: None,
             segment_cache: Arc::new(NoOpSegmentCache),
             initial_read_size: 0,
+            initial_read_segments: Default::default(),
             metrics: VortexMetrics::default(),
         }
     }
@@ -201,8 +201,7 @@ impl<F: FileType> VortexOpenOptions<F> {
 
         // If the initial read happened to cover any segments, then we can populate the
         // segment cache
-        self.populate_segments(initial_offset, &initial_read, &footer)
-            .await?;
+        self.populate_initial_segments(initial_offset, &initial_read, &footer);
 
         Ok(footer)
     }
@@ -265,30 +264,31 @@ impl<F: FileType> VortexOpenOptions<F> {
     }
 
     /// Populate segments in the cache that were covered by the initial read.
-    async fn populate_segments(
+    fn populate_initial_segments(
         &self,
         initial_offset: u64,
         initial_read: &ByteBuffer,
         footer: &Footer,
-    ) -> VortexResult<()> {
-        stream::iter(
-            footer
-                .segment_map()
-                .iter()
-                .enumerate()
-                .filter(|(_, segment)| segment.offset > initial_offset)
-                .map(|(idx, segment)| async move {
-                    let segment_id = SegmentId::from(u32::try_from(idx)?);
-                    let offset = usize::try_from(segment.offset - initial_offset)?;
-                    let buffer = initial_read
-                        .slice(offset..offset + (segment.length as usize))
-                        .aligned(segment.alignment);
-                    self.segment_cache.put(segment_id, buffer).await
-                }),
-        )
-        .collect::<FuturesUnordered<_>>()
-        .await
-        .try_collect::<()>()
-        .await
+    ) {
+        let first_idx = footer
+            .segment_map()
+            .partition_point(|segment| segment.offset < initial_offset);
+
+        let mut initial_segments = self
+            .initial_read_segments
+            .write()
+            .vortex_expect("poisoned lock");
+
+        for idx in first_idx..footer.segment_map().len() {
+            let segment = &footer.segment_map()[idx];
+            let segment_id =
+                SegmentId::from(u32::try_from(idx).vortex_expect("Invalid segment ID"));
+            let offset =
+                usize::try_from(segment.offset - initial_offset).vortex_expect("Invalid offset");
+            let buffer = initial_read
+                .slice(offset..offset + (segment.length as usize))
+                .aligned(segment.alignment);
+            initial_segments.insert(segment_id, buffer);
+        }
     }
 }
diff --git a/vortex-file/src/segments/cache.rs b/vortex-file/src/segments/cache.rs
@@ -1,17 +1,21 @@
+use std::sync::{Arc, RwLock};
+
 use async_trait::async_trait;
+use futures::FutureExt;
 use moka::future::{Cache, CacheBuilder};
 use moka::policy::EvictionPolicy;
 use rustc_hash::FxBuildHasher;
+use vortex_array::aliases::hash_map::HashMap;
 use vortex_buffer::ByteBuffer;
 use vortex_error::{VortexExpect, VortexResult};
-use vortex_layout::segments::SegmentId;
+use vortex_layout::segments::{SegmentFuture, SegmentId, SegmentSource};
+use vortex_metrics::{Counter, VortexMetrics};
 
 /// A cache for storing and retrieving individual segment data.
 #[async_trait]
 pub trait SegmentCache: Send + Sync {
     async fn get(&self, id: SegmentId) -> VortexResult<Option<ByteBuffer>>;
     async fn put(&self, id: SegmentId, buffer: ByteBuffer) -> VortexResult<()>;
-    async fn remove(&self, id: SegmentId) -> VortexResult<()>;
 }
 
 pub(crate) struct NoOpSegmentCache;
@@ -25,32 +29,31 @@ impl SegmentCache for NoOpSegmentCache {
     async fn put(&self, _id: SegmentId, _buffer: ByteBuffer) -> VortexResult<()> {
         Ok(())
     }
-
-    async fn remove(&self, _id: SegmentId) -> VortexResult<()> {
-        Ok(())
-    }
 }
 
-pub(crate) struct InMemorySegmentCache(Cache<SegmentId, ByteBuffer, FxBuildHasher>);
+/// A [`SegmentCache`] based around an in-memory Moka cache.
+pub struct MokaSegmentCache(Cache<SegmentId, ByteBuffer, FxBuildHasher>);
 
-impl InMemorySegmentCache {
-    pub fn new(builder: CacheBuilder<SegmentId, ByteBuffer, Cache<SegmentId, ByteBuffer>>) -> Self {
+impl MokaSegmentCache {
+    pub fn new(max_capacity_bytes: u64) -> Self {
         Self(
-            builder
+            CacheBuilder::new(max_capacity_bytes)
+                .name("vortex-segment-cache")
                 // Weight each segment by the number of bytes in the buffer.
-                .weigher(|_, buffer| {
+                .weigher(|_, buffer: &ByteBuffer| {
                     u32::try_from(buffer.len().min(u32::MAX as usize)).vortex_expect("must fit")
                 })
-                // We configure LRU instead of LFU since we're likely to re-read segments between
-                // filter and projection.
-                .eviction_policy(EvictionPolicy::lru())
+                // We configure LFU (vs LRU) since the cache is mostly used when re-reading the
+                // same file - it is _not_ used when reading the same segments during a single
+                // scan.
+                .eviction_policy(EvictionPolicy::tiny_lfu())
                 .build_with_hasher(FxBuildHasher),
         )
     }
 }
 
 #[async_trait]
-impl SegmentCache for InMemorySegmentCache {
+impl SegmentCache for MokaSegmentCache {
     async fn get(&self, id: SegmentId) -> VortexResult<Option<ByteBuffer>> {
         Ok(self.0.get(&id).await)
     }
@@ -59,9 +62,99 @@ impl SegmentCache for InMemorySegmentCache {
         self.0.insert(id, buffer).await;
         Ok(())
     }
+}
+
+/// Segment cache containing the initial read segments.
+pub(crate) struct InitialReadSegmentCache {
+    pub(crate) initial: RwLock<HashMap<SegmentId, ByteBuffer>>,
+    pub(crate) fallback: Arc<dyn SegmentCache>,
+}
+
+#[async_trait]
+impl SegmentCache for InitialReadSegmentCache {
+    async fn get(&self, id: SegmentId) -> VortexResult<Option<ByteBuffer>> {
+        if let Some(buffer) = self.initial.read().vortex_expect("poisoned lock").get(&id) {
+            return Ok(Some(buffer.clone()));
+        }
+        self.fallback.get(id).await
+    }
+
+    async fn put(&self, id: SegmentId, buffer: ByteBuffer) -> VortexResult<()> {
+        self.fallback.put(id, buffer).await
+    }
+}
+
+pub struct SegmentCacheMetrics<C> {
+    segment_cache: C,
+
+    hits: Arc<Counter>,
+    misses: Arc<Counter>,
+    stores: Arc<Counter>,
+}
 
-    async fn remove(&self, id: SegmentId) -> VortexResult<()> {
-        self.0.invalidate(&id).await;
+impl<C: SegmentCache> SegmentCacheMetrics<C> {
+    pub fn new(segment_cache: C, metrics: VortexMetrics) -> Self {
+        Self {
+            segment_cache,
+            hits: metrics.counter("vortex.file.segments.cache.hits"),
+            misses: metrics.counter("vortex.file.segments.cache.misses"),
+            stores: metrics.counter("vortex.file.segments.cache.stores"),
+        }
+    }
+}
+
+#[async_trait]
+impl<C: SegmentCache> SegmentCache for SegmentCacheMetrics<C> {
+    async fn get(&self, id: SegmentId) -> VortexResult<Option<ByteBuffer>> {
+        let result = self.segment_cache.get(id).await?;
+        if result.is_some() {
+            self.hits.inc()
+        } else {
+            self.misses.inc()
+        }
+        Ok(result)
+    }
+
+    async fn put(&self, id: SegmentId, buffer: ByteBuffer) -> VortexResult<()> {
+        self.segment_cache.put(id, buffer).await?;
+        self.stores.inc();
         Ok(())
     }
 }
+
+pub struct SegmentCacheSourceAdapter {
+    cache: Arc<dyn SegmentCache>,
+    source: Arc<dyn SegmentSource>,
+}
+
+impl SegmentCacheSourceAdapter {
+    pub fn new(cache: Arc<dyn SegmentCache>, source: Arc<dyn SegmentSource>) -> Self {
+        Self { cache, source }
+    }
+}
+
+impl SegmentSource for SegmentCacheSourceAdapter {
+    fn request(&self, id: SegmentId, for_whom: &Arc<str>) -> SegmentFuture {
+        let cache = self.cache.clone();
+        let delegate = self.source.request(id, for_whom);
+        let for_whom = for_whom.clone();
+
+        async move {
+            if let Ok(Some(segment)) = cache.get(id).await {
+                log::debug!("Resolved segment {} for {} from cache", id, &for_whom);
+                return Ok(segment);
+            }
+            let result = delegate.await?;
+            if let Err(e) = cache.put(id, result.clone()).await {
+                log::warn!(
+                    "Failed to store segment {} for {} in cache: {}",
+                    id,
+                    &for_whom,
+                    e
+                );
+            }
+            Ok(result)
+        }
+        .boxed()
+    }
+}
diff --git a/vortex-file/src/segments/cached.rs b/vortex-file/src/segments/cached.rs
diff --git a/vortex-file/src/segments/mod.rs b/vortex-file/src/segments/mod.rs
@@ -1,6 +1,4 @@
 mod cache;
-mod cached;
 pub(crate) mod writer;
 
 pub use cache::*;
-pub use cached::*;
