vortex_layout/segments/
cache.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use async_trait::async_trait;
7use futures::FutureExt;
8use moka::future::Cache;
9use moka::future::CacheBuilder;
10use moka::policy::EvictionPolicy;
11use rustc_hash::FxBuildHasher;
12use vortex_array::buffer::BufferHandle;
13use vortex_buffer::ByteBuffer;
14use vortex_error::VortexExpect;
15use vortex_error::VortexResult;
16use vortex_metrics::Counter;
17use vortex_metrics::Label;
18use vortex_metrics::MetricBuilder;
19use vortex_metrics::MetricsRegistry;
20
21use crate::segments::SegmentFuture;
22use crate::segments::SegmentId;
23use crate::segments::SegmentSource;
24
25/// Cache for individual segment byte buffers.
26///
27/// Caches are optional and operate above a [`SegmentSource`]. They should only store host buffers:
28/// device buffers and other non-host handles should be passed through uncached.
29///
30/// To bring your own cache, implement this trait and give it to the file with
31/// `VortexOpenOptions::with_segment_cache`. A [`SegmentId`] is unique only within one file, so a
32/// cache that many files share must also key its entries by file. Give each file its own
33/// implementation of this trait that holds the [`SegmentSourceId`] of that file and adds it to
34/// the key, as [`MokaSegmentCache::for_file`] does with [`FileSegmentCache`].
35///
36/// The Python bindings accept only a [`MokaSegmentCache`]. A cache written in Python can use the
37/// same design as `PyReadable` in `vortex-python`:
38///
39/// - The Python object has the methods `get(key: str, segment_id: int) -> Buffer | None` and
40///   `put(key: str, segment_id: int, data: Buffer) -> None`. `vortex.open` and
41///   `vortex.open_readable` accept it with the `cache_key` that they accept now, and the stubs
42///   declare it as a `typing.Protocol`. Then users can use `cachetools`, `diskcache`, Redis, or a
43///   cache that many processes share.
44/// - A Rust adapter holds the Python object and the [`SegmentSourceId`] of one file, and
45///   implements this trait. `get` calls Python through `spawn_blocking` and `Python::attach`, as
46///   `PyReadable::read_at` does, and wraps the buffer that Python returns without a copy.
47///   `put` gives Python a read-only object that owns a clone of the [`ByteBuffer`] and exports
48///   it through `__getbuffer__`, as `ReadBuffer` does, so Python can keep it without a copy.
49#[async_trait]
50pub trait SegmentCache: Send + Sync {
51    /// Return a cached segment, or `None` on cache miss.
52    async fn get(&self, id: SegmentId) -> VortexResult<Option<ByteBuffer>>;
53    /// Store a segment in the cache.
54    async fn put(&self, id: SegmentId, buffer: ByteBuffer) -> VortexResult<()>;
55}
56
57/// Segment cache implementation that never stores anything.
58pub struct NoOpSegmentCache;
59
60#[async_trait]
61impl SegmentCache for NoOpSegmentCache {
62    async fn get(&self, _id: SegmentId) -> VortexResult<Option<ByteBuffer>> {
63        Ok(None)
64    }
65
66    async fn put(&self, _id: SegmentId, _buffer: ByteBuffer) -> VortexResult<()> {
67        Ok(())
68    }
69}
70
71/// Which segments a [`MokaSegmentCache`] evicts when it is full.
72#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
73pub enum SegmentEviction {
74    /// Admit every new segment and evict the least recently used. Suits a cache shared by readers
75    /// that move from file to file, where a newly opened file's segments must displace older ones.
76    #[default]
77    Lru,
78    /// Admit a new segment only if it is likely to be used more often than the one it would
79    /// evict. Suits repeated scans of one file larger than the cache, which LRU would evict
80    /// entirely on every pass, but rejects a new file's segments while older ones are popular.
81    TinyLfu,
82}
83
84/// Identifies the contents of one segment source, usually a file, in a cache that many sources
85/// share.
86///
87/// Sources with the same ID share cached segments, so an ID must identify the contents of the
88/// source, not only its location. A file that has changed must get a new ID.
89#[derive(Clone, Debug, PartialEq, Eq, Hash)]
90pub struct SegmentSourceId(Arc<str>);
91
92impl SegmentSourceId {
93    /// Construct an ID from a string that identifies the contents of a source.
94    pub fn new(id: impl Into<Arc<str>>) -> Self {
95        Self(id.into())
96    }
97}
98
99impl From<&str> for SegmentSourceId {
100    fn from(id: &str) -> Self {
101        Self::new(id)
102    }
103}
104
105impl From<String> for SegmentSourceId {
106    fn from(id: String) -> Self {
107        Self::new(id)
108    }
109}
110
111impl From<Arc<str>> for SegmentSourceId {
112    fn from(id: Arc<str>) -> Self {
113        Self(id)
114    }
115}
116
117/// An in-memory Moka cache of segments, capped by total buffer bytes, that any number of files can
118/// share.
119///
120/// A [`SegmentId`] is unique only within one file, so files use the cache through a
121/// [`FileSegmentCache`] from [`Self::for_file`], which adds a [`SegmentSourceId`]. Opening the
122/// same file again with the same ID reuses the segments an earlier open read.
123#[derive(Clone)]
124pub struct MokaSegmentCache(Cache<(SegmentSourceId, SegmentId), ByteBuffer, FxBuildHasher>);
125
126impl MokaSegmentCache {
127    /// Construct a Moka-backed cache capped by total buffer bytes.
128    pub fn new(max_capacity_bytes: u64, eviction: SegmentEviction) -> Self {
129        let eviction_policy = match eviction {
130            SegmentEviction::Lru => EvictionPolicy::lru(),
131            SegmentEviction::TinyLfu => EvictionPolicy::tiny_lfu(),
132        };
133        Self(
134            CacheBuilder::new(max_capacity_bytes)
135                .name("vortex-segment-cache")
136                // Weight each segment by the number of bytes in the buffer.
137                .weigher(|_, buffer: &ByteBuffer| {
138                    u32::try_from(buffer.len().min(u32::MAX as usize)).vortex_expect("must fit")
139                })
140                .eviction_policy(eviction_policy)
141                .build_with_hasher(FxBuildHasher),
142        )
143    }
144
145    /// The view of this cache for the file identified by `source_id`.
146    pub fn for_file(&self, source_id: impl Into<SegmentSourceId>) -> FileSegmentCache {
147        FileSegmentCache {
148            cache: self.clone(),
149            source_id: source_id.into(),
150        }
151    }
152
153    /// The total bytes of the cached segments.
154    ///
155    /// Recent inserts and evictions may not be counted yet; call [`Self::run_pending_tasks`] first
156    /// for an exact value.
157    pub fn weighted_size(&self) -> u64 {
158        self.0.weighted_size()
159    }
160
161    /// The number of cached segments, with the same caveat as [`Self::weighted_size`].
162    pub fn entry_count(&self) -> u64 {
163        self.0.entry_count()
164    }
165
166    /// Apply pending inserts and evictions, so that the counts are exact.
167    pub async fn run_pending_tasks(&self) {
168        self.0.run_pending_tasks().await;
169    }
170
171    /// Remove every cached segment.
172    pub fn invalidate_all(&self) {
173        self.0.invalidate_all();
174    }
175}
176
177/// One file's view of a [`MokaSegmentCache`]; see [`MokaSegmentCache::for_file`].
178#[derive(Clone)]
179pub struct FileSegmentCache {
180    cache: MokaSegmentCache,
181    source_id: SegmentSourceId,
182}
183
184#[async_trait]
185impl SegmentCache for FileSegmentCache {
186    async fn get(&self, id: SegmentId) -> VortexResult<Option<ByteBuffer>> {
187        Ok(self.cache.0.get(&(self.source_id.clone(), id)).await)
188    }
189
190    async fn put(&self, id: SegmentId, buffer: ByteBuffer) -> VortexResult<()> {
191        self.cache
192            .0
193            .insert((self.source_id.clone(), id), buffer)
194            .await;
195        Ok(())
196    }
197}
198
199/// Wrapper for [`SegmentCache`] that tracks its hit rate.
200pub struct InstrumentedSegmentCache<C> {
201    segment_cache: C,
202
203    hits: Counter,
204    misses: Counter,
205    stores: Counter,
206}
207
208impl<C: SegmentCache> InstrumentedSegmentCache<C> {
209    /// Wrap a segment cache and record hit/miss/store metrics with the supplied labels.
210    pub fn new(
211        segment_cache: C,
212        metrics_registry: &dyn MetricsRegistry,
213        labels: Vec<Label>,
214    ) -> Self {
215        Self {
216            segment_cache,
217            hits: MetricBuilder::new(metrics_registry)
218                .add_labels(labels.clone())
219                .counter("vortex.file.segments.cache.hits"),
220            misses: MetricBuilder::new(metrics_registry)
221                .add_labels(labels.clone())
222                .counter("vortex.file.segments.cache.misses"),
223            stores: MetricBuilder::new(metrics_registry)
224                .add_labels(labels)
225                .counter("vortex.file.segments.cache.stores"),
226        }
227    }
228}
229
230#[async_trait]
231impl<C: SegmentCache> SegmentCache for InstrumentedSegmentCache<C> {
232    async fn get(&self, id: SegmentId) -> VortexResult<Option<ByteBuffer>> {
233        let result = self.segment_cache.get(id).await?;
234        if result.is_some() {
235            self.hits.add(1);
236        } else {
237            self.misses.add(1);
238        }
239        Ok(result)
240    }
241
242    async fn put(&self, id: SegmentId, buffer: ByteBuffer) -> VortexResult<()> {
243        self.segment_cache.put(id, buffer).await?;
244        self.stores.add(1);
245        Ok(())
246    }
247}
248
249/// [`SegmentSource`] wrapper that consults a [`SegmentCache`] before the underlying source.
250pub struct SegmentCacheSourceAdapter {
251    cache: Arc<dyn SegmentCache>,
252    source: Arc<dyn SegmentSource>,
253}
254
255impl SegmentCacheSourceAdapter {
256    /// Construct a cache-fronted source.
257    pub fn new(cache: Arc<dyn SegmentCache>, source: Arc<dyn SegmentSource>) -> Self {
258        Self { cache, source }
259    }
260}
261
262impl SegmentSource for SegmentCacheSourceAdapter {
263    fn request(&self, id: SegmentId) -> SegmentFuture {
264        let cache = Arc::clone(&self.cache);
265        let delegate = self.source.request(id);
266
267        async move {
268            if let Ok(Some(segment)) = cache.get(id).await {
269                tracing::debug!("Resolved segment {} from cache", id);
270                return Ok(BufferHandle::new_host(segment));
271            }
272            let result = delegate.await?;
273            // Cache only CPU buffers; device buffers are not cached.
274            if let Some(buffer) = result.as_host_opt()
275                && let Err(e) = cache.put(id, buffer.clone()).await
276            {
277                tracing::warn!("Failed to store segment {} in cache: {}", id, e);
278            }
279            Ok(result)
280        }
281        .boxed()
282    }
283}
284
285#[cfg(test)]
286mod tests {
287    use rstest::rstest;
288    use vortex_buffer::ByteBuffer;
289
290    use super::*;
291
292    #[rstest]
293    #[case::lru(SegmentEviction::Lru)]
294    #[case::tiny_lfu(SegmentEviction::TinyLfu)]
295    #[tokio::test]
296    async fn shared_cache_separates_files(#[case] eviction: SegmentEviction) -> VortexResult<()> {
297        let shared = MokaSegmentCache::new(1 << 20, eviction);
298        let a = shared.for_file("a");
299        let b = shared.for_file("b");
300        let id = SegmentId::from(0);
301
302        a.put(id, ByteBuffer::copy_from(b"from a")).await?;
303        assert_eq!(a.get(id).await?.as_deref(), Some(b"from a".as_slice()));
304        assert!(b.get(id).await?.is_none());
305
306        // A later view with the same source ID sees what the earlier one stored.
307        let a_again = shared.for_file("a");
308        assert_eq!(
309            a_again.get(id).await?.as_deref(),
310            Some(b"from a".as_slice())
311        );
312        Ok(())
313    }
314
315    #[rstest]
316    #[case::lru(SegmentEviction::Lru)]
317    #[case::tiny_lfu(SegmentEviction::TinyLfu)]
318    #[tokio::test]
319    async fn shared_cache_is_capped_by_bytes(
320        #[case] eviction: SegmentEviction,
321    ) -> VortexResult<()> {
322        let shared = MokaSegmentCache::new(1000, eviction);
323        let file = shared.for_file("file");
324        for i in 0..10u32 {
325            file.put(SegmentId::from(i), ByteBuffer::copy_from(vec![0u8; 400]))
326                .await?;
327        }
328
329        shared.run_pending_tasks().await;
330        assert!(shared.weighted_size() <= 1000);
331        assert!(shared.entry_count() <= 2);
332
333        // LRU admits every segment, so the most recently stored one survives.
334        if eviction == SegmentEviction::Lru {
335            assert!(file.get(SegmentId::from(9)).await?.is_some());
336        }
337        Ok(())
338    }
339}
vortex_layout/segments/cache.rs

vortex_layout/segments/
cache.rs
