feat: Layout VTable (#3331)

Similar refactor as #3273 but for layouts.

Again, this introduces the idea that a `LayoutEncoding` can be
independently heap allocated, and is instantiated as a `Layout`. Slight
difference to arrays is that the `Layout` tree is lazy in terms of
parsing children, so child accessors are fallible.

---------

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

Cargo.lock

@@ -96,6 +96,7 @@ impl<V: VTable> Encoding for EncodingAdapter<V> {
             &self.0, dtype, len, &metadata, buffers, children,
         )?;
         assert_eq!(array.len(), len, "Array length mismatch after building");
+        assert_eq!(array.dtype(), dtype, "Array dtype mismatch after building");
         Ok(array.to_array())
     }
 

vortex-array/src/encoding.rs

@@ -39,7 +39,7 @@ use crate::{Array, Encoding, EncodingId, EncodingRef, IntoArray};
 /// out of bounds). Post-conditions are validated after invocation of the vtable function and will
 /// panic if violated.
 pub trait VTable: 'static + Sized + Send + Sync + Debug {
-    type Array: 'static + Send + Sync + Clone + Deref<Target = dyn Array> + IntoArray;
+    type Array: 'static + Send + Sync + Clone + Debug + Deref<Target = dyn Array> + IntoArray;
     type Encoding: 'static + Send + Sync + Clone + Deref<Target = dyn Encoding>;
 
     type ArrayVTable: ArrayVTable<Self>;

vortex-array/src/vtable/mod.rs

@@ -30,6 +30,11 @@ impl<T, const A: usize> ConstBuffer<T, A> {
         self.0.as_slice()
     }
 
+    /// Unwrap the inner buffer.
+    pub fn inner(&self) -> &Buffer<T> {
+        &self.0
+    }
+
     /// Unwrap the inner buffer.
     pub fn into_inner(self) -> Buffer<T> {
         self.0

vortex-buffer/src/const.rs

@@ -162,8 +162,8 @@ fn estimate_layout_size(footer: &Footer) -> usize {
 
     let root_layout = footer.layout();
     let layout_size = size_of::<DType>()
-        + root_layout.metadata().map(|b| b.len()).unwrap_or_default()
-        + root_layout.nsegments() * size_of::<SegmentId>();
+        + root_layout.metadata().len()
+        + root_layout.segment_ids().len() * size_of::<SegmentId>();
 
     segments_size + stats_size + layout_size
 }

vortex-datafusion/src/persistent/cache.rs

@@ -1,4 +1,5 @@
 //! Field mask represents a field projection, which leads to a set of field paths under a given layout.
+// TODO(ngates): this API needs work. It could be made a lot easier to use.
 
 use vortex_error::{VortexResult, vortex_bail};
 

vortex-dtype/src/field_mask.rs

@@ -72,7 +72,8 @@ impl VortexFile {
         let segment_source = self.segment_source();
         self.footer
             .layout()
-            .reader(&segment_source, self.footer().ctx())
+            // TODO(ngates): we may want to allow the user pass in a name here?
+            .new_reader(&"".into(), &segment_source, self.footer().ctx())
     }
 
     /// Initiate a scan of the file, returning a builder for configuring the scan.

vortex-file/src/file.rs

@@ -23,15 +23,15 @@ 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, layout as fbl};
-use vortex_layout::{Layout, LayoutContext, LayoutRegistry};
+use vortex_flatbuffers::{FlatBuffer, footer as fb};
+use vortex_layout::{LayoutContext, LayoutRef, LayoutRegistry, layout_from_flatbuffer};
 
 /// Captures the layout information of a Vortex file.
 #[derive(Debug, Clone)]
 pub struct Footer {
     array_ctx: ArrayContext,
     layout_ctx: LayoutContext,
-    root_layout: Layout,
+    root_layout: LayoutRef,
     segments: Arc<[SegmentSpec]>,
     statistics: Option<FileStatistics>,
 }
