Skip to main content

vortex_file/footer/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Vortex file footer metadata.
5//!
6//! A footer contains the root layout, file-level statistics, the segment map, and the read contexts
7//! needed to resolve array/layout encoding ids during deserialization.
8//!
9//! The byte-level footer and postscript layout is part of the file-format spec; this module exposes
10//! the structured Rust representation and serializer/deserializer state machine.
11mod field_sizes;
12mod file_layout;
13mod file_statistics;
14mod postscript;
15mod segment;
16
17use std::sync::Arc;
18
19mod serializer;
20pub use serializer::*;
21mod deserializer;
22pub use deserializer::*;
23pub use field_sizes::CompressedFieldSizes;
24pub use file_statistics::FileStatistics;
25use flatbuffers::root;
26use itertools::Itertools;
27pub use segment::*;
28use vortex_array::ArrayId;
29use vortex_array::dtype::DType;
30use vortex_array::flatbuffers::FlatBuffer;
31use vortex_buffer::ByteBuffer;
32use vortex_error::VortexResult;
33use vortex_error::vortex_bail;
34use vortex_error::vortex_err;
35use vortex_layout::LayoutEncodingId;
36use vortex_layout::LayoutRef;
37use vortex_layout::layout_from_flatbuffer_with_options;
38use vortex_session::VortexSession;
39use vortex_session::registry::ReadContext;
40
41use crate::flatbuffers::footer as fb;
42
43/// Maximum number of user-defined metadata segments. Keeps postscript bookkeeping small so the
44/// footer and required segments still fit the initial tail read.
45pub(crate) const MAX_METADATA_SEGMENTS: usize = 16;
46
47/// Maximum length, in UTF-8 bytes, of a user-defined metadata key (keys live in the postscript).
48///
49/// 64 bytes covers reverse-DNS query-engine keys, not just short Iceberg-style keys: e.g.
50/// `org.apache.spark.sql.parquet.row.metadata` (41 bytes), which Spark writes into every Parquet
51/// file. With [`MAX_METADATA_SEGMENTS`] keys this bounds the postscript key budget at 1 KiB.
52pub(crate) const MAX_METADATA_KEY_BYTES: usize = 64;
53
54/// User-defined metadata segment locators stored as `(key, locator)` pairs.
55pub(crate) type MetadataSegments = Arc<[(String, SegmentSpec)]>;
56
57/// Captures the layout information of a Vortex file.
58#[derive(Debug, Clone)]
59pub struct Footer {
60    root_layout: LayoutRef,
61    segments: Arc<[SegmentSpec]>,
62    statistics: Option<FileStatistics>,
63    metadata: Arc<[(String, SegmentSpec)]>,
64    // The specific arrays used within the file, in the order they were registered.
65    array_read_ctx: ReadContext,
66    // The approximate size of the footer in bytes, used for caching and memory management.
67    approx_byte_size: Option<usize>,
68}
69
70impl Footer {
71    pub fn new(
72        root_layout: LayoutRef,
73        segments: Arc<[SegmentSpec]>,
74        statistics: Option<FileStatistics>,
75        array_read_ctx: ReadContext,
76    ) -> Self {
77        Self {
78            root_layout,
79            segments,
80            statistics,
81            metadata: Arc::from([]),
82            array_read_ctx,
83            approx_byte_size: None,
84        }
85    }
86
87    pub(crate) fn with_approx_byte_size(mut self, approx_byte_size: usize) -> Self {
88        self.approx_byte_size = Some(approx_byte_size);
89        self
90    }
91
92    /// Read the [`Footer`] from a flatbuffer.
93    pub(crate) fn from_flatbuffer(
94        footer_bytes: &[u8],
95        layout_bytes: FlatBuffer,
96        dtype: DType,
97        statistics: Option<FileStatistics>,
98        metadata: Arc<[(String, SegmentSpec)]>,
99        session: &VortexSession,
100    ) -> VortexResult<Self> {
101        let metadata_bytes: usize = metadata
102            .iter()
103            .map(|(key, _segment)| key.len() + size_of::<SegmentSpec>())
104            .sum();
105        let approx_byte_size = footer_bytes.len() + layout_bytes.len() + metadata_bytes;
106        let fb_footer = root::<fb::Footer>(footer_bytes)?;
107
108        // Create a LayoutContext from the registry.
109        let layout_specs = fb_footer.layout_specs();
110        #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")]
111        let layout_ids: Arc<[_]> = layout_specs
112            .iter()
113            .flat_map(|e| e.iter())
114            .map(|encoding| LayoutEncodingId::new(encoding.id()))
115            .collect();
116        let layout_read_ctx = ReadContext::new(layout_ids);
117
118        // Create an ArrayContext from the registry.
119        let array_specs = fb_footer.array_specs();
120        #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")]
121        let array_ids: Arc<[_]> = array_specs
122            .iter()
123            .flat_map(|e| e.iter())
124            .map(|encoding| ArrayId::new(encoding.id()))
125            .collect();
126        let array_read_ctx = ReadContext::new(array_ids);
127
128        let root_layout = layout_from_flatbuffer_with_options(
129            layout_bytes,
130            &dtype,
131            &layout_read_ctx,
132            &array_read_ctx,
133            session,
134            session.allows_unknown(),
135        )?;
136
137        let segments: Arc<[SegmentSpec]> = fb_footer
138            .segment_specs()
139            .ok_or_else(|| vortex_err!("FileLayout missing segment specs"))?
140            .iter()
141            .map(SegmentSpec::try_from)
142            .try_collect()?;
143
144        // Note this assertion is `<=` since we allow zero-length segments
145        if !segments.is_sorted_by_key(|segment| segment.offset) {
146            vortex_bail!("Segment offsets are not ordered");
147        }
148
149        Ok(Self {
150            root_layout,
151            segments,
152            statistics,
153            metadata,
154            array_read_ctx,
155            approx_byte_size: Some(approx_byte_size),
156        })
157    }
158
159    /// Returns the root [`LayoutRef`] of the file.
160    pub fn layout(&self) -> &LayoutRef {
161        &self.root_layout
162    }
163
164    /// Returns the segment map of the file.
165    pub fn segment_map(&self) -> &Arc<[SegmentSpec]> {
166        &self.segments
167    }
168
169    /// Returns the statistics of the file.
170    pub fn statistics(&self) -> Option<&FileStatistics> {
171        self.statistics.as_ref()
172    }
173
174    /// Returns the user-defined metadata segment locators stored in the postscript.
175    pub fn metadata_segments(&self) -> impl Iterator<Item = (&str, &SegmentSpec)> {
176        self.metadata
177            .iter()
178            .map(|(key, segment)| (key.as_str(), segment))
179    }
180
181    /// Returns the user-defined metadata segment locator for the given key.
182    pub fn metadata_segment(&self, key: &str) -> Option<&SegmentSpec> {
183        self.metadata
184            .iter()
185            .find_map(|(candidate, segment)| (candidate == key).then_some(segment))
186    }
187
188    pub(crate) fn with_metadata_segments(mut self, metadata: Arc<[(String, SegmentSpec)]>) -> Self {
189        self.metadata = metadata;
190        self
191    }
192
193    pub(crate) fn segment_specs_with_metadata(&self) -> Arc<[SegmentSpec]> {
194        self.segments
195            .iter()
196            .copied()
197            .chain(self.metadata.iter().map(|(_, segment)| *segment))
198            .collect()
199    }
200
201    /// Computes the compressed size in bytes of every field in the file, keyed by field path.
202    ///
203    /// Sizes are derived by attributing each segment in the [segment map][Self::segment_map] to a
204    /// field in the [layout tree][Self::layout]; see [`CompressedFieldSizes`] for the exact
205    /// attribution semantics. No IO is performed.
206    pub fn compressed_field_sizes(&self) -> VortexResult<CompressedFieldSizes> {
207        CompressedFieldSizes::try_new(&self.root_layout, &self.segments)
208    }
209
210    /// Returns the [`DType`] of the file.
211    pub fn dtype(&self) -> &DType {
212        self.root_layout.dtype()
213    }
214
215    /// Returns the approximate size of the footer in bytes, used for caching and memory management.
216    pub fn approx_byte_size(&self) -> Option<usize> {
217        self.approx_byte_size
218    }
219
220    /// Returns the number of rows in the file.
221    pub fn row_count(&self) -> u64 {
222        self.root_layout.row_count()
223    }
224
225    /// Validate that every segment declared in the footer lies within a file of `file_size` bytes.
226    pub(crate) fn validate_file_size(&self, file_size: u64) -> VortexResult<()> {
227        validate_segments_within_file(&self.segments, file_size)
228    }
229
230    /// Returns a serializer for this footer.
231    pub fn into_serializer(self) -> FooterSerializer {
232        FooterSerializer::new(self)
233    }
234
235    /// Create a deserializer for a Vortex file footer.
236    pub fn deserializer(eof_buffer: ByteBuffer, session: VortexSession) -> FooterDeserializer {
237        FooterDeserializer::new(eof_buffer, session)
238    }
239}
240
241/// Validate that every segment declared in the footer lies within a file of `file_size` bytes.
242///
243/// A corrupt or malicious file can declare a segment whose offset or length extends past the end
244/// of the file. Rejecting such files up front ensures that later slicing of the backing buffer
245/// returns a [`VortexError`](vortex_error::VortexError) rather than panicking (see issue #8819).
246fn validate_segments_within_file(segments: &[SegmentSpec], file_size: u64) -> VortexResult<()> {
247    for segment in segments {
248        let within_file = segment
249            .offset
250            .checked_add(segment.length as u64)
251            .is_some_and(|end| end <= file_size);
252        if !within_file {
253            vortex_bail!(
254                "Segment at offset {} with length {} extends past the end of the \
255                 {file_size}-byte file",
256                segment.offset,
257                segment.length,
258            );
259        }
260    }
261    Ok(())
262}
263
264#[cfg(test)]
265mod tests {
266    use vortex_buffer::Alignment;
267
268    use super::*;
269
270    fn segment(offset: u64, length: u32) -> SegmentSpec {
271        SegmentSpec {
272            offset,
273            length,
274            alignment: Alignment::none(),
275        }
276    }
277
278    #[test]
279    fn accepts_segments_within_file() -> VortexResult<()> {
280        validate_segments_within_file(&[segment(0, 100), segment(100, 50)], 150)?;
281        Ok(())
282    }
283
284    #[test]
285    fn rejects_segment_extending_past_end_of_file() {
286        let err =
287            validate_segments_within_file(&[segment(0, 100), segment(100, 51)], 150).unwrap_err();
288        assert!(err.to_string().contains("past the end"), "{err}");
289    }
290
291    #[test]
292    fn rejects_segment_offset_length_overflow() {
293        let err = validate_segments_within_file(&[segment(u64::MAX, 1)], u64::MAX).unwrap_err();
294        assert!(err.to_string().contains("past the end"), "{err}");
295    }
296}