vortex_layout/segments/
test.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 parking_lot::Mutex;
9use vortex_buffer::{ByteBuffer, ByteBufferMut};
10use vortex_error::{VortexExpect, VortexResult, vortex_err};
11
12use crate::segments::{SegmentFuture, SegmentId, SegmentSink, SegmentSource};
13use crate::sequence::SequenceId;
14
15/// A dummy in-memory implementation of a segment reader and writer.
16#[derive(Default, Clone)]
17pub struct TestSegments {
18    segments: Arc<Mutex<Vec<ByteBuffer>>>,
19}
20
21impl SegmentSource for TestSegments {
22    fn request(&self, id: SegmentId) -> SegmentFuture {
23        let buffer = self.segments.lock().get(*id as usize).cloned();
24        async move { buffer.ok_or_else(|| vortex_err!("Segment not found")) }.boxed()
25    }
26}
27
28#[async_trait]
29impl SegmentSink for TestSegments {
30    async fn write(
31        &self,
32        _sequence_id: SequenceId,
33        buffers: Vec<ByteBuffer>,
34    ) -> VortexResult<SegmentId> {
35        // Combine all the buffers since we're only a test implementation
36        let mut buffer = ByteBufferMut::empty();
37        for segment in buffers {
38            buffer.extend_from_slice(segment.as_ref());
39        }
40
41        let mut segments = self.segments.lock();
42        let segment_id =
43            SegmentId::from(u32::try_from(segments.len()).vortex_expect("Too many segments"));
44        segments.push(buffer.freeze());
45
46        Ok(segment_id)
47    }
48}
vortex_layout/segments/test.rs

vortex_layout/segments/
test.rs
