vortex_layout/segments/
test.rs

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