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#[async_trait]
30pub trait SegmentCache: Send + Sync {
31    /// Return a cached segment, or `None` on cache miss.
32    async fn get(&self, id: SegmentId) -> VortexResult<Option<ByteBuffer>>;
33    /// Store a segment in the cache.
34    async fn put(&self, id: SegmentId, buffer: ByteBuffer) -> VortexResult<()>;
35}
36
37/// Segment cache implementation that never stores anything.
38pub struct NoOpSegmentCache;
39
40#[async_trait]
41impl SegmentCache for NoOpSegmentCache {
42    async fn get(&self, _id: SegmentId) -> VortexResult<Option<ByteBuffer>> {
43        Ok(None)
44    }
45
46    async fn put(&self, _id: SegmentId, _buffer: ByteBuffer) -> VortexResult<()> {
47        Ok(())
48    }
49}
50
51/// A [`SegmentCache`] based around an in-memory Moka cache.
52pub struct MokaSegmentCache(Cache<SegmentId, ByteBuffer, FxBuildHasher>);
53
54impl MokaSegmentCache {
55    /// Construct a Moka-backed cache capped by total buffer bytes.
56    pub fn new(max_capacity_bytes: u64) -> Self {
57        Self(
58            CacheBuilder::new(max_capacity_bytes)
59                .name("vortex-segment-cache")
60                // Weight each segment by the number of bytes in the buffer.
61                .weigher(|_, buffer: &ByteBuffer| {
62                    u32::try_from(buffer.len().min(u32::MAX as usize)).vortex_expect("must fit")
63                })
64                // We configure LFU (vs LRU) since the cache is mostly used when re-reading the
65                // same file - it is _not_ used when reading the same segments during a single
66                // scan.
67                .eviction_policy(EvictionPolicy::tiny_lfu())
68                .build_with_hasher(FxBuildHasher),
69        )
70    }
71}
72
73#[async_trait]
74impl SegmentCache for MokaSegmentCache {
75    async fn get(&self, id: SegmentId) -> VortexResult<Option<ByteBuffer>> {
76        Ok(self.0.get(&id).await)
77    }
78
79    async fn put(&self, id: SegmentId, buffer: ByteBuffer) -> VortexResult<()> {
80        self.0.insert(id, buffer).await;
81        Ok(())
82    }
83}
84
85/// Wrapper for [`SegmentCache`] that tracks its hit rate.
86pub struct InstrumentedSegmentCache<C> {
87    segment_cache: C,
88
89    hits: Counter,
90    misses: Counter,
91    stores: Counter,
92}
93
94impl<C: SegmentCache> InstrumentedSegmentCache<C> {
95    /// Wrap a segment cache and record hit/miss/store metrics with the supplied labels.
96    pub fn new(
97        segment_cache: C,
98        metrics_registry: &dyn MetricsRegistry,
99        labels: Vec<Label>,
100    ) -> Self {
101        Self {
102            segment_cache,
103            hits: MetricBuilder::new(metrics_registry)
104                .add_labels(labels.clone())
105                .counter("vortex.file.segments.cache.hits"),
106            misses: MetricBuilder::new(metrics_registry)
107                .add_labels(labels.clone())
108                .counter("vortex.file.segments.cache.misses"),
109            stores: MetricBuilder::new(metrics_registry)
110                .add_labels(labels)
111                .counter("vortex.file.segments.cache.stores"),
112        }
113    }
114}
115
116#[async_trait]
117impl<C: SegmentCache> SegmentCache for InstrumentedSegmentCache<C> {
118    async fn get(&self, id: SegmentId) -> VortexResult<Option<ByteBuffer>> {
119        let result = self.segment_cache.get(id).await?;
120        if result.is_some() {
121            self.hits.add(1);
122        } else {
123            self.misses.add(1);
124        }
125        Ok(result)
126    }
127
128    async fn put(&self, id: SegmentId, buffer: ByteBuffer) -> VortexResult<()> {
129        self.segment_cache.put(id, buffer).await?;
130        self.stores.add(1);
131        Ok(())
132    }
133}
134
135/// [`SegmentSource`] wrapper that consults a [`SegmentCache`] before the underlying source.
136pub struct SegmentCacheSourceAdapter {
137    cache: Arc<dyn SegmentCache>,
138    source: Arc<dyn SegmentSource>,
139}
140
141impl SegmentCacheSourceAdapter {
142    /// Construct a cache-fronted source.
143    pub fn new(cache: Arc<dyn SegmentCache>, source: Arc<dyn SegmentSource>) -> Self {
144        Self { cache, source }
145    }
146}
147
148impl SegmentSource for SegmentCacheSourceAdapter {
149    fn request(&self, id: SegmentId) -> SegmentFuture {
150        let cache = Arc::clone(&self.cache);
151        let delegate = self.source.request(id);
152
153        async move {
154            if let Ok(Some(segment)) = cache.get(id).await {
155                tracing::debug!("Resolved segment {} from cache", id);
156                return Ok(BufferHandle::new_host(segment));
157            }
158            let result = delegate.await?;
159            // Cache only CPU buffers; device buffers are not cached.
160            if let Some(buffer) = result.as_host_opt()
161                && let Err(e) = cache.put(id, buffer.clone()).await
162            {
163                tracing::warn!("Failed to store segment {} in cache: {}", id, e);
164            }
165            Ok(result)
166        }
167        .boxed()
168    }
169}
vortex_layout/segments/cache.rs

vortex_layout/segments/
cache.rs
