Skip to main content

vortex_file/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 parking_lot::RwLock;
8use vortex_buffer::ByteBuffer;
9use vortex_error::VortexResult;
10use vortex_layout::segments::SegmentCache;
11use vortex_layout::segments::SegmentId;
12use vortex_utils::aliases::hash_map::HashMap;
13
14/// Segment cache containing the initial read segments.
15pub struct InitialReadSegmentCache {
16    /// Segments that were already covered by the footer initial read.
17    pub initial: RwLock<HashMap<SegmentId, ByteBuffer>>,
18    /// Delegate cache used for all misses and stores.
19    pub fallback: Arc<dyn SegmentCache>,
20}
21
22#[async_trait]
23impl SegmentCache for InitialReadSegmentCache {
24    async fn get(&self, id: SegmentId) -> VortexResult<Option<ByteBuffer>> {
25        if let Some(buffer) = self.initial.read().get(&id) {
26            return Ok(Some(buffer.clone()));
27        }
28        self.fallback.get(id).await
29    }
30
31    async fn put(&self, id: SegmentId, buffer: ByteBuffer) -> VortexResult<()> {
32        self.fallback.put(id, buffer).await
33    }
34}