Pull out layout flatbuffer (#3127)

For queries that exclusively use column-level stats, in theory we could
defer loading the layout flatbuffer entirely...

In any case, this makes it easier to build custom logic for deferred
flatbuffer verification over layouts (when it's not included as part of
footer validation). We could build similar logic into DType flatbuffer
if that also gets expensive.

Author: gatesn

vortex-file/src/footer/file_layout.rs

@@ -9,36 +9,32 @@ use std::sync::Arc;
 use flatbuffers::{FlatBufferBuilder, WIPOffset};
 use vortex_array::ArrayContext;
 use vortex_flatbuffers::{FlatBufferRoot, WriteFlatBuffer, footer as fb};
-use vortex_layout::{Layout, LayoutContext};
+use vortex_layout::LayoutContext;
 
 use crate::footer::segment::SegmentSpec;
 
 /// A writer for serializing a file layout to a FlatBuffer.
 ///
 /// This struct is used to write the layout component of a Vortex file footer,
 /// which describes the structure of the data in the file.
-pub(crate) struct FileLayoutFlatBufferWriter {
+pub(crate) struct FooterFlatBufferWriter {
     /// The array context containing encodings used in the file.
     pub(crate) ctx: ArrayContext,
-    /// The root layout of the file.
-    pub(crate) layout: Layout,
+    /// The layout context containing the layouts used in the file.
+    pub(crate) layout_ctx: LayoutContext,
     /// Specifications for all segments in the file.
     pub(crate) segment_specs: Arc<[SegmentSpec]>,
 }
 
-impl FlatBufferRoot for FileLayoutFlatBufferWriter {}
+impl FlatBufferRoot for FooterFlatBufferWriter {}
 
-impl WriteFlatBuffer for FileLayoutFlatBufferWriter {
-    type Target<'a> = fb::FileLayout<'a>;
+impl WriteFlatBuffer for FooterFlatBufferWriter {
+    type Target<'a> = fb::Footer<'a>;
 
     fn write_flatbuffer<'fb>(
         &self,
         fbb: &mut FlatBufferBuilder<'fb>,
     ) -> WIPOffset<Self::Target<'fb>> {
-        // Set up a layout context to capture the layouts used in the file.
-        let layout_ctx = LayoutContext::empty();
-        let layout = self.layout.write_flatbuffer(fbb, &layout_ctx);
-
         let segment_specs =
             fbb.create_vector_from_iter(self.segment_specs.iter().map(fb::SegmentSpec::from));
 
@@ -53,7 +49,8 @@ impl WriteFlatBuffer for FileLayoutFlatBufferWriter {
             .collect::<Vec<_>>();
         let array_specs = fbb.create_vector(array_specs.as_slice());
 
-        let layout_specs = layout_ctx
+        let layout_specs = self
+            .layout_ctx
             .encodings()
             .iter()
             .map(|e| {
@@ -63,10 +60,9 @@ impl WriteFlatBuffer for FileLayoutFlatBufferWriter {
             .collect::<Vec<_>>();
         let layout_specs = fbb.create_vector(layout_specs.as_slice());
 
-        fb::FileLayout::create(
+        fb::Footer::create(
             fbb,
-            &fb::FileLayoutArgs {
-                layout: Some(layout),
+            &fb::FooterArgs {
                 segment_specs: Some(segment_specs),
                 array_specs: Some(array_specs),
                 layout_specs: Some(layout_specs),

vortex-file/src/footer/mod.rs

@@ -23,7 +23,7 @@ use vortex_array::stats::StatsSet;
 use vortex_array::{ArrayContext, ArrayRegistry};
 use vortex_dtype::DType;
 use vortex_error::{VortexResult, vortex_bail, vortex_err};
-use vortex_flatbuffers::{FlatBuffer, footer as fb};
+use vortex_flatbuffers::{FlatBuffer, footer as fb, layout as fbl};
 use vortex_layout::{Layout, LayoutContext, LayoutRegistry};
 
 /// Captures the layout information of a Vortex file.
@@ -39,39 +39,38 @@ pub struct Footer {
 impl Footer {
     /// Read the [`Footer`] from a flatbuffer.
     pub(crate) fn from_flatbuffer(
-        flatbuffer: FlatBuffer,
+        footer_bytes: FlatBuffer,
+        layout_bytes: FlatBuffer,
         dtype: DType,
         statistics: Option<FileStatistics>,
         array_registry: &ArrayRegistry,
         layout_registry: &LayoutRegistry,
     ) -> VortexResult<Self> {
-        let fb = root::<fb::FileLayout>(&flatbuffer)?;
-        let fb_root_layout = fb
-            .layout()
-            .ok_or_else(|| vortex_err!("Footer missing root layout"))?;
+        let fb_footer = root::<fb::Footer>(&footer_bytes)?;
+        let fb_layout = root::<fbl::Layout>(&layout_bytes)?;
 
         // Create a LayoutContext from the registry.
-        let layout_specs = fb.layout_specs();
+        let layout_specs = fb_footer.layout_specs();
         let layout_ids = layout_specs
             .iter()
             .flat_map(|e| e.iter())
             .map(|encoding| encoding.id());
         let layout_ctx = layout_registry.new_context(layout_ids)?;
 
         // Create an ArrayContext from the registry.
-        let array_specs = fb.array_specs();
+        let array_specs = fb_footer.array_specs();
         let array_ids = array_specs
             .iter()
             .flat_map(|e| e.iter())
             .map(|encoding| encoding.id());
         let array_ctx = array_registry.new_context(array_ids)?;
 
         let root_encoding = layout_ctx
-            .lookup_encoding(fb_root_layout.encoding())
+            .lookup_encoding(fb_layout.encoding())
             .ok_or_else(|| {
                 vortex_err!(
                     "Footer root layout encoding {} not found",
-                    fb_root_layout.encoding()
+                    fb_layout.encoding()
                 )
             })?
             .clone();
@@ -82,13 +81,13 @@ impl Footer {
                 "".into(),
                 root_encoding,
                 dtype,
-                flatbuffer.clone(),
-                fb_root_layout._tab.loc(),
+                layout_bytes.clone(),
+                fb_layout._tab.loc(),
                 layout_ctx.clone(),
             )
         };
 
-        let segments: Arc<[SegmentSpec]> = fb
+        let segments: Arc<[SegmentSpec]> = fb_footer
             .segment_specs()
             .ok_or_else(|| vortex_err!("FileLayout missing segment specs"))?
             .iter()

vortex-file/src/footer/postscript.rs

@@ -7,8 +7,9 @@ use vortex_flatbuffers::{FlatBufferRoot, ReadFlatBuffer, WriteFlatBuffer, footer
 /// reading a Vortex file.
 pub(crate) struct Postscript {
     pub(crate) dtype: Option<PostscriptSegment>,
-    pub(crate) statistics: Option<PostscriptSegment>,
     pub(crate) layout: PostscriptSegment,
+    pub(crate) statistics: Option<PostscriptSegment>,
+    pub(crate) footer: PostscriptSegment,
 }
 
 impl FlatBufferRoot for Postscript {}
@@ -21,14 +22,16 @@ impl WriteFlatBuffer for Postscript {
         fbb: &mut FlatBufferBuilder<'fb>,
     ) -> WIPOffset<Self::Target<'fb>> {
         let dtype = self.dtype.as_ref().map(|ps| ps.write_flatbuffer(fbb));
-        let statistics = self.statistics.as_ref().map(|ps| ps.write_flatbuffer(fbb));
         let layout = self.layout.write_flatbuffer(fbb);
+        let statistics = self.statistics.as_ref().map(|ps| ps.write_flatbuffer(fbb));
+        let footer = self.footer.write_flatbuffer(fbb);
         fb::Postscript::create(
             fbb,
             &fb::PostscriptArgs {
                 dtype,
-                statistics,
                 layout: Some(layout),
+                statistics,
+                footer: Some(footer),
             },
         )
     }
@@ -46,13 +49,17 @@ impl ReadFlatBuffer for Postscript {
                 .dtype()
                 .map(|ps| PostscriptSegment::read_flatbuffer(&ps))
                 .transpose()?,
+            layout: PostscriptSegment::read_flatbuffer(
+                &fb.layout()
+                    .ok_or_else(|| vortex_err!("Postscript missing layout segment"))?,
+            )?,
             statistics: fb
                 .statistics()
                 .map(|ps| PostscriptSegment::read_flatbuffer(&ps))
                 .transpose()?,
-            layout: PostscriptSegment::read_flatbuffer(
-                &fb.layout()
-                    .ok_or_else(|| vortex_err!("Postscript missing layout segment"))?,
+            footer: PostscriptSegment::read_flatbuffer(
+                &fb.footer()
+                    .ok_or_else(|| vortex_err!("Postscript missing footer segment"))?,
             )?,
         })
     }

vortex-file/src/generic.rs

@@ -154,6 +154,7 @@ impl VortexOpenOptions<GenericVortexFile> {
             read_more_offset = read_more_offset.min(stats_segment.offset);
         }
         read_more_offset = read_more_offset.min(postscript.layout.offset);
+        read_more_offset = read_more_offset.min(postscript.footer.offset);
 
         // Read more bytes if necessary.
         if read_more_offset < initial_offset {
@@ -188,6 +189,7 @@ impl VortexOpenOptions<GenericVortexFile> {
         let footer = self.parse_footer(
             initial_offset,
             &initial_read,
+            &postscript.footer,
             &postscript.layout,
             dtype,
             file_stats,

vortex-file/src/memory.rs

@@ -47,7 +47,14 @@ impl VortexOpenOptions<InMemoryFileType> {
             .map(|segment| self.parse_file_statistics(0, &buffer, &segment))
             .transpose()?;
 
-        let footer = self.parse_footer(0, &buffer, &postscript.layout, dtype, file_stats)?;
+        let footer = self.parse_footer(
+            0,
+            &buffer,
+            &postscript.footer,
+            &postscript.layout,
+            dtype,
+            file_stats,
+        )?;
 
         let segment_source_factory = Arc::new(InMemorySegmentReader {
             buffer,

vortex-file/src/open.rs

@@ -170,15 +170,24 @@ impl<F: FileType> VortexOpenOptions<F> {
         &self,
         initial_offset: u64,
         initial_read: &[u8],
+        footer_segment: &PostscriptSegment,
         layout_segment: &PostscriptSegment,
         dtype: DType,
         file_stats: Option<FileStatistics>,
     ) -> VortexResult<Footer> {
-        let offset = usize::try_from(layout_segment.offset - initial_offset)?;
-        let bytes =
-            FlatBuffer::copy_from(&initial_read[offset..offset + (layout_segment.length as usize)]);
+        let footer_offset = usize::try_from(footer_segment.offset - initial_offset)?;
+        let footer_bytes = FlatBuffer::copy_from(
+            &initial_read[footer_offset..footer_offset + (footer_segment.length as usize)],
+        );
+
+        let layout_offset = usize::try_from(layout_segment.offset - initial_offset)?;
+        let layout_bytes = FlatBuffer::copy_from(
+            &initial_read[layout_offset..layout_offset + (layout_segment.length as usize)],
+        );
+
         Footer::from_flatbuffer(
-            bytes,
+            footer_bytes,
+            layout_bytes,
             dtype,
             file_stats,
             &self.registry,

vortex-file/src/writer.rs

@@ -6,9 +6,9 @@ use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
 use vortex_flatbuffers::{FlatBuffer, FlatBufferRoot, WriteFlatBuffer, WriteFlatBufferExt};
 use vortex_io::VortexWrite;
 use vortex_layout::layouts::file_stats::FileStatsLayoutWriter;
-use vortex_layout::{LayoutStrategy, LayoutWriter};
+use vortex_layout::{LayoutContext, LayoutStrategy, LayoutWriter};
 
-use crate::footer::{FileLayoutFlatBufferWriter, FileStatistics, Postscript, PostscriptSegment};
+use crate::footer::{FileStatistics, FooterFlatBufferWriter, Postscript, PostscriptSegment};
 use crate::segments::writer::BufferedSegmentWriter;
 use crate::strategy::VortexLayoutStrategy;
 use crate::{EOF_SIZE, MAGIC_BYTES, MAX_FOOTER_SIZE, VERSION};
@@ -112,19 +112,24 @@ impl VortexWriteOptions {
             Some(self.write_flatbuffer(&mut write, stream.dtype()).await?)
         };
 
+        let layout_ctx = LayoutContext::empty();
+        let layout_segment = self
+            .write_flatbuffer(&mut write, &layout.flatbuffer_writer(&layout_ctx))
+            .await?;
+
         let statistics_segment = if self.file_statistics.is_empty() {
             None
         } else {
             let file_statistics = FileStatistics(layout_writer.into_stats_sets().into());
             Some(self.write_flatbuffer(&mut write, &file_statistics).await?)
         };
 
-        let layout_segment = self
+        let footer_segment = self
             .write_flatbuffer(
                 &mut write,
-                &FileLayoutFlatBufferWriter {
+                &FooterFlatBufferWriter {
                     ctx,
-                    layout,
+                    layout_ctx,
                     segment_specs: segment_specs.into(),
                 },
             )
@@ -133,8 +138,9 @@ impl VortexWriteOptions {
         // Assemble the postscript, and write it manually to avoid any framing.
         let postscript = Postscript {
             dtype: dtype_segment,
-            statistics: statistics_segment,
             layout: layout_segment,
+            statistics: statistics_segment,
+            footer: footer_segment,
         };
         let postscript_buffer = postscript.write_flatbuffer_bytes();
         if postscript_buffer.len() > MAX_FOOTER_SIZE as usize {

vortex-flatbuffers/flatbuffers/vortex-file/footer.fbs

@@ -20,10 +20,12 @@ include "vortex-layout/layout.fbs";
 table Postscript {
     /// Segment containing the root `DType` flatbuffer.
     dtype: PostscriptSegment;
+    /// Segment containing the root `Layout` flatbuffer (required).
+    layout: PostscriptSegment;
     /// Segment containing the file-level `Statistics` flatbuffer.
     statistics: PostscriptSegment;
-    /// Segment containing the `FileLayout` flatbuffer (required).
-    layout: PostscriptSegment;
+    /// Segment containing the 'Footer' flatbuffer (required)
+    footer: PostscriptSegment;
 }
 
 /// A `PostscriptSegment` describes the location of a segment in the file without referencing any
@@ -47,10 +49,7 @@ table FileStatistics {
 
 /// The `Registry` object stores dictionary-encoded configuration for segments,
 /// compression schemes, encryption schemes, etc.
-table FileLayout {
-    // The root [`Layout`] of the file.
-    layout: Layout;
-
+table Footer {
     // Dictionary-encoded array specs, up to u16::MAX.
     array_specs: [ArraySpec];
     // Dictionary-encoded layout specs, up to u16::MAX.
@@ -111,5 +110,5 @@ table EncryptionSpec {
 // [footer]
 
 root_type FileStatistics;
-root_type FileLayout;
+root_type Footer;
 root_type Postscript;