vortex_file/footer/
deserializer.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use flatbuffers::root;
5use vortex_array::dtype::DType;
6use vortex_buffer::ByteBuffer;
7use vortex_buffer::ByteBufferMut;
8use vortex_error::VortexExpect;
9use vortex_error::VortexResult;
10use vortex_error::vortex_bail;
11use vortex_error::vortex_err;
12use vortex_flatbuffers::FlatBuffer;
13use vortex_flatbuffers::ReadFlatBuffer;
14use vortex_session::VortexSession;
15
16use crate::EOF_SIZE;
17use crate::Footer;
18use crate::MAGIC_BYTES;
19use crate::VERSION;
20use crate::footer::FileStatistics;
21use crate::footer::postscript::Postscript;
22use crate::footer::postscript::PostscriptSegment;
23
24/// Deserialize a footer from the end of a Vortex file or created from a
25/// [`crate::footer::FooterSerializer`].
26///
27/// The deserializer is incremental because callers may initially read only the tail of a file. Call
28/// [`deserialize`](Self::deserialize) until it returns [`DeserializeStep::Done`]. If it asks for
29/// [`DeserializeStep::NeedMoreData`], prefix the requested bytes with [`prefix_data`](Self::prefix_data).
30/// If it asks for [`DeserializeStep::NeedFileSize`], call [`with_size`](Self::with_size) and retry.
31pub struct FooterDeserializer {
32    // A buffer representing the end of a Vortex file.
33    // During deserialization, we may need to expand this buffer by requesting more data from
34    // the caller.
35    buffer: ByteBuffer,
36    // The session to use for deserialization.
37    session: VortexSession,
38    // The DType, if provided externally.
39    dtype: Option<DType>,
40
41    // Internal state that we accumulate
42
43    // The file size, possibly provided externally.
44    file_size: Option<u64>,
45    // The postscript, once we've parsed it.
46    postscript: Option<Postscript>,
47}
48
49impl FooterDeserializer {
50    pub(super) fn new(initial_read: ByteBuffer, session: VortexSession) -> Self {
51        Self {
52            buffer: initial_read,
53            session,
54            dtype: None,
55            file_size: None,
56            postscript: None,
57        }
58    }
59
60    /// Provide the file dtype externally.
61    ///
62    /// This is required for files written with [`VortexWriteOptions::exclude_dtype`](crate::VortexWriteOptions::exclude_dtype).
63    pub fn with_dtype(mut self, dtype: DType) -> Self {
64        self.dtype = Some(dtype);
65        self
66    }
67
68    /// Provide or clear the externally known file dtype.
69    pub fn with_some_dtype(mut self, dtype: Option<DType>) -> Self {
70        self.dtype = dtype;
71        self
72    }
73
74    /// Provide the total file size.
75    pub fn with_size(mut self, file_size: u64) -> Self {
76        self.file_size = Some(file_size);
77        self
78    }
79
80    /// Provide or clear the total file size.
81    pub fn with_some_size(mut self, file_size: Option<u64>) -> Self {
82        self.file_size = file_size;
83        self
84    }
85
86    /// Prefix more data to the existing buffer when requested by the deserializer.
87    pub fn prefix_data(&mut self, more_data: ByteBuffer) {
88        let mut buffer = ByteBufferMut::with_capacity(self.buffer.len() + more_data.len());
89        buffer.extend_from_slice(&more_data);
90        buffer.extend_from_slice(&self.buffer);
91        self.buffer = buffer.freeze();
92    }
93
94    /// Advance footer deserialization.
95    ///
96    /// Returns the next missing input requirement or the finished [`Footer`].
97    pub fn deserialize(&mut self) -> VortexResult<DeserializeStep> {
98        let postscript = if let Some(postscript) = &self.postscript {
99            postscript
100        } else {
101            self.postscript = Some(self.parse_postscript(&self.buffer)?);
102            self.postscript
103                .as_ref()
104                .vortex_expect("Just set postscript")
105        };
106
107        // If we haven't been provided a DType, we must read one from the file.
108        let dtype_segment = self
109            .dtype
110            .is_none()
111            .then(|| {
112                postscript.dtype.as_ref().ok_or_else(|| {
113                    vortex_err!(
114                        "Vortex file doesn't embed a DType and none provided to VortexOpenOptions"
115                    )
116                })
117            })
118            .transpose()?;
119
120        // The other postscript segments are required, so now we figure out our the offset that
121        // contains all the required segments.
122
123        // The initial offset is the file size - the size of our initial read.
124        let Some(file_size) = self.file_size else {
125            return Ok(DeserializeStep::NeedFileSize);
126        };
127        let initial_offset = file_size - (self.buffer.len() as u64);
128
129        let mut read_more_offset = initial_offset;
130        if let Some(dtype_segment) = &dtype_segment {
131            read_more_offset = read_more_offset.min(dtype_segment.offset);
132        }
133        if let Some(stats_segment) = &postscript.statistics {
134            read_more_offset = read_more_offset.min(stats_segment.offset);
135        }
136        read_more_offset = read_more_offset.min(postscript.layout.offset);
137        read_more_offset = read_more_offset.min(postscript.footer.offset);
138
139        // Read more bytes if necessary.
140        if read_more_offset < initial_offset {
141            tracing::trace!(
142                "Initial read from {initial_offset} did not cover all footer segments, reading from {read_more_offset}"
143            );
144            return Ok(DeserializeStep::NeedMoreData {
145                offset: read_more_offset,
146                len: usize::try_from(initial_offset - read_more_offset)?,
147            });
148        }
149
150        // Now we read our initial segments.
151        let dtype = dtype_segment
152            .map(|segment| self.parse_dtype(initial_offset, &self.buffer, segment))
153            .transpose()?
154            .unwrap_or_else(|| self.dtype.clone().vortex_expect("DType was provided"));
155        let file_stats = postscript
156            .statistics
157            .as_ref()
158            .map(|segment| {
159                self.parse_file_statistics(
160                    initial_offset,
161                    &self.buffer,
162                    segment,
163                    &dtype,
164                    &self.session,
165                )
166            })
167            .transpose()?;
168
169        Ok(DeserializeStep::Done(self.parse_footer(
170            initial_offset,
171            &self.buffer,
172            &postscript.footer,
173            &postscript.layout,
174            dtype,
175            file_stats,
176        )?))
177    }
178
179    /// The current buffer being used for deserialization.
180    pub fn buffer(&self) -> &ByteBuffer {
181        &self.buffer
182    }
183
184    /// Parse the postscript from the initial read.
185    fn parse_postscript(&self, initial_read: &[u8]) -> VortexResult<Postscript> {
186        if initial_read.len() < EOF_SIZE {
187            vortex_bail!(
188                "Initial read must be at least EOF_SIZE ({}) bytes",
189                EOF_SIZE
190            );
191        }
192        let eof_loc = initial_read.len() - EOF_SIZE;
193        let magic_bytes_loc = eof_loc + (EOF_SIZE - MAGIC_BYTES.len());
194
195        let magic_number = &initial_read[magic_bytes_loc..];
196        if magic_number != MAGIC_BYTES {
197            vortex_bail!("Malformed file, invalid magic bytes, got {magic_number:?}")
198        }
199
200        let version = u16::from_le_bytes(
201            initial_read[eof_loc..eof_loc + 2]
202                .try_into()
203                .map_err(|e| vortex_err!("Version was not a u16 {e}"))?,
204        );
205        if version != VERSION {
206            vortex_bail!("Malformed file, unsupported version {version}")
207        }
208
209        let ps_size = u16::from_le_bytes(
210            initial_read[eof_loc + 2..eof_loc + 4]
211                .try_into()
212                .map_err(|e| vortex_err!("Postscript size was not a u16 {e}"))?,
213        ) as usize;
214
215        if initial_read.len() < ps_size + EOF_SIZE {
216            vortex_bail!(
217                "Initial read must be at least {} bytes to include the Postscript",
218                ps_size + EOF_SIZE
219            );
220        }
221
222        Postscript::read_flatbuffer_bytes(&initial_read[eof_loc - ps_size..eof_loc])
223    }
224
225    /// Parse the DType from the initial read.
226    fn parse_dtype(
227        &self,
228        initial_offset: u64,
229        initial_read: &[u8],
230        segment: &PostscriptSegment,
231    ) -> VortexResult<DType> {
232        let offset = usize::try_from(segment.offset - initial_offset)?;
233        let sliced_buffer =
234            FlatBuffer::copy_from(&initial_read[offset..offset + (segment.length as usize)]);
235        DType::from_flatbuffer(sliced_buffer, &self.session)
236    }
237
238    /// Parse the [`FileStatistics`] from the initial read buffer.
239    fn parse_file_statistics(
240        &self,
241        initial_offset: u64,
242        initial_read: &[u8],
243        segment: &PostscriptSegment,
244        dtype: &DType,
245        session: &VortexSession,
246    ) -> VortexResult<FileStatistics> {
247        let offset = usize::try_from(segment.offset - initial_offset)?;
248        let sliced_buffer =
249            FlatBuffer::copy_from(&initial_read[offset..offset + (segment.length as usize)]);
250
251        let fb = root::<vortex_flatbuffers::footer::FileStatistics>(&sliced_buffer)?;
252        FileStatistics::from_flatbuffer(&fb, dtype, session)
253    }
254
255    /// Parse the rest of the footer from the initial read.
256    fn parse_footer(
257        &self,
258        initial_offset: u64,
259        initial_read: &[u8],
260        footer_segment: &PostscriptSegment,
261        layout_segment: &PostscriptSegment,
262        dtype: DType,
263        file_stats: Option<FileStatistics>,
264    ) -> VortexResult<Footer> {
265        let footer_offset = usize::try_from(footer_segment.offset - initial_offset)?;
266        let footer_bytes = FlatBuffer::copy_from(
267            &initial_read[footer_offset..footer_offset + (footer_segment.length as usize)],
268        );
269
270        let layout_offset = usize::try_from(layout_segment.offset - initial_offset)?;
271        let layout_bytes = FlatBuffer::copy_from(
272            &initial_read[layout_offset..layout_offset + (layout_segment.length as usize)],
273        );
274
275        Footer::from_flatbuffer(footer_bytes, layout_bytes, dtype, file_stats, &self.session)
276    }
277}
278
279#[derive(Debug)]
280/// Result of one [`FooterDeserializer::deserialize`] step.
281pub enum DeserializeStep {
282    /// Additional data needed to continue deserialization.
283    NeedMoreData {
284        /// Absolute file offset to read from.
285        offset: u64,
286        /// Number of bytes to read and prefix into the deserializer.
287        len: usize,
288    },
289    /// The total file size is required before offsets can be resolved.
290    NeedFileSize,
291    /// Footer deserialization is complete.
292    Done(Footer),
293}
vortex_file/footer/deserializer.rs

vortex_file/footer/
deserializer.rs
