vortex-file 0.89.0

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

//! Vortex file footer metadata.
//!
//! A footer contains the root layout, file-level statistics, the segment map, and the read contexts
//! needed to resolve array/layout encoding ids during deserialization.
//!
//! The byte-level footer and postscript layout is part of the file-format spec; this module exposes
//! the structured Rust representation and serializer/deserializer state machine.
mod field_sizes;
mod file_layout;
mod file_statistics;
mod postscript;
mod segment;

use std::sync::Arc;

mod serializer;
pub use serializer::*;
mod deserializer;
pub use deserializer::*;
pub use field_sizes::CompressedFieldSizes;
pub use file_statistics::FileStatistics;
use flatbuffers::root;
use itertools::Itertools;
pub use segment::*;
use vortex_array::ArrayId;
use vortex_array::dtype::DType;
use vortex_array::flatbuffers::FlatBuffer;
use vortex_buffer::ByteBuffer;
use vortex_error::VortexResult;
use vortex_error::vortex_bail;
use vortex_error::vortex_err;
use vortex_layout::LayoutEncodingId;
use vortex_layout::LayoutRef;
use vortex_layout::layout_from_flatbuffer_with_options;
use vortex_session::VortexSession;
use vortex_session::registry::ReadContext;

use crate::flatbuffers::footer as fb;

/// Maximum number of user-defined metadata segments. Keeps postscript bookkeeping small so the
/// footer and required segments still fit the initial tail read.
pub(crate) const MAX_METADATA_SEGMENTS: usize = 16;

/// Maximum length, in UTF-8 bytes, of a user-defined metadata key (keys live in the postscript).
///
/// 64 bytes covers reverse-DNS query-engine keys, not just short Iceberg-style keys: e.g.
/// `org.apache.spark.sql.parquet.row.metadata` (41 bytes), which Spark writes into every Parquet
/// file. With [`MAX_METADATA_SEGMENTS`] keys this bounds the postscript key budget at 1 KiB.
pub(crate) const MAX_METADATA_KEY_BYTES: usize = 64;

/// User-defined metadata segment locators stored as `(key, locator)` pairs.
pub(crate) type MetadataSegments = Arc<[(String, SegmentSpec)]>;

/// Captures the layout information of a Vortex file.
#[derive(Debug, Clone)]
pub struct Footer {
    root_layout: LayoutRef,
    segments: Arc<[SegmentSpec]>,
    statistics: Option<FileStatistics>,
    metadata: Arc<[(String, SegmentSpec)]>,
    // `segments` followed by the `metadata` segments, which is the segment table a file is opened
    // with. Precomputed so opening a file with a cached footer does not copy every segment spec.
    segments_with_metadata: Arc<[SegmentSpec]>,
    // The largest end offset of any segment in `segments_with_metadata`, or `None` if one
    // overflows `u64`. Precomputed so validating a cached footer against a file size does not walk
    // every segment.
    // Writers never overlap segments, so in a valid file this is the last segment's end, but only
    // their order is checked when reading, and this guards against malformed footers.
    segments_end: Option<u64>,
    // The specific arrays used within the file, in the order they were registered.
    array_read_ctx: ReadContext,
    // The approximate size of the footer in bytes, used for caching and memory management.
    approx_byte_size: Option<usize>,
}

impl Footer {
    pub fn new(
        root_layout: LayoutRef,
        segments: Arc<[SegmentSpec]>,
        statistics: Option<FileStatistics>,
        array_read_ctx: ReadContext,
    ) -> Self {
        Self {
            root_layout,
            segments_with_metadata: Arc::clone(&segments),
            segments_end: segments_end(&segments),
            segments,
            statistics,
            metadata: Arc::from([]),
            array_read_ctx,
            approx_byte_size: None,
        }
    }

    pub(crate) fn with_approx_byte_size(mut self, approx_byte_size: usize) -> Self {
        self.approx_byte_size = Some(approx_byte_size);
        self
    }

