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 futures::FutureExt;
7use parking_lot::Mutex;
8use vortex_buffer::{ByteBuffer, ByteBufferMut};
9use vortex_error::{VortexResult, vortex_err};
10
11use crate::segments::sink::SegmentWriter;
12use crate::segments::{SegmentFuture, SegmentId, SegmentSource};
13
14/// A dummy in-memory implementation of a segment reader and writer.
15#[derive(Default, Clone)]
16pub struct TestSegments {
17    segments: Arc<Mutex<Vec<ByteBuffer>>>,
18}
19
20impl SegmentSource for TestSegments {
21    fn request(&self, id: SegmentId, _for_whom: &Arc<str>) -> SegmentFuture {
22        let buffer = self.segments.lock().get(*id as usize).cloned();
23        async move { buffer.ok_or_else(|| vortex_err!("Segment not found")) }.boxed()
24    }
25}
26
27impl SegmentWriter for TestSegments {
28    fn put(&mut self, _segment_id: SegmentId, data: Vec<ByteBuffer>) -> VortexResult<()> {
29        // Combine all the buffers since we're only a test implementation
30        let mut buffer = ByteBufferMut::empty();
31        for segment in data {
32            buffer.extend_from_slice(segment.as_ref());
33        }
34        self.segments.lock().push(buffer.freeze());
35        Ok(())
36    }
37}
vortex_layout/segments/test.rs

vortex_layout/segments/
test.rs
