vortex_file/footer/
serializer.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use vortex_buffer::ByteBuffer;
7use vortex_error::VortexExpect;
8use vortex_error::VortexResult;
9use vortex_error::vortex_err;
10use vortex_flatbuffers::FlatBuffer;
11use vortex_flatbuffers::FlatBufferRoot;
12use vortex_flatbuffers::WriteFlatBuffer;
13use vortex_flatbuffers::WriteFlatBufferExt;
14use vortex_layout::LayoutContext;
15use vortex_session::registry::ReadContext;
16
17use crate::EOF_SIZE;
18use crate::Footer;
19use crate::MAGIC_BYTES;
20use crate::MAX_POSTSCRIPT_SIZE;
21use crate::VERSION;
22use crate::footer::file_layout::FooterFlatBufferWriter;
23use crate::footer::postscript::Postscript;
24use crate::footer::postscript::PostscriptSegment;
25
26pub struct FooterSerializer {
27    footer: Footer,
28    exclude_dtype: bool,
29    offset: u64,
30}
31
32impl FooterSerializer {
33    pub(super) fn new(footer: Footer) -> Self {
34        Self {
35            footer,
36            exclude_dtype: false,
37            offset: 0,
38        }
39    }
40
41    /// Update the offset used to generate absolute segment locations.
42    ///
43    /// This represents the byte position that the first buffer emitted by this serializer will be
44    /// written to.
45    pub fn with_offset(mut self, offset: u64) -> Self {
46        self.offset = offset;
47        self
48    }
49
50    /// Exclude the DType from the serialized footer.
51    /// If excluded, the reader must be provided the DType from an external source.
52    pub fn exclude_dtype(mut self) -> Self {
53        self.exclude_dtype = true;
54        self
55    }
56
57    /// Whether to exclude the DType from the serialized footer.
58    /// If excluded, the reader must be provided the DType from an external source.
59    pub fn with_exclude_dtype(mut self, exclude_dtype: bool) -> Self {
60        self.exclude_dtype = exclude_dtype;
61        self
62    }
63
64    /// Serialize the footer into a byte buffer that can later be deserialized as a [`Footer`].
65    /// This can be helpful for storing some footer data out-of-band to accelerate opening a file.
66    pub fn serialize(mut self) -> VortexResult<Vec<ByteBuffer>> {
67        let mut buffers = vec![];
68
69        let dtype_segment = if self.exclude_dtype {
70            None
71        } else {
72            let (buffer, dtype_segment) = write_flatbuffer(&mut self.offset, self.footer.dtype())?;
73            buffers.push(buffer);
74            Some(dtype_segment)
75        };
76
77        // TODO(ngates): we should separate the read/write side of Context since the write side
78        //  doesn't need to look anything up in the registry.
79        let layout_ctx = LayoutContext::default();
80
81        let (buffer, layout_segment) = write_flatbuffer(
82            &mut self.offset,
83            &self.footer.layout().flatbuffer_writer(&layout_ctx),
84        )?;
85        buffers.push(buffer);
86
87        let statistics_segment = match self.footer.statistics() {
88            None => None,
89            Some(stats) if stats.stats_sets().is_empty() => None,
90            Some(stats) => {
91                let (buffer, stats_segment) = write_flatbuffer(&mut self.offset, stats)?;
92                buffers.push(buffer);
93                Some(stats_segment)
94            }
95        };
96
97        let (buffer, footer_segment) = write_flatbuffer(
98            &mut self.offset,
99            &FooterFlatBufferWriter {
100                ctx: self.footer.array_read_ctx.clone(),
101                layout_ctx: ReadContext::new(layout_ctx.to_ids()),
102                segment_specs: Arc::clone(&self.footer.segments),
103            },
104        )?;
105        buffers.push(buffer);
106
107        // Assemble the postscript, and write it manually to avoid any framing.
108        let postscript = Postscript {
109            dtype: dtype_segment,
110            layout: layout_segment,
111            statistics: statistics_segment,
112            footer: footer_segment,
113        };
114        let postscript_buffer = postscript.write_flatbuffer_bytes()?;
115        if postscript_buffer.len() > MAX_POSTSCRIPT_SIZE as usize {
116            Err(vortex_err!(
117                "Postscript is too large ({} bytes); max postscript size is {}",
118                postscript_buffer.len(),
119                MAX_POSTSCRIPT_SIZE
120            ))?;
121        }
122
123        let postscript_len = u16::try_from(postscript_buffer.len())
124            .vortex_expect("Postscript already verified to fit into u16");
125        buffers.push(postscript_buffer.into_inner());
126
127        // And finally, the EOF 8-byte footer.
128        let mut eof = [0u8; EOF_SIZE];
129        eof[0..2].copy_from_slice(&VERSION.to_le_bytes());
130        eof[2..4].copy_from_slice(&postscript_len.to_le_bytes());
131        eof[4..8].copy_from_slice(&MAGIC_BYTES);
132        buffers.push(ByteBuffer::copy_from(eof));
133
134        Ok(buffers)
135    }
136}
137
138fn write_flatbuffer<F: FlatBufferRoot + WriteFlatBuffer>(
139    offset: &mut u64,
140    flatbuffer: &F,
141) -> VortexResult<(ByteBuffer, PostscriptSegment)> {
142    let buffer = flatbuffer.write_flatbuffer_bytes()?;
143    let length = u32::try_from(buffer.len())
144        .map_err(|_| vortex_err!("flatbuffer length exceeds maximum u32"))?;
145
146    let segment = PostscriptSegment {
147        offset: *offset,
148        length,
149        alignment: FlatBuffer::alignment(),
150    };
151
152    *offset += u64::from(length);
153
154    Ok((buffer.into_inner(), segment))
155}
vortex_file/footer/serializer.rs

vortex_file/footer/
serializer.rs
