vortex_file/
memory.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 vortex_buffer::ByteBuffer;
8use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
9use vortex_layout::segments::{SegmentFuture, SegmentId, SegmentSource};
10
11use crate::{FileType, Footer, VortexFile, VortexOpenOptions};
12
13/// A Vortex file that is backed by an in-memory buffer.
14///
15/// This type of file reader performs no coalescing or other clever orchestration, simply
16/// zero-copy slicing the segments from the buffer.
17pub struct InMemoryFileType;
18
19impl FileType for InMemoryFileType {
20    type Options = ();
21}
22
23impl VortexOpenOptions<InMemoryFileType> {
24    /// Create open options for an in-memory Vortex file.
25    pub fn in_memory() -> Self {
26        Self::new(())
27    }
28
29    /// Open an in-memory file contained in the provided buffer.
30    pub fn open<B: Into<ByteBuffer>>(self, buffer: B) -> VortexResult<VortexFile> {
31        let buffer = buffer.into();
32
33        let postscript = self.parse_postscript(&buffer)?;
34
35        // If we haven't been provided a DType, we must read one from the file.
36        let dtype = self.dtype
37            .clone()
38            .map(Ok)
39            .unwrap_or_else(|| {
40                let dtype_segment = postscript
41                    .dtype
42                    .ok_or_else(|| vortex_err!("Vortex file doesn't embed a DType and one has not been provided to VortexOpenOptions"))?;
43                self.parse_dtype(0, &buffer, &dtype_segment)
44            })?;
45
46        let file_stats = postscript
47            .statistics
48            .map(|segment| self.parse_file_statistics(0, &buffer, &segment))
49            .transpose()?;
50
51        let footer = self.parse_footer(
52            0,
53            &buffer,
54            &postscript.footer,
55            &postscript.layout,
56            dtype,
57            file_stats,
58        )?;
59
60        let segment_source = Arc::new(InMemorySegmentReader {
61            buffer,
62            footer: footer.clone(),
63        });
64
65        Ok(VortexFile {
66            footer,
67            segment_source,
68            metrics: self.metrics,
69        })
70    }
71}
72
73#[derive(Clone)]
74struct InMemorySegmentReader {
75    buffer: ByteBuffer,
76    footer: Footer,
77}
78
79impl SegmentSource for InMemorySegmentReader {
80    fn request(&self, id: SegmentId) -> SegmentFuture {
81        let Some(spec) = self.footer.segment_map().get(*id as usize) else {
82            return async move { vortex_bail!("segment not found {id}") }.boxed();
83        };
84
85        let start = usize::try_from(spec.offset).vortex_expect("segment offset larger than usize");
86        let end = start + spec.length as usize;
87        let buffer = self.buffer.slice(start..end);
88
89        async move { Ok(buffer) }.boxed()
90    }
91}
vortex_file/memory.rs

vortex_file/
memory.rs
