From a2b696c822f9a4e0691ede3c92372363a469ab33 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Wed, 22 Jul 2026 18:03:25 +0100 Subject: [PATCH 1/8] wip Signed-off-by: Joe Isaacs --- vortex-cuda/src/file.rs | 12 +- vortex-cuda/src/layout.rs | 147 ++-- vortex-file/src/tests.rs | 14 +- vortex-layout/src/children.rs | 5 + vortex-layout/src/display.rs | 3 +- vortex-layout/src/encoding.rs | 133 ++- vortex-layout/src/flatbuffers.rs | 51 +- vortex-layout/src/layout.rs | 883 ++++++++++---------- vortex-layout/src/layouts/chunked/mod.rs | 165 ++-- vortex-layout/src/layouts/chunked/reader.rs | 5 +- vortex-layout/src/layouts/chunked/writer.rs | 1 - vortex-layout/src/layouts/dict/mod.rs | 214 ++--- vortex-layout/src/layouts/dict/reader.rs | 8 +- vortex-layout/src/layouts/dict/writer.rs | 1 - vortex-layout/src/layouts/flat/mod.rs | 175 ++-- vortex-layout/src/layouts/flat/reader.rs | 2 +- vortex-layout/src/layouts/flat/writer.rs | 1 - vortex-layout/src/layouts/foreign/mod.rs | 36 +- vortex-layout/src/layouts/list/mod.rs | 288 +++---- vortex-layout/src/layouts/list/reader.rs | 6 +- vortex-layout/src/layouts/list/writer.rs | 1 - vortex-layout/src/layouts/struct_/mod.rs | 213 ++--- vortex-layout/src/layouts/struct_/reader.rs | 4 +- vortex-layout/src/layouts/struct_/writer.rs | 1 - vortex-layout/src/layouts/zoned/mod.rs | 340 ++++---- vortex-layout/src/layouts/zoned/reader.rs | 5 +- vortex-layout/src/layouts/zoned/writer.rs | 1 - vortex-layout/src/lib.rs | 6 +- vortex-layout/src/session.rs | 32 +- vortex-layout/src/vtable.rs | 182 ---- vortex-tui/src/segment_tree.rs | 6 +- vortex-tui/src/tree.rs | 2 +- vortex-web/crate/src/wasm.rs | 2 +- 33 files changed, 1239 insertions(+), 1706 deletions(-) delete mode 100644 vortex-layout/src/vtable.rs diff --git a/vortex-cuda/src/file.rs b/vortex-cuda/src/file.rs index 9a41f51f9b9..fac2abdff0a 100644 --- a/vortex-cuda/src/file.rs +++ b/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> { +fn control_plane_segments(layout: &dyn DynLayout, segment_count: usize) -> VortexResult> { 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::() || layout.is::() { 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; diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index 6bd81b59f03..c97f7b737cb 100644 --- a/vortex-cuda/src/layout.rs +++ b/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,12 +35,12 @@ 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; @@ -49,13 +48,13 @@ 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,15 +83,16 @@ pub struct CudaFlatLayoutMetadata { pub host_buffers: Vec, } -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, @@ -100,7 +100,10 @@ pub struct CudaFlatLayout { host_buffers: Arc>, } -impl CudaFlatLayout { +/// A CUDA-optimized terminal layout. +pub type CudaFlatLayout = Layout; + +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; - 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::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 { - vec![layout.segment_id] - } - - fn nchildren(_layout: &Self::Layout) -> usize { - 0 + fn deserialize( + &self, + args: &LayoutDeserializeArgs<'_>, + metadata: &CudaFlatLayoutMetadata, + ) -> VortexResult { + 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 { - vortex_bail!("CudaFlatLayout has no children"); + fn child_dtype(_layout: &Layout, idx: usize) -> VortexResult { + vortex_bail!("CudaFlatLayout has no child {idx}"); } - fn child_type(_layout: &Self::Layout, _idx: usize) -> LayoutChildType { + fn child_type(_layout: &Layout, _idx: usize) -> LayoutChildType { vortex_panic!("CudaFlatLayout has no children"); } fn new_reader( - layout: &Self::Layout, + layout: &Layout, name: Arc, segment_source: Arc, session: &VortexSession, @@ -189,40 +195,6 @@ impl VTable for CudaFlat { array: Default::default(), })) } - - fn build( - _encoding: &Self::Encoding, - dtype: &DType, - row_count: u64, - metadata: &::Output, - segment_ids: Vec, - _children: &dyn LayoutChildren, - build_ctx: &LayoutBuildContext<'_>, - ) -> VortexResult { - if segment_ids.len() != 1 { - vortex_bail!("CudaFlatLayout must have exactly one segment ID"); - } - let host_buffers: HashMap = 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) -> 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)); } diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index e665ab18d50..0729318cf89 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -67,7 +67,7 @@ use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; use vortex_error::VortexResult; 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; @@ -1860,14 +1860,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 { +fn collect_segment_offsets(layout: &dyn DynLayout, segment_specs: &[SegmentSpec]) -> Vec { 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, ) { @@ -1891,7 +1891,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() @@ -2061,7 +2061,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 = @@ -2183,7 +2183,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::() || layout.is::() { // child 0 = data, child 1 = zones let data_offsets = @@ -2211,7 +2211,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, all_zones: &mut Vec, diff --git a/vortex-layout/src/children.rs b/vortex-layout/src/children.rs index 4903265890e..fa9b3ae47f8 100644 --- a/vortex-layout/src/children.rs +++ b/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) -> Arc { + OwnedLayoutChildren::layout_children(children) +} + /// In-memory implementation of [`LayoutChildren`]. impl LayoutChildren for OwnedLayoutChildren { fn to_arc(&self) -> Arc { diff --git a/vortex-layout/src/display.rs b/vortex-layout/src/display.rs index 5acf685e946..56f9691c722 100644 --- a/vortex-layout/src/display.rs +++ b/vortex-layout/src/display.rs @@ -127,7 +127,7 @@ impl DisplayLayoutTree { fn make_tree(&self, layout: LayoutRef) -> VortexResult> { // Build the node label with encoding, dtype, and metadata let mut node_parts = vec![ - format!("{}", layout.encoding()), + format!("{}", layout.encoding_id()), format!("dtype: {}", layout.dtype()), ]; @@ -232,7 +232,6 @@ mod tests { use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSessionExt; - use crate::IntoLayout; use crate::OwnedLayoutChildren; use crate::layouts::chunked::ChunkedLayout; use crate::layouts::flat::Flat; diff --git a/vortex-layout/src/encoding.rs b/vortex-layout/src/encoding.rs index b6e10b3b38a..3a32ea82899 100644 --- a/vortex-layout/src/encoding.rs +++ b/vortex-layout/src/encoding.rs @@ -11,37 +11,53 @@ use vortex_array::DeserializeMetadata; use vortex_array::dtype::DType; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_session::registry::Id; +use vortex_session::VortexSession; +use vortex_session::registry::ReadContext; -use crate::IntoLayout; -use crate::LayoutBuildContext; -use crate::LayoutChildren; +use crate::LayoutId; use crate::LayoutRef; use crate::VTable; +use crate::children::LayoutChildren; use crate::segments::SegmentId; -/// A unique identifier for a layout encoding. -pub type LayoutEncodingId = Id; -/// Shared reference to a registered layout encoding. -pub type LayoutEncodingRef = ArcRef; - -/// Object-safe layout encoding registered in a [`LayoutSession`](crate::session::LayoutSession). -/// -/// Encoding instances deserialize serialized layout metadata into concrete [`LayoutRef`] nodes. -/// New in-tree encodings usually implement [`VTable`] and use [`LayoutEncodingAdapter`], while -/// foreign encodings can provide an object-safe implementation directly. -pub trait LayoutEncoding: 'static + Send + Sync + Debug + private::Sealed { - /// Returns this encoding as [`Any`] for downcasting. - fn as_any(&self) -> &dyn Any; +/// Backwards-compatible name for a layout ID. +pub type LayoutEncodingId = LayoutId; +/// Shared reference to a registered layout-vtable plugin. +pub type LayoutVTableRef = ArcRef; +/// Backwards-compatible name for a layout-vtable reference. +pub type LayoutEncodingRef = LayoutVTableRef; + +/// Common fields available while deserializing layout-specific data. +pub struct LayoutDeserializeArgs<'a> { + /// Session used to resolve plugin-owned metadata. + pub session: &'a VortexSession, + /// Array read context referenced by serialized array metadata. + pub array_read_ctx: &'a ReadContext, + /// Logical dtype of this layout. + pub dtype: &'a DType, + /// Number of rows in this layout. + pub row_count: u64, + /// Directly referenced segments. + pub segment_ids: Vec, + /// Lazy child access. + pub children: &'a dyn LayoutChildren, +} - /// Returns the globally unique encoding id. - fn id(&self) -> LayoutEncodingId; +/// Context shared while recursively deserializing layouts. +pub struct LayoutBuildContext<'a> { + /// Session used to resolve plugin-owned metadata. + pub session: &'a VortexSession, + /// Array read context referenced by serialized array metadata. + pub array_read_ctx: &'a ReadContext, +} - /// Build a layout from serialized metadata, segment ids, and children. - /// - /// Implementations must use `build_ctx` for session-scoped plugin resolution instead of global - /// registries. +/// Object-safe plugin registered for a layout ID. +pub trait LayoutVTablePlugin: 'static + Send + Sync + Debug { + /// Returns this plugin as [`Any`]. + fn as_any(&self) -> &dyn Any; + /// Returns the globally unique layout ID. + fn id(&self) -> LayoutEncodingId; + /// Deserializes a layout node. fn build( &self, dtype: &DType, @@ -53,17 +69,16 @@ pub trait LayoutEncoding: 'static + Send + Sync + Debug + private::Sealed { ) -> VortexResult; } -/// Object-safe adapter from a typed layout [`VTable`] to [`LayoutEncoding`]. -#[repr(transparent)] -pub struct LayoutEncodingAdapter(V::Encoding); +/// Backwards-compatible name for the object-safe layout-vtable plugin. +pub use LayoutVTablePlugin as LayoutEncoding; -impl LayoutEncoding for LayoutEncodingAdapter { +impl LayoutVTablePlugin for V { fn as_any(&self) -> &dyn Any { self } fn id(&self) -> LayoutEncodingId { - V::id(&self.0) + VTable::id(self) } fn build( @@ -76,77 +91,47 @@ impl LayoutEncoding for LayoutEncodingAdapter { build_ctx: &LayoutBuildContext<'_>, ) -> VortexResult { let metadata = ::deserialize(metadata)?; - let layout = V::build( - &self.0, + Ok(V::build( + self, dtype, row_count, &metadata, segment_ids, children, build_ctx, - )?; - - // Validate that the builder function returned the expected values. - if layout.row_count() != row_count { - vortex_bail!( - "Layout row count mismatch: {} != {}", - layout.row_count(), - row_count - ); - } - if layout.dtype() != dtype { - vortex_bail!("Layout dtype mismatch: {} != {}", layout.dtype(), dtype); - } - - Ok(layout.into_layout()) - } -} - -impl Debug for LayoutEncodingAdapter { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - f.debug_struct("LayoutEncoding") - .field("id", &self.id()) - .finish() + )? + .into_layout()) } } -impl Display for dyn LayoutEncoding + '_ { +impl Display for dyn LayoutVTablePlugin + '_ { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { write!(f, "{}", self.id()) } } -impl PartialEq for dyn LayoutEncoding + '_ { +impl PartialEq for dyn LayoutVTablePlugin + '_ { fn eq(&self, other: &Self) -> bool { self.id() == other.id() } } -impl Eq for dyn LayoutEncoding + '_ {} +impl Eq for dyn LayoutVTablePlugin + '_ {} -impl dyn LayoutEncoding + '_ { +impl dyn LayoutVTablePlugin + '_ { + /// Returns whether this plugin is vtable `V`. pub fn is(&self) -> bool { self.as_opt::().is_some() } - pub fn as_(&self) -> &V::Encoding { + /// Downcasts this plugin to vtable `V`. + pub fn as_(&self) -> &V { self.as_opt::() - .vortex_expect("LayoutEncoding is not of the expected type") + .vortex_expect("layout encoding type mismatch") } - pub fn as_opt(&self) -> Option<&V::Encoding> { - self.as_any() - .downcast_ref::>() - .map(|e| &e.0) + /// Attempts to downcast this plugin to vtable `V`. + pub fn as_opt(&self) -> Option<&V> { + self.as_any().downcast_ref() } } - -mod private { - use super::*; - use crate::layouts::foreign::ForeignLayoutEncoding; - - pub trait Sealed {} - - impl Sealed for LayoutEncodingAdapter {} - impl Sealed for ForeignLayoutEncoding {} -} diff --git a/vortex-layout/src/flatbuffers.rs b/vortex-layout/src/flatbuffers.rs index 54ce0450df8..8e461d32d1f 100644 --- a/vortex-layout/src/flatbuffers.rs +++ b/vortex-layout/src/flatbuffers.rs @@ -18,7 +18,7 @@ use vortex_flatbuffers::layout; use vortex_session::VortexSession; use vortex_session::registry::ReadContext; -use crate::Layout; +use crate::DynLayout; use crate::LayoutBuildContext; use crate::LayoutContext; use crate::LayoutRef; @@ -149,7 +149,7 @@ fn foreign_layout_from_fb( )) } -impl dyn Layout + '_ { +impl dyn DynLayout + '_ { /// Serialize the layout into a [`FlatBufferBuilder`]. pub fn flatbuffer_writer<'a>( &'a self, @@ -161,7 +161,7 @@ impl dyn Layout + '_ { /// An adapter struct for writing a layout to a FlatBuffer. struct LayoutFlatBufferWriter<'a> { - layout: &'a dyn Layout, + layout: &'a dyn DynLayout, ctx: &'a LayoutContext, } @@ -224,15 +224,58 @@ impl WriteFlatBuffer for LayoutFlatBufferWriter<'_> { #[cfg(test)] mod tests { use flatbuffers::FlatBufferBuilder; + use vortex_array::array_session; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_flatbuffers::WriteFlatBufferExt; use vortex_flatbuffers::layout as fbl; use vortex_session::registry::ReadContext; + use super::layout_from_flatbuffer; use super::layout_from_flatbuffer_with_options; + use crate::LayoutContext; use crate::LayoutEncodingId; + use crate::layouts::flat::Flat; + use crate::layouts::flat::FlatLayout; + use crate::segments::SegmentId; use crate::session::LayoutSession; + #[test] + fn flat_layout_preserves_wire_fields_and_deserializes() { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let array_ctx = ReadContext::new([]); + let layout = FlatLayout::new_with_metadata( + 3, + dtype.clone(), + SegmentId::from(7), + array_ctx.clone(), + Some(vec![1, 2].into()), + ); + let layout_ctx = LayoutContext::default(); + let layout = layout.to_layout(); + let buffer = layout + .flatbuffer_writer(&layout_ctx) + .write_flatbuffer_bytes() + .unwrap(); + + let wire = flatbuffers::root::(&buffer).unwrap(); + assert_eq!(wire.encoding(), 0); + assert_eq!(wire.row_count(), 3); + assert_eq!(wire.metadata().unwrap().bytes(), &[0x0a, 0x02, 1, 2]); + assert_eq!(wire.segments().unwrap().iter().collect::>(), [7]); + assert!(wire.children().is_none()); + + let layout_read_ctx = ReadContext::new(layout_ctx.to_ids()); + let session = array_session().with::(); + let decoded = + layout_from_flatbuffer(buffer, &dtype, &layout_read_ctx, &array_ctx, &session).unwrap(); + let decoded = decoded.as_::(); + assert_eq!(decoded.row_count(), 3); + assert_eq!(*decoded.segment_id(), 7); + assert_eq!(decoded.array_tree().unwrap().as_ref(), &[1, 2]); + } + #[expect(clippy::disallowed_methods, reason = "test-only id")] #[test] fn unknown_layout_encoding_allow_unknown() { @@ -274,7 +317,7 @@ mod tests { LayoutEncodingId::new("vortex.test.foreign_child_layout"), ]); let array_ctx = ReadContext::new([]); - let session = vortex_array::array_session().with::(); + let session = array_session().with::(); let layout = layout_from_flatbuffer_with_options( layout_buffer, diff --git a/vortex-layout/src/layout.rs b/vortex-layout/src/layout.rs index 17b542217d5..df665ecd44e 100644 --- a/vortex-layout/src/layout.rs +++ b/vortex-layout/src/layout.rs @@ -5,6 +5,7 @@ use std::any::Any; use std::fmt::Debug; use std::fmt::Display; use std::fmt::Formatter; +use std::ops::Deref; use std::sync::Arc; use itertools::Itertools; @@ -13,15 +14,14 @@ use vortex_array::dtype::DType; use vortex_array::dtype::FieldName; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_err; use vortex_session::VortexSession; use vortex_session::registry::Id; -use crate::LayoutEncodingId; -use crate::LayoutEncodingRef; +use crate::LayoutBuildContext; +use crate::LayoutDeserializeArgs; use crate::LayoutReaderContext; use crate::LayoutReaderRef; -use crate::VTable; +use crate::children::LayoutChildren; use crate::display::DisplayLayoutTree; use crate::display::display_tree_with_segment_sizes; use crate::segments::SegmentId; @@ -31,66 +31,277 @@ use crate::segments::SegmentSource; pub type LayoutId = Id; /// Shared, erased handle to a layout tree node. -pub type LayoutRef = Arc; +pub type LayoutRef = Arc; -/// Erased layout tree node. +/// Layout-specific behavior for a typed [`Layout`]. /// -/// A layout is immutable metadata describing how row data is organized into child layouts and -/// physical byte segments. It is self-describing only when paired with the session registries used -/// to deserialize its encoding ids and the segment source used to read bytes. -pub trait Layout: 'static + Send + Sync + Debug + private::Sealed { - /// Returns this layout as [`Any`] for downcasting through the `dyn Layout` helpers. - fn as_any(&self) -> &dyn Any; +/// Common serialized fields are stored by [`Layout`]. Implementations own only their +/// layout-specific data, metadata codec, child typing, and reader construction. +pub trait VTable: 'static + Clone + Send + Sync + Debug { + /// Layout-specific data. + type LayoutData: 'static + Clone + Send + Sync + Debug; + /// Serialized metadata type. + type Metadata: SerializeMetadata + vortex_array::DeserializeMetadata + Debug; + + /// Returns the globally unique layout ID. + fn id(&self) -> LayoutId; + + /// Returns the serializable metadata for a layout. + fn metadata(layout: &Layout) -> Self::Metadata; + + /// Deserialize and validate layout-specific data. + fn deserialize( + &self, + args: &LayoutDeserializeArgs<'_>, + metadata: &::Output, + ) -> VortexResult; + + /// Construct a typed layout from deserialized common fields. + fn build( + vtable: &Self, + dtype: &DType, + row_count: u64, + metadata: &::Output, + segment_ids: Vec, + children: &dyn LayoutChildren, + build_ctx: &LayoutBuildContext<'_>, + ) -> VortexResult> { + let args = LayoutDeserializeArgs { + session: build_ctx.session, + array_read_ctx: build_ctx.array_read_ctx, + dtype, + row_count, + segment_ids, + children, + }; + let data = vtable.deserialize(&args, metadata)?; + Ok(LayoutParts::new( + vtable.clone(), + dtype.clone(), + row_count, + args.segment_ids, + children.to_arc(), + data, + ) + .into_typed()) + } - /// Converts an owned layout reference into erased [`Any`] for `Arc` downcasting. - fn as_any_arc(self: Arc) -> Arc; - - /// Clone this layout as an erased [`LayoutRef`]. - fn to_layout(&self) -> LayoutRef; - - /// Returns the [`crate::LayoutEncoding`] for this layout. - fn encoding(&self) -> LayoutEncodingRef; - - /// The number of rows in this layout. - fn row_count(&self) -> u64; - - /// The logical dtype of this layout when evaluated at the root scope. - fn dtype(&self) -> &DType; - - /// The number of children in this layout. - fn nchildren(&self) -> usize; - - /// Get the child at the given index. - /// - /// Child ordering and meaning are encoding-specific and described by - /// [`child_type`](Self::child_type). - fn child(&self, idx: usize) -> VortexResult; - - /// Get the relative row offset of the child at the given index, returning `None` for - /// any auxiliary children, e.g. dictionary values, zone maps, etc. - fn child_type(&self, idx: usize) -> LayoutChildType; - - /// Serialize this layout's encoding metadata. - fn metadata(&self) -> Vec; - - /// Get the segment IDs directly referenced by this layout node. - fn segment_ids(&self) -> Vec; - - /// Construct a new reader for this layout. - /// - /// - `name` — human-readable label for this reader, propagated to child readers - /// (typically by appending a path component) and surfaced in tracing/debug output. - /// - `segment_source` — source of segment bytes for this and any descendant readers - /// constructed from the returned reader; recursive callers should pass the same - /// source through. - /// - `session` — the [`VortexSession`] hosting the encoding/scalar/layout registries - /// and execution context the reader needs at evaluation time. - /// - `ctx` — id-keyed dependency registry threaded through reader construction (see - /// [`LayoutReaderContext`]). Top-level callers (file open, tests) typically pass - /// `&LayoutReaderContext::new()`; recursive callers inside layout implementations - /// must propagate the `ctx` they were handed so ancestor-published values reach - /// descendants. + /// Returns the expected dtype of child `idx`. + fn child_dtype(layout: &Layout, idx: usize) -> VortexResult; + + /// Returns the relationship between child `idx` and its parent. + fn child_type(layout: &Layout, idx: usize) -> LayoutChildType; + + /// Construct a reader for this layout. fn new_reader( + layout: &Layout, + name: Arc, + segment_source: Arc, + session: &VortexSession, + ctx: &LayoutReaderContext, + ) -> VortexResult; +} + +/// Pieces used to construct a typed layout. +pub struct LayoutParts { + vtable: V, + dtype: DType, + row_count: u64, + segment_ids: Vec, + children: Arc, + data: V::LayoutData, +} + +impl LayoutParts { + /// Create layout parts from common fields and vtable-specific data. + pub fn new( + vtable: V, + dtype: DType, + row_count: u64, + segment_ids: Vec, + children: Arc, + data: V::LayoutData, + ) -> Self { + Self { + vtable, + dtype, + row_count, + segment_ids, + children, + data, + } + } + + /// Convert these parts into a typed layout. + pub fn into_typed(self) -> Layout { + Layout::from_parts(self) + } + + /// Erase these parts into a layout reference. + pub fn into_layout(self) -> LayoutRef { + self.into_typed().into_layout() + } +} + +/// A typed layout node. +pub struct Layout { + inner: Arc>, +} + +struct LayoutInner { + vtable: V, + dtype: DType, + row_count: u64, + segment_ids: Vec, + children: Arc, + data: V::LayoutData, +} + +impl Layout { + /// Construct a layout from explicit parts. + pub fn from_parts(parts: LayoutParts) -> Self { + Self { + inner: Arc::new(LayoutInner { + vtable: parts.vtable, + dtype: parts.dtype, + row_count: parts.row_count, + segment_ids: parts.segment_ids, + children: parts.children, + data: parts.data, + }), + } + } + + /// Returns the vtable. + pub fn vtable(&self) -> &V { + &self.inner.vtable + } + + /// Returns layout-specific data. + pub fn data(&self) -> &V::LayoutData { + &self.inner.data + } + + /// Returns the logical dtype. + pub fn dtype(&self) -> &DType { + &self.inner.dtype + } + + /// Returns the number of rows. + pub fn row_count(&self) -> u64 { + self.inner.row_count + } + + /// Returns directly referenced segment IDs. + pub fn segment_ids(&self) -> &[SegmentId] { + &self.inner.segment_ids + } + + /// Returns the child adapter. + pub fn children(&self) -> &Arc { + &self.inner.children + } + + /// Returns the number of children. + pub fn nchildren(&self) -> usize { + self.inner.children.nchildren() + } + + /// Materialize child `idx` with its expected dtype. + pub fn child(&self, idx: usize) -> VortexResult { + self.inner.children.child(idx, &V::child_dtype(self, idx)?) + } + + /// Returns a child's serialized row count without materializing it. + pub fn child_row_count(&self, idx: usize) -> u64 { + self.inner.children.child_row_count(idx) + } + + /// Returns a child's relationship to this layout. + pub fn child_type(&self, idx: usize) -> LayoutChildType { + V::child_type(self, idx) + } + + /// Erase this typed layout into a shared layout reference. + pub fn to_layout(&self) -> LayoutRef { + self.clone().into_layout() + } + + /// Erase this typed layout into a shared layout reference. + pub fn into_layout(self) -> LayoutRef { + Arc::new(self) + } + + /// Construct a reader for this layout. + pub fn new_reader( + &self, + name: Arc, + segment_source: Arc, + session: &VortexSession, + ctx: &LayoutReaderContext, + ) -> VortexResult { + V::new_reader(self, name, segment_source, session, ctx) + } +} + +impl Clone for Layout { + fn clone(&self) -> Self { + Self { + inner: Arc::clone(&self.inner), + } + } +} + +impl Debug for Layout { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Layout") + .field("encoding_id", &self.vtable().id()) + .field("dtype", &self.inner.dtype) + .field("row_count", &self.inner.row_count) + .field("segment_ids", &self.inner.segment_ids) + .field("data", &self.inner.data) + .finish() + } +} + +impl Deref for Layout { + type Target = V::LayoutData; + + fn deref(&self) -> &Self::Target { + self.data() + } +} + +impl From> for LayoutRef { + fn from(value: Layout) -> Self { + value.into_layout() + } +} + +/// Erased layout behavior used by [`LayoutRef`]. +pub trait DynLayout: 'static + Send + Sync + Debug { + /// Returns this layout as [`Any`] for downcasting. + fn as_any(&self) -> &dyn Any; + /// Clone this layout as an erased reference. + fn dyn_to_layout(&self) -> LayoutRef; + /// Returns the layout ID. + fn dyn_encoding_id(&self) -> LayoutId; + /// Returns the row count. + fn dyn_row_count(&self) -> u64; + /// Returns the logical dtype. + fn dyn_dtype(&self) -> &DType; + /// Returns the number of children. + fn dyn_nchildren(&self) -> usize; + /// Returns child `idx`. + fn dyn_child(&self, idx: usize) -> VortexResult; + /// Returns the relationship of child `idx`. + fn dyn_child_type(&self, idx: usize) -> LayoutChildType; + /// Serializes layout-specific metadata. + fn dyn_metadata(&self) -> Vec; + /// Returns directly referenced segment IDs. + fn dyn_segment_ids(&self) -> Vec; + /// Constructs a reader. + fn dyn_new_reader( &self, name: Arc, segment_source: Arc, @@ -99,111 +310,187 @@ pub trait Layout: 'static + Send + Sync + Debug + private::Sealed { ) -> VortexResult; } -/// Conversion into an erased [`LayoutRef`]. -pub trait IntoLayout { - /// Converts this type into a [`LayoutRef`]. - fn into_layout(self) -> LayoutRef; +impl DynLayout for Layout { + fn as_any(&self) -> &dyn Any { + self + } + + fn dyn_to_layout(&self) -> LayoutRef { + Layout::to_layout(self) + } + + fn dyn_encoding_id(&self) -> LayoutId { + self.vtable().id() + } + + fn dyn_row_count(&self) -> u64 { + Layout::row_count(self) + } + + fn dyn_dtype(&self) -> &DType { + Layout::dtype(self) + } + + fn dyn_nchildren(&self) -> usize { + Layout::nchildren(self) + } + + fn dyn_child(&self, idx: usize) -> VortexResult { + Layout::child(self, idx) + } + + fn dyn_child_type(&self, idx: usize) -> LayoutChildType { + Layout::child_type(self, idx) + } + + fn dyn_metadata(&self) -> Vec { + V::metadata(self).serialize() + } + + fn dyn_segment_ids(&self) -> Vec { + self.inner.segment_ids.clone() + } + + fn dyn_new_reader( + &self, + name: Arc, + segment_source: Arc, + session: &VortexSession, + ctx: &LayoutReaderContext, + ) -> VortexResult { + Layout::new_reader(self, name, segment_source, session, ctx) + } } -/// A type that allows us to identify how a layout child relates to its parent. +/// Identifies how a layout child relates to its parent. #[derive(Debug, Clone, PartialEq, Eq)] pub enum LayoutChildType { - /// A layout child that retains the same schema and row offset position in the dataset. + /// A child retaining the parent's schema and row offset. Transparent(Arc), - /// A layout child that provides auxiliary data, e.g. dictionary values, zone maps, etc. - /// Contains a human-readable name of the child. + /// Auxiliary data, such as dictionary values or zone maps. Auxiliary(Arc), - /// A layout child that represents a row-based chunk of data. - /// Contains the chunk index and relative row offset of the child. + /// A row-based chunk with its relative row offset. Chunk((usize, u64)), - /// A layout child that represents a single field of data. - /// Contains the field name of the child. + /// A single field of a struct. Field(FieldName), } impl LayoutChildType { - /// Returns the name of this child. + /// Returns the child name. pub fn name(&self) -> Arc { match self { - LayoutChildType::Chunk((idx, _offset)) => format!("[{idx}]").into(), - LayoutChildType::Auxiliary(name) => Arc::clone(name), - LayoutChildType::Transparent(name) => Arc::clone(name), - LayoutChildType::Field(name) => name.clone().into(), + Self::Chunk((idx, _)) => format!("[{idx}]").into(), + Self::Auxiliary(name) | Self::Transparent(name) => Arc::clone(name), + Self::Field(name) => name.clone().into(), } } - /// Returns the relative row offset of this child. - /// For auxiliary children, this is `None`. + /// Returns the relative row offset, or `None` for auxiliary children. pub fn row_offset(&self) -> Option { match self { - LayoutChildType::Chunk((_idx, offset)) => Some(*offset), - LayoutChildType::Auxiliary(_) => None, - LayoutChildType::Transparent(_) => Some(0), - LayoutChildType::Field(_) => Some(0), + Self::Chunk((_, offset)) => Some(*offset), + Self::Auxiliary(_) => None, + Self::Transparent(_) | Self::Field(_) => Some(0), } } } -impl dyn Layout + '_ { - /// The ID of the encoding for this layout. - pub fn encoding_id(&self) -> LayoutEncodingId { - self.encoding().id() +impl dyn DynLayout + '_ { + /// Returns a cloned erased layout reference. + pub fn to_layout(&self) -> LayoutRef { + self.dyn_to_layout() + } + + /// Returns the layout ID. + pub fn encoding_id(&self) -> LayoutId { + self.dyn_encoding_id() + } + + /// Returns the logical dtype. + pub fn dtype(&self) -> &DType { + self.dyn_dtype() + } + + /// Returns the number of rows. + pub fn row_count(&self) -> u64 { + self.dyn_row_count() + } + + /// Returns the number of children. + pub fn nchildren(&self) -> usize { + self.dyn_nchildren() + } + + /// Returns child `idx`. + pub fn child(&self, idx: usize) -> VortexResult { + self.dyn_child(idx) + } + + /// Returns the relationship of child `idx`. + pub fn child_type(&self, idx: usize) -> LayoutChildType { + self.dyn_child_type(idx) + } + + /// Returns serialized layout-specific metadata. + pub fn metadata(&self) -> Vec { + self.dyn_metadata() + } + + /// Returns directly referenced segment IDs. + pub fn segment_ids(&self) -> Vec { + self.dyn_segment_ids() + } + + /// Constructs a reader for this layout. + pub fn new_reader( + &self, + name: Arc, + segment_source: Arc, + session: &VortexSession, + ctx: &LayoutReaderContext, + ) -> VortexResult { + self.dyn_new_reader(name, segment_source, session, ctx) } - /// The children of this layout. + /// Returns all children. pub fn children(&self) -> VortexResult> { - (0..self.nchildren()).map(|i| self.child(i)).try_collect() + (0..self.nchildren()) + .map(|idx| self.child(idx)) + .try_collect() } - /// The child types of this layout. - pub fn child_types(&self) -> impl Iterator { - (0..self.nchildren()).map(|i| self.child_type(i)) + /// Returns all child types. + pub fn child_types(&self) -> impl Iterator + '_ { + (0..self.nchildren()).map(|idx| self.child_type(idx)) } - /// The names of the children of this layout. - pub fn child_names(&self) -> impl Iterator> { + /// Returns all child names. + pub fn child_names(&self) -> impl Iterator> + '_ { self.child_types().map(|child| child.name()) } - /// The row offsets of the children of this layout, where `None` indicates an auxiliary child. - pub fn child_row_offsets(&self) -> impl Iterator> { + /// Returns all child row offsets. + pub fn child_row_offsets(&self) -> impl Iterator> + '_ { self.child_types().map(|child| child.row_offset()) } + /// Returns whether this layout uses vtable `V`. pub fn is(&self) -> bool { self.as_opt::().is_some() } - /// Downcast a layout to a specific type. - pub fn as_(&self) -> &V::Layout { + /// Downcasts this layout to vtable `V`. + pub fn as_(&self) -> &Layout { self.as_opt::().vortex_expect("Failed to downcast") } - /// Downcast a layout to a specific type. - pub fn as_opt(&self) -> Option<&V::Layout> { - self.as_any() - .downcast_ref::>() - .map(|adapter| &adapter.0) + /// Attempts to downcast this layout to vtable `V`. + pub fn as_opt(&self) -> Option<&Layout> { + self.as_any().downcast_ref() } - /// Downcast a layout to a specific type. - pub fn into(self: Arc) -> Arc { - let layout_adapter = self - .as_any_arc() - .downcast::>() - .map_err(|_| vortex_err!("Invalid layout type")) - .vortex_expect("Invalid layout type"); - - // SAFETY: LayoutAdapter is #[repr(transparent)] (see line 192) which guarantees - // it has the same memory layout as V::Layout. The downcast above ensures we have - // the correct type. This transmute is safe because both Arc types point to data - // with identical layout and alignment. - unsafe { std::mem::transmute::>, Arc>(layout_adapter) } - } - - /// Depth-first traversal of the layout and its children. + /// Returns a depth-first pre-order traversal. pub fn depth_first_traversal(&self) -> impl Iterator> { - /// A depth-first pre-order iterator over a layout. struct ChildrenIterator { stack: Vec, } @@ -216,9 +503,7 @@ impl dyn Layout + '_ { let Ok(children) = next.children() else { return Some(Ok(next)); }; - for child in children.into_iter().rev() { - self.stack.push(child); - } + self.stack.extend(children.into_iter().rev()); Some(Ok(next)) } } @@ -228,22 +513,17 @@ impl dyn Layout + '_ { } } - /// Display the layout as a tree structure. + /// Displays the layout as a tree. pub fn display_tree(&self) -> DisplayLayoutTree { DisplayLayoutTree::new(self.to_layout(), false) } - /// Display the layout as a tree structure with optional verbose metadata. + /// Displays the layout as a tree with optional verbose metadata. pub fn display_tree_verbose(&self, verbose: bool) -> DisplayLayoutTree { DisplayLayoutTree::new(self.to_layout(), verbose) } - /// Display the layout as a tree structure, fetching segment buffer sizes from the segment source. - /// - /// # Warning - /// - /// This function performs IO to fetch each segment's buffer. For layouts with - /// many segments, this may result in significant IO overhead. + /// Displays the tree after fetching segment sizes. pub async fn display_tree_with_segments( &self, segment_source: Arc, @@ -252,11 +532,10 @@ impl dyn Layout + '_ { } } -/// Display the encoding, dtype, row count, and segment IDs of this layout. -impl Display for dyn Layout + '_ { +impl Display for dyn DynLayout + '_ { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - let segment_ids = self.segment_ids(); - if segment_ids.is_empty() { + let segments = self.segment_ids(); + if segments.is_empty() { write!( f, "{}({}, rows={})", @@ -271,351 +550,29 @@ impl Display for dyn Layout + '_ { self.encoding_id(), self.dtype(), self.row_count(), - segment_ids - .iter() - .map(|s| format!("{}", **s)) - .collect::>() - .join(", ") + segments.iter().map(|s| format!("{}", **s)).join(", ") ) } } } -#[repr(transparent)] -pub struct LayoutAdapter(V::Layout); - -impl Debug for LayoutAdapter { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - self.0.fmt(f) - } -} - -impl Layout for LayoutAdapter { - fn as_any(&self) -> &dyn Any { - self - } - - fn as_any_arc(self: Arc) -> Arc { - self - } - - fn to_layout(&self) -> LayoutRef { - Arc::new(LayoutAdapter::(self.0.clone())) - } - - fn encoding(&self) -> LayoutEncodingRef { - V::encoding(&self.0) - } - - fn row_count(&self) -> u64 { - V::row_count(&self.0) - } - - fn dtype(&self) -> &DType { - V::dtype(&self.0) - } - - fn nchildren(&self) -> usize { - V::nchildren(&self.0) - } - - fn child(&self, idx: usize) -> VortexResult { - V::child(&self.0, idx) - } - - fn child_type(&self, idx: usize) -> LayoutChildType { - V::child_type(&self.0, idx) - } - - fn metadata(&self) -> Vec { - V::metadata(&self.0).serialize() - } - - fn segment_ids(&self) -> Vec { - V::segment_ids(&self.0) - } - - fn new_reader( - &self, - name: Arc, - segment_source: Arc, - session: &VortexSession, - ctx: &LayoutReaderContext, - ) -> VortexResult { - V::new_reader(&self.0, name, segment_source, session, ctx) - } -} - -mod private { - use super::*; - use crate::layouts::foreign::ForeignLayout; - - pub trait Sealed {} - - impl Sealed for LayoutAdapter {} - impl Sealed for ForeignLayout {} -} - #[cfg(test)] mod tests { - use rstest::rstest; - use vortex_session::registry::ReadContext; - use super::*; #[test] - fn test_layout_child_type_name() { - // Test Chunk variant + fn layout_child_type_names_and_offsets() { let chunk = LayoutChildType::Chunk((5, 100)); assert_eq!(chunk.name().as_ref(), "[5]"); + assert_eq!(chunk.row_offset(), Some(100)); - // Test Field variant let field = LayoutChildType::Field(FieldName::from("customer_id")); assert_eq!(field.name().as_ref(), "customer_id"); - - // Test Auxiliary variant - let aux = LayoutChildType::Auxiliary(Arc::from("zone_map")); - assert_eq!(aux.name().as_ref(), "zone_map"); - - // Test Transparent variant - let transparent = LayoutChildType::Transparent(Arc::from("compressed")); - assert_eq!(transparent.name().as_ref(), "compressed"); - } - - #[test] - fn test_layout_child_type_row_offset() { - // Chunk should return the offset - let chunk = LayoutChildType::Chunk((0, 42)); - assert_eq!(chunk.row_offset(), Some(42)); - - // Field should return 0 - let field = LayoutChildType::Field(FieldName::from("field1")); assert_eq!(field.row_offset(), Some(0)); - // Auxiliary should return None - let aux = LayoutChildType::Auxiliary(Arc::from("metadata")); - assert_eq!(aux.row_offset(), None); - - // Transparent should return 0 - let transparent = LayoutChildType::Transparent(Arc::from("wrapper")); + let auxiliary = LayoutChildType::Auxiliary("zone_map".into()); + assert_eq!(auxiliary.row_offset(), None); + let transparent = LayoutChildType::Transparent("compressed".into()); assert_eq!(transparent.row_offset(), Some(0)); } - - #[test] - fn test_layout_child_type_equality() { - // Test Chunk equality - let chunk1 = LayoutChildType::Chunk((1, 100)); - let chunk2 = LayoutChildType::Chunk((1, 100)); - let chunk3 = LayoutChildType::Chunk((2, 100)); - let chunk4 = LayoutChildType::Chunk((1, 200)); - - assert_eq!(chunk1, chunk2); - assert_ne!(chunk1, chunk3); - assert_ne!(chunk1, chunk4); - - // Test Field equality - let field1 = LayoutChildType::Field(FieldName::from("name")); - let field2 = LayoutChildType::Field(FieldName::from("name")); - let field3 = LayoutChildType::Field(FieldName::from("age")); - - assert_eq!(field1, field2); - assert_ne!(field1, field3); - - // Test Auxiliary equality - let aux1 = LayoutChildType::Auxiliary(Arc::from("stats")); - let aux2 = LayoutChildType::Auxiliary(Arc::from("stats")); - let aux3 = LayoutChildType::Auxiliary(Arc::from("index")); - - assert_eq!(aux1, aux2); - assert_ne!(aux1, aux3); - - // Test Transparent equality - let trans1 = LayoutChildType::Transparent(Arc::from("enc")); - let trans2 = LayoutChildType::Transparent(Arc::from("enc")); - let trans3 = LayoutChildType::Transparent(Arc::from("dec")); - - assert_eq!(trans1, trans2); - assert_ne!(trans1, trans3); - - // Test cross-variant inequality - assert_ne!(chunk1, field1); - assert_ne!(field1, aux1); - assert_ne!(aux1, trans1); - } - - #[rstest] - #[case(LayoutChildType::Chunk((0, 0)), "[0]", Some(0))] - #[case(LayoutChildType::Chunk((999, 1000000)), "[999]", Some(1000000))] - #[case(LayoutChildType::Field(FieldName::from("")), "", Some(0))] - #[case( - LayoutChildType::Field(FieldName::from("very_long_field_name_that_is_quite_lengthy")), - "very_long_field_name_that_is_quite_lengthy", - Some(0) - )] - #[case(LayoutChildType::Auxiliary(Arc::from("aux")), "aux", None)] - #[case(LayoutChildType::Transparent(Arc::from("t")), "t", Some(0))] - fn test_layout_child_type_parameterized( - #[case] child_type: LayoutChildType, - #[case] expected_name: &str, - #[case] expected_offset: Option, - ) { - assert_eq!(child_type.name().as_ref(), expected_name); - assert_eq!(child_type.row_offset(), expected_offset); - } - - #[test] - fn test_chunk_with_different_indices_and_offsets() { - let chunks = [ - LayoutChildType::Chunk((0, 0)), - LayoutChildType::Chunk((1, 100)), - LayoutChildType::Chunk((2, 200)), - LayoutChildType::Chunk((100, 10000)), - ]; - - for chunk in chunks.iter() { - let name = chunk.name(); - assert!(name.starts_with('[')); - assert!(name.ends_with(']')); - - if let LayoutChildType::Chunk((idx, offset)) = chunk { - assert_eq!(name.as_ref(), format!("[{}]", idx)); - assert_eq!(chunk.row_offset(), Some(*offset)); - } - } - } - - #[test] - fn test_field_names_with_special_characters() { - let special_fields: Vec> = vec![ - Arc::from("field-with-dashes"), - Arc::from("field_with_underscores"), - Arc::from("field.with.dots"), - Arc::from("field::with::colons"), - Arc::from("field/with/slashes"), - Arc::from("field@with#symbols"), - ]; - - for field_name in special_fields { - let field = LayoutChildType::Field(Arc::clone(&field_name).into()); - assert_eq!(field.name(), field_name); - assert_eq!(field.row_offset(), Some(0)); - } - } - - #[test] - fn test_struct_layout_display() { - use vortex_array::dtype::Nullability::NonNullable; - use vortex_array::dtype::PType; - use vortex_array::dtype::StructFields; - - use crate::IntoLayout; - use crate::layouts::chunked::ChunkedLayout; - use crate::layouts::dict::DictLayout; - use crate::layouts::flat::FlatLayout; - use crate::layouts::struct_::StructLayout; - use crate::segments::SegmentId; - - let ctx = ReadContext::new([]); - - // Create a flat layout for dict values (utf8 strings) - let dict_values = - FlatLayout::new(3, DType::Utf8(NonNullable), SegmentId::from(0), ctx.clone()) - .into_layout(); - - // Test flat layout display shows segment - assert_eq!( - format!("{}", dict_values), - "vortex.flat(utf8, rows=3, segments=[0])" - ); - - // Create a flat layout for dict codes - let dict_codes = FlatLayout::new( - 10, - DType::Primitive(PType::U16, NonNullable), - SegmentId::from(1), - ctx.clone(), - ) - .into_layout(); - - // Test flat layout display shows segment - assert_eq!( - format!("{}", dict_codes), - "vortex.flat(u16, rows=10, segments=[1])" - ); - - // Create dict layout (column "name") - let dict_layout = - DictLayout::new(Arc::clone(&dict_values), Arc::clone(&dict_codes)).into_layout(); - - // Test dict layout display (no direct segments) - assert_eq!(format!("{}", dict_layout), "vortex.dict(utf8, rows=10)"); - - // Create flat layouts for chunks - let chunk1 = FlatLayout::new( - 5, - DType::Primitive(PType::I64, NonNullable), - SegmentId::from(2), - ctx.clone(), - ) - .into_layout(); - - let chunk2 = FlatLayout::new( - 5, - DType::Primitive(PType::I64, NonNullable), - SegmentId::from(3), - ctx, - ) - .into_layout(); - - // Create chunked layout (column "value") - let chunked_layout = ChunkedLayout::new( - 10, - DType::Primitive(PType::I64, NonNullable), - crate::OwnedLayoutChildren::layout_children(vec![ - Arc::clone(&chunk1), - Arc::clone(&chunk2), - ]), - ) - .into_layout(); - - // Test chunked layout display (no direct segments) - assert_eq!( - format!("{}", chunked_layout), - "vortex.chunked(i64, rows=10)" - ); - - // Test chunk displays show segments - assert_eq!( - format!("{}", chunk1), - "vortex.flat(i64, rows=5, segments=[2])" - ); - assert_eq!( - format!("{}", chunk2), - "vortex.flat(i64, rows=5, segments=[3])" - ); - - // Create struct layout with two fields - let field_names: Vec> = vec!["name".into(), "value".into()]; - let struct_dtype = DType::Struct( - StructFields::new( - field_names.into(), - vec![ - DType::Utf8(NonNullable), - DType::Primitive(PType::I64, NonNullable), - ], - ), - NonNullable, - ); - - let struct_layout = - StructLayout::new(10, struct_dtype, vec![dict_layout, chunked_layout]).into_layout(); - - println!("{}", struct_layout.display_tree_verbose(true)); - - // Test Display impl for struct (no direct segments) - assert_eq!( - format!("{}", struct_layout), - "vortex.struct({name=utf8, value=i64}, rows=10)" - ); - } } diff --git a/vortex-layout/src/layouts/chunked/mod.rs b/vortex-layout/src/layouts/chunked/mod.rs index 336de3235b1..459429046c2 100644 --- a/vortex-layout/src/layouts/chunked/mod.rs +++ b/vortex-layout/src/layouts/chunked/mod.rs @@ -6,74 +6,80 @@ pub mod writer; use std::sync::Arc; -use vortex_array::DeserializeMetadata; use vortex_array::EmptyMetadata; use vortex_array::dtype::DType; +use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::LayoutBuildContext; +use crate::Layout; use crate::LayoutChildType; -use crate::LayoutEncodingRef; +use crate::LayoutChildren; +use crate::LayoutDeserializeArgs; use crate::LayoutId; +use crate::LayoutParts; use crate::LayoutReaderContext; use crate::LayoutReaderRef; use crate::LayoutRef; use crate::VTable; -use crate::children::LayoutChildren; use crate::children::OwnedLayoutChildren; use crate::layouts::chunked::reader::ChunkedReader; -use crate::segments::SegmentId; use crate::segments::SegmentSource; -use crate::vtable; -vtable!(Chunked); +/// Chunked layout vtable. +#[derive(Clone, Debug)] +pub struct Chunked; + +/// Backwards-compatible name for the chunked layout plugin. +pub use Chunked as ChunkedLayoutEncoding; + +/// Chunked-layout-specific data. +#[derive(Clone, Debug)] +pub struct ChunkedData { + chunk_offsets: Vec, +} + +/// A layout partitioned into independently readable row chunks. +pub type ChunkedLayout = Layout; impl VTable for Chunked { - type Layout = ChunkedLayout; - type Encoding = ChunkedLayoutEncoding; + type LayoutData = ChunkedData; type Metadata = EmptyMetadata; - fn id(_encoding: &Self::Encoding) -> LayoutId { + fn id(&self) -> LayoutId { static ID: CachedId = CachedId::new("vortex.chunked"); *ID } - fn encoding(_layout: &Self::Layout) -> LayoutEncodingRef { - LayoutEncodingRef::new_ref(ChunkedLayoutEncoding.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::Metadata { EmptyMetadata } - fn segment_ids(_layout: &Self::Layout) -> Vec { - vec![] - } - - fn nchildren(layout: &Self::Layout) -> usize { - layout.children.nchildren() + fn deserialize( + &self, + args: &LayoutDeserializeArgs<'_>, + _metadata: &EmptyMetadata, + ) -> VortexResult { + let chunk_offsets = chunk_offsets(args.children)?; + if chunk_offsets.last().copied() != Some(args.row_count) { + vortex_bail!("Chunked child row counts do not add up to parent row count"); + } + Ok(ChunkedData { chunk_offsets }) } - fn child(layout: &Self::Layout, idx: usize) -> VortexResult { - layout.children.child(idx, Self::dtype(layout)) + fn child_dtype(layout: &Layout, _idx: usize) -> VortexResult { + Ok(layout.dtype().clone()) } - fn child_type(layout: &Self::Layout, idx: usize) -> LayoutChildType { + fn child_type(layout: &Layout, idx: usize) -> LayoutChildType { LayoutChildType::Chunk((idx, layout.chunk_offsets[idx])) } fn new_reader( - layout: &Self::Layout, + layout: &Layout, name: Arc, segment_source: Arc, session: &VortexSession, @@ -87,74 +93,49 @@ impl VTable for Chunked { ctx.clone(), ))) } - - fn build( - _encoding: &Self::Encoding, - dtype: &DType, - row_count: u64, - _metadata: &::Output, - _segment_ids: Vec, - children: &dyn LayoutChildren, - _build_ctx: &LayoutBuildContext<'_>, - ) -> VortexResult { - Ok(ChunkedLayout::new( - row_count, - dtype.clone(), - children.to_arc(), - )) - } - - fn with_children(layout: &mut Self::Layout, children: Vec) -> VortexResult<()> { - let new_children = OwnedLayoutChildren::layout_children(children); - - // Recalculate chunk offsets based on new children - let mut chunk_offsets = vec![0; new_children.nchildren() + 1]; - for i in 0..new_children.nchildren() { - chunk_offsets[i + 1] = chunk_offsets[i] + new_children.child_row_count(i); - } - - layout.children = new_children; - layout.chunk_offsets = chunk_offsets; - Ok(()) - } -} - -#[derive(Debug)] -pub struct ChunkedLayoutEncoding; - -/// Partitions a column into row-based chunks so that each chunk can be read independently. -/// -/// Used to break large columns into smaller pieces for parallel I/O and to limit memory -/// usage when scanning. -#[derive(Clone, Debug)] -pub struct ChunkedLayout { - row_count: u64, - dtype: DType, - children: Arc, - chunk_offsets: Vec, } -impl ChunkedLayout { +impl Layout { + /// Construct a chunked layout. pub fn new(row_count: u64, dtype: DType, children: Arc) -> Self { - let mut chunk_offsets = vec![0; children.nchildren() + 1]; - for i in 0..children.nchildren() { - chunk_offsets[i + 1] = chunk_offsets[i] + children.child_row_count(i); - } - + let offsets = chunk_offsets(children.as_ref()).vortex_expect("chunk row counts overflow"); assert_eq!( - chunk_offsets[children.nchildren()], - row_count, + offsets.last().copied(), + Some(row_count), "Row count mismatch" ); - Self { - row_count, + LayoutParts::new( + Chunked, dtype, + row_count, + Vec::new(), children, - chunk_offsets, - } + ChunkedData { + chunk_offsets: offsets, + }, + ) + .into_typed() } - pub fn children(&self) -> &Arc { - &self.children + /// Rebuild this layout with owned children. + pub fn with_children(&self, children: Vec) -> Self { + Self::new( + self.row_count(), + self.dtype().clone(), + OwnedLayoutChildren::layout_children(children), + ) + } +} + +fn chunk_offsets(children: &dyn LayoutChildren) -> VortexResult> { + let mut offsets = Vec::with_capacity(children.nchildren() + 1); + offsets.push(0u64); + for idx in 0..children.nchildren() { + offsets.push( + offsets[idx] + .checked_add(children.child_row_count(idx)) + .ok_or_else(|| vortex_err!("Chunked child row counts overflow"))?, + ); } + Ok(offsets) } diff --git a/vortex-layout/src/layouts/chunked/reader.rs b/vortex-layout/src/layouts/chunked/reader.rs index a9870229792..171bb138c6c 100644 --- a/vortex-layout/src/layouts/chunked/reader.rs +++ b/vortex-layout/src/layouts/chunked/reader.rs @@ -53,7 +53,7 @@ impl ChunkedReader { ctx: LayoutReaderContext, ) -> Self { let nchildren = layout.nchildren(); - let dtypes = vec![layout.dtype.clone(); nchildren]; + let dtypes = vec![layout.dtype().clone(); nchildren]; // format!() has non-marginal overhead for short queries like random // access benchmarks @@ -66,7 +66,7 @@ impl ChunkedReader { }; let lazy_children = LazyReaderChildren::new( - Arc::clone(&layout.children), + Arc::clone(layout.children()), dtypes, names, segment_source, @@ -365,7 +365,6 @@ mod test { use vortex_io::session::RuntimeSessionExt; use vortex_session::registry::ReadContext; - use crate::IntoLayout; use crate::LayoutRef; use crate::LayoutStrategy; use crate::OwnedLayoutChildren; diff --git a/vortex-layout/src/layouts/chunked/writer.rs b/vortex-layout/src/layouts/chunked/writer.rs index 29596fadff0..8cc4536588d 100644 --- a/vortex-layout/src/layouts/chunked/writer.rs +++ b/vortex-layout/src/layouts/chunked/writer.rs @@ -14,7 +14,6 @@ use vortex_error::VortexResult; use vortex_io::session::RuntimeSessionExt; use vortex_session::VortexSession; -use crate::IntoLayout; use crate::LayoutRef; use crate::LayoutStrategy; use crate::children::OwnedLayoutChildren; diff --git a/vortex-layout/src/layouts/dict/mod.rs b/vortex-layout/src/layouts/dict/mod.rs index 7e9125428c6..84771304690 100644 --- a/vortex-layout/src/layouts/dict/mod.rs +++ b/vortex-layout/src/layouts/dict/mod.rs @@ -7,7 +7,6 @@ pub mod writer; use std::sync::Arc; use reader::DictReader; -use vortex_array::DeserializeMetadata; use vortex_array::ProstMetadata; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -16,85 +15,105 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; -use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::LayoutBuildContext; +use crate::Layout; use crate::LayoutChildType; -use crate::LayoutEncodingRef; +use crate::LayoutDeserializeArgs; use crate::LayoutId; +use crate::LayoutParts; +use crate::LayoutReaderContext; use crate::LayoutReaderRef; use crate::LayoutRef; use crate::VTable; -use crate::children::LayoutChildren; -use crate::segments::SegmentId; +use crate::children::OwnedLayoutChildren; use crate::segments::SegmentSource; -use crate::vtable; -vtable!(Dict); +/// Dictionary layout vtable. +#[derive(Clone, Debug)] +pub struct Dict; + +/// Backwards-compatible name for the dictionary layout plugin. +pub use Dict as DictLayoutEncoding; + +/// Dictionary-layout-specific data. +#[derive(Clone, Debug)] +pub struct DictData { + codes_dtype: DType, + all_values_referenced: bool, +} + +/// A dictionary values child paired with a codes child. +pub type DictLayout = Layout; impl VTable for Dict { - type Layout = DictLayout; - type Encoding = DictLayoutEncoding; + type LayoutData = DictData; type Metadata = ProstMetadata; - fn id(_encoding: &Self::Encoding) -> LayoutId { + fn id(&self) -> LayoutId { static ID: CachedId = CachedId::new("vortex.dict"); *ID } - fn encoding(_layout: &Self::Layout) -> LayoutEncodingRef { - LayoutEncodingRef::new_ref(DictLayoutEncoding.as_ref()) - } - - fn row_count(layout: &Self::Layout) -> u64 { - layout.codes.row_count() - } - - fn dtype(layout: &Self::Layout) -> &DType { - layout.values.dtype() - } - - fn metadata(layout: &Self::Layout) -> Self::Metadata { - let mut metadata = - DictLayoutMetadata::new(PType::try_from(layout.codes.dtype()).vortex_expect("ptype")); - metadata.is_nullable_codes = Some(layout.codes.dtype().is_nullable()); + fn metadata(layout: &Layout) -> Self::Metadata { + let mut metadata = DictLayoutMetadata::new( + PType::try_from(&layout.codes_dtype).vortex_expect("codes ptype"), + ); + metadata.is_nullable_codes = Some(layout.codes_dtype.is_nullable()); metadata.all_values_referenced = Some(layout.all_values_referenced); ProstMetadata(metadata) } - fn segment_ids(_layout: &Self::Layout) -> Vec { - vec![] - } - - fn nchildren(_layout: &Self::Layout) -> usize { - 2 + fn deserialize( + &self, + args: &LayoutDeserializeArgs<'_>, + metadata: &DictLayoutMetadata, + ) -> VortexResult { + vortex_ensure!( + args.children.nchildren() == 2, + "DictLayout expects exactly 2 children" + ); + let codes_nullable = metadata + .is_nullable_codes + .map(Nullability::from) + .unwrap_or_else(|| args.dtype.nullability()); + let codes_dtype = DType::Primitive(metadata.codes_ptype(), codes_nullable); + args.children.child(0, args.dtype)?; + let codes = args.children.child(1, &codes_dtype)?; + vortex_ensure!( + codes.row_count() == args.row_count, + "Dictionary codes row count does not match parent" + ); + Ok(DictData { + codes_dtype, + all_values_referenced: metadata.all_values_referenced.unwrap_or(false), + }) } - fn child(layout: &Self::Layout, idx: usize) -> VortexResult { + fn child_dtype(layout: &Layout, idx: usize) -> VortexResult { match idx { - 0 => Ok(Arc::clone(&layout.values)), - 1 => Ok(Arc::clone(&layout.codes)), - _ => vortex_bail!("Unreachable child index: {}", idx), + 0 => Ok(layout.dtype().clone()), + 1 => Ok(layout.codes_dtype.clone()), + _ => vortex_bail!("Dict child index out of bounds: {idx}"), } } - fn child_type(_layout: &Self::Layout, idx: usize) -> LayoutChildType { + fn child_type(_layout: &Layout, idx: usize) -> LayoutChildType { match idx { 0 => LayoutChildType::Auxiliary("values".into()), 1 => LayoutChildType::Transparent("codes".into()), - _ => vortex_panic!("Unreachable child index: {}", idx), + _ => vortex_panic!("Dict child index out of bounds: {idx}"), } } fn new_reader( - layout: &Self::Layout, + layout: &Layout, name: Arc, segment_source: Arc, session: &VortexSession, - ctx: &crate::LayoutReaderContext, + ctx: &LayoutReaderContext, ) -> VortexResult { Ok(Arc::new(DictReader::try_new( layout.clone(), @@ -104,89 +123,36 @@ impl VTable for Dict { ctx.clone(), )?)) } - - fn build( - _encoding: &Self::Encoding, - dtype: &DType, - _row_count: u64, - metadata: &::Output, - _segment_ids: Vec, - children: &dyn LayoutChildren, - _build_ctx: &LayoutBuildContext<'_>, - ) -> VortexResult { - let values = children.child(0, dtype)?; - let codes_nullable = metadata - .is_nullable_codes - .map(Nullability::from) - // The old behaviour (without `is_nullable_codes` metadata) used the nullability - // of the values (and whole array). - // see [`SerdeVTable::build`]. - .unwrap_or_else(|| dtype.nullability()); - let codes = children.child(1, &DType::Primitive(metadata.codes_ptype(), codes_nullable))?; - Ok(unsafe { - DictLayout::new(values, codes) - .set_all_values_referenced(metadata.all_values_referenced.unwrap_or(false)) - }) - } - - fn with_children(layout: &mut Self::Layout, children: Vec) -> VortexResult<()> { - vortex_ensure!( - children.len() == 2, - "DictLayout expects exactly 2 children (values, codes), got {}", - children.len() - ); - let mut children_iter = children.into_iter(); - layout.values = children_iter - .next() - .ok_or_else(|| vortex_err!("Missing values child"))?; - layout.codes = children_iter - .next() - .ok_or_else(|| vortex_err!("Missing codes child"))?; - Ok(()) - } -} - -#[derive(Debug)] -pub struct DictLayoutEncoding; - -/// Stores a shared dictionary of values alongside compact integer codes that index into it. -/// -/// Useful for columns with many repeated values, where storing each value once and -/// referencing it by index saves significant space. -#[derive(Clone, Debug)] -pub struct DictLayout { - values: LayoutRef, - codes: LayoutRef, - /// Indicates whether all dictionary values are definitely referenced by at least one code. - /// `true` = all values are referenced (computed during encoding). - /// `false` = unknown/might have unreferenced values. - all_values_referenced: bool, } -impl DictLayout { +impl Layout { pub(crate) fn new(values: LayoutRef, codes: LayoutRef) -> Self { - Self { - values, - codes, - all_values_referenced: false, - } - } - - /// Set whether all dictionary values are definitely referenced. - /// - /// # Safety - /// The caller must ensure that when setting `all_values_referenced = true`, ALL dictionary - /// values are actually referenced by at least one valid code. Setting this incorrectly can - /// lead to incorrect query results in operations like min/max. - /// - /// This is typically only set to `true` during dictionary encoding when we know for certain - /// that all values are referenced. - /// See `DictArray::set_all_values_referenced`. - pub unsafe fn set_all_values_referenced(mut self, all_values_referenced: bool) -> Self { - self.all_values_referenced = all_values_referenced; - self - } - + Self::new_with_all_values_referenced(values, codes, false) + } + + fn new_with_all_values_referenced( + values: LayoutRef, + codes: LayoutRef, + all_values_referenced: bool, + ) -> Self { + let dtype = values.dtype().clone(); + let row_count = codes.row_count(); + let codes_dtype = codes.dtype().clone(); + LayoutParts::new( + Dict, + dtype, + row_count, + Vec::new(), + OwnedLayoutChildren::layout_children(vec![values, codes]), + DictData { + codes_dtype, + all_values_referenced, + }, + ) + .into_typed() + } + + /// Returns whether every dictionary value is known to be referenced. pub fn has_all_values_referenced(&self) -> bool { self.all_values_referenced } @@ -195,15 +161,9 @@ impl DictLayout { #[derive(prost::Message)] pub struct DictLayoutMetadata { #[prost(enumeration = "PType", tag = "1")] - // i32 is required for proto, use the generated getter to read this field. codes_ptype: i32, - // nullable codes are optional since they were added after stabilisation #[prost(optional, bool, tag = "2")] is_nullable_codes: Option, - // all_values_referenced is optional for backward compatibility - // true = all dictionary values are definitely referenced by at least one code - // false/None = unknown whether all values are referenced (conservative default) - // see `DictArray::all_values_referenced` #[prost(optional, bool, tag = "3")] pub(crate) all_values_referenced: Option, } diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 30172eb45d2..84c484dd5fa 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -67,14 +67,16 @@ impl DictReader { session: VortexSession, ctx: crate::LayoutReaderContext, ) -> VortexResult { - let values_len = usize::try_from(layout.values.row_count())?; - let values = layout.values.new_reader( + let values_layout = layout.child(0)?; + let codes_layout = layout.child(1)?; + let values_len = usize::try_from(values_layout.row_count())?; + let values = values_layout.new_reader( format!("{name}.values").into(), Arc::clone(&segment_source), &session, &ctx, )?; - let codes = layout.codes.new_reader( + let codes = codes_layout.new_reader( format!("{name}.codes").into(), segment_source, &session, diff --git a/vortex-layout/src/layouts/dict/writer.rs b/vortex-layout/src/layouts/dict/writer.rs index 83483c92662..c93142d2f15 100644 --- a/vortex-layout/src/layouts/dict/writer.rs +++ b/vortex-layout/src/layouts/dict/writer.rs @@ -38,7 +38,6 @@ use vortex_io::kanal_ext::KanalExt; use vortex_io::session::RuntimeSessionExt; use vortex_session::VortexSession; -use crate::IntoLayout; use crate::LayoutRef; use crate::LayoutStrategy; use crate::OwnedLayoutChildren; diff --git a/vortex-layout/src/layouts/flat/mod.rs b/vortex-layout/src/layouts/flat/mod.rs index aa75a3b67a8..f7589a4c77f 100644 --- a/vortex-layout/src/layouts/flat/mod.rs +++ b/vortex-layout/src/layouts/flat/mod.rs @@ -8,7 +8,6 @@ use std::env; use std::sync::Arc; use std::sync::LazyLock; -use vortex_array::DeserializeMetadata; use vortex_array::ProstMetadata; use vortex_array::dtype::DType; use vortex_buffer::ByteBuffer; @@ -19,18 +18,18 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use vortex_session::registry::ReadContext; -use crate::LayoutBuildContext; +use crate::Layout; use crate::LayoutChildType; -use crate::LayoutEncodingRef; +use crate::LayoutDeserializeArgs; use crate::LayoutId; +use crate::LayoutParts; +use crate::LayoutReaderContext; use crate::LayoutReaderRef; -use crate::LayoutRef; use crate::VTable; -use crate::children::LayoutChildren; +use crate::children::OwnedLayoutChildren; use crate::layouts::flat::reader::FlatReader; use crate::segments::SegmentId; use crate::segments::SegmentSource; -use crate::vtable; /// Check if inline array node is enabled. pub(super) fn flat_layout_inline_array_node() -> bool { @@ -39,58 +38,74 @@ pub(super) fn flat_layout_inline_array_node() -> bool { *FLAT_LAYOUT_INLINE_ARRAY_NODE } -vtable!(Flat); +/// Flat layout vtable. +#[derive(Clone, Debug)] +pub struct Flat; + +/// Backwards-compatible name for the flat layout plugin. +pub use Flat as FlatLayoutEncoding; + +/// Flat-layout-specific data. +#[derive(Clone, Debug)] +pub struct FlatData { + segment_id: SegmentId, + ctx: ReadContext, + array_tree: Option, +} + +/// A terminal layout storing one serialized array segment. +pub type FlatLayout = Layout; impl VTable for Flat { - type Layout = FlatLayout; - type Encoding = FlatLayoutEncoding; + type LayoutData = FlatData; type Metadata = ProstMetadata; - fn id(_encoding: &Self::Encoding) -> LayoutId { + fn id(&self) -> LayoutId { static ID: CachedId = CachedId::new("vortex.flat"); *ID } - fn encoding(_layout: &Self::Layout) -> LayoutEncodingRef { - LayoutEncodingRef::new_ref(FlatLayoutEncoding.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::Metadata { ProstMetadata(FlatLayoutMetadata { array_encoding_tree: layout.array_tree.as_ref().map(|bytes| bytes.to_vec()), }) } - fn segment_ids(layout: &Self::Layout) -> Vec { - vec![layout.segment_id] - } - - fn nchildren(_layout: &Self::Layout) -> usize { - 0 + fn deserialize( + &self, + args: &LayoutDeserializeArgs<'_>, + metadata: &FlatLayoutMetadata, + ) -> VortexResult { + if args.segment_ids.len() != 1 { + vortex_bail!("Flat layout must have exactly one segment ID"); + } + if args.children.nchildren() != 0 { + vortex_bail!("Flat layout must not have children"); + } + Ok(FlatData { + segment_id: args.segment_ids[0], + ctx: args.array_read_ctx.clone(), + array_tree: metadata + .array_encoding_tree + .as_ref() + .map(|bytes| ByteBuffer::from(bytes.clone())), + }) } - fn child(_layout: &Self::Layout, _idx: usize) -> VortexResult { - vortex_bail!("Flat layout has no children"); + fn child_dtype(_layout: &Layout, idx: usize) -> VortexResult { + vortex_bail!("Flat layout has no child {idx}") } - fn child_type(_layout: &Self::Layout, _idx: usize) -> LayoutChildType { - vortex_panic!("Flat layout has no children"); + fn child_type(_layout: &Layout, idx: usize) -> LayoutChildType { + vortex_panic!("Flat layout has no child {idx}") } fn new_reader( - layout: &Self::Layout, + layout: &Layout, name: Arc, segment_source: Arc, session: &VortexSession, - _ctx: &crate::LayoutReaderContext, + _ctx: &LayoutReaderContext, ) -> VortexResult { Ok(Arc::new(FlatReader::new( layout.clone(), @@ -99,91 +114,48 @@ impl VTable for Flat { session.clone(), ))) } - - fn build( - _encoding: &Self::Encoding, - dtype: &DType, - row_count: u64, - metadata: &::Output, - segment_ids: Vec, - _children: &dyn LayoutChildren, - build_ctx: &LayoutBuildContext<'_>, - ) -> VortexResult { - if segment_ids.len() != 1 { - vortex_bail!("Flat layout must have exactly one segment ID"); - } - Ok(FlatLayout::new_with_metadata( - row_count, - dtype.clone(), - segment_ids[0], - build_ctx.array_read_ctx.clone(), - metadata - .array_encoding_tree - .as_ref() - .map(|v| ByteBuffer::from(v.clone())), - )) - } - - fn with_children(_layout: &mut Self::Layout, children: Vec) -> VortexResult<()> { - if !children.is_empty() { - vortex_bail!("Flat layout has no children, got {}", children.len()); - } - Ok(()) - } -} - -#[derive(Debug)] -pub struct FlatLayoutEncoding; - -/// The terminal node of a layout tree. Stores a single chunk of array data as one serialized -/// segment on disk. -#[derive(Clone, Debug)] -pub struct FlatLayout { - row_count: u64, - dtype: DType, - segment_id: SegmentId, - ctx: ReadContext, - array_tree: Option, } -impl FlatLayout { +impl Layout { + /// Construct a flat layout without an inline array encoding tree. pub fn new(row_count: u64, dtype: DType, segment_id: SegmentId, ctx: ReadContext) -> Self { - Self { - row_count, - dtype, - segment_id, - ctx, - array_tree: None, - } + Self::new_with_metadata(row_count, dtype, segment_id, ctx, None) } + /// Construct a flat layout with optional inline array metadata. pub fn new_with_metadata( row_count: u64, dtype: DType, segment_id: SegmentId, ctx: ReadContext, - metadata: Option, + array_tree: Option, ) -> Self { - Self { - row_count, + LayoutParts::new( + Flat, dtype, - segment_id, - ctx, - array_tree: metadata, - } - } - - #[inline] + row_count, + vec![segment_id], + OwnedLayoutChildren::layout_children(Vec::new()), + FlatData { + segment_id, + ctx, + array_tree, + }, + ) + .into_typed() + } + + /// Returns the serialized array segment ID. pub fn segment_id(&self) -> SegmentId { self.segment_id } - #[inline] + /// Returns the array read context. pub fn array_ctx(&self) -> &ReadContext { &self.ctx } - #[inline] + /// Returns the optional inline array encoding tree. pub fn array_tree(&self) -> Option<&ByteBuffer> { self.array_tree.as_ref() } @@ -191,9 +163,6 @@ impl FlatLayout { #[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>, } diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index 50ec634a61a..42e1221af45 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -107,7 +107,7 @@ impl LayoutReader for FlatReader { 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(()) } diff --git a/vortex-layout/src/layouts/flat/writer.rs b/vortex-layout/src/layouts/flat/writer.rs index 55c07275301..723556b5f49 100644 --- a/vortex-layout/src/layouts/flat/writer.rs +++ b/vortex-layout/src/layouts/flat/writer.rs @@ -22,7 +22,6 @@ use vortex_error::vortex_bail; use vortex_session::VortexSession; use vortex_session::registry::ReadContext; -use crate::IntoLayout; use crate::LayoutRef; use crate::LayoutStrategy; use crate::children::OwnedLayoutChildren; diff --git a/vortex-layout/src/layouts/foreign/mod.rs b/vortex-layout/src/layouts/foreign/mod.rs index 491b8744262..1cee18c2eda 100644 --- a/vortex-layout/src/layouts/foreign/mod.rs +++ b/vortex-layout/src/layouts/foreign/mod.rs @@ -10,7 +10,7 @@ use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_session::VortexSession; -use crate::Layout; +use crate::DynLayout; use crate::LayoutBuildContext; use crate::LayoutChildType; use crate::LayoutChildren; @@ -119,54 +119,54 @@ pub(crate) fn new_foreign_layout( )) } -impl Layout for ForeignLayout { +impl DynLayout for ForeignLayout { fn as_any(&self) -> &dyn Any { self } - fn as_any_arc(self: Arc) -> Arc { - self - } - - fn to_layout(&self) -> LayoutRef { + fn dyn_to_layout(&self) -> LayoutRef { Arc::new(self.clone()) } - fn encoding(&self) -> LayoutEncodingRef { - self.encoding.clone() + fn dyn_encoding_id(&self) -> LayoutEncodingId { + self.encoding.id() } - fn row_count(&self) -> u64 { + fn dyn_row_count(&self) -> u64 { self.row_count } - fn dtype(&self) -> &DType { + fn dyn_dtype(&self) -> &DType { &self.dtype } - fn nchildren(&self) -> usize { + fn dyn_nchildren(&self) -> usize { self.children.len() } - fn child(&self, idx: usize) -> VortexResult { + fn dyn_child(&self, idx: usize) -> VortexResult { self.children.get(idx).cloned().ok_or_else(|| { - vortex_err!("Child index out of bounds: {} of {}", idx, self.nchildren()) + vortex_err!( + "Child index out of bounds: {} of {}", + idx, + self.children.len() + ) }) } - fn child_type(&self, idx: usize) -> LayoutChildType { + fn dyn_child_type(&self, idx: usize) -> LayoutChildType { LayoutChildType::Auxiliary(format!("[{idx}]").into()) } - fn metadata(&self) -> Vec { + fn dyn_metadata(&self) -> Vec { self.metadata.clone() } - fn segment_ids(&self) -> Vec { + fn dyn_segment_ids(&self) -> Vec { self.segment_ids.clone() } - fn new_reader( + fn dyn_new_reader( &self, _name: Arc, _segment_source: Arc, diff --git a/vortex-layout/src/layouts/list/mod.rs b/vortex-layout/src/layouts/list/mod.rs index 0d47ad5266d..f77d4c325b3 100644 --- a/vortex-layout/src/layouts/list/mod.rs +++ b/vortex-layout/src/layouts/list/mod.rs @@ -1,18 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! An experimental structural layout for list-typed columns. Note that this is expected to change. -//! -//! [`ListLayout`] decomposes a list column into independently configurable child layouts: -//! `elements`, `offsets`, and, for nullable lists, `validity`. Keeping the children independent allows -//! each child to use its own configurable layout and lets nested list elements (e.g. with `List`) be decomposed recursively. -//! -//! This provides benefits such as: -//! * Reading only the children needed to evaluate an expression. For example, `ListLength` does -//! not need to read list elements. -//! * Restricting element reads to the range covered by the selected outer rows, avoiding elements -//! belonging exclusively to unselected leading or trailing lists. -//! * Allowing each child to use its own compression, chunking, and pruning strategy. +//! An experimental structural layout for list-typed columns. mod expr; mod reader; @@ -21,7 +10,6 @@ pub mod writer; use std::sync::Arc; use reader::ListReader; -use vortex_array::DeserializeMetadata; use vortex_array::ProstMetadata; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -35,89 +23,118 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::LayoutBuildContext; +use crate::Layout; use crate::LayoutChildType; -use crate::LayoutEncodingRef; +use crate::LayoutDeserializeArgs; use crate::LayoutId; +use crate::LayoutParts; use crate::LayoutReaderContext; use crate::LayoutReaderRef; use crate::LayoutRef; use crate::VTable; -use crate::children::LayoutChildren; -use crate::segments::SegmentId; +use crate::children::OwnedLayoutChildren; use crate::segments::SegmentSource; -use crate::vtable; -/// Child index of the `elements` layout. +/// Child index of the elements layout. pub const ELEMENTS_CHILD_INDEX: usize = 0; -/// Child index of the `offsets` layout. +/// Child index of the offsets layout. pub const OFFSETS_CHILD_INDEX: usize = 1; -/// Child index of the `validity` layout (only present when the list dtype is nullable). +/// Child index of the optional validity layout. pub const VALIDITY_CHILD_INDEX: usize = 2; - -/// Number of children when the list dtype is non-nullable. +/// Number of children for a non-nullable list. pub const NUM_CHILDREN_NON_NULLABLE: usize = 2; -vtable!(List); +/// List layout vtable. +#[derive(Clone, Debug)] +pub struct List; + +/// Backwards-compatible name for the list layout plugin. +pub use List as ListLayoutEncoding; + +/// List-layout-specific data. +#[derive(Clone, Debug)] +pub struct ListData { + offsets_ptype: PType, +} + +/// A list layout shredded into elements, offsets, and optional validity children. +pub type ListLayout = Layout; impl VTable for List { - type Layout = ListLayout; - type Encoding = ListLayoutEncoding; + type LayoutData = ListData; type Metadata = ProstMetadata; - fn id(_encoding: &Self::Encoding) -> LayoutId { + fn id(&self) -> LayoutId { static ID: CachedId = CachedId::new("vortex.list"); *ID } - fn encoding(_layout: &Self::Layout) -> LayoutEncodingRef { - LayoutEncodingRef::new_ref(ListLayoutEncoding.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 { - ProstMetadata(ListLayoutMetadata::new(layout.offsets_ptype())) + fn metadata(layout: &Layout) -> Self::Metadata { + ProstMetadata(ListLayoutMetadata::new(layout.offsets_ptype)) } - fn segment_ids(_layout: &Self::Layout) -> Vec { - vec![] - } - - fn nchildren(layout: &Self::Layout) -> usize { - if layout.dtype.is_nullable() { - NUM_CHILDREN_NON_NULLABLE + 1 - } else { - NUM_CHILDREN_NON_NULLABLE + fn deserialize( + &self, + args: &LayoutDeserializeArgs<'_>, + metadata: &ListLayoutMetadata, + ) -> VortexResult { + validate_children(args.dtype, args.children.nchildren())?; + let elements_dtype = args + .dtype + .as_list_element_opt() + .ok_or_else(|| vortex_err!("ListLayout requires a List dtype, got {}", args.dtype))?; + args.children.child(ELEMENTS_CHILD_INDEX, elements_dtype)?; + let offsets_dtype = DType::Primitive(metadata.offsets_ptype(), Nullability::NonNullable); + let offsets = args.children.child(OFFSETS_CHILD_INDEX, &offsets_dtype)?; + vortex_error::vortex_ensure!( + offsets.row_count().saturating_sub(1) == args.row_count, + "List offsets row count does not match parent" + ); + if args.dtype.is_nullable() { + let validity = args + .children + .child(VALIDITY_CHILD_INDEX, &DType::Bool(Nullability::NonNullable))?; + vortex_error::vortex_ensure!( + validity.row_count() == args.row_count, + "List validity row count does not match parent" + ); } + Ok(ListData { + offsets_ptype: metadata.offsets_ptype(), + }) } - fn child(layout: &Self::Layout, idx: usize) -> VortexResult { - match (idx, layout.validity.as_ref()) { - (ELEMENTS_CHILD_INDEX, _) => Ok(Arc::clone(&layout.elements)), - (OFFSETS_CHILD_INDEX, _) => Ok(Arc::clone(&layout.offsets)), - (VALIDITY_CHILD_INDEX, Some(validity)) => Ok(Arc::clone(validity)), + fn child_dtype(layout: &Layout, idx: usize) -> VortexResult { + match idx { + ELEMENTS_CHILD_INDEX => layout + .dtype() + .as_list_element_opt() + .map(|dtype| dtype.as_ref().clone()) + .ok_or_else(|| vortex_err!("ListLayout requires a List dtype")), + OFFSETS_CHILD_INDEX => Ok(DType::Primitive( + layout.offsets_ptype, + Nullability::NonNullable, + )), + VALIDITY_CHILD_INDEX if layout.dtype().is_nullable() => { + Ok(DType::Bool(Nullability::NonNullable)) + } _ => vortex_bail!("Invalid child index {idx} for ListLayout"), } } - fn child_type(layout: &Self::Layout, idx: usize) -> LayoutChildType { - match (idx, layout.validity.is_some()) { - (ELEMENTS_CHILD_INDEX, _) => LayoutChildType::Auxiliary("elements".into()), - (OFFSETS_CHILD_INDEX, _) => LayoutChildType::Auxiliary("offsets".into()), - (VALIDITY_CHILD_INDEX, true) => LayoutChildType::Auxiliary("validity".into()), + fn child_type(layout: &Layout, idx: usize) -> LayoutChildType { + match idx { + ELEMENTS_CHILD_INDEX => LayoutChildType::Auxiliary("elements".into()), + OFFSETS_CHILD_INDEX => LayoutChildType::Auxiliary("offsets".into()), + VALIDITY_CHILD_INDEX if layout.dtype().is_nullable() => { + LayoutChildType::Auxiliary("validity".into()) + } _ => vortex_panic!("Invalid child index {idx} for ListLayout"), } } fn new_reader( - layout: &Self::Layout, + layout: &Layout, name: Arc, segment_source: Arc, session: &VortexSession, @@ -131,144 +148,69 @@ impl VTable for List { ctx, )?)) } - - fn build( - _encoding: &Self::Encoding, - dtype: &DType, - _row_count: u64, - metadata: &::Output, - _segment_ids: Vec, - children: &dyn LayoutChildren, - _ctx: &LayoutBuildContext<'_>, - ) -> VortexResult { - validate_children(dtype, children.nchildren())?; - - let elements_dtype = dtype - .as_list_element_opt() - .ok_or_else(|| vortex_err!("ListLayout requires a List dtype, got {dtype}"))?; - let elements = children.child(ELEMENTS_CHILD_INDEX, elements_dtype.as_ref())?; - - let offsets_dtype = DType::Primitive(metadata.offsets_ptype(), Nullability::NonNullable); - let offsets = children.child(OFFSETS_CHILD_INDEX, &offsets_dtype)?; - - let validity = dtype - .is_nullable() - .then(|| children.child(VALIDITY_CHILD_INDEX, &DType::Bool(Nullability::NonNullable))) - .transpose()?; - - Ok(ListLayout { - dtype: dtype.clone(), - elements, - offsets, - validity, - }) - } - - fn with_children(layout: &mut Self::Layout, children: Vec) -> VortexResult<()> { - validate_children(layout.dtype(), children.len())?; - - let mut iter = children.into_iter(); - layout.elements = iter - .next() - .ok_or_else(|| vortex_err!("missing elements child"))?; - layout.offsets = iter - .next() - .ok_or_else(|| vortex_err!("missing offsets child"))?; - layout.validity = layout - .dtype - .is_nullable() - .then(|| { - iter.next() - .ok_or_else(|| vortex_err!("missing validity child")) - }) - .transpose()?; - Ok(()) - } -} - -/// Validates expected number of children based on `dtype` -fn validate_children(dtype: &DType, n_children: usize) -> VortexResult<()> { - let expected = if dtype.is_nullable() { - NUM_CHILDREN_NON_NULLABLE + 1 - } else { - NUM_CHILDREN_NON_NULLABLE - }; - - vortex_ensure_eq!(n_children, expected); - Ok(()) -} - -#[derive(Debug)] -pub struct ListLayoutEncoding; - -/// Stores a list-typed array by shredding `elements`, `offsets`, and optional `validity` children. -#[derive(Clone, Debug)] -pub struct ListLayout { - dtype: DType, - elements: LayoutRef, - offsets: LayoutRef, - validity: Option, } -impl ListLayout { - /// Construct a new `ListLayout` from its components. - /// - /// # Invariants - /// - /// - `dtype` must be a [`DType::List`]. - /// - `validity` must be `Some` iff `dtype.is_nullable()`. - /// - `offsets.dtype()` must be a non-nullable integer. - /// - `offsets.row_count()` is the Arrow-canonical `n+1` for `n` lists (or `0` for empty). - /// - When present, `validity.row_count() == offsets.row_count().saturating_sub(1)`. +impl Layout { + /// Construct a list layout from its children. pub fn new( dtype: DType, elements: LayoutRef, offsets: LayoutRef, validity: Option, ) -> Self { - Self { + let row_count = offsets.row_count().saturating_sub(1); + let offsets_ptype = offsets.dtype().as_ptype(); + let mut children = vec![elements, offsets]; + children.extend(validity); + validate_children(&dtype, children.len()).vortex_expect("invalid list children"); + LayoutParts::new( + List, dtype, - elements, - offsets, - validity, - } + row_count, + Vec::new(), + OwnedLayoutChildren::layout_children(children), + ListData { offsets_ptype }, + ) + .into_typed() } - /// Number of lists in this layout. - #[inline] - pub fn row_count(&self) -> u64 { - self.offsets.row_count().saturating_sub(1) + /// Returns the elements child. + pub fn elements(&self) -> VortexResult { + self.child(ELEMENTS_CHILD_INDEX) } - #[inline] - pub fn elements(&self) -> &LayoutRef { - &self.elements + /// Returns the offsets child. + pub fn offsets(&self) -> VortexResult { + self.child(OFFSETS_CHILD_INDEX) } - #[inline] - pub fn offsets(&self) -> &LayoutRef { - &self.offsets - } - - #[inline] - pub fn validity(&self) -> Option<&LayoutRef> { - self.validity.as_ref() + /// Returns the optional validity child. + pub fn validity(&self) -> VortexResult> { + self.dtype() + .is_nullable() + .then(|| self.child(VALIDITY_CHILD_INDEX)) + .transpose() } - /// The integer type used for the `offsets` child layout. - #[inline] + /// Returns the integer ptype used by offsets. pub fn offsets_ptype(&self) -> PType { - self.offsets.dtype().as_ptype() + self.offsets_ptype } - /// The dtype of the inner elements column. + /// Returns the list element dtype. pub fn elements_dtype(&self) -> &DType { - self.dtype + self.dtype() .as_list_element_opt() .vortex_expect("ListLayout dtype must be a List") } } +fn validate_children(dtype: &DType, nchildren: usize) -> VortexResult<()> { + let expected = NUM_CHILDREN_NON_NULLABLE + usize::from(dtype.is_nullable()); + vortex_ensure_eq!(nchildren, expected); + Ok(()) +} + #[derive(prost::Message)] pub struct ListLayoutMetadata { #[prost(enumeration = "PType", tag = "1")] diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index 0d33f1f4ac5..1e21714fdd0 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -69,20 +69,20 @@ impl ListReader { session: VortexSession, ctx: &LayoutReaderContext, ) -> VortexResult { - let elements = layout.elements().new_reader( + let elements = layout.elements()?.new_reader( format!("{name}.elements").into(), Arc::clone(&segment_source), &session, ctx, )?; - let offsets = layout.offsets().new_reader( + let offsets = layout.offsets()?.new_reader( format!("{name}.offsets").into(), Arc::clone(&segment_source), &session, ctx, )?; let validity = layout - .validity() + .validity()? .map(|v| { v.new_reader( format!("{name}.validity").into(), diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 9f33c1eee3e..f9f0516df23 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -31,7 +31,6 @@ use vortex_io::kanal_ext::KanalExt; use vortex_io::session::RuntimeSessionExt; use vortex_session::VortexSession; -use crate::IntoLayout; use crate::LayoutRef; use crate::LayoutStrategy; use crate::layouts::flat::writer::FlatLayoutStrategy; diff --git a/vortex-layout/src/layouts/struct_/mod.rs b/vortex-layout/src/layouts/struct_/mod.rs index 853a6793f4c..f65b93fa6a8 100644 --- a/vortex-layout/src/layouts/struct_/mod.rs +++ b/vortex-layout/src/layouts/struct_/mod.rs @@ -7,7 +7,6 @@ pub mod writer; use std::sync::Arc; use reader::StructReader; -use vortex_array::DeserializeMetadata; use vortex_array::EmptyMetadata; use vortex_array::dtype::DType; use vortex_array::dtype::Field; @@ -24,83 +23,69 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; pub use writer::StructStrategy; -use crate::LayoutBuildContext; +use crate::Layout; use crate::LayoutChildType; -use crate::LayoutEncodingRef; +use crate::LayoutDeserializeArgs; use crate::LayoutId; +use crate::LayoutParts; +use crate::LayoutReaderContext; use crate::LayoutReaderRef; use crate::LayoutRef; use crate::VTable; -use crate::children::LayoutChildren; use crate::children::OwnedLayoutChildren; -use crate::segments::SegmentId; use crate::segments::SegmentSource; -use crate::vtable; -vtable!(Struct); +/// Struct layout vtable. +#[derive(Clone, Debug)] +pub struct Struct; + +/// Backwards-compatible name for the struct layout plugin. +pub use Struct as StructLayoutEncoding; + +/// A layout decomposing a struct into one child per field and optional validity. +pub type StructLayout = Layout; impl VTable for Struct { - type Layout = StructLayout; - type Encoding = StructLayoutEncoding; + type LayoutData = (); type Metadata = EmptyMetadata; - fn id(_encoding: &Self::Encoding) -> LayoutId { + fn id(&self) -> LayoutId { static ID: CachedId = CachedId::new("vortex.struct"); *ID } - fn encoding(_layout: &Self::Layout) -> LayoutEncodingRef { - LayoutEncodingRef::new_ref(StructLayoutEncoding.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::Metadata { EmptyMetadata } - fn segment_ids(_layout: &Self::Layout) -> Vec { - vec![] - } - - fn nchildren(layout: &Self::Layout) -> usize { - let validity_children = if layout.dtype.is_nullable() { 1 } else { 0 }; - layout.struct_fields().nfields() + validity_children + fn deserialize( + &self, + args: &LayoutDeserializeArgs<'_>, + _metadata: &EmptyMetadata, + ) -> VortexResult { + validate_children(args.dtype, args.children.nchildren())?; + for idx in 0..args.children.nchildren() { + let child_dtype = struct_child_dtype(args.dtype, idx)?; + let child = args.children.child(idx, &child_dtype)?; + vortex_ensure!( + child.row_count() == args.row_count, + "Struct child {idx} row count does not match parent" + ); + } + Ok(()) } - fn child(layout: &Self::Layout, index: usize) -> VortexResult { - let schema_index = if layout.dtype.is_nullable() { - index.saturating_sub(1) - } else { - index - }; - - let child_dtype = if index == 0 && layout.dtype.is_nullable() { - DType::Bool(Nullability::NonNullable) - } else { - layout - .struct_fields() - .field_by_index(schema_index) - .ok_or_else(|| vortex_err!("Missing field {schema_index}"))? - }; - - layout.children.child(index, &child_dtype) + fn child_dtype(layout: &Layout, index: usize) -> VortexResult { + struct_child_dtype(layout.dtype(), index) } - fn child_type(layout: &Self::Layout, idx: usize) -> LayoutChildType { - let schema_index = if layout.dtype.is_nullable() { + fn child_type(layout: &Layout, idx: usize) -> LayoutChildType { + let schema_index = if layout.dtype().is_nullable() { idx.saturating_sub(1) } else { idx }; - - if idx == 0 && layout.dtype.is_nullable() { + if idx == 0 && layout.dtype().is_nullable() { LayoutChildType::Auxiliary("validity".into()) } else { LayoutChildType::Field( @@ -114,11 +99,11 @@ impl VTable for Struct { } fn new_reader( - layout: &Self::Layout, + layout: &Layout, name: Arc, segment_source: Arc, session: &VortexSession, - ctx: &crate::LayoutReaderContext, + ctx: &LayoutReaderContext, ) -> VortexResult { Ok(Arc::new(StructReader::try_new( layout.clone(), @@ -128,97 +113,35 @@ impl VTable for Struct { ctx.clone(), )?)) } - - fn build( - _encoding: &Self::Encoding, - dtype: &DType, - row_count: u64, - _metadata: &::Output, - _segment_ids: Vec, - children: &dyn LayoutChildren, - _build_ctx: &LayoutBuildContext<'_>, - ) -> VortexResult { - let struct_dt = dtype - .as_struct_fields_opt() - .ok_or_else(|| vortex_err!("Expected struct dtype"))?; - - let expected_children = struct_dt.nfields() + (dtype.is_nullable() as usize); - vortex_ensure!( - children.nchildren() == expected_children, - "Struct layout has {} children, but dtype has {} fields", - children.nchildren(), - struct_dt.nfields() - ); - - Ok(StructLayout { - row_count, - dtype: dtype.clone(), - children: children.to_arc(), - }) - } - - fn with_children(layout: &mut Self::Layout, children: Vec) -> VortexResult<()> { - let struct_dt = layout - .dtype - .as_struct_fields_opt() - .ok_or_else(|| vortex_err!("Expected struct dtype"))?; - - let expected_children = struct_dt.nfields() + (layout.dtype.is_nullable() as usize); - vortex_ensure!( - children.len() == expected_children, - "StructLayout expects {} children, got {}", - expected_children, - children.len() - ); - - layout.children = OwnedLayoutChildren::layout_children(children); - Ok(()) - } } -#[derive(Debug)] -pub struct StructLayoutEncoding; - -/// Decomposes a struct-typed column into one child per field, enabling columnar projection. -/// -/// Queries that only need a subset of fields can skip reading the rest entirely. -#[derive(Clone, Debug)] -pub struct StructLayout { - row_count: u64, - dtype: DType, - children: Arc, -} - -impl StructLayout { +impl Layout { + /// Construct a struct layout from owned children. pub fn new(row_count: u64, dtype: DType, children: Vec) -> Self { - Self { - row_count, + validate_children(&dtype, children.len()).vortex_expect("invalid struct children"); + LayoutParts::new( + Struct, dtype, - children: OwnedLayoutChildren::layout_children(children), - } + row_count, + Vec::new(), + OwnedLayoutChildren::layout_children(children), + (), + ) + .into_typed() } + /// Returns the struct fields. pub fn struct_fields(&self) -> &StructFields { - self.dtype + self.dtype() .as_struct_fields_opt() .vortex_expect("Struct layout dtype must be a struct") } - #[inline] - pub fn row_count(&self) -> u64 { - self.row_count - } - - #[inline] - pub fn children(&self) -> &Arc { - &self.children - } - + /// Invokes `per_child` for fields selected by `field_mask`. pub fn matching_fields(&self, field_mask: &[FieldMask], mut per_child: F) -> VortexResult<()> where F: FnMut(FieldMask, usize) -> VortexResult<()>, { - // If the field mask contains an `All` fields, then enumerate all fields. if field_mask.iter().any(|mask| mask.matches_all()) { for idx in 0..self.struct_fields().nfields() { per_child(FieldMask::All, idx)?; @@ -226,10 +149,8 @@ impl StructLayout { return Ok(()); } - // Enumerate each field in the mask for path in field_mask { let Some(field) = path.starting_field()? else { - // skip fields not in mask continue; }; let Field::Name(field_name) = field else { @@ -239,10 +160,36 @@ impl StructLayout { .struct_fields() .find(field_name) .ok_or_else(|| vortex_err!("Field not found: {field_name}"))?; - per_child(path.clone().step_into()?, idx)?; } - Ok(()) } } + +fn validate_children(dtype: &DType, nchildren: usize) -> VortexResult<()> { + let fields = dtype + .as_struct_fields_opt() + .ok_or_else(|| vortex_err!("Expected struct dtype"))?; + let expected = fields.nfields() + usize::from(dtype.is_nullable()); + vortex_ensure!( + nchildren == expected, + "Struct layout has {nchildren} children, expected {expected}" + ); + Ok(()) +} + +fn struct_child_dtype(dtype: &DType, index: usize) -> VortexResult { + let schema_index = if dtype.is_nullable() { + index.saturating_sub(1) + } else { + index + }; + if index == 0 && dtype.is_nullable() { + Ok(DType::Bool(Nullability::NonNullable)) + } else { + dtype + .as_struct_fields_opt() + .and_then(|fields| fields.field_by_index(schema_index)) + .ok_or_else(|| vortex_err!("Missing field {schema_index}")) + } +} diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 18a14a2c97b..a9623e61f8e 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -82,7 +82,7 @@ impl StructReader { .collect() }); - let nullable = layout.dtype.is_nullable(); + let nullable = layout.dtype().is_nullable(); let extra = nullable as usize; let mut dtypes: Vec = Vec::with_capacity(struct_dt.nfields() + extra); @@ -95,7 +95,7 @@ impl StructReader { names.extend(struct_dt.names().iter().map(|x| Arc::clone(x.inner()))); let lazy_children = LazyReaderChildren::new( - Arc::clone(&layout.children), + Arc::clone(layout.children()), dtypes, names, Arc::clone(&segment_source), diff --git a/vortex-layout/src/layouts/struct_/writer.rs b/vortex-layout/src/layouts/struct_/writer.rs index 3076c3ed39c..cc2e239144f 100644 --- a/vortex-layout/src/layouts/struct_/writer.rs +++ b/vortex-layout/src/layouts/struct_/writer.rs @@ -39,7 +39,6 @@ use vortex_utils::aliases::DefaultHashBuilder; use vortex_utils::aliases::hash_map::HashMap; use vortex_utils::aliases::hash_set::HashSet; -use crate::IntoLayout; use crate::LayoutRef; use crate::LayoutStrategy; use crate::layouts::struct_::StructLayout; diff --git a/vortex-layout/src/layouts/zoned/mod.rs b/vortex-layout/src/layouts/zoned/mod.rs index ce6d869acd8..3ecbf9c0dde 100644 --- a/vortex-layout/src/layouts/zoned/mod.rs +++ b/vortex-layout/src/layouts/zoned/mod.rs @@ -43,14 +43,15 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::LayoutBuildContext; +use crate::Layout; use crate::LayoutChildType; -use crate::LayoutEncodingRef; +use crate::LayoutDeserializeArgs; use crate::LayoutId; +use crate::LayoutParts; +use crate::LayoutReaderContext; use crate::LayoutReaderRef; use crate::LayoutRef; use crate::VTable; -use crate::children::LayoutChildren; use crate::children::OwnedLayoutChildren; use crate::layouts::zoned::reader::ZonedReader; use crate::layouts::zoned::schema::AggregateSpecProto; @@ -58,36 +59,43 @@ use crate::layouts::zoned::schema::aggregate_fns_from_specs; use crate::layouts::zoned::schema::aggregate_specs_from_fns; use crate::layouts::zoned::schema::aggregate_stats_table_dtype; use crate::layouts::zoned::schema::legacy_stats_table_dtype; -use crate::segments::SegmentId; use crate::segments::SegmentSource; -use crate::vtable; -vtable!(Zoned); -vtable!(LegacyStats); +/// Zoned layout vtable. +#[derive(Clone, Debug)] +pub struct Zoned; + +/// Legacy `vortex.stats` layout vtable. +#[derive(Clone, Debug)] +pub struct LegacyStats; + +/// Backwards-compatible plugin names. +pub use LegacyStats as LegacyStatsLayoutEncoding; +pub use Zoned as ZonedLayoutEncoding; + +/// Zoned-layout-specific data. +#[derive(Clone, Debug)] +pub struct ZonedData { + zone_len: usize, + zone_map_schema: ZoneMapSchema, + stats_table_dtype: DType, +} + +/// A layout annotating a data child with per-zone statistics. +pub type ZonedLayout = Layout; +/// A legacy `vortex.stats` layout using the same runtime data. +pub type LegacyStatsLayout = Layout; impl VTable for Zoned { - type Layout = ZonedLayout; - type Encoding = ZonedLayoutEncoding; + type LayoutData = ZonedData; type Metadata = ZonedMetadata; - fn id(_encoding: &Self::Encoding) -> LayoutId { + fn id(&self) -> LayoutId { static ID: CachedId = CachedId::new("vortex.zoned"); *ID } - fn encoding(_layout: &Self::Layout) -> LayoutEncodingRef { - LayoutEncodingRef::new_ref(ZonedLayoutEncoding.as_ref()) - } - - fn row_count(layout: &Self::Layout) -> u64 { - layout.children.child_row_count(0) - } - - fn dtype(layout: &Self::Layout) -> &DType { - &layout.dtype - } - - fn metadata(layout: &Self::Layout) -> Self::Metadata { + fn metadata(layout: &Layout) -> Self::Metadata { ZonedMetadata { zone_len: u32::try_from(layout.zone_len).vortex_expect("Invalid zone length"), aggregate_specs: match &layout.zone_map_schema { @@ -103,23 +111,38 @@ impl VTable for Zoned { } } - fn segment_ids(_layout: &Self::Layout) -> Vec { - vec![] - } - - fn nchildren(_layout: &Self::Layout) -> usize { - 2 + fn deserialize( + &self, + args: &LayoutDeserializeArgs<'_>, + metadata: &ZonedMetadata, + ) -> VortexResult { + vortex_ensure_eq!( + args.children.nchildren(), + 2, + "ZonedLayout expects exactly 2 children (data, zones)" + ); + vortex_ensure_eq!(args.children.child_row_count(0), args.row_count); + let aggregate_fns = aggregate_fns_from_specs(&metadata.aggregate_specs, args.session)?; + aggregate_specs_from_fns(&aggregate_fns)?; + let stats_table_dtype = aggregate_stats_table_dtype(args.dtype, &aggregate_fns); + args.children.child(0, args.dtype)?; + args.children.child(1, &stats_table_dtype)?; + Ok(ZonedData { + zone_len: metadata.zone_len as usize, + zone_map_schema: ZoneMapSchema::AggregateFns(aggregate_fns), + stats_table_dtype, + }) } - fn child(layout: &Self::Layout, idx: usize) -> VortexResult { + fn child_dtype(layout: &Layout, idx: usize) -> VortexResult { match idx { - 0 => layout.children.child(0, layout.dtype()), - 1 => layout.children.child(1, &layout.stats_table_dtype), - _ => vortex_bail!("Invalid child index: {}", idx), + 0 => Ok(layout.dtype().clone()), + 1 => Ok(layout.stats_table_dtype.clone()), + _ => vortex_bail!("Invalid child index: {idx}"), } } - fn child_type(_layout: &Self::Layout, idx: usize) -> LayoutChildType { + fn child_type(_layout: &Layout, idx: usize) -> LayoutChildType { match idx { 0 => LayoutChildType::Transparent("data".into()), 1 => LayoutChildType::Auxiliary("zones".into()), @@ -128,11 +151,11 @@ impl VTable for Zoned { } fn new_reader( - layout: &Self::Layout, + layout: &Layout, name: Arc, segment_source: Arc, session: &VortexSession, - ctx: &crate::LayoutReaderContext, + ctx: &LayoutReaderContext, ) -> VortexResult { Ok(Arc::new(ZonedReader::try_new( layout.clone(), @@ -142,172 +165,95 @@ impl VTable for Zoned { ctx.clone(), )?)) } - - fn build( - _encoding: &Self::Encoding, - dtype: &DType, - _row_count: u64, - metadata: &ZonedMetadata, - _segment_ids: Vec, - children: &dyn LayoutChildren, - build_ctx: &LayoutBuildContext<'_>, - ) -> VortexResult { - vortex_ensure_eq!( - children.nchildren(), - 2, - "ZonedLayout expects exactly 2 children (data, zones)" - ); - let aggregate_fns = aggregate_fns_from_specs(&metadata.aggregate_specs, build_ctx.session)?; - aggregate_specs_from_fns(&aggregate_fns)?; - let stats_table_dtype = aggregate_stats_table_dtype(dtype, &aggregate_fns); - Ok(ZonedLayout { - dtype: dtype.clone(), - children: children.to_arc(), - zone_len: metadata.zone_len as usize, - zone_map_schema: ZoneMapSchema::AggregateFns(aggregate_fns), - stats_table_dtype, - }) - } - - fn with_children(layout: &mut Self::Layout, children: Vec) -> VortexResult<()> { - if children.len() != 2 { - vortex_bail!( - "ZonedLayout expects exactly 2 children (data, zones), got {}", - children.len() - ); - } - layout.children = OwnedLayoutChildren::layout_children(children); - Ok(()) - } } -// TODO: This legacy vtable is only needed until layouts move onto the new vtable structure, where -// a LayoutPlugin can deserialize legacy `vortex.stats` metadata directly into `vortex.zoned`. impl VTable for LegacyStats { - type Layout = LegacyStatsLayout; - type Encoding = LegacyStatsLayoutEncoding; + type LayoutData = ZonedData; type Metadata = LegacyStatsMetadata; - fn id(_encoding: &Self::Encoding) -> LayoutId { + fn id(&self) -> LayoutId { static ID: CachedId = CachedId::new("vortex.stats"); *ID } - fn encoding(_layout: &Self::Layout) -> LayoutEncodingRef { - LayoutEncodingRef::new_ref(LegacyStatsLayoutEncoding.as_ref()) - } - - fn row_count(layout: &Self::Layout) -> u64 { - ::row_count(&layout.0) - } - - fn dtype(layout: &Self::Layout) -> &DType { - ::dtype(&layout.0) - } - - fn metadata(layout: &Self::Layout) -> Self::Metadata { + fn metadata(layout: &Layout) -> Self::Metadata { LegacyStatsMetadata { - zone_len: u32::try_from(layout.0.zone_len).vortex_expect("Invalid zone length"), - zone_map_schema: layout.0.zone_map_schema.clone(), + zone_len: u32::try_from(layout.zone_len).vortex_expect("Invalid zone length"), + zone_map_schema: layout.zone_map_schema.clone(), } } - fn segment_ids(layout: &Self::Layout) -> Vec { - ::segment_ids(&layout.0) - } - - fn nchildren(layout: &Self::Layout) -> usize { - ::nchildren(&layout.0) + fn deserialize( + &self, + args: &LayoutDeserializeArgs<'_>, + metadata: &LegacyStatsMetadata, + ) -> VortexResult { + vortex_ensure_eq!( + args.children.nchildren(), + 2, + "LegacyStatsLayout expects exactly 2 children (data, zones)" + ); + let stats_table_dtype = match &metadata.zone_map_schema { + ZoneMapSchema::LegacyStats(stats) => legacy_stats_table_dtype(args.dtype, stats), + ZoneMapSchema::AggregateFns(aggregate_fns) => { + aggregate_stats_table_dtype(args.dtype, aggregate_fns) + } + }; + args.children.child(0, args.dtype)?; + args.children.child(1, &stats_table_dtype)?; + Ok(ZonedData { + zone_len: metadata.zone_len as usize, + zone_map_schema: metadata.zone_map_schema.clone(), + stats_table_dtype, + }) } - fn child(layout: &Self::Layout, idx: usize) -> VortexResult { - ::child(&layout.0, idx) + fn child_dtype(layout: &Layout, idx: usize) -> VortexResult { + match idx { + 0 => Ok(layout.dtype().clone()), + 1 => Ok(layout.stats_table_dtype.clone()), + _ => vortex_bail!("Invalid child index: {idx}"), + } } - fn child_type(layout: &Self::Layout, idx: usize) -> LayoutChildType { - ::child_type(&layout.0, idx) + fn child_type(_layout: &Layout, idx: usize) -> LayoutChildType { + match idx { + 0 => LayoutChildType::Transparent("data".into()), + 1 => LayoutChildType::Auxiliary("zones".into()), + _ => vortex_panic!("Invalid child index: {idx}"), + } } fn new_reader( - layout: &Self::Layout, + layout: &Layout, name: Arc, segment_source: Arc, session: &VortexSession, - ctx: &crate::LayoutReaderContext, + ctx: &LayoutReaderContext, ) -> VortexResult { + let zoned = LayoutParts::new( + Zoned, + layout.dtype().clone(), + layout.row_count(), + layout.segment_ids().to_vec(), + Arc::clone(layout.children()), + layout.data().clone(), + ) + .into_typed(); Ok(Arc::new(ZonedReader::try_new( - layout.0.clone(), + zoned, name, segment_source, session.clone(), ctx.clone(), )?)) } - - fn build( - _encoding: &Self::Encoding, - dtype: &DType, - _row_count: u64, - metadata: &LegacyStatsMetadata, - _segment_ids: Vec, - children: &dyn LayoutChildren, - _build_ctx: &LayoutBuildContext<'_>, - ) -> VortexResult { - vortex_ensure_eq!( - children.nchildren(), - 2, - "LegacyStatsLayout expects exactly 2 children (data, zones)" - ); - let stats_table_dtype = match &metadata.zone_map_schema { - ZoneMapSchema::LegacyStats(stats) => legacy_stats_table_dtype(dtype, stats), - ZoneMapSchema::AggregateFns(aggregate_fns) => { - aggregate_stats_table_dtype(dtype, aggregate_fns) - } - }; - Ok(LegacyStatsLayout(ZonedLayout { - dtype: dtype.clone(), - children: children.to_arc(), - zone_len: metadata.zone_len as usize, - zone_map_schema: metadata.zone_map_schema.clone(), - stats_table_dtype, - })) - } - - fn with_children(layout: &mut Self::Layout, children: Vec) -> VortexResult<()> { - ::with_children(&mut layout.0, children) - } } -/// Encoding marker for the zoned layout. -#[derive(Debug)] -pub struct ZonedLayoutEncoding; - -/// Encoding marker for the legacy `vortex.stats` zoned layout. -#[derive(Debug)] -pub struct LegacyStatsLayoutEncoding; - -/// A layout that annotates a data child with one row of aggregate statistics per zone. -/// -/// The first child is the underlying data layout. The second child is an auxiliary stats table -/// whose rows align with logical row zones of length `zone_len`, except for the final partial zone. -/// During reads, pruning uses the stats table to skip zones whose rows cannot satisfy a filter. -#[derive(Clone, Debug)] -pub struct ZonedLayout { - dtype: DType, - children: Arc, - zone_len: usize, - zone_map_schema: ZoneMapSchema, - stats_table_dtype: DType, -} - -/// A legacy `vortex.stats` layout backed by the shared zoned runtime implementation. -#[derive(Clone, Debug)] -pub struct LegacyStatsLayout(ZonedLayout); - impl LegacyStatsLayout { /// Returns display names for the zone-map aggregates stored by this layout. pub fn present_aggregates(&self) -> Arc<[String]> { - self.0.present_aggregates() + present_aggregates(&self.zone_map_schema) } } @@ -332,17 +278,25 @@ impl ZonedLayout { } aggregate_specs_from_fns(&aggregate_fns)?; - Ok(Self { - dtype: data.dtype().clone(), - children: OwnedLayoutChildren::layout_children(vec![data, zones]), - zone_len: zone_len.get(), - zone_map_schema: ZoneMapSchema::AggregateFns(aggregate_fns), - stats_table_dtype: expected_dtype, - }) + let dtype = data.dtype().clone(); + let row_count = data.row_count(); + Ok(LayoutParts::new( + Zoned, + dtype, + row_count, + Vec::new(), + OwnedLayoutChildren::layout_children(vec![data, zones]), + ZonedData { + zone_len: zone_len.get(), + zone_map_schema: ZoneMapSchema::AggregateFns(aggregate_fns), + stats_table_dtype: expected_dtype, + }, + ) + .into_typed()) } pub fn nzones(&self) -> usize { - usize::try_from(self.children.child_row_count(1)).vortex_expect("Invalid number of zones") + usize::try_from(self.child_row_count(1)).vortex_expect("Invalid number of zones") } pub fn zone_len(&self) -> usize { @@ -351,19 +305,7 @@ impl ZonedLayout { /// Returns display names for the zone-map aggregates stored by this layout. pub fn present_aggregates(&self) -> Arc<[String]> { - match &self.zone_map_schema { - ZoneMapSchema::LegacyStats(stats) => stats - .iter() - .filter_map(Stat::aggregate_fn) - .map(|aggregate_fn| aggregate_fn.to_string()) - .collect::>() - .into(), - ZoneMapSchema::AggregateFns(aggregate_fns) => aggregate_fns - .iter() - .map(ToString::to_string) - .collect::>() - .into(), - } + present_aggregates(&self.zone_map_schema) } pub(super) fn aggregate_fns( @@ -382,10 +324,26 @@ impl ZonedLayout { pub(super) fn stats_table_dtype_for(&self, aggregate_fns: &[AggregateFnRef]) -> DType { if let ZoneMapSchema::LegacyStats(stats) = &self.zone_map_schema { - return legacy_stats_table_dtype(&self.dtype, stats); + return legacy_stats_table_dtype(self.dtype(), stats); } - aggregate_stats_table_dtype(&self.dtype, aggregate_fns) + aggregate_stats_table_dtype(self.dtype(), aggregate_fns) + } +} + +fn present_aggregates(schema: &ZoneMapSchema) -> Arc<[String]> { + match schema { + ZoneMapSchema::LegacyStats(stats) => stats + .iter() + .filter_map(Stat::aggregate_fn) + .map(|aggregate_fn| aggregate_fn.to_string()) + .collect::>() + .into(), + ZoneMapSchema::AggregateFns(aggregate_fns) => aggregate_fns + .iter() + .map(ToString::to_string) + .collect::>() + .into(), } } @@ -510,7 +468,7 @@ mod tests { use vortex_session::registry::ReadContext; use super::*; - use crate::IntoLayout; + use crate::LayoutBuildContext; use crate::children::OwnedLayoutChildren; use crate::layouts::flat::FlatLayout; use crate::segments::SegmentId; @@ -641,7 +599,7 @@ mod tests { &build_ctx, )?; - assert_eq!(layout.0.zone_len, 0); + assert_eq!(layout.zone_len, 0); Ok(()) } diff --git a/vortex-layout/src/layouts/zoned/reader.rs b/vortex-layout/src/layouts/zoned/reader.rs index 162742b6f22..51b4ad4ac33 100644 --- a/vortex-layout/src/layouts/zoned/reader.rs +++ b/vortex-layout/src/layouts/zoned/reader.rs @@ -45,12 +45,12 @@ impl ZonedReader { ) -> VortexResult { let aggregate_fns = layout.aggregate_fns(&session)?; let dtypes = vec![ - layout.dtype.clone(), + layout.dtype().clone(), layout.stats_table_dtype_for(&aggregate_fns), ]; let names = vec![Arc::clone(&name), format!("{}.zones", name).into()]; let lazy_children = Arc::new(LazyReaderChildren::new( - Arc::clone(&layout.children), + Arc::clone(layout.children()), dtypes, names, Arc::clone(&segment_source), @@ -248,7 +248,6 @@ mod test { use vortex_session::VortexSession; use vortex_session::registry::ReadContext; - use crate::IntoLayout; use crate::LayoutBuildContext; use crate::LayoutRef; use crate::LayoutStrategy; diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index 048905dade6..77d7c770319 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -34,7 +34,6 @@ use vortex_io::session::RuntimeSessionExt; use vortex_session::VortexSession; use vortex_utils::parallelism::get_available_parallelism; -use crate::IntoLayout; use crate::LayoutRef; use crate::LayoutStrategy; use crate::layouts::zoned::AggregateStatsAccumulator; diff --git a/vortex-layout/src/lib.rs b/vortex-layout/src/lib.rs index 6612f5e0350..226dcefcf0a 100644 --- a/vortex-layout/src/lib.rs +++ b/vortex-layout/src/lib.rs @@ -8,8 +8,8 @@ //! segment bytes. A [`LayoutReader`] pairs a layout with a [`SegmentSource`](segments::SegmentSource) //! and session so scans can evaluate projections and filters. //! -//! Most users enter this crate through file APIs, but extension authors implement [`VTable`], -//! [`LayoutEncoding`], and [`LayoutStrategy`] to add new on-disk organizations. +//! Most users enter this crate through file APIs, but extension authors implement [`VTable`] and +//! [`LayoutStrategy`] to add new on-disk organizations. //! //! Scanning is built with [`scan::scan_builder::ScanBuilder`]. It accepts a projection expression, //! optional filter, optional row range, [`Selection`](vortex_scan::selection::Selection), split @@ -24,7 +24,6 @@ pub use reader::*; pub use reader_context::*; pub use strategy::*; use vortex_session::registry::Context; -pub use vtable::*; pub mod aliases; mod children; pub mod display; @@ -40,6 +39,5 @@ pub mod session; mod strategy; #[cfg(test)] mod test; -pub mod vtable; pub type LayoutContext = Context; diff --git a/vortex-layout/src/session.rs b/vortex-layout/src/session.rs index 6bb52b1d1e8..53ced66e955 100644 --- a/vortex-layout/src/session.rs +++ b/vortex-layout/src/session.rs @@ -8,14 +8,15 @@ use vortex_session::SessionGuard; use vortex_session::SessionVar; use vortex_session::registry::Registry; +use crate::LayoutEncoding; use crate::LayoutEncodingRef; -use crate::layouts::chunked::ChunkedLayoutEncoding; -use crate::layouts::dict::DictLayoutEncoding; -use crate::layouts::flat::FlatLayoutEncoding; -use crate::layouts::list::ListLayoutEncoding; -use crate::layouts::struct_::StructLayoutEncoding; -use crate::layouts::zoned::LegacyStatsLayoutEncoding; -use crate::layouts::zoned::ZonedLayoutEncoding; +use crate::layouts::chunked::Chunked; +use crate::layouts::dict::Dict; +use crate::layouts::flat::Flat; +use crate::layouts::list::List; +use crate::layouts::struct_::Struct; +use crate::layouts::zoned::LegacyStats; +use crate::layouts::zoned::Zoned; pub type LayoutRegistry = Registry; @@ -49,16 +50,13 @@ impl Default for LayoutSession { let layouts = LayoutRegistry::default(); // Register the built-in layout encodings. - layouts.register(ChunkedLayoutEncoding.id(), ChunkedLayoutEncoding.as_ref()); - layouts.register(FlatLayoutEncoding.id(), FlatLayoutEncoding.as_ref()); - layouts.register(StructLayoutEncoding.id(), StructLayoutEncoding.as_ref()); - layouts.register(ZonedLayoutEncoding.id(), ZonedLayoutEncoding.as_ref()); - layouts.register( - LegacyStatsLayoutEncoding.id(), - LegacyStatsLayoutEncoding.as_ref(), - ); - layouts.register(DictLayoutEncoding.id(), DictLayoutEncoding.as_ref()); - layouts.register(ListLayoutEncoding.id(), ListLayoutEncoding.as_ref()); + layouts.register(Chunked.id(), &Chunked as &dyn LayoutEncoding); + layouts.register(Flat.id(), &Flat as &dyn LayoutEncoding); + layouts.register(Struct.id(), &Struct as &dyn LayoutEncoding); + layouts.register(Zoned.id(), &Zoned as &dyn LayoutEncoding); + layouts.register(LegacyStats.id(), &LegacyStats as &dyn LayoutEncoding); + layouts.register(Dict.id(), &Dict as &dyn LayoutEncoding); + layouts.register(List.id(), &List as &dyn LayoutEncoding); Self { registry: layouts } } diff --git a/vortex-layout/src/vtable.rs b/vortex-layout/src/vtable.rs deleted file mode 100644 index d584dd629a3..00000000000 --- a/vortex-layout/src/vtable.rs +++ /dev/null @@ -1,182 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::fmt::Debug; -use std::ops::Deref; -use std::sync::Arc; - -use vortex_array::DeserializeMetadata; -use vortex_array::SerializeMetadata; -use vortex_array::dtype::DType; -use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_session::VortexSession; -use vortex_session::registry::ReadContext; - -use crate::IntoLayout; -use crate::Layout; -use crate::LayoutChildType; -use crate::LayoutEncoding; -use crate::LayoutEncodingRef; -use crate::LayoutId; -use crate::LayoutReaderContext; -use crate::LayoutReaderRef; -use crate::LayoutRef; -use crate::children::LayoutChildren; -use crate::segments::SegmentId; -use crate::segments::SegmentSource; - -/// Context available while constructing a layout from serialized metadata. -pub struct LayoutBuildContext<'a> { - /// The session used to resolve plugin-owned metadata such as aggregate function options. - pub session: &'a VortexSession, - /// The array read context referenced by serialized array metadata in descendant layouts. - pub array_read_ctx: &'a ReadContext, -} - -/// Typed implementation contract for a layout encoding. -/// -/// A layout vtable connects a concrete layout node type, its registered encoding object, and the -/// metadata representation used for serialization. The object-safe [`Layout`] and -/// [`LayoutEncoding`] APIs delegate to this trait through adapters. -pub trait VTable: 'static + Sized + Send + Sync + Debug { - /// Concrete layout node type for this encoding. - type Layout: 'static + Send + Sync + Clone + Debug + Deref + IntoLayout; - /// Concrete encoding object registered in the session. - type Encoding: 'static + Send + Sync + Deref; - /// Serialized layout metadata type. - type Metadata: SerializeMetadata + DeserializeMetadata + Debug; - - /// Returns the ID of the layout encoding. - fn id(encoding: &Self::Encoding) -> LayoutId; - - /// Returns the encoding for the layout. - fn encoding(layout: &Self::Layout) -> LayoutEncodingRef; - - /// Returns the row count for the layout. - fn row_count(layout: &Self::Layout) -> u64; - - /// Returns the dtype for the layout reader. - fn dtype(layout: &Self::Layout) -> &DType; - - /// Returns the metadata for the layout. - fn metadata(layout: &Self::Layout) -> Self::Metadata; - - /// Returns the segment IDs for the layout. - fn segment_ids(layout: &Self::Layout) -> Vec; - - /// Returns the number of children for the layout. - fn nchildren(layout: &Self::Layout) -> usize; - - /// Return the child at the given index. - fn child(layout: &Self::Layout, idx: usize) -> VortexResult; - - /// Return the type of the child at the given index. - fn child_type(layout: &Self::Layout, idx: usize) -> LayoutChildType; - - /// Create a new reader for the layout. - /// - /// **Layouts with children MUST propagate `ctx` to descendants** by passing it - /// through `Layout::new_reader` (or `LazyReaderChildren::new`) when constructing - /// child readers. If `ctx` is dropped at any link in the chain, ancestor-published - /// values won't reach affected descendants — a silent runtime regression for any - /// descendant that looked up an ancestor-published value via `ctx.get::()`. - /// There is no compile-time check that catches this; reviewer discipline + the - /// integration tests in `vortex-layout` are the only safety net. - fn new_reader( - layout: &Self::Layout, - name: Arc, - segment_source: Arc, - session: &VortexSession, - ctx: &LayoutReaderContext, - ) -> VortexResult; - - /// Construct a new [`Layout`] from deserialized parts. - /// - /// Implementations should validate child count, child types, row counts, segment references, - /// and dtype consistency for their encoding. The generic adapter checks the returned layout's - /// top-level dtype and row count, but encoding-specific invariants belong here. - fn build( - encoding: &Self::Encoding, - dtype: &DType, - row_count: u64, - metadata: &::Output, - segment_ids: Vec, - children: &dyn LayoutChildren, - build_ctx: &LayoutBuildContext<'_>, - ) -> VortexResult; - - /// Replaces the children of the layout with the given layout references. - /// - /// The count and types of children must match the layout's requirements. - /// This method is used for transforming layout trees by replacing child layouts. - fn with_children(_layout: &mut Self::Layout, _children: Vec) -> VortexResult<()> { - vortex_bail!("with_children not implemented for this layout") - } -} - -#[macro_export] -macro_rules! vtable { - ($V:ident) => { - $crate::aliases::paste::paste! { - #[derive(Debug)] - pub struct $V; - - impl AsRef for [<$V Layout>] { - fn as_ref(&self) -> &dyn $crate::Layout { - // SAFETY: LayoutAdapter is #[repr(transparent)] over the Layout type, - // which guarantees identical memory layout. This cast is safe because - // we're only changing the type metadata, not the actual data. - unsafe { &*(self as *const [<$V Layout>] as *const $crate::LayoutAdapter<$V>) } - } - } - - impl std::ops::Deref for [<$V Layout>] { - type Target = dyn $crate::Layout; - - fn deref(&self) -> &Self::Target { - // SAFETY: LayoutAdapter is #[repr(transparent)] over the Layout type, - // which guarantees identical memory layout. This cast is safe because - // we're only changing the type metadata, not the actual data. - unsafe { &*(self as *const [<$V Layout>] as *const $crate::LayoutAdapter<$V>) } - } - } - - impl $crate::IntoLayout for [<$V Layout>] { - fn into_layout(self) -> $crate::LayoutRef { - // SAFETY: LayoutAdapter is #[repr(transparent)] over the Layout type, - // guaranteeing identical memory layout and alignment. The transmute is safe - // because both types have the same size and representation. - std::sync::Arc::new(unsafe { std::mem::transmute::<[<$V Layout>], $crate::LayoutAdapter::<$V>>(self) }) - } - } - - impl From<[<$V Layout>]> for $crate::LayoutRef { - fn from(value: [<$V Layout>]) -> $crate::LayoutRef { - use $crate::IntoLayout; - value.into_layout() - } - } - - impl AsRef for [<$V LayoutEncoding>] { - fn as_ref(&self) -> &dyn $crate::LayoutEncoding { - // SAFETY: LayoutEncodingAdapter is #[repr(transparent)] over the LayoutEncoding type, - // which guarantees identical memory layout. This cast is safe because - // we're only changing the type metadata, not the actual data. - unsafe { &*(self as *const [<$V LayoutEncoding>] as *const $crate::LayoutEncodingAdapter<$V>) } - } - } - - impl std::ops::Deref for [<$V LayoutEncoding>] { - type Target = dyn $crate::LayoutEncoding; - - fn deref(&self) -> &Self::Target { - // SAFETY: LayoutEncodingAdapter is #[repr(transparent)] over the LayoutEncoding type, - // which guarantees identical memory layout. This cast is safe because - // we're only changing the type metadata, not the actual data. - unsafe { &*(self as *const [<$V LayoutEncoding>] as *const $crate::LayoutEncodingAdapter<$V>) } - } - } - } - }; -} diff --git a/vortex-tui/src/segment_tree.rs b/vortex-tui/src/segment_tree.rs index fd88650f2fd..40918806cdd 100644 --- a/vortex-tui/src/segment_tree.rs +++ b/vortex-tui/src/segment_tree.rs @@ -8,7 +8,7 @@ use std::sync::Arc; use vortex::dtype::FieldName; use vortex::error::VortexResult; use vortex::file::SegmentSpec; -use vortex::layout::Layout; +use vortex::layout::DynLayout; use vortex::layout::LayoutChildType; use vortex::utils::aliases::hash_map::HashMap; @@ -34,7 +34,7 @@ pub struct SegmentTree { /// Collect segment tree from a layout and segment map. pub fn collect_segment_tree( - root_layout: &dyn Layout, + root_layout: &dyn DynLayout, segments: &Arc<[SegmentSpec]>, ) -> SegmentTree { let mut tree = SegmentTree { @@ -54,7 +54,7 @@ pub fn collect_segment_tree( } fn segments_by_name_impl( - root: &dyn Layout, + root: &dyn DynLayout, group_name: Option, name: Option, row_offset: Option, diff --git a/vortex-tui/src/tree.rs b/vortex-tui/src/tree.rs index 27218d1d8db..281af5d7472 100644 --- a/vortex-tui/src/tree.rs +++ b/vortex-tui/src/tree.rs @@ -152,7 +152,7 @@ fn layout_to_json(layout: LayoutRef) -> VortexResult { .collect::>>()?; Ok(LayoutTreeNode { - encoding: layout.encoding().to_string(), + encoding: layout.encoding_id().to_string(), dtype: layout.dtype().to_string(), row_count: layout.row_count(), metadata_bytes: layout.metadata().len(), diff --git a/vortex-web/crate/src/wasm.rs b/vortex-web/crate/src/wasm.rs index 0323a98c271..3c89ede5cb3 100644 --- a/vortex-web/crate/src/wasm.rs +++ b/vortex-web/crate/src/wasm.rs @@ -591,7 +591,7 @@ fn build_layout_tree( Ok(LayoutTreeNodeJson { id, - encoding: layout.encoding().id().to_string(), + encoding: layout.encoding_id().to_string(), dtype: layout.dtype().to_string(), row_count, row_offset: parent_row_offset, From 1c381298f826b73e149153b50d6ae3d84741b9f4 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 23 Jul 2026 11:45:02 +0100 Subject: [PATCH 2/8] wip Signed-off-by: Joe Isaacs --- vortex-layout/src/flatbuffers.rs | 60 ++++-------------------- vortex-layout/src/layout.rs | 10 ++++ vortex-layout/src/layouts/list/mod.rs | 16 ++++--- vortex-layout/src/layouts/struct_/mod.rs | 56 ++++++++++++---------- 4 files changed, 58 insertions(+), 84 deletions(-) diff --git a/vortex-layout/src/flatbuffers.rs b/vortex-layout/src/flatbuffers.rs index 8e461d32d1f..4f57d6e92d7 100644 --- a/vortex-layout/src/flatbuffers.rs +++ b/vortex-layout/src/flatbuffers.rs @@ -4,28 +4,28 @@ use std::env; use std::sync::LazyLock; +use flatbuffers::root_with_opts; use flatbuffers::FlatBufferBuilder; use flatbuffers::VerifierOptions; use flatbuffers::WIPOffset; -use flatbuffers::root_with_opts; use vortex_array::dtype::DType; -use vortex_error::VortexResult; use vortex_error::vortex_err; +use vortex_error::VortexResult; +use vortex_flatbuffers::layout; use vortex_flatbuffers::FlatBuffer; use vortex_flatbuffers::FlatBufferRoot; use vortex_flatbuffers::WriteFlatBuffer; -use vortex_flatbuffers::layout; -use vortex_session::VortexSession; use vortex_session::registry::ReadContext; +use vortex_session::VortexSession; -use crate::DynLayout; -use crate::LayoutBuildContext; -use crate::LayoutContext; -use crate::LayoutRef; use crate::children::ViewedLayoutChildren; use crate::layouts::foreign::new_foreign_layout; use crate::segments::SegmentId; use crate::session::LayoutSessionExt; +use crate::DynLayout; +use crate::LayoutBuildContext; +use crate::LayoutContext; +use crate::LayoutRef; static LAYOUT_VERIFIER: LazyLock = LazyLock::new(|| { VerifierOptions { @@ -227,54 +227,12 @@ mod tests { use vortex_array::array_session; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; - use vortex_array::dtype::PType; - use vortex_flatbuffers::WriteFlatBufferExt; use vortex_flatbuffers::layout as fbl; use vortex_session::registry::ReadContext; - use super::layout_from_flatbuffer; use super::layout_from_flatbuffer_with_options; - use crate::LayoutContext; - use crate::LayoutEncodingId; - use crate::layouts::flat::Flat; - use crate::layouts::flat::FlatLayout; - use crate::segments::SegmentId; use crate::session::LayoutSession; - - #[test] - fn flat_layout_preserves_wire_fields_and_deserializes() { - let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let array_ctx = ReadContext::new([]); - let layout = FlatLayout::new_with_metadata( - 3, - dtype.clone(), - SegmentId::from(7), - array_ctx.clone(), - Some(vec![1, 2].into()), - ); - let layout_ctx = LayoutContext::default(); - let layout = layout.to_layout(); - let buffer = layout - .flatbuffer_writer(&layout_ctx) - .write_flatbuffer_bytes() - .unwrap(); - - let wire = flatbuffers::root::(&buffer).unwrap(); - assert_eq!(wire.encoding(), 0); - assert_eq!(wire.row_count(), 3); - assert_eq!(wire.metadata().unwrap().bytes(), &[0x0a, 0x02, 1, 2]); - assert_eq!(wire.segments().unwrap().iter().collect::>(), [7]); - assert!(wire.children().is_none()); - - let layout_read_ctx = ReadContext::new(layout_ctx.to_ids()); - let session = array_session().with::(); - let decoded = - layout_from_flatbuffer(buffer, &dtype, &layout_read_ctx, &array_ctx, &session).unwrap(); - let decoded = decoded.as_::(); - assert_eq!(decoded.row_count(), 3); - assert_eq!(*decoded.segment_id(), 7); - assert_eq!(decoded.array_tree().unwrap().as_ref(), &[1, 2]); - } + use crate::LayoutEncodingId; #[expect(clippy::disallowed_methods, reason = "test-only id")] #[test] diff --git a/vortex-layout/src/layout.rs b/vortex-layout/src/layout.rs index df665ecd44e..f4cd4ecca02 100644 --- a/vortex-layout/src/layout.rs +++ b/vortex-layout/src/layout.rs @@ -282,24 +282,34 @@ impl From> for LayoutRef { pub trait DynLayout: 'static + Send + Sync + Debug { /// Returns this layout as [`Any`] for downcasting. fn as_any(&self) -> &dyn Any; + /// Clone this layout as an erased reference. fn dyn_to_layout(&self) -> LayoutRef; + /// Returns the layout ID. fn dyn_encoding_id(&self) -> LayoutId; + /// Returns the row count. fn dyn_row_count(&self) -> u64; + /// Returns the logical dtype. fn dyn_dtype(&self) -> &DType; + /// Returns the number of children. fn dyn_nchildren(&self) -> usize; + /// Returns child `idx`. fn dyn_child(&self, idx: usize) -> VortexResult; + /// Returns the relationship of child `idx`. fn dyn_child_type(&self, idx: usize) -> LayoutChildType; + /// Serializes layout-specific metadata. fn dyn_metadata(&self) -> Vec; + /// Returns directly referenced segment IDs. fn dyn_segment_ids(&self) -> Vec; + /// Constructs a reader. fn dyn_new_reader( &self, diff --git a/vortex-layout/src/layouts/list/mod.rs b/vortex-layout/src/layouts/list/mod.rs index f77d4c325b3..165d2525a01 100644 --- a/vortex-layout/src/layouts/list/mod.rs +++ b/vortex-layout/src/layouts/list/mod.rs @@ -78,7 +78,7 @@ impl VTable for List { args: &LayoutDeserializeArgs<'_>, metadata: &ListLayoutMetadata, ) -> VortexResult { - validate_children(args.dtype, args.children.nchildren())?; + ListLayout::validate_children(args.dtype, args.children.nchildren())?; let elements_dtype = args .dtype .as_list_element_opt() @@ -162,7 +162,7 @@ impl Layout { let offsets_ptype = offsets.dtype().as_ptype(); let mut children = vec![elements, offsets]; children.extend(validity); - validate_children(&dtype, children.len()).vortex_expect("invalid list children"); + Self::validate_children(&dtype, children.len()).vortex_expect("invalid list children"); LayoutParts::new( List, dtype, @@ -203,14 +203,16 @@ impl Layout { .as_list_element_opt() .vortex_expect("ListLayout dtype must be a List") } -} -fn validate_children(dtype: &DType, nchildren: usize) -> VortexResult<()> { - let expected = NUM_CHILDREN_NON_NULLABLE + usize::from(dtype.is_nullable()); - vortex_ensure_eq!(nchildren, expected); - Ok(()) + fn validate_children(dtype: &DType, nchildren: usize) -> VortexResult<()> { + let expected = NUM_CHILDREN_NON_NULLABLE + usize::from(dtype.is_nullable()); + vortex_ensure_eq!(nchildren, expected); + Ok(()) + } } + + #[derive(prost::Message)] pub struct ListLayoutMetadata { #[prost(enumeration = "PType", tag = "1")] diff --git a/vortex-layout/src/layouts/struct_/mod.rs b/vortex-layout/src/layouts/struct_/mod.rs index f65b93fa6a8..3d8de3e0e34 100644 --- a/vortex-layout/src/layouts/struct_/mod.rs +++ b/vortex-layout/src/layouts/struct_/mod.rs @@ -63,9 +63,10 @@ impl VTable for Struct { args: &LayoutDeserializeArgs<'_>, _metadata: &EmptyMetadata, ) -> VortexResult { - validate_children(args.dtype, args.children.nchildren())?; + Layout::::validate_children(args.dtype, args.children.nchildren())?; + for idx in 0..args.children.nchildren() { - let child_dtype = struct_child_dtype(args.dtype, idx)?; + let child_dtype = StructLayout::child_dtype(args.dtype, idx)?; let child = args.children.child(idx, &child_dtype)?; vortex_ensure!( child.row_count() == args.row_count, @@ -76,7 +77,7 @@ impl VTable for Struct { } fn child_dtype(layout: &Layout, index: usize) -> VortexResult { - struct_child_dtype(layout.dtype(), index) + StructLayout::child_dtype(layout.dtype(), index) } fn child_type(layout: &Layout, idx: usize) -> LayoutChildType { @@ -118,7 +119,7 @@ impl VTable for Struct { impl Layout { /// Construct a struct layout from owned children. pub fn new(row_count: u64, dtype: DType, children: Vec) -> Self { - validate_children(&dtype, children.len()).vortex_expect("invalid struct children"); + Self::validate_children(&dtype, children.len()).vortex_expect("invalid struct children"); LayoutParts::new( Struct, dtype, @@ -164,32 +165,35 @@ impl Layout { } Ok(()) } -} -fn validate_children(dtype: &DType, nchildren: usize) -> VortexResult<()> { - let fields = dtype - .as_struct_fields_opt() - .ok_or_else(|| vortex_err!("Expected struct dtype"))?; - let expected = fields.nfields() + usize::from(dtype.is_nullable()); - vortex_ensure!( + fn validate_children(dtype: &DType, nchildren: usize) -> VortexResult<()> { + let fields = dtype + .as_struct_fields_opt() + .ok_or_else(|| vortex_err!("Expected struct dtype"))?; + let expected = fields.nfields() + usize::from(dtype.is_nullable()); + vortex_ensure!( nchildren == expected, "Struct layout has {nchildren} children, expected {expected}" ); - Ok(()) -} + Ok(()) + } -fn struct_child_dtype(dtype: &DType, index: usize) -> VortexResult { - let schema_index = if dtype.is_nullable() { - index.saturating_sub(1) - } else { - index - }; - if index == 0 && dtype.is_nullable() { - Ok(DType::Bool(Nullability::NonNullable)) - } else { - dtype - .as_struct_fields_opt() - .and_then(|fields| fields.field_by_index(schema_index)) - .ok_or_else(|| vortex_err!("Missing field {schema_index}")) + fn child_dtype(dtype: &DType, index: usize) -> VortexResult { + let schema_index = if dtype.is_nullable() { + index.saturating_sub(1) + } else { + index + }; + if index == 0 && dtype.is_nullable() { + Ok(DType::Bool(Nullability::NonNullable)) + } else { + dtype + .as_struct_fields_opt() + .and_then(|fields| fields.field_by_index(schema_index)) + .ok_or_else(|| vortex_err!("Missing field {schema_index}")) + } } } + + + From 795dc416ee7e25d7e2c67a5834265e6329b644d8 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 23 Jul 2026 11:50:38 +0100 Subject: [PATCH 3/8] wip Signed-off-by: Joe Isaacs --- vortex-layout/src/display.rs | 10 +++---- vortex-layout/src/layouts/chunked/reader.rs | 3 ++- vortex-layout/src/layouts/flat/reader.rs | 8 +++--- vortex-layout/src/layouts/flat/writer.rs | 11 ++++---- vortex-layout/src/layouts/repartition.rs | 6 ++--- vortex-layout/src/layouts/row_idx/mod.rs | 8 +++--- vortex-layout/src/layouts/struct_/reader.rs | 11 ++++---- vortex-layout/src/layouts/table.rs | 29 +++++++++------------ vortex-layout/src/scan/split_by.rs | 4 +-- vortex-layout/src/test.rs | 7 +++-- 10 files changed, 49 insertions(+), 48 deletions(-) diff --git a/vortex-layout/src/display.rs b/vortex-layout/src/display.rs index 56f9691c722..20b13b7afbc 100644 --- a/vortex-layout/src/display.rs +++ b/vortex-layout/src/display.rs @@ -241,7 +241,7 @@ mod tests { use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; use crate::strategy::LayoutStrategy; - use crate::test::SESSION; + use crate::test::new_session; /// Test display_tree with inline array_tree metadata (no segment source needed). #[test] @@ -250,7 +250,7 @@ mod tests { if std::env::var("NEXTEST_RUN_ID").is_ok() { temp_env::with_var("FLAT_LAYOUT_INLINE_ARRAY_NODE", Some("1"), || { block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); @@ -348,7 +348,7 @@ vortex.struct, dtype: {numbers=i64?, strings=utf8}, children: 2, rows: 5 if std::env::var("NEXTEST_RUN_ID").is_ok() { temp_env::with_var("FLAT_LAYOUT_INLINE_ARRAY_NODE", None::<&str>, || { block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); @@ -418,7 +418,7 @@ vortex.chunked, dtype: i32, children: 2, rows: 10 // Create a simple primitive array let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5], Validity::AllValid); let layout = block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); FlatLayoutStrategy::default() .write_stream( ctx.clone(), @@ -463,7 +463,7 @@ vortex.flat, dtype: i32?, segment 0, buffers=[20B], total=20B // Create a simple primitive array let array = PrimitiveArray::new(buffer![10i64, 20, 30], Validity::NonNullable); let layout = block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); FlatLayoutStrategy::default() .write_stream( ctx, diff --git a/vortex-layout/src/layouts/chunked/reader.rs b/vortex-layout/src/layouts/chunked/reader.rs index 171bb138c6c..0e33f39df0c 100644 --- a/vortex-layout/src/layouts/chunked/reader.rs +++ b/vortex-layout/src/layouts/chunked/reader.rs @@ -380,6 +380,7 @@ mod test { use crate::sequence::SequentialStreamAdapter; use crate::sequence::SequentialStreamExt as _; use crate::test::SESSION; + use crate::test::new_session; #[fixture] /// Create a chunked layout with three chunks of primitive arrays. @@ -391,7 +392,7 @@ mod test { let (mut sequence_id, eof) = SequenceId::root().split(); let segments2 = Arc::::clone(&segments); let layout = block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); strategy .write_stream( ctx, diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index 42e1221af45..42e707155a8 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -251,12 +251,12 @@ mod test { use crate::segments::TestSegments; use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; - use crate::test::SESSION; + use crate::test::new_session; #[test] fn flat_identity() -> VortexResult<()> { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); let mut ctx = session.create_execution_ctx(); let array_ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); @@ -296,7 +296,7 @@ mod test { #[test] fn flat_expr() { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); let mut ctx = session.create_execution_ctx(); let array_ctx = ArrayContext::empty(); @@ -336,7 +336,7 @@ mod test { #[test] fn flat_unaligned_row_mask() { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); let mut ctx = session.create_execution_ctx(); let array_ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); diff --git a/vortex-layout/src/layouts/flat/writer.rs b/vortex-layout/src/layouts/flat/writer.rs index 723556b5f49..efb2c11cbe3 100644 --- a/vortex-layout/src/layouts/flat/writer.rs +++ b/vortex-layout/src/layouts/flat/writer.rs @@ -221,6 +221,7 @@ mod tests { use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; use crate::test::SESSION; + use crate::test::new_session; // Currently, flat layouts do not force compute stats during write, they only retain // pre-computed stats. @@ -228,7 +229,7 @@ mod tests { #[test] fn flat_stats() { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); @@ -266,7 +267,7 @@ mod tests { #[test] fn truncates_variable_size_stats() { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); @@ -328,7 +329,7 @@ mod tests { fn struct_array_round_trip() { block_on(|handle| async { let mut ctx_exec = array_session().create_execution_ctx(); - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); let mut validity_builder = BitBufferMut::with_capacity(2); validity_builder.append(true); validity_builder.append(false); @@ -414,7 +415,7 @@ mod tests { #[test] fn flat_invalid_array_fails() -> VortexResult<()> { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); let prim: PrimitiveArray = (0..10).collect(); let filter = prim.filter(Mask::from_indices(10, vec![2, 3]))?; @@ -454,7 +455,7 @@ mod tests { #[test] fn flat_valid_array_writes() -> VortexResult<()> { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); let codes: PrimitiveArray = (0u32..10).collect(); let values: PrimitiveArray = (0..10).collect(); let dict = DictArray::new(codes.into_array(), values.into_array()); diff --git a/vortex-layout/src/layouts/repartition.rs b/vortex-layout/src/layouts/repartition.rs index 89a1317d240..963bfb36f40 100644 --- a/vortex-layout/src/layouts/repartition.rs +++ b/vortex-layout/src/layouts/repartition.rs @@ -297,7 +297,7 @@ mod tests { use crate::segments::TestSegments; use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; - use crate::test::SESSION; + use crate::test::new_session; const ONE_MEG: u64 = 1 << 20; @@ -398,7 +398,7 @@ mod tests { let stream = fsl.into_array().to_array_stream().sequenced(ptr); let layout = block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); strategy .write_stream( ctx, @@ -463,7 +463,7 @@ mod tests { let stream = elements.into_array().to_array_stream().sequenced(ptr); let layout = block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); strategy .write_stream( ctx, diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index 4e5e8e494fa..fe94d1b9e77 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -347,12 +347,12 @@ mod tests { use crate::segments::TestSegments; use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; - use crate::test::SESSION; + use crate::test::new_session; #[test] fn flat_expr_no_row_id() { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); let mut ctx = session.create_execution_ctx(); let array_ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); @@ -397,7 +397,7 @@ mod tests { #[test] fn flat_expr_row_id() { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); let mut ctx = session.create_execution_ctx(); let array_ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); @@ -442,7 +442,7 @@ mod tests { #[test] fn flat_expr_or() { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); let mut ctx = session.create_execution_ctx(); let array_ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index a9623e61f8e..3d85e754cf4 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -434,6 +434,7 @@ mod tests { use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; use crate::test::SESSION; + use crate::test::new_session; #[fixture] fn empty_struct() -> (Arc, LayoutRef) { @@ -447,7 +448,7 @@ mod tests { ); let segments2 = Arc::::clone(&segments); let layout = block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); strategy .write_stream( ctx, @@ -484,7 +485,7 @@ mod tests { ); let segments2 = Arc::::clone(&segments); let layout = block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); strategy .write_stream( ctx, @@ -524,7 +525,7 @@ mod tests { ); let segments2 = Arc::::clone(&segments); let layout = block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); strategy .write_stream( ctx, @@ -569,7 +570,7 @@ mod tests { ); let segments2 = Arc::::clone(&segments); let layout = block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); strategy .write_stream( ctx, @@ -848,7 +849,7 @@ mod tests { ); let segments2 = Arc::::clone(&segments); let layout = block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); strategy .write_stream( ctx, diff --git a/vortex-layout/src/layouts/table.rs b/vortex-layout/src/layouts/table.rs index 0853478c14a..049051c4272 100644 --- a/vortex-layout/src/layouts/table.rs +++ b/vortex-layout/src/layouts/table.rs @@ -349,8 +349,6 @@ mod tests { use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; - use vortex_io::runtime::single::block_on; - use vortex_io::session::RuntimeSessionExt; use crate::LayoutRef; use crate::LayoutStrategy; @@ -658,9 +656,9 @@ mod tests { .with_field_writer(FieldPath::root(), flat); } - #[test] + #[tokio::test] #[should_panic(expected = "panic while transposing table stream")] - fn table_fanout_panic_propagates() { + async fn table_fanout_panic_propagates() { let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); let (_, eof) = SequenceId::root().split(); @@ -680,18 +678,15 @@ mod tests { Arc::new(FlatLayoutStrategy::default()), ); - block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); - strategy - .write_stream( - ctx, - segments, - SequentialStreamAdapter::new(dtype, stream).sendable(), - eof, - &session, - ) - .await - .unwrap(); - }); + strategy + .write_stream( + ctx, + segments, + SequentialStreamAdapter::new(dtype, stream).sendable(), + eof, + &SESSION, + ) + .await + .unwrap(); } } diff --git a/vortex-layout/src/scan/split_by.rs b/vortex-layout/src/scan/split_by.rs index 2af312f7537..6c2cde11b72 100644 --- a/vortex-layout/src/scan/split_by.rs +++ b/vortex-layout/src/scan/split_by.rs @@ -136,7 +136,6 @@ mod test { use vortex_array::expr::Expression; use vortex_buffer::buffer; use vortex_io::runtime::single::block_on; - use vortex_io::session::RuntimeSessionExt; use vortex_mask::Mask; use super::*; @@ -145,6 +144,7 @@ mod test { use crate::RowSplits; use crate::layouts::flat::writer::FlatLayoutStrategy; use crate::scan::test::SCAN_SESSION; + use crate::scan::test::session_with_handle; use crate::segments::TestSegments; use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; @@ -154,7 +154,7 @@ mod test { let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); let layout = block_on(|handle| async { - let session = SCAN_SESSION.clone().with_handle(handle); + let session = session_with_handle(handle); FlatLayoutStrategy::default() .write_stream( ctx, diff --git a/vortex-layout/src/test.rs b/vortex-layout/src/test.rs index 72258a1161e..15bb6eff543 100644 --- a/vortex-layout/src/test.rs +++ b/vortex-layout/src/test.rs @@ -4,12 +4,15 @@ use std::sync::LazyLock; use vortex_io::session::RuntimeSession; +use vortex_io::session::RuntimeSessionExt; use vortex_session::VortexSession; use crate::session::LayoutSession; -pub static SESSION: LazyLock = LazyLock::new(|| { +pub static SESSION: LazyLock = LazyLock::new(|| new_session().with_tokio()); + +pub fn new_session() -> VortexSession { vortex_array::array_session() .with::() .with::() -}); +} From 12fbe6725fb0cf9e78a61f67eadf53e9000f864c Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 23 Jul 2026 17:23:41 +0100 Subject: [PATCH 4/8] fix Signed-off-by: Joe Isaacs --- vortex-layout/src/layouts/struct_/mod.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/vortex-layout/src/layouts/struct_/mod.rs b/vortex-layout/src/layouts/struct_/mod.rs index d6e38844455..7e7809330c3 100644 --- a/vortex-layout/src/layouts/struct_/mod.rs +++ b/vortex-layout/src/layouts/struct_/mod.rs @@ -66,10 +66,9 @@ impl VTable for Struct { Layout::::validate_children(args.dtype, args.children.nchildren())?; for idx in 0..args.children.nchildren() { - let child_dtype = StructLayout::child_dtype(args.dtype, idx)?; - let child = args.children.child(idx, &child_dtype)?; + let child_row_count = args.children.child_row_count(idx); vortex_ensure!( - child.row_count() == args.row_count, + child_row_count == args.row_count, "Struct child {idx} row count does not match parent" ); } From 2b1f514503c4765662b2cbf4c3e8094c5c1d05ab Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 24 Jul 2026 11:08:29 +0100 Subject: [PATCH 5/8] fix Signed-off-by: Joe Isaacs --- vortex-layout/src/layouts/zoned/mod.rs | 116 ++++++++++++--------- vortex-layout/src/layouts/zoned/pruning.rs | 19 ++-- vortex-layout/src/layouts/zoned/reader.rs | 66 ++++++++---- 3 files changed, 122 insertions(+), 79 deletions(-) diff --git a/vortex-layout/src/layouts/zoned/mod.rs b/vortex-layout/src/layouts/zoned/mod.rs index 3a223e3d9b6..4b6707a43ce 100644 --- a/vortex-layout/src/layouts/zoned/mod.rs +++ b/vortex-layout/src/layouts/zoned/mod.rs @@ -86,35 +86,6 @@ pub type ZonedLayout = Layout; /// A legacy `vortex.stats` layout using the same runtime data. pub type LegacyStatsLayout = Layout; -/// Builds a reader for a zoned layout, bypassing [`ZonedReader`] when the zone map is empty -/// (`zone_len == 0`) and reading the data child directly, since there is nothing to prune with. -/// This covers both legacy zero-length zones and layouts whose aggregates the session cannot -/// reconstruct. -fn zoned_reader( - layout: &ZonedLayout, - name: Arc, - segment_source: Arc, - session: &VortexSession, - ctx: &crate::LayoutReaderContext, -) -> VortexResult { - if layout.zone_len == 0 { - return layout.children.child(0, &layout.dtype)?.new_reader( - name, - segment_source, - session, - ctx, - ); - } - - Ok(Arc::new(ZonedReader::try_new( - layout.clone(), - name, - segment_source, - session.clone(), - ctx.clone(), - )?)) -} - impl VTable for Zoned { type LayoutData = ZonedData; type Metadata = ZonedMetadata; @@ -195,7 +166,7 @@ impl VTable for Zoned { session: &VortexSession, ctx: &LayoutReaderContext, ) -> VortexResult { - zoned_reader(layout, name, segment_source, session, ctx) + layout.zoned_reader(name, segment_source, session, ctx) } } @@ -263,11 +234,39 @@ impl VTable for LegacyStats { session: &VortexSession, ctx: &LayoutReaderContext, ) -> VortexResult { - zoned_reader(&layout.0, name, segment_source, session, ctx) + layout.zoned_reader(name, segment_source, session, ctx) } } impl LegacyStatsLayout { + fn zoned_reader( + &self, + name: Arc, + segment_source: Arc, + session: &VortexSession, + ctx: &LayoutReaderContext, + ) -> VortexResult { + if self.zone_len == 0 { + return self + .child(0)? + .new_reader(name, segment_source, session, ctx); + } + + Ok(Arc::new(ZonedReader::try_new( + self.clone(), + self.nzones(), + name, + segment_source, + session.clone(), + ctx.clone(), + )?)) + } + + pub fn nzones(&self) -> usize { + usize::try_from(self.children().child_row_count(1)) + .vortex_expect("Invalid number of zones, cannot handle more than usize zones") + } + /// Returns display names for the zone-map aggregates stored by this layout. pub fn present_aggregates(&self) -> Arc<[String]> { present_aggregates(&self.zone_map_schema) @@ -314,10 +313,6 @@ impl ZonedLayout { .into_typed()) } - pub fn nzones(&self) -> usize { - usize::try_from(self.child_row_count(1)).vortex_expect("Invalid number of zones") - } - pub fn zone_len(&self) -> usize { self.zone_len } @@ -327,27 +322,50 @@ impl ZonedLayout { present_aggregates(&self.zone_map_schema) } - pub(super) fn aggregate_fns( + /// Builds a reader for a zoned layout, bypassing [`ZonedReader`] when the zone map is empty + /// (`zone_len == 0`) and reading the data child directly, since there is nothing to prune with. + /// This covers both legacy zero-length zones and layouts whose aggregates the session cannot + /// reconstruct. + fn zoned_reader( &self, - _session: &VortexSession, - ) -> VortexResult> { + name: Arc, + segment_source: Arc, + session: &VortexSession, + ctx: &LayoutReaderContext, + ) -> VortexResult { + if self.zone_len == 0 { + return self + .child(0)? + .new_reader(name, segment_source, session, ctx); + } + + Ok(Arc::new(ZonedReader::try_new( + self.clone(), + self.nzones(), + name, + segment_source, + session.clone(), + ctx.clone(), + )?)) + } + + pub fn nzones(&self) -> usize { + usize::try_from(self.children().child_row_count(1)) + .vortex_expect("Invalid number of zones, cannot handle more than usize zones") + } +} + +impl ZonedData { + fn aggregate_fns(&self) -> Arc<[AggregateFnRef]> { match &self.zone_map_schema { - ZoneMapSchema::LegacyStats(stats) => Ok(stats + ZoneMapSchema::LegacyStats(stats) => stats .iter() .filter_map(Stat::aggregate_fn) .collect::>() - .into()), - ZoneMapSchema::AggregateFns(aggregate_fns) => Ok(Arc::clone(aggregate_fns)), + .into(), + ZoneMapSchema::AggregateFns(aggregate_fns) => Arc::clone(aggregate_fns), } } - - pub(super) fn stats_table_dtype_for(&self, aggregate_fns: &[AggregateFnRef]) -> DType { - if let ZoneMapSchema::LegacyStats(stats) = &self.zone_map_schema { - return legacy_stats_table_dtype(self.dtype(), stats); - } - - aggregate_stats_table_dtype(self.dtype(), aggregate_fns) - } } fn present_aggregates(schema: &ZoneMapSchema) -> Arc<[String]> { diff --git a/vortex-layout/src/layouts/zoned/pruning.rs b/vortex-layout/src/layouts/zoned/pruning.rs index b517df985c9..aaf11150bde 100644 --- a/vortex-layout/src/layouts/zoned/pruning.rs +++ b/vortex-layout/src/layouts/zoned/pruning.rs @@ -28,9 +28,9 @@ use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_utils::aliases::dash_map::DashMap; -use crate::LazyReaderChildren; -use crate::layouts::zoned::ZonedLayout; +use crate::layouts::zoned::ZonedData; use crate::layouts::zoned::zone_map::ZoneMap; +use crate::{Layout, LazyReaderChildren, VTable}; type SharedZoneMap = Shared>>; pub(super) type SharedPruningResult = @@ -45,23 +45,26 @@ pub(super) struct PruningState { aggregate_fns: Arc<[AggregateFnRef]>, lazy_children: Arc, session: VortexSession, - pruning_result: LazyLock>>, zone_map: OnceLock, pruning_predicates: LazyLock>>, } impl PruningState { - pub(super) fn new( - layout: &ZonedLayout, + pub(super) fn new( + layout: Layout, + zone_count: usize, aggregate_fns: Arc<[AggregateFnRef]>, lazy_children: Arc, session: VortexSession, - ) -> Self { + ) -> Self + where + V: VTable, + { Self { - zone_count: layout.nzones(), + zone_count, row_count: layout.row_count(), - zone_len: layout.zone_len() as u64, + zone_len: layout.zone_len as u64, dtype: layout.dtype().clone(), aggregate_fns, lazy_children, diff --git a/vortex-layout/src/layouts/zoned/reader.rs b/vortex-layout/src/layouts/zoned/reader.rs index d1585b0d1cc..4c0e3bc73e4 100644 --- a/vortex-layout/src/layouts/zoned/reader.rs +++ b/vortex-layout/src/layouts/zoned/reader.rs @@ -19,35 +19,40 @@ use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; +use crate::Layout; use crate::LayoutReader; use crate::LayoutReaderRef; use crate::LazyReaderChildren; use crate::RowSplits; use crate::SplitRange; -use crate::layouts::zoned::ZonedLayout; +use crate::VTable; +use crate::layouts::zoned::ZonedData; use crate::layouts::zoned::pruning::PruningState; use crate::segments::SegmentSource; pub struct ZonedReader { - layout: ZonedLayout, + dtype: DType, + row_count: u64, + zone_len: usize, name: Arc, lazy_children: Arc, pruning: PruningState, } impl ZonedReader { - pub(super) fn try_new( - layout: ZonedLayout, + pub(super) fn try_new( + layout: Layout, + zone_count: usize, name: Arc, segment_source: Arc, session: VortexSession, ctx: crate::LayoutReaderContext, - ) -> VortexResult { - let aggregate_fns = layout.aggregate_fns(&session)?; - let dtypes = vec![ - layout.dtype().clone(), - layout.stats_table_dtype_for(&aggregate_fns), - ]; + ) -> VortexResult + where + V: VTable, + { + let aggregate_fns = layout.aggregate_fns(); + let dtypes = vec![layout.dtype().clone(), layout.stats_table_dtype.clone()]; let names = vec![Arc::clone(&name), format!("{}.zones", name).into()]; let lazy_children = Arc::new(LazyReaderChildren::new( Arc::clone(layout.children()), @@ -57,10 +62,21 @@ impl ZonedReader { session.clone(), ctx, )); + let dtype = layout.dtype().clone(); + let row_count = layout.row_count(); + let zone_len = layout.zone_len; Ok(Self { - pruning: PruningState::new(&layout, aggregate_fns, Arc::clone(&lazy_children), session), - layout, + pruning: PruningState::new( + layout, + zone_count, + aggregate_fns, + Arc::clone(&lazy_children), + session, + ), + dtype, + row_count, + zone_len, name, lazy_children, }) @@ -74,9 +90,9 @@ impl ZonedReader { pub(crate) fn zone_range(&self, row_range: &Range) -> Range { // Callers rely on `zone_len > 0`. `new_reader` never constructs a `ZonedReader` for a // zero-length zone map (it reads the data child directly), so this holds by construction. - debug_assert!(self.layout.zone_len > 0, "zone_len must be > 0"); + debug_assert!(self.zone_len > 0, "zone_len must be > 0"); - let zone_len_u64 = self.layout.zone_len as u64; + let zone_len_u64 = self.zone_len as u64; let zone_start = row_range.start / zone_len_u64; let zone_end = row_range.end.div_ceil(zone_len_u64); zone_start..zone_end @@ -85,8 +101,8 @@ impl ZonedReader { /// Get the row index for the first row in a zone with the given `zone_index`. pub(crate) fn first_row_offset(&self, zone_idx: u64) -> u64 { zone_idx - .saturating_mul(self.layout.zone_len as u64) - .min(self.layout.row_count()) + .saturating_mul(self.zone_len as u64) + .min(self.row_count) } } @@ -100,11 +116,11 @@ impl LayoutReader for ZonedReader { } fn dtype(&self) -> &DType { - self.layout.dtype() + &self.dtype } fn row_count(&self) -> u64 { - self.layout.row_count() + self.row_count } fn register_splits( @@ -422,8 +438,15 @@ mod test { } #[rstest] - fn test_legacy_zero_zone_len_skips_zoned_pruning( + #[case::zero_zone_len(0, [true; 9])] + #[case::zoned_reader( + 3, + [false, false, false, false, false, false, true, true, true] + )] + fn test_legacy_zoned_reader( #[from(stats_layout)] (segments, layout): (Arc, LayoutRef), + #[case] zone_len: u32, + #[case] expected: [bool; 9], ) -> VortexResult<()> { let zoned_layout = layout.as_::(); let children = @@ -439,7 +462,7 @@ mod test { layout.dtype(), layout.row_count(), &LegacyStatsMetadata { - zone_len: 0, + zone_len, zone_map_schema: zoned_layout.zone_map_schema.clone(), }, vec![], @@ -462,9 +485,8 @@ mod test { )? .await?; - assert!(result.all_true()); + assert_eq!(result, Mask::from_iter(expected)); - // The bypass returns the data child's reader, so projection reads the underlying data. let projected = reader .projection_evaluation( &(0..row_count), From e3b803eae037ce98443c9d6f8ccaf875206b3fa3 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 24 Jul 2026 11:36:54 +0100 Subject: [PATCH 6/8] fix Signed-off-by: Joe Isaacs --- vortex-layout/src/layouts/zoned/pruning.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vortex-layout/src/layouts/zoned/pruning.rs b/vortex-layout/src/layouts/zoned/pruning.rs index aaf11150bde..7764378eb05 100644 --- a/vortex-layout/src/layouts/zoned/pruning.rs +++ b/vortex-layout/src/layouts/zoned/pruning.rs @@ -28,9 +28,11 @@ use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_utils::aliases::dash_map::DashMap; +use crate::Layout; +use crate::LazyReaderChildren; +use crate::VTable; use crate::layouts::zoned::ZonedData; use crate::layouts::zoned::zone_map::ZoneMap; -use crate::{Layout, LazyReaderChildren, VTable}; type SharedZoneMap = Shared>>; pub(super) type SharedPruningResult = From a4e6ca6ed328b385cfaafdda5a52ddbe7b3306e8 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 24 Jul 2026 13:40:33 +0100 Subject: [PATCH 7/8] update Signed-off-by: Joe Isaacs --- vortex-layout/src/layout.rs | 75 +------------------------------- vortex-layout/src/lib.rs | 2 + vortex-layout/src/vtable.rs | 86 +++++++++++++++++++++++++++++++++++++ 3 files changed, 89 insertions(+), 74 deletions(-) create mode 100644 vortex-layout/src/vtable.rs diff --git a/vortex-layout/src/layout.rs b/vortex-layout/src/layout.rs index f4cd4ecca02..fadec7b93c5 100644 --- a/vortex-layout/src/layout.rs +++ b/vortex-layout/src/layout.rs @@ -17,8 +17,6 @@ use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::Id; -use crate::LayoutBuildContext; -use crate::LayoutDeserializeArgs; use crate::LayoutReaderContext; use crate::LayoutReaderRef; use crate::children::LayoutChildren; @@ -26,82 +24,11 @@ use crate::display::DisplayLayoutTree; use crate::display::display_tree_with_segment_sizes; use crate::segments::SegmentId; use crate::segments::SegmentSource; +use crate::vtable::{LayoutRef, VTable}; /// A unique identifier for a layout encoding. pub type LayoutId = Id; -/// Shared, erased handle to a layout tree node. -pub type LayoutRef = Arc; - -/// Layout-specific behavior for a typed [`Layout`]. -/// -/// Common serialized fields are stored by [`Layout`]. Implementations own only their -/// layout-specific data, metadata codec, child typing, and reader construction. -pub trait VTable: 'static + Clone + Send + Sync + Debug { - /// Layout-specific data. - type LayoutData: 'static + Clone + Send + Sync + Debug; - /// Serialized metadata type. - type Metadata: SerializeMetadata + vortex_array::DeserializeMetadata + Debug; - - /// Returns the globally unique layout ID. - fn id(&self) -> LayoutId; - - /// Returns the serializable metadata for a layout. - fn metadata(layout: &Layout) -> Self::Metadata; - - /// Deserialize and validate layout-specific data. - fn deserialize( - &self, - args: &LayoutDeserializeArgs<'_>, - metadata: &::Output, - ) -> VortexResult; - - /// Construct a typed layout from deserialized common fields. - fn build( - vtable: &Self, - dtype: &DType, - row_count: u64, - metadata: &::Output, - segment_ids: Vec, - children: &dyn LayoutChildren, - build_ctx: &LayoutBuildContext<'_>, - ) -> VortexResult> { - let args = LayoutDeserializeArgs { - session: build_ctx.session, - array_read_ctx: build_ctx.array_read_ctx, - dtype, - row_count, - segment_ids, - children, - }; - let data = vtable.deserialize(&args, metadata)?; - Ok(LayoutParts::new( - vtable.clone(), - dtype.clone(), - row_count, - args.segment_ids, - children.to_arc(), - data, - ) - .into_typed()) - } - - /// Returns the expected dtype of child `idx`. - fn child_dtype(layout: &Layout, idx: usize) -> VortexResult; - - /// Returns the relationship between child `idx` and its parent. - fn child_type(layout: &Layout, idx: usize) -> LayoutChildType; - - /// Construct a reader for this layout. - fn new_reader( - layout: &Layout, - name: Arc, - segment_source: Arc, - session: &VortexSession, - ctx: &LayoutReaderContext, - ) -> VortexResult; -} - /// Pieces used to construct a typed layout. pub struct LayoutParts { vtable: V, diff --git a/vortex-layout/src/lib.rs b/vortex-layout/src/lib.rs index fd19a6c6fa9..779076bddac 100644 --- a/vortex-layout/src/lib.rs +++ b/vortex-layout/src/lib.rs @@ -24,6 +24,7 @@ pub use reader::*; pub use reader_context::*; pub use strategy::*; use vortex_session::registry::Context; +pub use vtable::*; pub mod aliases; mod children; pub mod display; @@ -39,5 +40,6 @@ pub mod session; mod strategy; #[cfg(test)] mod test; +mod vtable; pub type LayoutContext = Context; diff --git a/vortex-layout/src/vtable.rs b/vortex-layout/src/vtable.rs new file mode 100644 index 00000000000..2eb0162b858 --- /dev/null +++ b/vortex-layout/src/vtable.rs @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use crate::segments::{SegmentId, SegmentSource}; +use crate::{ + DynLayout, Layout, LayoutBuildContext, LayoutChildType, LayoutChildren, LayoutDeserializeArgs, + LayoutId, LayoutParts, LayoutReaderContext, LayoutReaderRef, +}; +use std::fmt::Debug; +use std::sync::Arc; +use vortex_array::SerializeMetadata; +use vortex_array::dtype::DType; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +/// Shared, erased handle to a layout tree node. +pub type LayoutRef = Arc; + +/// Layout-specific behavior for a typed [`Layout`]. +/// +/// Common serialized fields are stored by [`Layout`]. Implementations own only their +/// layout-specific data, metadata codec, child typing, and reader construction. +pub trait VTable: 'static + Clone + Send + Sync + Debug { + /// Layout-specific data. + type LayoutData: 'static + Clone + Send + Sync + Debug; + /// Serialized metadata type. + type Metadata: SerializeMetadata + vortex_array::DeserializeMetadata + Debug; + + /// Returns the globally unique layout ID. + fn id(&self) -> LayoutId; + + /// Returns the serializable metadata for a layout. + fn metadata(layout: &Layout) -> Self::Metadata; + + /// Deserialize and validate layout-specific data. + fn deserialize( + &self, + args: &LayoutDeserializeArgs<'_>, + metadata: &::Output, + ) -> VortexResult; + + /// Construct a typed layout from deserialized common fields. + fn build( + vtable: &Self, + dtype: &DType, + row_count: u64, + metadata: &::Output, + segment_ids: Vec, + children: &dyn LayoutChildren, + build_ctx: &LayoutBuildContext<'_>, + ) -> VortexResult> { + let args = LayoutDeserializeArgs { + session: build_ctx.session, + array_read_ctx: build_ctx.array_read_ctx, + dtype, + row_count, + segment_ids, + children, + }; + let data = vtable.deserialize(&args, metadata)?; + Ok(LayoutParts::new( + vtable.clone(), + dtype.clone(), + row_count, + args.segment_ids, + children.to_arc(), + data, + ) + .into_typed()) + } + + /// Returns the expected dtype of child `idx`. + fn child_dtype(layout: &Layout, idx: usize) -> VortexResult; + + /// Returns the relationship between child `idx` and its parent. + fn child_type(layout: &Layout, idx: usize) -> LayoutChildType; + + /// Construct a reader for this layout. + fn new_reader( + layout: &Layout, + name: Arc, + segment_source: Arc, + session: &VortexSession, + ctx: &LayoutReaderContext, + ) -> VortexResult; +} From adc3f2150aedfaaef49de8263a45be63fe7500c8 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 24 Jul 2026 13:45:13 +0100 Subject: [PATCH 8/8] update Signed-off-by: Joe Isaacs --- vortex-layout/src/layout.rs | 3 ++- vortex-layout/src/vtable.rs | 19 ++++++++++++++----- 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/vortex-layout/src/layout.rs b/vortex-layout/src/layout.rs index fadec7b93c5..a51ff96e1ed 100644 --- a/vortex-layout/src/layout.rs +++ b/vortex-layout/src/layout.rs @@ -24,7 +24,8 @@ use crate::display::DisplayLayoutTree; use crate::display::display_tree_with_segment_sizes; use crate::segments::SegmentId; use crate::segments::SegmentSource; -use crate::vtable::{LayoutRef, VTable}; +use crate::vtable::LayoutRef; +use crate::vtable::VTable; /// A unique identifier for a layout encoding. pub type LayoutId = Id; diff --git a/vortex-layout/src/vtable.rs b/vortex-layout/src/vtable.rs index 2eb0162b858..369107a057f 100644 --- a/vortex-layout/src/vtable.rs +++ b/vortex-layout/src/vtable.rs @@ -1,18 +1,27 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use crate::segments::{SegmentId, SegmentSource}; -use crate::{ - DynLayout, Layout, LayoutBuildContext, LayoutChildType, LayoutChildren, LayoutDeserializeArgs, - LayoutId, LayoutParts, LayoutReaderContext, LayoutReaderRef, -}; use std::fmt::Debug; use std::sync::Arc; + use vortex_array::SerializeMetadata; use vortex_array::dtype::DType; use vortex_error::VortexResult; use vortex_session::VortexSession; +use crate::DynLayout; +use crate::Layout; +use crate::LayoutBuildContext; +use crate::LayoutChildType; +use crate::LayoutChildren; +use crate::LayoutDeserializeArgs; +use crate::LayoutId; +use crate::LayoutParts; +use crate::LayoutReaderContext; +use crate::LayoutReaderRef; +use crate::segments::SegmentId; +use crate::segments::SegmentSource; + /// Shared, erased handle to a layout tree node. pub type LayoutRef = Arc;