    /// Read the [`Footer`] from a flatbuffer.
    pub(crate) fn from_flatbuffer(
        footer_bytes: &[u8],
        layout_bytes: FlatBuffer,
        dtype: DType,
        statistics: Option<FileStatistics>,
        metadata: Arc<[(String, SegmentSpec)]>,
        session: &VortexSession,
    ) -> VortexResult<Self> {
        let metadata_bytes: usize = metadata
            .iter()
            .map(|(key, _segment)| key.len() + size_of::<SegmentSpec>())
            .sum();
        let approx_byte_size = footer_bytes.len() + layout_bytes.len() + metadata_bytes;
        let fb_footer = root::<fb::Footer>(footer_bytes)?;

        // Create a LayoutContext from the registry.
        let layout_specs = fb_footer.layout_specs();
        #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")]
        let layout_ids: Arc<[_]> = layout_specs
            .iter()
            .flat_map(|e| e.iter())
            .map(|encoding| LayoutEncodingId::new(encoding.id()))
            .collect();
        let layout_read_ctx = ReadContext::new(layout_ids);

        // Create an ArrayContext from the registry.
        let array_specs = fb_footer.array_specs();
        #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")]
        let array_ids: Arc<[_]> = array_specs
            .iter()
            .flat_map(|e| e.iter())
            .map(|encoding| ArrayId::new(encoding.id()))
            .collect();
        let array_read_ctx = ReadContext::new(array_ids);

        let root_layout = layout_from_flatbuffer_with_options(
            layout_bytes,
            &dtype,
            &layout_read_ctx,
            &array_read_ctx,
            session,
            session.allows_unknown(),
        )?;

        let segments: Arc<[SegmentSpec]> = fb_footer
            .segment_specs()
            .ok_or_else(|| vortex_err!("FileLayout missing segment specs"))?
            .iter()
            .map(SegmentSpec::try_from)
            .try_collect()?;

        // Note this assertion is `<=` since we allow zero-length segments
        if !segments.is_sorted_by_key(|segment| segment.offset) {
            vortex_bail!("Segment offsets are not ordered");
        }

        let segments_with_metadata = segments_with_metadata(&segments, &metadata);

        Ok(Self {
            root_layout,
            segments_end: segments_end(&segments_with_metadata),
            segments_with_metadata,
            segments,
            statistics,
            metadata,
            array_read_ctx,
            approx_byte_size: Some(approx_byte_size),
        })
    }

    /// Returns the root [`LayoutRef`] of the file.
    pub fn layout(&self) -> &LayoutRef {
        &self.root_layout
    }

    /// Returns the segment map of the file.
    pub fn segment_map(&self) -> &Arc<[SegmentSpec]> {
        &self.segments
    }

    /// Returns the statistics of the file.
    pub fn statistics(&self) -> Option<&FileStatistics> {
        self.statistics.as_ref()
    }

    /// Returns the user-defined metadata segment locators stored in the postscript.
    pub fn metadata_segments(&self) -> impl Iterator<Item = (&str, &SegmentSpec)> {
        self.metadata
            .iter()
            .map(|(key, segment)| (key.as_str(), segment))
    }

    /// Returns the user-defined metadata segment locator for the given key.
    pub fn metadata_segment(&self, key: &str) -> Option<&SegmentSpec> {
        self.metadata
            .iter()
            .find_map(|(candidate, segment)| (candidate == key).then_some(segment))
    }

    pub(crate) fn with_metadata_segments(mut self, metadata: Arc<[(String, SegmentSpec)]>) -> Self {
        self.segments_with_metadata = segments_with_metadata(&self.segments, &metadata);
        self.segments_end = segments_end(&self.segments_with_metadata);
        self.metadata = metadata;
        self
    }

    pub(crate) fn segment_specs_with_metadata(&self) -> Arc<[SegmentSpec]> {
        Arc::clone(&self.segments_with_metadata)
    }

    /// Computes the compressed size in bytes of every field in the file, keyed by field path.
    ///
    /// Sizes are derived by attributing each segment in the [segment map][Self::segment_map] to a
    /// field in the [layout tree][Self::layout]; see [`CompressedFieldSizes`] for the exact
    /// attribution semantics. No IO is performed.
    pub fn compressed_field_sizes(&self) -> VortexResult<CompressedFieldSizes> {
        CompressedFieldSizes::try_new(&self.root_layout, &self.segments)
    }

    /// Returns the [`DType`] of the file.
    pub fn dtype(&self) -> &DType {
        self.root_layout.dtype()
    }

    /// Returns the approximate size of the footer in bytes, used for caching and memory management.
    pub fn approx_byte_size(&self) -> Option<usize> {
        self.approx_byte_size
    }

    /// Returns the number of rows in the file.
    pub fn row_count(&self) -> u64 {
        self.root_layout.row_count()
    }

    /// Validate that every segment declared in the footer lies within a file of `file_size` bytes.
    pub(crate) fn validate_file_size(&self, file_size: u64) -> VortexResult<()> {
        if self.segments_end.is_some_and(|end| end <= file_size) {
            return Ok(());
        }
        // Walk the segments only to name the offending one in the error.
        validate_segments_within_file(&self.segments_with_metadata, file_size)
    }

    /// Returns a serializer for this footer.
    pub fn into_serializer(self) -> FooterSerializer {
        FooterSerializer::new(self)
    }

