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

vortex_file/footer/
serializer.rs