@@ -47,7 +47,6 @@ impl Footer {
         layout_registry: &LayoutRegistry,
     ) -> VortexResult<Self> {
         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_footer.layout_specs();
@@ -65,27 +64,7 @@ impl Footer {
             .map(|encoding| encoding.id());
         let array_ctx = array_registry.new_context(array_ids)?;
 
-        let root_encoding = layout_ctx
-            .lookup_encoding(fb_layout.encoding())
-            .ok_or_else(|| {
-                vortex_err!(
-                    "Footer root layout encoding {} not found",
-                    fb_layout.encoding()
-                )
-            })?
-            .clone();
-
-        // SAFETY: We have validated the fb_root_layout at the beginning of this function
-        let root_layout = unsafe {
-            Layout::new_viewed_unchecked(
-                "".into(),
-                root_encoding,
-                dtype,
-                layout_bytes.clone(),
-                fb_layout._tab.loc(),
-                layout_ctx.clone(),
-            )
-        };
+        let root_layout = layout_from_flatbuffer(layout_bytes, &dtype, &layout_ctx)?;
 
         let segments: Arc<[SegmentSpec]> = fb_footer
             .segment_specs()
@@ -118,8 +97,8 @@ impl Footer {
         &self.layout_ctx
     }
 
-    /// Returns the root [`Layout`] of the file.
-    pub fn layout(&self) -> &Layout {
+    /// Returns the root [`LayoutRef`] of the file.
+    pub fn layout(&self) -> &LayoutRef {
         &self.root_layout
     }
 

vortex-file/src/footer/mod.rs

@@ -17,7 +17,7 @@
 //!    offsets. Finding the chunks containing row range is an `Nlog(N)` operation of searching the
 //!    offsets.
 //!
-//! 4. The [`StatsLayout`](vortex_layout::layouts::stats::StatsLayout).
+//! 4. The [`ZonedLayout`](vortex_layout::layouts::zoned::ZonedLayout).
 //!
 //! A layout, alone, is _not_ a standalone Vortex file because layouts are not self-describing. They
 //! neither contain a description of the kind of layout (e.g. flat, column of flat, chunked of

vortex-file/src/lib.rs

@@ -17,10 +17,10 @@ use vortex_layout::layouts::flat::writer::FlatLayoutStrategy;
 use vortex_layout::layouts::repartition::{
     RepartitionStrategy, RepartitionWriter, RepartitionWriterOptions,
 };
-use vortex_layout::layouts::stats::writer::{StatsLayoutOptions, StatsLayoutWriter};
 use vortex_layout::layouts::struct_::writer::StructLayoutWriter;
+use vortex_layout::layouts::zoned::writer::{ZonedLayoutOptions, ZonedLayoutWriter};
 use vortex_layout::segments::SegmentWriter;
-use vortex_layout::{Layout, LayoutStrategy, LayoutWriter, LayoutWriterExt};
+use vortex_layout::{LayoutRef, LayoutStrategy, LayoutWriter, LayoutWriterExt};
 use vortex_scalar::Scalar;
 
 const ROW_BLOCK_SIZE: usize = 8192;
@@ -70,15 +70,15 @@ impl LayoutStrategy for VortexLayoutStrategy {
 
         let writer = dict_strategy.new_writer(ctx, dtype)?;
 
-        // Prior to repartitioning, we record statistics
-        let stats_writer = StatsLayoutWriter::new(
+        // Prior to repartitioning, we create a zone map
+        let zoned_writer = ZonedLayoutWriter::new(
             ctx.clone(),
             dtype,
             writer,
             ArcRef::new_arc(Arc::new(BtrBlocksCompressedStrategy {
                 child: ArcRef::new_arc(Arc::new(FlatLayoutStrategy::default())),
             })),
-            StatsLayoutOptions {
+            ZonedLayoutOptions {
                 block_size: ROW_BLOCK_SIZE,
                 stats: PRUNING_STATS.into(),
             },
@@ -87,7 +87,7 @@ impl LayoutStrategy for VortexLayoutStrategy {
 
         let writer = RepartitionWriter::new(
             dtype.clone(),
-            stats_writer,
+            zoned_writer,
             RepartitionWriterOptions {
                 // No minimum block size in bytes
                 block_size_minimum: 0,
@@ -211,7 +211,7 @@ impl LayoutWriter for BtrBlocksCompressedWriter {
         self.child.flush(segment_writer)
     }
 
-    fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout> {
+    fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<LayoutRef> {
         self.child.finish(segment_writer)
     }
 }
@@ -271,7 +271,7 @@ impl LayoutWriter for BufferedWriter {
         self.child.flush(segment_writer)
     }
 
-    fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout> {
+    fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<LayoutRef> {
         self.child.finish(segment_writer)
     }
 }