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;
16use vortex_utils::aliases::hash_map::HashMap;
17
18use crate::EOF_SIZE;
19use crate::Footer;
20use crate::MAGIC_BYTES;
21use crate::MAX_POSTSCRIPT_SIZE;
22use crate::VERSION;
23use crate::footer::SegmentSpec;
24use crate::footer::file_layout::FooterFlatBufferWriter;
25use crate::footer::postscript::Postscript;
26use crate::footer::postscript::PostscriptMetadata;
27use crate::footer::postscript::PostscriptSegment;
28
29/// Serializes a [`Footer`] into footer buffers and the trailing postscript/EOF marker.
30pub struct FooterSerializer {
31    footer: Footer,
32    metadata: HashMap<String, ByteBuffer>,
33    exclude_dtype: bool,
34    offset: u64,
35}
36
37impl FooterSerializer {
38    pub(super) fn new(footer: Footer) -> Self {
39        Self {
40            footer,
41            metadata: HashMap::default(),
42            exclude_dtype: false,
43            offset: 0,
44        }
45    }
46
47    /// Update the offset used to generate absolute segment locations.
48    ///
49    /// This represents the byte position that the first buffer emitted by this serializer will be
50    /// written to.
51    pub fn with_offset(mut self, offset: u64) -> Self {
52        self.offset = offset;
53        self
54    }
55
56    /// Exclude the DType from the serialized footer.
57    /// If excluded, the reader must be provided the DType from an external source.
58    pub fn exclude_dtype(mut self) -> Self {
59        self.exclude_dtype = true;
60        self
61    }
62
63    /// Whether to exclude the DType from the serialized footer.
64    /// If excluded, the reader must be provided the DType from an external source.
65    pub fn with_exclude_dtype(mut self, exclude_dtype: bool) -> Self {
66        self.exclude_dtype = exclude_dtype;
67        self
68    }
69
70    pub(crate) fn with_metadata_segments(mut self, metadata: HashMap<String, ByteBuffer>) -> Self {
71        self.metadata = metadata;
72        self
73    }
74
75    /// Serialize the footer into a byte buffer that can later be deserialized as a [`Footer`].
76    /// This can be helpful for storing some footer data out-of-band to accelerate opening a file.
77    pub fn serialize(self) -> VortexResult<Vec<ByteBuffer>> {
78        self.serialize_with_metadata()
79            .map(|(buffers, _metadata, _approx_byte_size)| buffers)
80    }
81
82    pub(crate) fn serialize_with_metadata(
83        mut self,
84    ) -> VortexResult<(Vec<ByteBuffer>, super::MetadataSegments, usize)> {
85        let mut buffers = vec![];
86
87        let (metadata_segments, metadata_locators) = if self.metadata.is_empty() {
88            let locators = self
89                .footer
90                .metadata_segments()
91                .map(|(key, segment)| (key.to_string(), *segment))
92                .collect::<Vec<_>>();
93            let segments = locators
94                .iter()
95                .map(|(key, segment)| PostscriptMetadata {
96                    key: key.clone(),
97                    segment: PostscriptSegment::from(segment),
98                })
99                .collect();
100            (segments, locators)
101        } else {
102            let mut metadata = self.metadata.into_iter().collect::<Vec<_>>();
103            // Metadata is stored in a HashMap, so sort keys before assigning segment offsets.
104            metadata.sort_unstable_by(|(left, _), (right, _)| left.cmp(right));
105
106            let mut segments = Vec::with_capacity(metadata.len());
107            let mut locators = Vec::with_capacity(metadata.len());
108            for (key, metadata) in metadata {
109                let (metadata_buffers, segment) = write_buffer(&mut self.offset, metadata)?;
110                buffers.extend(metadata_buffers);
111                locators.push((key.clone(), SegmentSpec::from(&segment)));
112                segments.push(PostscriptMetadata { key, segment });
113            }
114            (segments, locators)
115        };
116        let metadata_storage_bytes = buffers.iter().map(ByteBuffer::len).sum::<usize>();
117
118        let dtype_segment = if self.exclude_dtype {
119            None
120        } else {
121            let (buffer, dtype_segment) = write_flatbuffer(&mut self.offset, self.footer.dtype())?;
122            buffers.push(buffer);
123            Some(dtype_segment)
124        };
125
126        // TODO(ngates): we should separate the read/write side of Context since the write side
127        //  doesn't need to look anything up in the registry.
128        let layout_ctx = LayoutContext::default();
129
130        let (buffer, layout_segment) = write_flatbuffer(
131            &mut self.offset,
132            &self.footer.layout().flatbuffer_writer(&layout_ctx),
133        )?;
134        buffers.push(buffer);
135
136        let statistics_segment = match self.footer.statistics() {
137            None => None,
138            Some(stats) if stats.stats_sets().is_empty() => None,
139            Some(stats) => {
140                let (buffer, stats_segment) = write_flatbuffer(&mut self.offset, stats)?;
141                buffers.push(buffer);
142                Some(stats_segment)
143            }
144        };
145
146        let (buffer, footer_segment) = write_flatbuffer(
147            &mut self.offset,
148            &FooterFlatBufferWriter {
149                ctx: self.footer.array_read_ctx.clone(),
150                layout_ctx: ReadContext::new(layout_ctx.to_ids()),
151                segment_specs: Arc::clone(&self.footer.segments),
152            },
153        )?;
154        buffers.push(buffer);
155
156        // Assemble the postscript, and write it manually to avoid any framing.
157        let postscript = Postscript {
158            dtype: dtype_segment,
159            layout: layout_segment,
160            statistics: statistics_segment,
161            footer: footer_segment,
162            metadata: metadata_segments,
163        };
164        let postscript_buffer = postscript.write_flatbuffer_bytes()?;
165        if postscript_buffer.len() > MAX_POSTSCRIPT_SIZE as usize {
166            Err(vortex_err!(
167                "Postscript is too large ({} bytes); max postscript size is {}",
168                postscript_buffer.len(),
169                MAX_POSTSCRIPT_SIZE
170            ))?;
171        }
172
173        let postscript_len = u16::try_from(postscript_buffer.len())
174            .vortex_expect("Postscript already verified to fit into u16");
175        buffers.push(postscript_buffer.into_inner());
176
177        // And finally, the EOF 8-byte footer.
178        let mut eof = [0u8; EOF_SIZE];
179        eof[0..2].copy_from_slice(&VERSION.to_le_bytes());
180        eof[2..4].copy_from_slice(&postscript_len.to_le_bytes());
181        eof[4..8].copy_from_slice(&MAGIC_BYTES);
182        buffers.push(ByteBuffer::copy_from(eof));
183
184        let approx_byte_size =
185            buffers.iter().map(ByteBuffer::len).sum::<usize>() - metadata_storage_bytes;
186        Ok((buffers, metadata_locators.into(), approx_byte_size))
187    }
188}
189
190impl From<&PostscriptSegment> for SegmentSpec {
191    fn from(value: &PostscriptSegment) -> Self {
192        Self {
193            offset: value.offset,
194            length: value.length,
195            alignment: value.alignment,
196        }
197    }
198}
199
200impl From<&SegmentSpec> for PostscriptSegment {
201    fn from(value: &SegmentSpec) -> Self {
202        Self {
203            offset: value.offset,
204            length: value.length,
205            alignment: value.alignment,
206        }
207    }
208}
209
210fn write_flatbuffer<F: FlatBufferRoot + WriteFlatBuffer>(
211    offset: &mut u64,
212    flatbuffer: &F,
213) -> VortexResult<(ByteBuffer, PostscriptSegment)> {
214    let buffer = flatbuffer.write_flatbuffer_bytes()?;
215    let length = u32::try_from(buffer.len())
216        .map_err(|_| vortex_err!("flatbuffer length exceeds maximum u32"))?;
217
218    let segment = PostscriptSegment {
219        offset: *offset,
220        length,
221        alignment: FlatBuffer::alignment(),
222    };
223
224    *offset += u64::from(length);
225
226    Ok((buffer.into_inner(), segment))
227}
228
229fn write_buffer(
230    offset: &mut u64,
231    buffer: ByteBuffer,
232) -> VortexResult<(Vec<ByteBuffer>, PostscriptSegment)> {
233    let length = u32::try_from(buffer.len())
234        .map_err(|_| vortex_err!("metadata segment length exceeds maximum u32"))?;
235    let alignment = buffer.alignment();
236
237    let padding = offset.next_multiple_of(*alignment as u64) - *offset;
238    let segment_offset = *offset + padding;
239
240    let segment = PostscriptSegment {
241        offset: segment_offset,
242        length,
243        alignment,
244    };
245
246    *offset += padding + u64::from(length);
247
248    let mut buffers = Vec::with_capacity(if padding == 0 { 1 } else { 2 });
249    if padding > 0 {
250        buffers.push(ByteBuffer::zeroed(padding as usize));
251    }
252    buffers.push(buffer);
253
254    Ok((buffers, segment))
255}
vortex_file/footer/serializer.rs

vortex_file/footer/
serializer.rs
