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

vortex_file/footer/
serializer.rs
