break: new layout tree vtable

vortex-cuda/src/file.rs

@@ -11,7 +11,7 @@ use vortex::error::vortex_bail;
 use vortex::file::VortexFile;
 use vortex::file::VortexOpenOptions;
 use vortex::io::session::RuntimeSessionExt;
-use vortex::layout::Layout;
+use vortex::layout::DynLayout;
 use vortex::layout::layouts::zoned::LegacyStats;
 use vortex::layout::layouts::zoned::Zoned;
 use vortex::layout::segments::SegmentFuture;
@@ -107,13 +107,16 @@ impl SegmentSource for RoutingSegmentSource {
     }
 }
 
-fn control_plane_segments(layout: &dyn Layout, segment_count: usize) -> VortexResult<Vec<bool>> {
+fn control_plane_segments(layout: &dyn DynLayout, segment_count: usize) -> VortexResult<Vec<bool>> {
     let mut control_segments = vec![false; segment_count];
     mark_zone_map_segments(layout, &mut control_segments)?;
     Ok(control_segments)
 }
 
-fn mark_zone_map_segments(layout: &dyn Layout, control_segments: &mut [bool]) -> VortexResult<()> {
+fn mark_zone_map_segments(
+    layout: &dyn DynLayout,
+    control_segments: &mut [bool],
+) -> VortexResult<()> {
     if layout.is::<Zoned>() || layout.is::<LegacyStats>() {
         mark_layout_segments(layout.child(1)?.as_ref(), control_segments)?;
         mark_zone_map_segments(layout.child(0)?.as_ref(), control_segments)?;
@@ -126,7 +129,7 @@ fn mark_zone_map_segments(layout: &dyn Layout, control_segments: &mut [bool]) ->
     Ok(())
 }
 
-fn mark_layout_segments(layout: &dyn Layout, control_segments: &mut [bool]) -> VortexResult<()> {
+fn mark_layout_segments(layout: &dyn DynLayout, control_segments: &mut [bool]) -> VortexResult<()> {
     for id in layout.segment_ids() {
         let idx = usize::try_from(*id)?;
         let Some(is_control) = control_segments.get_mut(idx) else {
@@ -153,7 +156,6 @@ mod tests {
     use vortex::array::dtype::PType;
     use vortex::array::dtype::StructFields;
     use vortex::buffer::ByteBuffer;
-    use vortex::layout::IntoLayout;
     use vortex::layout::layouts::flat::FlatLayout;
     use vortex::layout::layouts::zoned::ZonedLayout;
     use vortex::session::registry::ReadContext;

vortex-cuda/src/layout.rs

@@ -16,7 +16,6 @@ use futures::future::BoxFuture;
 use vortex::array::ArrayContext;
 use vortex::array::ArrayRef;
 use vortex::array::ArrayVTable;
-use vortex::array::DeserializeMetadata;
 use vortex::array::MaskFuture;
 use vortex::array::ProstMetadata;
 use vortex::array::VortexSessionExecute;
@@ -36,26 +35,26 @@ use vortex::error::VortexExpect;
 use vortex::error::VortexResult;
 use vortex::error::vortex_bail;
 use vortex::error::vortex_panic;
-use vortex::layout::IntoLayout;
-use vortex::layout::LayoutBuildContext;
+use vortex::layout::Layout;
 use vortex::layout::LayoutChildType;
-use vortex::layout::LayoutChildren;
+use vortex::layout::LayoutDeserializeArgs;
 use vortex::layout::LayoutEncodingRef;
 use vortex::layout::LayoutId;
+use vortex::layout::LayoutParts;
 use vortex::layout::LayoutReader;
 use vortex::layout::LayoutReaderRef;
 use vortex::layout::LayoutRef;
 use vortex::layout::LayoutStrategy;
 use vortex::layout::RowSplits;
 use vortex::layout::SplitRange;
 use vortex::layout::VTable;
+use vortex::layout::layout_children;
 use vortex::layout::layouts::SharedArrayFuture;
 use vortex::layout::segments::SegmentId;
 use vortex::layout::segments::SegmentSinkRef;
 use vortex::layout::segments::SegmentSource;
 use vortex::layout::sequence::SendableSequentialStream;
 use vortex::layout::sequence::SequencePointer;
-use vortex::layout::vtable;
 use vortex::mask::Mask;
 use vortex::scalar::Scalar;
 use vortex::scalar::ScalarTruncation;
@@ -84,23 +83,27 @@ pub struct CudaFlatLayoutMetadata {
     pub host_buffers: Vec<InlinedBuffer>,
 }
 
-vtable!(CudaFlat);
+/// CUDA flat layout vtable.
+#[derive(Clone, Debug)]
+pub struct CudaFlat;
 
-#[derive(Debug)]
-pub struct CudaFlatLayoutEncoding;
+/// Backwards-compatible plugin name.
+pub use CudaFlat as CudaFlatLayoutEncoding;
 
+/// CUDA-flat-specific data.
 #[derive(Clone, Debug)]
-pub struct CudaFlatLayout {
-    row_count: u64,
-    dtype: DType,
+pub struct CudaFlatData {
     segment_id: SegmentId,
     ctx: ReadContext,
     array_tree: ByteBuffer,
     /// Small buffers kept on host, keyed by global buffer index.
     host_buffers: Arc<HashMap<u32, ByteBuffer>>,
 }
 
-impl CudaFlatLayout {
+/// A CUDA-optimized terminal layout.
+pub type CudaFlatLayout = Layout<CudaFlat>;
+
+impl CudaFlatData {
     #[inline]
     pub fn segment_id(&self) -> SegmentId {
         self.segment_id
@@ -123,28 +126,15 @@ impl CudaFlatLayout {
 }
 
 impl VTable for CudaFlat {
-    type Layout = CudaFlatLayout;
-    type Encoding = CudaFlatLayoutEncoding;
+    type LayoutData = CudaFlatData;
     type Metadata = ProstMetadata<CudaFlatLayoutMetadata>;
 
-    fn id(_encoding: &Self::Encoding) -> LayoutId {
+    fn id(&self) -> LayoutId {
         static ID: CachedId = CachedId::new("vortex.cuda_flat");
         *ID
     }
 
-    fn encoding(_layout: &Self::Layout) -> LayoutEncodingRef {
-        LayoutEncodingRef::new_ref(CudaFlatLayoutEncoding.as_ref())
-    }
-
-    fn row_count(layout: &Self::Layout) -> u64 {
-        layout.row_count
-    }
-
-    fn dtype(layout: &Self::Layout) -> &DType {
-        &layout.dtype
-    }
-
-    fn metadata(layout: &Self::Layout) -> Self::Metadata {
+    fn metadata(layout: &Layout<Self>) -> Self::Metadata {
         ProstMetadata(CudaFlatLayoutMetadata {
             array_encoding_tree: layout.array_tree.to_vec(),
             host_buffers: layout
@@ -158,24 +148,40 @@ impl VTable for CudaFlat {
         })
     }
 
-    fn segment_ids(layout: &Self::Layout) -> Vec<SegmentId> {
-        vec![layout.segment_id]
-    }
-
-    fn nchildren(_layout: &Self::Layout) -> usize {
-        0
+    fn deserialize(
+        &self,
+        args: &LayoutDeserializeArgs<'_>,
+        metadata: &CudaFlatLayoutMetadata,
+    ) -> VortexResult<Self::LayoutData> {
+        if args.segment_ids.len() != 1 {
+            vortex_bail!("CudaFlatLayout must have exactly one segment ID");
+        }
+        if args.children.nchildren() != 0 {
+            vortex_bail!("CudaFlatLayout must not have children");
+        }
+        let host_buffers = metadata
+            .host_buffers
+            .iter()
+            .map(|hb| (hb.buffer_index, ByteBuffer::from(hb.data.clone())))
+            .collect();
+        Ok(CudaFlatData {
+            segment_id: args.segment_ids[0],
+            ctx: args.array_read_ctx.clone(),
+            array_tree: ByteBuffer::from(metadata.array_encoding_tree.clone()),
+            host_buffers: Arc::new(host_buffers),
+        })
     }
 
-    fn child(_layout: &Self::Layout, _idx: usize) -> VortexResult<LayoutRef> {
-        vortex_bail!("CudaFlatLayout has no children");
+    fn child_dtype(_layout: &Layout<Self>, idx: usize) -> VortexResult<DType> {
+        vortex_bail!("CudaFlatLayout has no child {idx}");
     }
 
-    fn child_type(_layout: &Self::Layout, _idx: usize) -> LayoutChildType {
+    fn child_type(_layout: &Layout<Self>, _idx: usize) -> LayoutChildType {
         vortex_panic!("CudaFlatLayout has no children");
     }
 
     fn new_reader(
-        layout: &Self::Layout,
+        layout: &Layout<Self>,
         name: Arc<str>,
         segment_source: Arc<dyn SegmentSource>,
         session: &VortexSession,
@@ -189,40 +195,6 @@ impl VTable for CudaFlat {
             array: Default::default(),
         }))
     }
-
-    fn build(
-        _encoding: &Self::Encoding,
-        dtype: &DType,
-        row_count: u64,
-        metadata: &<Self::Metadata as DeserializeMetadata>::Output,
-        segment_ids: Vec<SegmentId>,
-        _children: &dyn LayoutChildren,
-        build_ctx: &LayoutBuildContext<'_>,
-    ) -> VortexResult<Self::Layout> {
-        if segment_ids.len() != 1 {
-            vortex_bail!("CudaFlatLayout must have exactly one segment ID");
-        }
-        let host_buffers: HashMap<u32, ByteBuffer> = metadata
-            .host_buffers
-            .iter()
-            .map(|hb| (hb.buffer_index, ByteBuffer::from(hb.data.clone())))
-            .collect();
-        Ok(CudaFlatLayout {
-            row_count,
-            dtype: dtype.clone(),
-            segment_id: segment_ids[0],
-            ctx: build_ctx.array_read_ctx.clone(),
-            array_tree: ByteBuffer::from(metadata.array_encoding_tree.clone()),
-            host_buffers: Arc::new(host_buffers),
-        })
-    }
-
-    fn with_children(_layout: &mut Self::Layout, children: Vec<LayoutRef>) -> VortexResult<()> {
-        if !children.is_empty() {
-            vortex_bail!("CudaFlatLayout has no children, got {}", children.len());
-        }
-        Ok(())
-    }
 }
 
 // Threshold to order filter and apply expression, copied from FlatLayout.
@@ -240,14 +212,14 @@ impl CudaFlatReader {
     fn array_future(&self) -> SharedArrayFuture {
         self.array
             .get_or_init(|| {
-                let row_count = usize::try_from(self.layout.row_count)
+                let row_count = usize::try_from(self.layout.row_count())
                     .vortex_expect("row count must fit in usize");
 
                 let segment_fut = self.segment_source.request(self.layout.segment_id);
 
                 let ctx = self.layout.ctx.clone();
                 let session = self.session.clone();
-                let dtype = self.layout.dtype.clone();
+                let dtype = self.layout.dtype().clone();
                 let array_tree = self.layout.array_tree.clone();
                 let host_buffers = Arc::clone(&self.layout.host_buffers);
 
@@ -275,11 +247,11 @@ impl LayoutReader for CudaFlatReader {
     }
 
     fn dtype(&self) -> &DType {
-        &self.layout.dtype
+        self.layout.dtype()
     }
 
     fn row_count(&self) -> u64 {
-        self.layout.row_count
+        self.layout.row_count()
     }
 
     fn register_splits(
@@ -288,7 +260,7 @@ impl LayoutReader for CudaFlatReader {
         split_range: &SplitRange,
         splits: &mut RowSplits,
     ) -> VortexResult<()> {
-        split_range.check_bounds(self.layout.row_count)?;
+        split_range.check_bounds(self.layout.row_count())?;
         splits.push(split_range.root_row_range().end);
         Ok(())
     }
@@ -523,14 +495,19 @@ impl LayoutStrategy for CudaFlatLayoutStrategy {
             .map(|hb| (hb.buffer_index, ByteBuffer::from(hb.data.clone())))
             .collect();
 
-        Ok(CudaFlatLayout {
+        Ok(LayoutParts::new(
+            CudaFlat,
+            stream.dtype().clone(),
             row_count,
-            dtype: stream.dtype().clone(),
-            segment_id,
-            ctx: ReadContext::new(ctx.to_ids()),
-            array_tree,
-            host_buffers: Arc::new(host_buffer_map),
-        }
+            vec![segment_id],
+            layout_children(Vec::new()),
+            CudaFlatData {
+                segment_id,
+                ctx: ReadContext::new(ctx.to_ids()),
+                array_tree,
+                host_buffers: Arc::new(host_buffer_map),
+            },
+        )
         .into_layout())
     }
 }
@@ -565,5 +542,5 @@ pub fn register_cuda_layout(session: &VortexSession) {
     use vortex::layout::session::LayoutSessionExt;
     session
         .layouts()
-        .register(LayoutEncodingRef::new_ref(CudaFlatLayoutEncoding.as_ref()));
+        .register(LayoutEncodingRef::new_ref(&CudaFlat));
 }

vortex-file/src/tests.rs

@@ -71,7 +71,7 @@ use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_flatbuffers::footer as fb;
 use vortex_io::session::RuntimeSession;
-use vortex_layout::Layout;
+use vortex_layout::DynLayout;
 use vortex_layout::layouts::flat::writer::FlatLayoutStrategy;
 use vortex_layout::layouts::zoned::LegacyStats;
 use vortex_layout::layouts::zoned::Zoned;
@@ -2036,14 +2036,14 @@ async fn timestamp_unit_mismatch_errors_with_constant_children()
 }
 
 /// Collect all segment byte offsets reachable from a layout node.
-fn collect_segment_offsets(layout: &dyn Layout, segment_specs: &[SegmentSpec]) -> Vec<u64> {
+fn collect_segment_offsets(layout: &dyn DynLayout, segment_specs: &[SegmentSpec]) -> Vec<u64> {
     let mut result = Vec::new();
     collect_segment_offsets_inner(layout, segment_specs, &mut result);
     result
 }
 
 fn collect_segment_offsets_inner(
-    layout: &dyn Layout,
+    layout: &dyn DynLayout,
     segment_specs: &[SegmentSpec],
     result: &mut Vec<u64>,
 ) {
@@ -2067,7 +2067,7 @@ fn assert_offsets_ordered(before: &[u64], after: &[u64], context: &str) {
 }
 
 /// Whether any node in the layout tree is a dict layout.
-fn layout_has_dict(layout: &dyn Layout) -> bool {
+fn layout_has_dict(layout: &dyn DynLayout) -> bool {
     layout.encoding_id().as_ref() == "vortex.dict"
         || layout
             .children()
@@ -2237,7 +2237,7 @@ async fn test_segment_ordering_dict_codes_before_values() -> VortexResult<()> {
 
     // Walk the layout tree and find all dict layouts.
     // Verify codes segments come before values segments in byte order within each run.
-    fn check_dict_ordering(layout: &dyn Layout, segment_specs: &[SegmentSpec]) {
+    fn check_dict_ordering(layout: &dyn DynLayout, segment_specs: &[SegmentSpec]) {
         if layout.encoding_id().as_ref() == "vortex.dict" {
             // child 0 = values, child 1 = codes
             let values_offsets =
@@ -2359,7 +2359,7 @@ async fn test_segment_ordering_zonemaps_after_data() -> VortexResult<()> {
     let root = footer.layout();
 
     // Find all zoned layouts and verify data segments come before zone map segments.
-    fn check_zoned_ordering(layout: &dyn Layout, segment_specs: &[SegmentSpec]) {
+    fn check_zoned_ordering(layout: &dyn DynLayout, segment_specs: &[SegmentSpec]) {
         if layout.is::<Zoned>() || layout.is::<LegacyStats>() {
             // child 0 = data, child 1 = zones
             let data_offsets =
@@ -2387,7 +2387,7 @@ async fn test_segment_ordering_zonemaps_after_data() -> VortexResult<()> {
     let mut all_zones_offsets = Vec::new();
 
     fn collect_all_zoned(
-        layout: &dyn Layout,
+        layout: &dyn DynLayout,
         segment_specs: &[SegmentSpec],
         all_data: &mut Vec<u64>,
         all_zones: &mut Vec<u64>,

vortex-layout/src/children.rs

@@ -73,6 +73,11 @@ impl OwnedLayoutChildren {
     }
 }
 
+/// Create an in-memory child adapter from owned layout references.
+pub fn layout_children(children: Vec<LayoutRef>) -> Arc<dyn LayoutChildren> {
+    OwnedLayoutChildren::layout_children(children)
+}
+
 /// In-memory implementation of [`LayoutChildren`].
 impl LayoutChildren for OwnedLayoutChildren {
     fn to_arc(&self) -> Arc<dyn LayoutChildren> {