feat[layout]: support array encoding outlined in flat layout (#5086)

set with `FLAT_LAYOUT_INLINE_ARRAY_NODE`.

This can be used to plan a scan without loading the layout segments

---------

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

vortex-layout/src/display.rs

@@ -7,11 +7,20 @@ use vortex_error::VortexResult;
 use crate::LayoutRef;
 
 /// Display wrapper for layout tree visualization
-pub struct DisplayLayoutTree(pub LayoutRef);
+pub struct DisplayLayoutTree {
+    layout: LayoutRef,
+    verbose: bool,
+}
+
+impl DisplayLayoutTree {
+    pub fn new(layout: LayoutRef, verbose: bool) -> Self {
+        Self { layout, verbose }
+    }
+}
 
 impl std::fmt::Display for DisplayLayoutTree {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        fn make_tree(layout: LayoutRef) -> VortexResult<Tree<String>> {
+        fn make_tree(layout: LayoutRef, verbose: bool) -> VortexResult<Tree<String>> {
             // Build the node label with encoding, dtype, and metadata
             let mut node_parts = vec![
                 format!("{}", layout.encoding()),
@@ -24,6 +33,30 @@ impl std::fmt::Display for DisplayLayoutTree {
                 node_parts.push(format!("children: {}", nchildren));
             }
 
+            // Add metadata information if verbose mode is enabled
+            if verbose {
+                let metadata = layout.metadata();
+                if !metadata.is_empty() {
+                    node_parts.push(format!("metadata: {} bytes", metadata.len()));
+                }
+
+                // Add segment IDs
+                let segment_ids = layout.segment_ids();
+                if !segment_ids.is_empty() {
+                    node_parts.push(format!(
+                        "segments: [{}]",
+                        segment_ids
+                            .iter()
+                            .map(|s| s.to_string())
+                            .collect::<Vec<_>>()
+                            .join(", ")
+                    ));
+                }
+
+                // Add row count
+                node_parts.push(format!("rows: {}", layout.row_count()));
+            }
+
             let node_name = node_parts.join(", ");
 
             // Get children and child names directly from the layout (not loading arrays)
@@ -38,14 +71,17 @@ impl std::fmt::Display for DisplayLayoutTree {
                         .into_iter()
                         .zip(child_names.iter())
                         .map(|(child, name)| {
-                            let child_tree = make_tree(child)?;
+                            let child_tree = make_tree(child, verbose)?;
                             Ok(Tree::new(format!("{}: {}", name, child_tree.root))
                                 .with_leaves(child_tree.leaves))
                         })
                         .collect()
                 } else if !children.is_empty() {
                     // No names available, just show children
-                    children.into_iter().map(make_tree).collect()
+                    children
+                        .into_iter()
+                        .map(|c| make_tree(c, verbose))
+                        .collect()
                 } else {
                     // Leaf node - no children
                     Ok(Vec::new())
@@ -54,7 +90,7 @@ impl std::fmt::Display for DisplayLayoutTree {
             Ok(Tree::new(node_name).with_leaves(child_trees?))
         }
 
-        match make_tree(self.0.clone()) {
+        match make_tree(self.layout.clone(), self.verbose) {
             Ok(tree) => write!(f, "{}", tree),
             Err(e) => write!(f, "Error building layout tree: {}", e),
         }

vortex-layout/src/layout.rs

@@ -188,7 +188,12 @@ impl dyn Layout + '_ {
 
     /// Display the layout as a tree structure.
     pub fn display_tree(&self) -> DisplayLayoutTree {
-        DisplayLayoutTree(self.to_layout())
+        DisplayLayoutTree::new(self.to_layout(), false)
+    }
+
+    /// Display the layout as a tree structure with optional verbose metadata.
+    pub fn display_tree_verbose(&self, verbose: bool) -> DisplayLayoutTree {
+        DisplayLayoutTree::new(self.to_layout(), verbose)
     }
 }
 

vortex-layout/src/layouts/flat/mod.rs

@@ -4,9 +4,11 @@
 mod reader;
 pub mod writer;
 
-use std::sync::Arc;
+use std::env;
+use std::sync::{Arc, LazyLock};
 
-use vortex_array::{ArrayContext, DeserializeMetadata, EmptyMetadata};
+use vortex_array::{ArrayContext, DeserializeMetadata, ProstMetadata};
+use vortex_buffer::ByteBuffer;
 use vortex_dtype::DType;
 use vortex_error::{VortexResult, vortex_bail, vortex_panic};
 
@@ -17,12 +19,15 @@ use crate::{
     LayoutChildType, LayoutEncodingRef, LayoutId, LayoutReaderRef, LayoutRef, VTable, vtable,
 };
 
+static FLAT_LAYOUT_INLINE_ARRAY_NODE: LazyLock<bool> =
+    LazyLock::new(|| env::var("FLAT_LAYOUT_INLINE_ARRAY_NODE").is_ok());
+
 vtable!(Flat);
 
 impl VTable for FlatVTable {
     type Layout = FlatLayout;
     type Encoding = FlatLayoutEncoding;
-    type Metadata = EmptyMetadata;
+    type Metadata = ProstMetadata<FlatLayoutMetadata>;
 
     fn id(_encoding: &Self::Encoding) -> LayoutId {
         LayoutId::new_ref("vortex.flat")
@@ -40,8 +45,10 @@ impl VTable for FlatVTable {
         &layout.dtype
     }
 
-    fn metadata(_layout: &Self::Layout) -> Self::Metadata {
-        EmptyMetadata
+    fn metadata(layout: &Self::Layout) -> Self::Metadata {
+        ProstMetadata(FlatLayoutMetadata {
+            array_encoding_tree: layout.array_tree.as_ref().map(|bytes| bytes.to_vec()),
+        })
     }
 
     fn segment_ids(layout: &Self::Layout) -> Vec<SegmentId> {
@@ -76,20 +83,24 @@ impl VTable for FlatVTable {
         _encoding: &Self::Encoding,
         dtype: &DType,
         row_count: u64,
-        _metadata: &<Self::Metadata as DeserializeMetadata>::Output,
+        metadata: &<Self::Metadata as DeserializeMetadata>::Output,
         segment_ids: Vec<SegmentId>,
         _children: &dyn LayoutChildren,
         ctx: ArrayContext,
     ) -> VortexResult<Self::Layout> {
         if segment_ids.len() != 1 {
             vortex_bail!("Flat layout must have exactly one segment ID");
         }
-        Ok(FlatLayout {
+        Ok(FlatLayout::new_with_metadata(
             row_count,
-            dtype: dtype.clone(),
-            segment_id: segment_ids[0],
+            dtype.clone(),
+            segment_ids[0],
             ctx,
-        })
+            metadata
+                .array_encoding_tree
+                .as_ref()
+                .map(|v| ByteBuffer::from(v.clone())),
+        ))
     }
 }
 
@@ -102,6 +113,7 @@ pub struct FlatLayout {
     dtype: DType,
     segment_id: SegmentId,
     ctx: ArrayContext,
+    array_tree: Option<ByteBuffer>,
 }
 
 impl FlatLayout {
@@ -111,10 +123,36 @@ impl FlatLayout {
             dtype,
             segment_id,
             ctx,
+            array_tree: None,
+        }
+    }
+
+    pub fn new_with_metadata(
+        row_count: u64,
+        dtype: DType,
+        segment_id: SegmentId,
+        ctx: ArrayContext,
+        metadata: Option<ByteBuffer>,
+    ) -> Self {
+        Self {
+            row_count,
+            dtype,
+            segment_id,
+            ctx,
+            array_tree: metadata,
         }
     }
 
     pub fn segment_id(&self) -> SegmentId {
         self.segment_id
     }
 }
+
+#[derive(prost::Message)]
+pub struct FlatLayoutMetadata {
+    // We can optionally store the array encoding tree here to avoid needing to fetch the segment
+    // to plan array deserialization.
+    // This will be a `ArrayNode`.
+    #[prost(optional, bytes, tag = "1")]
+    pub array_encoding_tree: Option<Vec<u8>>,
+}

vortex-layout/src/layouts/flat/writer.rs

@@ -10,7 +10,7 @@ use vortex_dtype::DType;
 use vortex_error::{VortexResult, vortex_bail};
 use vortex_io::runtime::Handle;
 
-use crate::layouts::flat::FlatLayout;
+use crate::layouts::flat::{FLAT_LAYOUT_INLINE_ARRAY_NODE, FlatLayout};
 use crate::layouts::zoned::{lower_bound, upper_bound};
 use crate::segments::SegmentSinkRef;
 use crate::sequence::{SendableSequentialStream, SequencePointer};
@@ -122,15 +122,23 @@ impl LayoutStrategy for FlatLayoutStrategy {
                 include_padding: options.include_padding,
             },
         )?;
+        // there is at least the flatbuffer and the length
+        assert!(buffers.len() >= 2);
+        let array_node =
+            (*FLAT_LAYOUT_INLINE_ARRAY_NODE).then(|| buffers[buffers.len() - 2].clone());
         let segment_id = segment_sink.write(sequence_id, buffers).await?;
 
         let None = stream.next().await else {
             vortex_bail!("flat layout received stream with more than a single chunk");
         };
-        Ok(
-            FlatLayout::new(row_count, stream.dtype().clone(), segment_id, ctx.clone())
-                .into_layout(),
+        Ok(FlatLayout::new_with_metadata(
+            row_count,
+            stream.dtype().clone(),
+            segment_id,
+            ctx.clone(),
+            array_node,
         )
+        .into_layout())
     }
 
     fn buffered_bytes(&self) -> u64 {

vortex-tui/src/browse/app.rs

@@ -10,10 +10,10 @@ use ratatui::widgets::ListState;
 use vortex::dtype::DType;
 use vortex::error::{VortexExpect, VortexResult, VortexUnwrap};
 use vortex::file::{Footer, SegmentSpec, VortexFile, VortexOpenOptions};
-use vortex::layout::LayoutRef;
 use vortex::layout::layouts::flat::FlatVTable;
 use vortex::layout::layouts::zoned::ZonedVTable;
 use vortex::layout::segments::{SegmentId, SegmentSource};
+use vortex::layout::{LayoutRef, VTable};
 use vortex::serde::ArrayParts;
 
 use crate::browse::ui::SegmentGridState;
@@ -103,6 +103,27 @@ impl LayoutCursor {
             .len()
     }
 
+    /// Get information about the flat layout metadata.
+    ///
+    /// NOTE: this is only safe to run against a FLAT layout.
+    pub fn flat_layout_metadata_info(&self) -> String {
+        let flat_layout = self.layout.as_::<FlatVTable>();
+        let metadata = FlatVTable::metadata(flat_layout);
+
+        // Check if array_encoding_tree is present and get its size
+        match metadata.0.array_encoding_tree.as_ref() {
+            Some(tree) => {
+                let size = tree.len();
+                // Truncate to a single line - show the size and presence
+                format!(
+                    "Flat Metadata: array_encoding_tree present ({} bytes)",
+                    size
+                )
+            }
+            None => "Flat Metadata: array_encoding_tree not present".to_string(),
+        }
+    }
+
     pub fn total_size(&self) -> usize {
         self.layout_segments()
             .iter()

vortex-tui/src/browse/ui/layouts.rs

@@ -64,6 +64,10 @@ fn render_layout_header(cursor: &LayoutCursor, area: Rect, buf: &mut Buffer) {
             "FlatBuffer Size: {}",
             size_formatter(cursor.flatbuffer_size())
         )));
+
+        // Display metadata info about the flat layout
+        let metadata_info = cursor.flat_layout_metadata_info();
+        rows.push(Text::from(metadata_info));
     }
 
     if let Some(layout) = cursor.layout().as_opt::<ZonedVTable>() {

vortex-tui/src/main.rs

@@ -39,7 +39,7 @@ impl Commands {
         match self {
             Commands::Tree(args) => match &args.mode {
                 tree::TreeMode::Array { file } => file,
-                tree::TreeMode::Layout { file } => file,
+                tree::TreeMode::Layout { file, .. } => file,
             },
             Commands::Browse { file } => file,
             Commands::Convert(flags) => &flags.file,

vortex-tui/src/tree.rs

@@ -25,13 +25,16 @@ pub enum TreeMode {
     Layout {
         /// Path to the Vortex file
         file: PathBuf,
+        /// Show additional metadata information
+        #[arg(short, long)]
+        verbose: bool,
     },
 }
 
 pub async fn exec_tree(args: TreeArgs) -> VortexResult<()> {
     match args.mode {
         TreeMode::Array { file } => exec_array_tree(&file).await?,
-        TreeMode::Layout { file } => exec_layout_tree(&file).await?,
+        TreeMode::Layout { file, verbose } => exec_layout_tree(&file, verbose).await?,
     }
 
     Ok(())
@@ -51,11 +54,11 @@ async fn exec_array_tree(file: &Path) -> VortexResult<()> {
     Ok(())
 }
 
-async fn exec_layout_tree(file: &Path) -> VortexResult<()> {
+async fn exec_layout_tree(file: &Path, verbose: bool) -> VortexResult<()> {
     let vxf = VortexOpenOptions::new().open(file).await?;
     let footer = vxf.footer();
 
-    println!("{}", footer.layout().display_tree());
+    println!("{}", footer.layout().display_tree_verbose(verbose));
 
     Ok(())
 }