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    pub initial: RwLock<HashMap<SegmentId, ByteBuffer>>,
17    pub fallback: Arc<dyn SegmentCache>,
18}
19
20#[async_trait]
21impl SegmentCache for InitialReadSegmentCache {
22    async fn get(&self, id: SegmentId) -> VortexResult<Option<ByteBuffer>> {
23        if let Some(buffer) = self.initial.read().get(&id) {
24            return Ok(Some(buffer.clone()));
25        }
26        self.fallback.get(id).await
27    }
28
29    async fn put(&self, id: SegmentId, buffer: ByteBuffer) -> VortexResult<()> {
30        self.fallback.put(id, buffer).await
31    }
32}