vortex_layout/segments/
test.rs1use 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#[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 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}