    /// Create a deserializer for a Vortex file footer.
    pub fn deserializer(eof_buffer: ByteBuffer, session: VortexSession) -> FooterDeserializer {
        FooterDeserializer::new(eof_buffer, session)
    }
}

fn segments_with_metadata(
    segments: &Arc<[SegmentSpec]>,
    metadata: &[(String, SegmentSpec)],
) -> Arc<[SegmentSpec]> {
    if metadata.is_empty() {
        return Arc::clone(segments);
    }
    segments
        .iter()
        .copied()
        .chain(metadata.iter().map(|(_, segment)| *segment))
        .collect()
}

fn segments_end(segments: &[SegmentSpec]) -> Option<u64> {
    segments.iter().try_fold(0u64, |end, segment| {
        Some(end.max(segment.offset.checked_add(u64::from(segment.length))?))
    })
}

/// Validate that every segment declared in the footer lies within a file of `file_size` bytes.
///
/// A corrupt or malicious file can declare a segment whose offset or length extends past the end
/// of the file. Rejecting such files up front ensures that later slicing of the backing buffer
/// returns a [`VortexError`](vortex_error::VortexError) rather than panicking (see issue #8819).
fn validate_segments_within_file(segments: &[SegmentSpec], file_size: u64) -> VortexResult<()> {
    for segment in segments {
        let within_file = segment
            .offset
            .checked_add(segment.length as u64)
            .is_some_and(|end| end <= file_size);
        if !within_file {
            vortex_bail!(
                "Segment at offset {} with length {} extends past the end of the \
                 {file_size}-byte file",
                segment.offset,
                segment.length,
            );
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use vortex_array::dtype::Nullability;
    use vortex_array::dtype::PType;
    use vortex_buffer::Alignment;
    use vortex_layout::layouts::flat::FlatLayout;
    use vortex_layout::segments::SegmentId;

    use super::*;

    fn segment(offset: u64, length: u32) -> SegmentSpec {
        SegmentSpec {
            offset,
            length,
            alignment: Alignment::none(),
        }
    }

    fn footer_with_metadata_segment(metadata: SegmentSpec) -> Footer {
        let layout = FlatLayout::new(
            1,
            DType::Primitive(PType::I32, Nullability::NonNullable),
            SegmentId::from(0),
            ReadContext::new([]),
        )
        .into_layout();
        Footer::new(
            layout,
            Arc::from([segment(0, 100)]),
            None,
            ReadContext::new([]),
        )
        .with_metadata_segments(Arc::from([("some.key".to_string(), metadata)]))
    }

    #[test]
    fn accepts_segments_within_file() -> VortexResult<()> {
        validate_segments_within_file(&[segment(0, 100), segment(100, 50)], 150)?;
        Ok(())
    }

    #[test]
    fn rejects_segment_extending_past_end_of_file() {
        let err =
            validate_segments_within_file(&[segment(0, 100), segment(100, 51)], 150).unwrap_err();
        assert!(err.to_string().contains("past the end"), "{err}");
    }

    #[test]
    fn rejects_segment_offset_length_overflow() {
        let err = validate_segments_within_file(&[segment(u64::MAX, 1)], u64::MAX).unwrap_err();
        assert!(err.to_string().contains("past the end"), "{err}");
    }

    #[test]
    fn segments_end_covers_overlapping_segments() {
        // Valid files never overlap segments, but a malformed footer may declare an earlier
        // segment that ends after a later one, and the end must cover it so that file size
        // validation still rejects it.
        assert_eq!(
            segments_end(&[segment(0, 300), segment(100, 50)]),
            Some(300)
        );
        assert_eq!(segments_end(&[]), Some(0));
        assert_eq!(segments_end(&[segment(0, 1), segment(u64::MAX, 1)]), None);
    }

    #[test]
    fn segments_with_metadata_appends_metadata_segments() {
        let segments: Arc<[SegmentSpec]> = Arc::from([segment(0, 10), segment(10, 20)]);

        let without = segments_with_metadata(&segments, &[]);
        assert!(Arc::ptr_eq(&without, &segments));

        let with = segments_with_metadata(&segments, &[("key".to_string(), segment(30, 5))]);
        let spans: Vec<_> = with.iter().map(|s| (s.offset, s.length)).collect();
        assert_eq!(spans, [(0, 10), (10, 20), (30, 5)]);
    }

    #[test]
    fn rejects_metadata_segment_extending_past_end_of_file() {
        let err = footer_with_metadata_segment(segment(100, 51))
            .validate_file_size(150)
            .unwrap_err();
        assert!(err.to_string().contains("past the end"), "{err}");
    }

    #[test]
    fn accepts_metadata_segment_within_file() -> VortexResult<()> {
        footer_with_metadata_segment(segment(100, 50)).validate_file_size(150)?;
        Ok(())
    }
}