layout builder overrides

Signed-off-by: Onur Satici <onur@spiraldb.com>

Author: onursatici

vortex-cuda/src/layout.rs

@@ -180,6 +180,7 @@ impl VTable for CudaFlat {
         name: Arc<str>,
         segment_source: Arc<dyn SegmentSource>,
         session: &VortexSession,
+        _ctx: &vortex::layout::LayoutReaderContext,
     ) -> VortexResult<LayoutReaderRef> {
         Ok(Arc::new(CudaFlatReader {
             layout: layout.clone(),

vortex-layout/public-api.lock

@@ -19,6 +19,7 @@ use vortex_session::registry::Id;
 
 use crate::LayoutEncodingId;
 use crate::LayoutEncodingRef;
+use crate::LayoutReaderContext;
 use crate::LayoutReaderRef;
 use crate::VTable;
 use crate::display::DisplayLayoutTree;
@@ -63,12 +64,30 @@ pub trait Layout: 'static + Send + Sync + Debug + private::Sealed {
     /// Get the segment IDs for this layout.
     fn segment_ids(&self) -> Vec<SegmentId>;
 
-    fn new_reader(
+    /// Construct a new reader for this layout, using the given dependency context.
+    ///
+    /// Implementations check `ctx` for an override registered against this layout's encoding
+    /// ID before dispatching to the layout's default [`crate::VTable::new_reader`].
+    fn new_reader_in_ctx(
         &self,
         name: Arc<str>,
         segment_source: Arc<dyn SegmentSource>,
         session: &VortexSession,
+        ctx: &LayoutReaderContext,
     ) -> VortexResult<LayoutReaderRef>;
+
+    /// Convenience: construct a new reader for this layout using a fresh, empty context.
+    ///
+    /// Top-level callers (file open, tests) typically use this. Recursive calls inside
+    /// layout implementations should use [`Self::new_reader_in_ctx`] to propagate `ctx`.
+    fn new_reader(
+        &self,
+        name: Arc<str>,
+        segment_source: Arc<dyn SegmentSource>,
+        session: &VortexSession,
+    ) -> VortexResult<LayoutReaderRef> {
+        self.new_reader_in_ctx(name, segment_source, session, &LayoutReaderContext::new())
+    }
 }
 
 pub trait IntoLayout {
@@ -306,13 +325,18 @@ impl<V: VTable> Layout for LayoutAdapter<V> {
         V::segment_ids(&self.0)
     }
 
-    fn new_reader(
+    fn new_reader_in_ctx(
         &self,
         name: Arc<str>,
         segment_source: Arc<dyn SegmentSource>,
         session: &VortexSession,
+        ctx: &LayoutReaderContext,
     ) -> VortexResult<LayoutReaderRef> {
-        V::new_reader(&self.0, name, segment_source, session)
+        let id = V::encoding(&self.0).id();
+        if let Some(builder) = ctx.get_override(&id) {
+            return builder(self, name, segment_source, session, ctx);
+        }
+        V::new_reader(&self.0, name, segment_source, session, ctx)
     }
 }
 

vortex-layout/src/layout.rs

@@ -2,14 +2,10 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use std::sync::Arc;
-use std::sync::OnceLock;
 
-use vortex_array::DeserializeMetadata;
-use vortex_array::SerializeMetadata;
+use vortex_array::EmptyMetadata;
 use vortex_array::dtype::DType;
-use vortex_array::dtype::TryFromBytes;
 use vortex_buffer::ByteBuffer;
-use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_panic;
@@ -19,12 +15,11 @@ use vortex_session::registry::ReadContext;
 use crate::LayoutChildType;
 use crate::LayoutEncodingRef;
 use crate::LayoutId;
+use crate::LayoutReaderContext;
 use crate::LayoutReaderRef;
 use crate::LayoutRef;
 use crate::VTable;
 use crate::children::LayoutChildren;
-use crate::layouts::array_tree::ArrayTreesSource;
-use crate::layouts::array_tree::reader::ArrayTreeFlatReader;
 use crate::layouts::flat::FlatLayout;
 use crate::segments::SegmentId;
 use crate::segments::SegmentSource;
@@ -36,44 +31,34 @@ vtable!(ArrayTreeFlat);
 #[derive(Debug)]
 pub struct ArrayTreeFlatLayoutEncoding;
 
-/// A flat layout variant that stores its compact encoding tree separately from the data segment.
+/// A flat layout variant that retrieves its compact encoding tree from a sibling layout's
+/// VarBin payload rather than from the data segment trailer.
 ///
 /// At write time, the compact flatbuffer (encoding tree + buffer descriptors, no stats) is
-/// stored in this layout and later collected by [`super::ArrayTreeLayout`] into a shared VarBin
-/// array.
+/// attached to this layout in memory and later collected by [`super::ArrayTreeLayout`] into
+/// a struct array keyed by segment ID.
 ///
-/// At read time, the compact flatbuffer is retrieved from the shared [`ArrayTreesSource`]
-/// (injected by the parent [`super::ArrayTreeLayout`]'s reader-construction walk) rather than
-/// being parsed from the data segment. This avoids fetching the segment for decode planning
-/// and prevents device-to-host copies for device-resident buffers.
+/// At read time, this layout's reader looks up its compact tree in a shared
+/// [`super::ArrayTreesSource`] using its own [`SegmentId`]. Construction requires that an
+/// ancestor [`super::ArrayTreeLayout`] has registered a reader-builder override against
+/// this encoding's ID — this layout has no useful default reader.
 #[derive(Clone, Debug)]
 pub struct ArrayTreeFlatLayout {
     inner: FlatLayout,
-    chunk_idx: usize,
     /// The compact flatbuffer produced at write time. Not persisted — only used to communicate
     /// between the leaf strategy and the collector strategy via the layout tree.
     compact_tree: Option<ByteBuffer>,
-    /// Shared source for compact flatbuffers, injected by the parent [`super::ArrayTreeLayout`]
-    /// during reader construction.
-    source: OnceLock<Arc<ArrayTreesSource>>,
 }
 
 impl ArrayTreeFlatLayout {
     /// Creates a new layout at write time with a compact flatbuffer.
-    pub fn new(inner: FlatLayout, chunk_idx: usize, compact_tree: ByteBuffer) -> Self {
+    pub fn new(inner: FlatLayout, compact_tree: ByteBuffer) -> Self {
         Self {
             inner,
-            chunk_idx,
             compact_tree: Some(compact_tree),
-            source: OnceLock::new(),
         }
     }
 
-    /// Returns the chunk index of this layout in the array trees VarBin.
-    pub fn chunk_idx(&self) -> usize {
-        self.chunk_idx
-    }
-
     /// Returns the compact flatbuffer, if available (write-time only).
     pub fn compact_tree(&self) -> Option<&ByteBuffer> {
         self.compact_tree.as_ref()
@@ -83,45 +68,12 @@ impl ArrayTreeFlatLayout {
     pub fn inner(&self) -> &FlatLayout {
         &self.inner
     }
-
-    /// Sets the shared array trees source. Called by the parent [`super::ArrayTreeLayout`]
-    /// during the reader-construction injection walk.
-    pub fn set_source(&self, source: Arc<ArrayTreesSource>) {
-        // Ignore if already set (e.g., in tests or double-init scenarios).
-        drop(self.source.set(source));
-    }
-
-    /// Returns the shared array trees source, if set.
-    pub fn source(&self) -> Option<&Arc<ArrayTreesSource>> {
-        self.source.get()
-    }
-}
-
-/// Metadata for [`ArrayTreeFlatLayout`]: stores the chunk index.
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub struct ArrayTreeFlatMetadata {
-    pub chunk_idx: u32,
-}
-
-impl SerializeMetadata for ArrayTreeFlatMetadata {
-    fn serialize(self) -> Vec<u8> {
-        self.chunk_idx.to_le_bytes().to_vec()
-    }
-}
-
-impl DeserializeMetadata for ArrayTreeFlatMetadata {
-    type Output = Self;
-
-    fn deserialize(metadata: &[u8]) -> VortexResult<Self::Output> {
-        let chunk_idx = u32::try_from_le_bytes(&metadata[0..4])?;
-        Ok(Self { chunk_idx })
-    }
 }
 
 impl VTable for ArrayTreeFlat {
     type Layout = ArrayTreeFlatLayout;
     type Encoding = ArrayTreeFlatLayoutEncoding;
-    type Metadata = ArrayTreeFlatMetadata;
+    type Metadata = EmptyMetadata;
 
     fn id(_encoding: &Self::Encoding) -> LayoutId {
         LayoutId::new_static("vortex.array_tree_flat")
@@ -139,10 +91,8 @@ impl VTable for ArrayTreeFlat {
         layout.inner.dtype()
     }
 
-    fn metadata(layout: &Self::Layout) -> Self::Metadata {
-        ArrayTreeFlatMetadata {
-            chunk_idx: u32::try_from(layout.chunk_idx).vortex_expect("chunk_idx must fit in u32"),
-        }
+    fn metadata(_layout: &Self::Layout) -> Self::Metadata {
+        EmptyMetadata
     }
 
     fn segment_ids(layout: &Self::Layout) -> Vec<SegmentId> {
@@ -162,24 +112,27 @@ impl VTable for ArrayTreeFlat {
     }
 
     fn new_reader(
-        layout: &Self::Layout,
-        name: Arc<str>,
-        segment_source: Arc<dyn SegmentSource>,
-        session: &VortexSession,
+        _layout: &Self::Layout,
+        _name: Arc<str>,
+        _segment_source: Arc<dyn SegmentSource>,
+        _session: &VortexSession,
+        _ctx: &LayoutReaderContext,
     ) -> VortexResult<LayoutReaderRef> {
-        Ok(Arc::new(ArrayTreeFlatReader::new(
-            layout.clone(),
-            name,
-            segment_source,
-            session.clone(),
-        )))
+        // ArrayTreeFlatLayout has no useful default reader. It exists to be intercepted by an
+        // ancestor ArrayTreeLayout that registers a reader-builder override carrying the
+        // shared ArrayTreesSource. If the dispatcher reached this method, no such ancestor
+        // was present in the layout tree.
+        vortex_bail!(
+            "ArrayTreeFlatLayout requires an ancestor ArrayTreeLayout to register a reader \
+             builder override; this layout cannot be read on its own"
+        )
     }
 
     fn build(
         _encoding: &Self::Encoding,
         dtype: &DType,
         row_count: u64,
-        metadata: &ArrayTreeFlatMetadata,
+        _metadata: &EmptyMetadata,
         segment_ids: Vec<SegmentId>,
         _children: &dyn LayoutChildren,
         ctx: &ReadContext,
@@ -189,9 +142,7 @@ impl VTable for ArrayTreeFlat {
         }
         Ok(ArrayTreeFlatLayout {
             inner: FlatLayout::new(row_count, dtype.clone(), segment_ids[0], ctx.clone()),
-            chunk_idx: metadata.chunk_idx as usize,
             compact_tree: None,
-            source: OnceLock::new(),
         })
     }
 }