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, SegmentId};
11use vortex_utils::aliases::hash_map::HashMap;
12
13/// Segment cache containing the initial read segments.
14pub struct InitialReadSegmentCache {
15    pub initial: RwLock<HashMap<SegmentId, ByteBuffer>>,
16    pub fallback: Arc<dyn SegmentCache>,
17}
18
19#[async_trait]
20impl SegmentCache for InitialReadSegmentCache {
21    async fn get(&self, id: SegmentId) -> VortexResult<Option<ByteBuffer>> {
22        if let Some(buffer) = self.initial.read().get(&id) {
23            return Ok(Some(buffer.clone()));
24        }
25        self.fallback.get(id).await
26    }
27
28    async fn put(&self, id: SegmentId, buffer: ByteBuffer) -> VortexResult<()> {
29        self.fallback.put(id, buffer).await
30    }
31}