From 48d2004bf2f060c4f69cec925ec52dc7f12dba4c Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 24 Jul 2026 16:09:13 +0100 Subject: [PATCH 1/3] wip Signed-off-by: Joe Isaacs --- vortex-array-macros/src/lib.rs | 253 ++++++++++++++++-- .../fns/uncompressed_size_in_bytes/union.rs | 1 + vortex-array/src/array/mod.rs | 58 ++++ vortex-array/src/arrays/chunked/array.rs | 46 ++-- .../src/arrays/chunked/compute/filter.rs | 4 +- .../src/arrays/chunked/compute/mask.rs | 2 +- .../src/arrays/chunked/compute/take.rs | 2 +- vortex-array/src/arrays/chunked/tests.rs | 2 +- .../src/arrays/filter/execute/union.rs | 1 + vortex-array/src/arrays/masked/execute.rs | 3 +- vortex-array/src/arrays/patched/array.rs | 67 +++++ vortex-array/src/arrays/struct_/array.rs | 22 +- vortex-array/src/arrays/union/array.rs | 48 ++-- vortex-array/src/arrays/union/compute/mask.rs | 3 +- .../src/arrays/union/compute/slice.rs | 1 + vortex-array/src/arrays/union/mod.rs | 3 + vortex-array/src/arrays/union/vtable/mod.rs | 23 +- .../src/arrays/union/vtable/operations.rs | 1 + .../src/arrays/union/vtable/validate.rs | 6 +- .../src/arrays/union/vtable/validity.rs | 2 +- vortex-compressor/src/compressor/cascade.rs | 1 + 21 files changed, 455 insertions(+), 94 deletions(-) diff --git a/vortex-array-macros/src/lib.rs b/vortex-array-macros/src/lib.rs index 3105dac4add..1522cc75ed9 100644 --- a/vortex-array-macros/src/lib.rs +++ b/vortex-array-macros/src/lib.rs @@ -19,7 +19,8 @@ use syn::spanned::Spanned; /// Generate slot index constants, a borrowed view struct, and a typed ext trait /// from a slot struct definition. /// -/// Fields must be `ArrayRef` (required slot) or `Option` (optional slot). +/// Fields must be `ArrayRef` (required slot), `Option` (optional slot), or — +/// for the final field only — `Vec` (variadic tail of required slots). /// Field declaration order determines slot indices. /// /// # Example @@ -94,6 +95,61 @@ use syn::spanned::Spanned; /// /// The underlying storage is always `ArraySlots` — the field type only /// controls whether the macro inserts a `.vortex_expect()` unwrap or not. +/// +/// # Variadic tail slots +/// +/// The final field may be `Vec`, declaring that every slot from that +/// position onward belongs to a homogeneous, variable-length run of required slots. +/// This supports encodings like `Chunked` (`[chunk_offsets, chunks...]`), +/// `Struct` (`[validity?, fields...]`), and `Union` (`[type_ids, children...]`). +/// +/// ```ignore +/// #[array_slots(Chunked)] +/// pub struct ChunkedSlots { +/// pub chunk_offsets: ArrayRef, +/// pub chunks: Vec, +/// } +/// ``` +/// +/// For a struct with a variadic tail, the macro generates a different set of +/// constants — slot count is no longer a compile-time constant: +/// +/// ```ignore +/// impl ChunkedSlots { +/// pub const CHUNK_OFFSETS: usize = 0; +/// /// Offset at which the `chunks` slots begin. +/// pub const CHUNKS_OFFSET: usize = 1; +/// /// Number of fixed (non-variadic) slots. +/// pub const FIXED_COUNT: usize = 1; +/// /// Names of the fixed slots in storage order. +/// pub const FIXED_NAMES: [&'static str; 1] = ["chunk_offsets"]; +/// +/// /// Name of the slot at `idx`, e.g. "chunk_offsets" or "chunks[3]". +/// pub fn slot_name(idx: usize) -> String { ... } +/// +/// pub fn from_slots(slots: ArraySlots) -> Self { ... } +/// pub fn into_slots(self) -> ArraySlots { ... } +/// } +/// ``` +/// +/// The view field and ext trait accessor for the tail are a [`SlotSlice`], a +/// borrowed run of required slots supporting `len()`, `get()`, `iter()`, and +/// indexing: +/// +/// ```ignore +/// pub struct ChunkedSlotsView<'a> { +/// pub chunk_offsets: &'a ArrayRef, +/// pub chunks: SlotSlice<'a>, +/// } +/// +/// pub trait ChunkedArraySlotsExt: TypedArrayRef { +/// fn chunk_offsets(&self) -> &ArrayRef { ... } +/// fn chunks(&self) -> SlotSlice<'_> { ... } +/// fn slots_view(&self) -> ChunkedSlotsView<'_> { ... } +/// } +/// ``` +/// +/// [`SlotSlice`]: ::vortex_array::SlotSlice #[proc_macro_attribute] pub fn array_slots(attr: TokenStream, item: TokenStream) -> TokenStream { let encoding = parse_macro_input!(attr as Path); @@ -143,15 +199,32 @@ fn expand_array_slots( .map(|(index, field)| SlotField::new(field, index, struct_ident)) .collect::>>()?; - let idx_consts = field_specs.iter().map(SlotField::idx_const); + // A variadic tail is only permitted as the final field, at most once. + for spec in field_specs.iter().rev().skip(1) { + if matches!(spec.slot_type, SlotFieldType::VariadicTail) { + return Err(syn::Error::new( + spec.field_ident.span(), + "#[array_slots] only the final field may be a variadic Vec tail", + )); + } + } + + let (fixed_specs, tail_spec) = match field_specs.split_last() { + Some((last, rest)) if matches!(last.slot_type, SlotFieldType::VariadicTail) => { + (rest, Some(last)) + } + _ => (field_specs.as_slice(), None), + }; + + let idx_consts = fixed_specs.iter().map(SlotField::idx_const); let view_fields = field_specs.iter().map(SlotField::view_field); let view_from_slots = field_specs.iter().map(SlotField::view_from_slots); let view_to_owned = field_specs.iter().map(SlotField::view_to_owned); - let owned_from_slots = field_specs.iter().map(SlotField::owned_from_slots); - let into_slots = field_specs.iter().map(SlotField::storage_slot); let ext_methods = field_specs.iter().map(SlotField::ext_method); - let slot_names = field_specs.iter().map(|field| field.slot_name.as_str()); - let slot_count = field_specs.len(); + + let counts = gen_counts(fixed_specs, tail_spec); + let from_slots = gen_from_slots(fixed_specs, tail_spec); + let into_slots = gen_into_slots(fixed_specs, tail_spec); Ok(quote! { #item_struct @@ -159,23 +232,13 @@ fn expand_array_slots( impl #struct_ident { #(#idx_consts)* - #[doc = "Total number of slots."] - pub const COUNT: usize = #slot_count; - - #[doc = "Slot names in storage order."] - pub const NAMES: [&'static str; #slot_count] = [#(#slot_names),*]; + #counts #[doc = "Convert owned slot storage into an owned slot struct."] - pub fn from_slots(mut slots: ::vortex_array::ArraySlots) -> Self { - Self { - #(#owned_from_slots,)* - } - } + #from_slots #[doc = "Convert this slot struct into storage order."] - pub fn into_slots(self) -> ::vortex_array::ArraySlots { - ::vortex_array::smallvec::smallvec![#(#into_slots),*] - } + #into_slots } #[derive(Clone, Copy, Debug)] @@ -214,6 +277,110 @@ fn expand_array_slots( }) } +fn gen_counts( + fixed_specs: &[SlotField], + tail_spec: Option<&SlotField>, +) -> proc_macro2::TokenStream { + let names = fixed_specs.iter().map(|field| field.slot_name.as_str()); + let fixed_count = fixed_specs.len(); + + match tail_spec { + None => quote! { + #[doc = "Total number of slots."] + pub const COUNT: usize = #fixed_count; + + #[doc = "Slot names in storage order."] + pub const NAMES: [&'static str; #fixed_count] = [#(#names),*]; + }, + Some(tail) => { + let offset_const = &tail.const_ident; + let tail_name = &tail.slot_name; + quote! { + #[doc = concat!("Offset at which the `", #tail_name, "` slots begin.")] + pub const #offset_const: usize = #fixed_count; + + #[doc = "Number of fixed (non-variadic) slots."] + pub const FIXED_COUNT: usize = #fixed_count; + + #[doc = "Names of the fixed slots in storage order."] + pub const FIXED_NAMES: [&'static str; #fixed_count] = [#(#names),*]; + + #[doc = "Name of the slot at the given index."] + pub fn slot_name(idx: usize) -> String { + if idx < Self::FIXED_COUNT { + Self::FIXED_NAMES[idx].to_string() + } else { + format!(concat!(#tail_name, "[{}]"), idx - Self::#offset_const) + } + } + } + } + } +} + +fn gen_from_slots( + fixed_specs: &[SlotField], + tail_spec: Option<&SlotField>, +) -> proc_macro2::TokenStream { + let owned_from_slots = fixed_specs.iter().map(SlotField::owned_from_slots); + + match tail_spec { + None => quote! { + pub fn from_slots(mut slots: ::vortex_array::ArraySlots) -> Self { + Self { + #(#owned_from_slots,)* + } + } + }, + Some(tail) => { + let tail_ident = &tail.field_ident; + let offset_const = &tail.const_ident; + let expect_message = &tail.expect_message; + quote! { + pub fn from_slots(mut slots: ::vortex_array::ArraySlots) -> Self { + let __variadic_tail: ::std::vec::Vec<::vortex_array::ArrayRef> = slots + .drain(Self::#offset_const..) + .map(|slot| ::vortex_error::VortexExpect::vortex_expect( + slot, + #expect_message, + )) + .collect(); + Self { + #(#owned_from_slots,)* + #tail_ident: __variadic_tail, + } + } + } + } + } +} + +fn gen_into_slots( + fixed_specs: &[SlotField], + tail_spec: Option<&SlotField>, +) -> proc_macro2::TokenStream { + let fixed_into_slots = fixed_specs.iter().map(SlotField::storage_slot); + + match tail_spec { + None => quote! { + pub fn into_slots(self) -> ::vortex_array::ArraySlots { + ::vortex_array::smallvec::smallvec![#(#fixed_into_slots),*] + } + }, + Some(tail) => { + let tail_ident = &tail.field_ident; + quote! { + pub fn into_slots(self) -> ::vortex_array::ArraySlots { + let mut slots: ::vortex_array::ArraySlots = + ::vortex_array::smallvec::smallvec![#(#fixed_into_slots),*]; + slots.extend(self.#tail_ident.into_iter().map(Some)); + slots + } + } + } + } +} + struct SlotField { field_ident: Ident, field_vis: Visibility, @@ -232,8 +399,13 @@ impl SlotField { .clone() .ok_or_else(|| syn::Error::new(field.span(), "slot fields must be named"))?; let field_name = ident_name(&field_ident); - let const_ident = format_ident!("{}", to_screaming_snake_case(&field_name)); let slot_type = SlotFieldType::from_syn_type(&field.ty)?; + let const_ident = match slot_type { + SlotFieldType::VariadicTail => { + format_ident!("{}_OFFSET", to_screaming_snake_case(&field_name)) + } + _ => format_ident!("{}", to_screaming_snake_case(&field_name)), + }; let expect_message = syn::LitStr::new( &format!("{} {} slot", ident_name(struct_ident), field_name), field.span(), @@ -288,6 +460,12 @@ impl SlotField { SlotFieldType::Optional => quote! { #field_ident: slots[#struct_ident::#const_ident].as_ref() }, + SlotFieldType::VariadicTail => quote! { + #field_ident: ::vortex_array::SlotSlice::new( + &slots[#struct_ident::#const_ident..], + #expect_message, + ) + }, } } @@ -301,6 +479,9 @@ impl SlotField { SlotFieldType::Optional => quote! { #field_ident: self.#field_ident.cloned() }, + SlotFieldType::VariadicTail => quote! { + #field_ident: self.#field_ident.to_vec() + }, } } @@ -320,6 +501,9 @@ impl SlotField { SlotFieldType::Optional => quote! { #field_ident: slots[#struct_ident::#const_ident].take() }, + SlotFieldType::VariadicTail => { + unreachable!("variadic tail is drained before fixed fields") + } } } @@ -333,6 +517,9 @@ impl SlotField { SlotFieldType::Optional => quote! { self.#field_ident }, + SlotFieldType::VariadicTail => { + unreachable!("variadic tail is appended after fixed fields") + } } } @@ -358,6 +545,15 @@ impl SlotField { self.as_ref().slots()[#struct_ident::#const_ident].as_ref() } }, + SlotFieldType::VariadicTail => quote! { + #[inline] + fn #field_ident(&self) -> ::vortex_array::SlotSlice<'_> { + ::vortex_array::SlotSlice::new( + &self.as_ref().slots()[#struct_ident::#const_ident..], + #expect_message, + ) + } + }, } } } @@ -366,6 +562,7 @@ impl SlotField { enum SlotFieldType { Required, Optional, + VariadicTail, } impl SlotFieldType { @@ -374,15 +571,22 @@ impl SlotFieldType { return Ok(Self::Required); } - if let Some(inner_ty) = option_inner_type(ty) + if let Some(inner_ty) = wrapper_inner_type(ty, "Option") && is_array_ref_type(inner_ty) { return Ok(Self::Optional); } + if let Some(inner_ty) = wrapper_inner_type(ty, "Vec") + && is_array_ref_type(inner_ty) + { + return Ok(Self::VariadicTail); + } + Err(syn::Error::new( ty.span(), - "#[array_slots] fields must be ArrayRef or Option", + "#[array_slots] fields must be ArrayRef, Option, or (final field only) \ + Vec", )) } @@ -390,6 +594,7 @@ impl SlotFieldType { match self { Self::Required => quote! { &'a ::vortex_array::ArrayRef }, Self::Optional => quote! { Option<&'a ::vortex_array::ArrayRef> }, + Self::VariadicTail => quote! { ::vortex_array::SlotSlice<'a> }, } } } @@ -407,12 +612,12 @@ fn is_array_ref_type(ty: &Type) -> bool { ) } -fn option_inner_type(ty: &Type) -> Option<&Type> { +fn wrapper_inner_type<'a>(ty: &'a Type, wrapper: &str) -> Option<&'a Type> { let Type::Path(type_path) = ty else { return None; }; let segment = type_path.path.segments.last()?; - if segment.ident != "Option" { + if segment.ident != wrapper { return None; } diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/union.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/union.rs index 2f13252de9d..cf43ed4e186 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/union.rs +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/union.rs @@ -8,6 +8,7 @@ use super::uncompressed_size_in_bytes_u64; use crate::ExecutionCtx; use crate::arrays::UnionArray; use crate::arrays::union::UnionArrayExt; +use crate::arrays::union::UnionArraySlotsExt; pub(super) fn union_uncompressed_size_in_bytes( array: &UnionArray, diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index e34accebfb2..a65c8ecf708 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -53,6 +53,64 @@ use crate::hash::ArrayHash; /// heap allocation in the common case. pub type ArraySlots = SmallVec<[Option; 4]>; +/// A borrowed run of required slots, e.g. the variadic tail of a slot layout. +/// +/// Wraps a `&[Option]` whose entries are guaranteed present by encoding +/// validation, exposing them as `&ArrayRef` without per-call-site unwrapping. +#[derive(Clone, Copy, Debug)] +pub struct SlotSlice<'a> { + slots: &'a [Option], + expect: &'static str, +} + +impl<'a> SlotSlice<'a> { + /// Wrap a slice of slots that validation guarantees are all present. + /// + /// `expect` names the slot run in the panic message if a slot is unexpectedly absent. + pub fn new(slots: &'a [Option], expect: &'static str) -> Self { + Self { slots, expect } + } + + /// The number of slots in the run. + #[allow(clippy::len_without_is_empty)] + pub fn len(&self) -> usize { + self.slots.len() + } + + /// Returns `true` if the run contains no slots. + pub fn is_empty(&self) -> bool { + self.slots.is_empty() + } + + /// Returns the slot at `idx`, or `None` if out of bounds. + pub fn get(&self, idx: usize) -> Option<&'a ArrayRef> { + self.slots + .get(idx) + .map(|slot| slot.as_ref().vortex_expect(self.expect)) + } + + /// Iterate the slots in order. + pub fn iter(&self) -> impl ExactSizeIterator + use<'a> { + let expect = self.expect; + self.slots + .iter() + .map(move |slot| slot.as_ref().vortex_expect(expect)) + } + + /// Clone every slot into an owned `Vec`. + pub fn to_vec(&self) -> Vec { + self.iter().cloned().collect() + } +} + +impl std::ops::Index for SlotSlice<'_> { + type Output = ArrayRef; + + fn index(&self, idx: usize) -> &Self::Output { + self.slots[idx].as_ref().vortex_expect(self.expect) + } +} + /// The public API trait for all Vortex arrays. /// /// This trait is sealed and cannot be implemented outside of `vortex-array`. diff --git a/vortex-array/src/arrays/chunked/array.rs b/vortex-array/src/arrays/chunked/array.rs index 3963d6ba950..eb79c567f1b 100644 --- a/vortex-array/src/arrays/chunked/array.rs +++ b/vortex-array/src/arrays/chunked/array.rs @@ -10,7 +10,6 @@ use std::fmt::Display; use std::fmt::Formatter; use futures::stream; -use smallvec::SmallVec; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -24,6 +23,7 @@ use crate::IntoArray; use crate::array::Array; use crate::array::ArrayParts; use crate::array::TypedArrayRef; +use crate::array_slots; use crate::arrays::Chunked; use crate::arrays::PrimitiveArray; use crate::dtype::DType; @@ -35,8 +35,17 @@ use crate::stream::ArrayStream; use crate::stream::ArrayStreamAdapter; use crate::validity::Validity; -pub(super) const CHUNK_OFFSETS_SLOT: usize = 0; -pub(super) const CHUNKS_OFFSET: usize = 1; +/// Slot layout of a [`Chunked`] array: `[chunk_offsets, chunks...]`. +#[array_slots(Chunked)] +pub struct ChunkedSlots { + /// The non-nullable `u64` array of cumulative chunk offsets. + pub chunk_offsets: ArrayRef, + /// The chunk arrays, each sharing the outer dtype. + pub chunks: Vec, +} + +pub(super) const CHUNK_OFFSETS_SLOT: usize = ChunkedSlots::CHUNK_OFFSETS; +pub(super) const CHUNKS_OFFSET: usize = ChunkedSlots::CHUNKS_OFFSET; #[derive(Clone, Debug)] pub struct ChunkedData { @@ -84,7 +93,8 @@ pub trait ChunkedArrayExt: TypedArrayRef { Box::new(self.iter_chunks().filter(|chunk| !chunk.is_empty())) } - fn chunk_offsets(&self) -> &[usize] { + /// Returns the cached chunk boundary offsets. + fn chunk_offset_values(&self) -> &[usize] { &self.chunk_offsets } @@ -93,12 +103,12 @@ pub trait ChunkedArrayExt: TypedArrayRef { index <= self.as_ref().len(), "Index out of bounds of the array" ); - let chunk_offsets = self.chunk_offsets(); - let index_chunk = chunk_offsets + let chunk_offset_values = self.chunk_offset_values(); + let index_chunk = chunk_offset_values .search_sorted(&index, SearchSortedSide::Right)? .to_ends_index(self.nchunks() + 1) .saturating_sub(1); - let chunk_start = chunk_offsets[index_chunk]; + let chunk_start = chunk_offset_values[index_chunk]; let index_in_chunk = index - chunk_start; Ok((index_chunk, index_in_chunk)) } @@ -138,20 +148,14 @@ impl ChunkedData { chunk_offsets } - pub(super) fn make_slots(chunk_offsets: &[usize], chunks: &[ArrayRef]) -> ArraySlots { + pub(super) fn make_chunk_offsets_array(chunk_offsets: &[usize]) -> ArrayRef { let mut chunk_offsets_buf = BufferMut::::with_capacity(chunk_offsets.len()); for &offset in chunk_offsets { let offset = u64::try_from(offset) .vortex_expect("chunk offset must fit in u64 for serialization"); unsafe { chunk_offsets_buf.push_unchecked(offset) } } - - let mut slots = SmallVec::with_capacity(1 + chunks.len()); - slots.push(Some( - PrimitiveArray::new(chunk_offsets_buf.freeze(), Validity::NonNullable).into_array(), - )); - slots.extend(chunks.iter().map(|c| Some(c.clone()))); - slots + PrimitiveArray::new(chunk_offsets_buf.freeze(), Validity::NonNullable).into_array() } /// Validates the components that would be used to create a `ChunkedArray`. @@ -256,10 +260,18 @@ impl Array { pub unsafe fn new_unchecked(chunks: Vec, dtype: DType) -> Self { let len = chunks.iter().map(|chunk| chunk.len()).sum(); let chunk_offsets = ChunkedData::compute_chunk_offsets(&chunks); + // Move `chunks` into the slot storage rather than cloning each `ArrayRef`. The generated + // `into_slots` pushes the offsets array then `extend`s the tail, which reserves from the + // `Vec`'s exact size hint (a single allocation). + let slots = ChunkedSlots { + chunk_offsets: ChunkedData::make_chunk_offsets_array(&chunk_offsets), + chunks, + } + .into_slots(); unsafe { Array::from_parts_unchecked( - ArrayParts::new(Chunked, dtype, len, ChunkedData::new(chunk_offsets.clone())) - .with_slots(ChunkedData::make_slots(&chunk_offsets, &chunks)), + ArrayParts::new(Chunked, dtype, len, ChunkedData::new(chunk_offsets)) + .with_slots(slots), ) } } diff --git a/vortex-array/src/arrays/chunked/compute/filter.rs b/vortex-array/src/arrays/chunked/compute/filter.rs index 5f50ce18387..4642cb08dbe 100644 --- a/vortex-array/src/arrays/chunked/compute/filter.rs +++ b/vortex-array/src/arrays/chunked/compute/filter.rs @@ -91,7 +91,7 @@ pub(crate) fn chunk_filters( array: ArrayView<'_, Chunked>, slices: impl Iterator, ) -> VortexResult> { - let chunk_offsets = array.chunk_offsets(); + let chunk_offsets = array.chunk_offset_values(); let mut chunk_filters = vec![ChunkFilter::None; array.nchunks()]; @@ -153,7 +153,7 @@ fn filter_indices( let mut current_chunk_id = 0; let mut chunk_indices = BufferMut::with_capacity(array.nchunks()); - let chunk_offsets = array.chunk_offsets(); + let chunk_offsets = array.chunk_offset_values(); for set_index in indices { let (chunk_id, index) = find_chunk_idx(set_index, chunk_offsets)?; diff --git a/vortex-array/src/arrays/chunked/compute/mask.rs b/vortex-array/src/arrays/chunked/compute/mask.rs index 01170a232bb..2944235e522 100644 --- a/vortex-array/src/arrays/chunked/compute/mask.rs +++ b/vortex-array/src/arrays/chunked/compute/mask.rs @@ -21,7 +21,7 @@ impl MaskKernel for Chunked { mask: &ArrayRef, _ctx: &mut ExecutionCtx, ) -> VortexResult> { - let chunk_offsets = array.chunk_offsets(); + let chunk_offsets = array.chunk_offset_values(); let new_chunks: Vec = array .iter_chunks() .enumerate() diff --git a/vortex-array/src/arrays/chunked/compute/take.rs b/vortex-array/src/arrays/chunked/compute/take.rs index 1b867981236..05fc1183563 100644 --- a/vortex-array/src/arrays/chunked/compute/take.rs +++ b/vortex-array/src/arrays/chunked/compute/take.rs @@ -51,7 +51,7 @@ fn take_chunked( // 2. Fused pass: walk sorted pairs against chunk boundaries. // - Dedup inline → build per-chunk filter masks // - Scatter final_take[orig_pos] = dedup_idx for every pair - let chunk_offsets = array.chunk_offsets(); + let chunk_offsets = array.chunk_offset_values(); let nchunks = array.nchunks(); let mut chunks = Vec::with_capacity(nchunks); let mut final_take = BufferMut::::with_capacity(n); diff --git a/vortex-array/src/arrays/chunked/tests.rs b/vortex-array/src/arrays/chunked/tests.rs index a32d03e207b..dc83b73aa80 100644 --- a/vortex-array/src/arrays/chunked/tests.rs +++ b/vortex-array/src/arrays/chunked/tests.rs @@ -213,7 +213,7 @@ fn with_slot_rewrites_chunk_and_offsets() { let array = array.as_::(); assert_eq!(array.nchunks(), 3); - assert_eq!(array.chunk_offsets(), [0, 3, 6, 9]); + assert_eq!(array.chunk_offset_values(), [0, 3, 6, 9]); assert_arrays_eq!( array.chunk(0).clone(), PrimitiveArray::from_iter([1u64, 2, 3]), diff --git a/vortex-array/src/arrays/filter/execute/union.rs b/vortex-array/src/arrays/filter/execute/union.rs index c00e8fba71e..0e911c8102f 100644 --- a/vortex-array/src/arrays/filter/execute/union.rs +++ b/vortex-array/src/arrays/filter/execute/union.rs @@ -10,6 +10,7 @@ use vortex_mask::MaskValues; use crate::ArrayRef; use crate::arrays::UnionArray; use crate::arrays::union::UnionArrayExt; +use crate::arrays::union::UnionArraySlotsExt; pub fn filter_union(array: &UnionArray, mask: &Arc) -> UnionArray { let filter_mask = Mask::Values(Arc::clone(mask)); diff --git a/vortex-array/src/arrays/masked/execute.rs b/vortex-array/src/arrays/masked/execute.rs index 14adcef59d6..f0d61b05454 100644 --- a/vortex-array/src/arrays/masked/execute.rs +++ b/vortex-array/src/arrays/masked/execute.rs @@ -26,6 +26,7 @@ use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::listview::ListViewArrayExt; use crate::arrays::struct_::StructArrayExt; use crate::arrays::union::UnionArrayExt; +use crate::arrays::union::UnionArraySlotsExt; use crate::arrays::variant::VariantArrayExt; use crate::builtins::ArrayBuiltins; use crate::executor::ExecutionCtx; @@ -153,7 +154,7 @@ fn mask_validity_union(array: UnionArray, validity: Validity) -> VortexResult, } + #[array_slots(Chunked)] + struct VariadicSlots { + offsets: ArrayRef, + maybe_validity: Option, + chunks: Vec, + } + #[test] fn generated_slots_round_trip() { let required = PrimitiveArray::new(buffer![1u8, 2, 3], Validity::NonNullable).into_array(); @@ -405,4 +413,63 @@ mod tests { assert_eq!(rebuilt.inner.len(), required.len()); assert_eq!(rebuilt.patch_values.len(), optional.len()); } + + #[test] + fn variadic_slots_round_trip() { + let offsets = PrimitiveArray::new(buffer![0u64, 3, 5], Validity::NonNullable).into_array(); + let chunk0 = PrimitiveArray::new(buffer![1u8, 2, 3], Validity::NonNullable).into_array(); + let chunk1 = PrimitiveArray::new(buffer![4u8, 5], Validity::NonNullable).into_array(); + + assert_eq!(VariadicSlots::OFFSETS, 0); + assert_eq!(VariadicSlots::MAYBE_VALIDITY, 1); + assert_eq!(VariadicSlots::CHUNKS_OFFSET, 2); + assert_eq!(VariadicSlots::FIXED_COUNT, 2); + assert_eq!(VariadicSlots::slot_name(0), "offsets"); + assert_eq!(VariadicSlots::slot_name(3), "chunks[1]"); + + let slot_vec = vec![Some(offsets.clone()), None, Some(chunk0), Some(chunk1)]; + + let view = VariadicSlotsView::from_slots(&slot_vec); + assert_eq!(view.offsets.len(), 3); + assert!(view.maybe_validity.is_none()); + assert_eq!(view.chunks.len(), 2); + assert_eq!(view.chunks[0].len(), 3); + assert_eq!(view.chunks.get(1).map(|c| c.len()), Some(2)); + assert!(view.chunks.get(2).is_none()); + assert_eq!( + view.chunks.iter().map(|c| c.len()).collect::>(), + vec![3, 2] + ); + + let owned = view.to_owned(); + assert_eq!(owned.chunks.len(), 2); + + let owned = VariadicSlots::from_slots(slot_vec.into()); + assert_eq!(owned.offsets.len(), offsets.len()); + assert!(owned.maybe_validity.is_none()); + assert_eq!(owned.chunks.len(), 2); + + let slots = owned.into_slots(); + assert_eq!(slots.len(), 4); + assert!(slots[1].is_none()); + assert_eq!( + slots[VariadicSlots::CHUNKS_OFFSET] + .as_ref() + .map(|c| c.len()), + Some(3) + ); + } + + #[test] + fn variadic_slots_empty_tail() { + let offsets = PrimitiveArray::new(buffer![0u64], Validity::NonNullable).into_array(); + let slot_vec = vec![Some(offsets), None]; + + let view = VariadicSlotsView::from_slots(&slot_vec); + assert!(view.chunks.is_empty()); + + let owned = VariadicSlots::from_slots(slot_vec.into()); + assert!(owned.chunks.is_empty()); + assert_eq!(owned.into_slots().len(), 2); + } } diff --git a/vortex-array/src/arrays/struct_/array.rs b/vortex-array/src/arrays/struct_/array.rs index 8b939deb911..11a8212f7ac 100644 --- a/vortex-array/src/arrays/struct_/array.rs +++ b/vortex-array/src/arrays/struct_/array.rs @@ -6,6 +6,7 @@ use std::iter::once; use std::sync::Arc; use itertools::Itertools; +use vortex_array_macros::array_slots; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -30,10 +31,19 @@ use crate::dtype::StructFields; use crate::validity::Validity; // StructArray has a variable number of slots: [validity?, field_0, ..., field_N] +/// Slot layout of a [`Struct`] array: `[validity?, fields...]`. +#[array_slots(Struct)] +pub struct StructSlots { + /// The optional row-level validity child. + pub validity: Option, + /// The field arrays, one per struct field, all sharing the outer length. + pub fields: Vec, +} + /// The validity bitmap indicating which struct elements are non-null. -pub(super) const VALIDITY_SLOT: usize = 0; +pub(super) const VALIDITY_SLOT: usize = StructSlots::VALIDITY; /// The offset at which the struct field arrays begin in the slots vector. -pub(super) const FIELDS_OFFSET: usize = 1; +pub(super) const FIELDS_OFFSET: usize = StructSlots::FIELDS_OFFSET; /// A struct array that stores multiple named fields as columns, similar to a database row. /// @@ -168,9 +178,11 @@ pub(super) fn make_struct_slots( validity: &Validity, length: usize, ) -> ArraySlots { - once(validity_to_child(validity, length)) - .chain(fields.iter().cloned().map(Some)) - .collect() + StructSlots { + validity: validity_to_child(validity, length), + fields: fields.to_vec(), + } + .into_slots() } pub trait StructArrayExt: TypedArrayRef { diff --git a/vortex-array/src/arrays/union/array.rs b/vortex-array/src/arrays/union/array.rs index 3a7603c056f..7df432214f8 100644 --- a/vortex-array/src/arrays/union/array.rs +++ b/vortex-array/src/arrays/union/array.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::iter::once; use std::sync::Arc; use vortex_error::VortexExpect; @@ -10,12 +9,12 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_err; use crate::ArrayRef; -use crate::ArraySlots; use crate::IntoArray; use crate::array::Array; use crate::array::ArrayParts; use crate::array::EmptyArrayData; use crate::array::TypedArrayRef; +use crate::array_slots; use crate::arrays::PrimitiveArray; use crate::arrays::Union; use crate::dtype::DType; @@ -23,10 +22,14 @@ use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::UnionVariants; -/// The row-aligned array of type IDs selecting a union child. -pub(super) const TYPE_IDS_SLOT: usize = 0; -/// The offset at which the sparse child arrays begin in the slots vector. -pub(super) const CHILDREN_OFFSET: usize = 1; +/// Slot layout of a canonical sparse union array. +#[array_slots(Union)] +pub struct UnionSlots { + /// The row-aligned array of type IDs selecting a union child. + pub type_ids: ArrayRef, + /// The row-aligned sparse children in variant order. + pub children: Vec, +} pub(super) fn make_union_parts( type_ids: ArrayRef, @@ -35,9 +38,11 @@ pub(super) fn make_union_parts( ) -> ArrayParts { let len = type_ids.len(); let nullability = type_ids.dtype().nullability(); - let slots: ArraySlots = once(Some(type_ids)) - .chain(children.iter().cloned().map(Some)) - .collect(); + let slots = UnionSlots { + type_ids, + children: children.to_vec(), + } + .into_slots(); ArrayParts::new( Union, @@ -59,7 +64,10 @@ pub struct UnionDataParts { } /// Accessors for a canonical sparse union array. -pub trait UnionArrayExt: TypedArrayRef { +/// +/// Slot accessors (`type_ids`, `children`, `slots_view`) live on the generated +/// [`UnionArraySlotsExt`] supertrait; this trait layers union-specific lookups on top. +pub trait UnionArrayExt: UnionArraySlotsExt { /// The union's variant schema. fn variants(&self) -> &UnionVariants { match self.as_ref().dtype() { @@ -68,28 +76,14 @@ pub trait UnionArrayExt: TypedArrayRef { } } - /// The row-aligned `u8` type IDs whose nulls represent outer union nulls. - fn type_ids(&self) -> &ArrayRef { - self.as_ref().slots()[TYPE_IDS_SLOT] - .as_ref() - .vortex_expect("UnionArray type_ids slot") - } - /// Iterate over sparse children in variant order. fn iter_children(&self) -> impl ExactSizeIterator + '_ { - self.as_ref().slots()[CHILDREN_OFFSET..] - .iter() - .map(|slot| slot.as_ref().vortex_expect("UnionArray child slot")) - } - - /// Return the sparse children in variant order. - fn children(&self) -> Arc<[ArrayRef]> { - self.iter_children().cloned().collect() + self.children().iter() } /// Return a sparse child by its variant index. fn child(&self, index: usize) -> Option<&ArrayRef> { - self.as_ref().slots().get(CHILDREN_OFFSET + index)?.as_ref() + self.children().get(index) } /// Return a sparse child selected by a data-level type ID. @@ -173,7 +167,7 @@ impl Array { pub fn into_data_parts(self) -> UnionDataParts { let variants = self.variants().clone(); let type_ids = self.type_ids().clone(); - let children = self.children(); + let children = self.children().to_vec().into(); UnionDataParts { variants, type_ids, diff --git a/vortex-array/src/arrays/union/compute/mask.rs b/vortex-array/src/arrays/union/compute/mask.rs index 0f1cbf618b7..03ee226fedf 100644 --- a/vortex-array/src/arrays/union/compute/mask.rs +++ b/vortex-array/src/arrays/union/compute/mask.rs @@ -9,6 +9,7 @@ use crate::array::ArrayView; use crate::arrays::Union; use crate::arrays::UnionArray; use crate::arrays::union::UnionArrayExt; +use crate::arrays::union::UnionArraySlotsExt; use crate::builtins::ArrayBuiltins; use crate::scalar_fn::fns::mask::MaskReduce; @@ -17,7 +18,7 @@ impl MaskReduce for Union { UnionArray::try_new( array.type_ids().clone().mask(mask.clone())?, array.variants().clone(), - array.children(), + array.children().to_vec(), ) .map(|a| Some(a.into_array())) } diff --git a/vortex-array/src/arrays/union/compute/slice.rs b/vortex-array/src/arrays/union/compute/slice.rs index 5a03cbebc46..be1336f5154 100644 --- a/vortex-array/src/arrays/union/compute/slice.rs +++ b/vortex-array/src/arrays/union/compute/slice.rs @@ -13,6 +13,7 @@ use crate::arrays::Union; use crate::arrays::UnionArray; use crate::arrays::slice::SliceReduce; use crate::arrays::union::UnionArrayExt; +use crate::arrays::union::UnionArraySlotsExt; impl SliceReduce for Union { fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { diff --git a/vortex-array/src/arrays/union/mod.rs b/vortex-array/src/arrays/union/mod.rs index aaeab88f44d..33a906d5cae 100644 --- a/vortex-array/src/arrays/union/mod.rs +++ b/vortex-array/src/arrays/union/mod.rs @@ -17,7 +17,10 @@ use crate::dtype::PType; mod array; pub use array::UnionArrayExt; +pub use array::UnionArraySlotsExt; pub use array::UnionDataParts; +pub use array::UnionSlots; +pub use array::UnionSlotsView; pub use vtable::UnionArray; pub(crate) mod compute; diff --git a/vortex-array/src/arrays/union/vtable/mod.rs b/vortex-array/src/arrays/union/vtable/mod.rs index 5357b0d4aaa..ea461d1ae09 100644 --- a/vortex-array/src/arrays/union/vtable/mod.rs +++ b/vortex-array/src/arrays/union/vtable/mod.rs @@ -21,8 +21,7 @@ use crate::array::EmptyArrayData; use crate::array::VTable; use crate::array::with_empty_buffers; use crate::arrays::union::UnionArrayExt; -use crate::arrays::union::array::CHILDREN_OFFSET; -use crate::arrays::union::array::TYPE_IDS_SLOT; +use crate::arrays::union::UnionSlots; use crate::arrays::union::array::make_union_parts; use crate::arrays::union::compute::rules::PARENT_RULES; use crate::arrays::union::union_type_ids_dtype; @@ -64,11 +63,11 @@ impl VTable for Union { slots: &[Option], ) -> VortexResult<()> { let type_ids = slots - .get(TYPE_IDS_SLOT) + .get(UnionSlots::TYPE_IDS) .and_then(Option::as_ref) .ok_or_else(|| vortex_err!("UnionArray is missing its type_ids slot"))?; let variant_arrays = slots - .get(CHILDREN_OFFSET..) + .get(UnionSlots::CHILDREN_OFFSET..) .unwrap_or_default() .iter() .enumerate() @@ -124,17 +123,21 @@ impl VTable for Union { }; vortex_ensure_eq!( children.len(), - CHILDREN_OFFSET + variants.len(), + UnionSlots::CHILDREN_OFFSET + variants.len(), "UnionArray expected {} children, found {}", - CHILDREN_OFFSET + variants.len(), + UnionSlots::CHILDREN_OFFSET + variants.len(), children.len() ); - let type_ids = children.get(TYPE_IDS_SLOT, &union_type_ids_dtype(*nullability), len)?; + let type_ids = children.get( + UnionSlots::TYPE_IDS, + &union_type_ids_dtype(*nullability), + len, + )?; let sparse_children = variants .variants() .enumerate() - .map(|(index, dtype)| children.get(CHILDREN_OFFSET + index, &dtype, len)) + .map(|(index, dtype)| children.get(UnionSlots::CHILDREN_OFFSET + index, &dtype, len)) .collect::>>()?; Ok(make_union_parts( @@ -145,10 +148,10 @@ impl VTable for Union { } fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String { - if idx == TYPE_IDS_SLOT { + if idx == UnionSlots::TYPE_IDS { "type_ids".to_string() } else { - array.variants().names()[idx - CHILDREN_OFFSET].to_string() + array.variants().names()[idx - UnionSlots::CHILDREN_OFFSET].to_string() } } diff --git a/vortex-array/src/arrays/union/vtable/operations.rs b/vortex-array/src/arrays/union/vtable/operations.rs index d8adb83992d..39ba95ed4ee 100644 --- a/vortex-array/src/arrays/union/vtable/operations.rs +++ b/vortex-array/src/arrays/union/vtable/operations.rs @@ -10,6 +10,7 @@ use crate::array::ArrayView; use crate::array::OperationsVTable; use crate::arrays::Union; use crate::arrays::union::UnionArrayExt; +use crate::arrays::union::UnionArraySlotsExt; use crate::scalar::Scalar; impl OperationsVTable for Union { diff --git a/vortex-array/src/arrays/union/vtable/validate.rs b/vortex-array/src/arrays/union/vtable/validate.rs index 95a97bb016a..e01c5486a94 100644 --- a/vortex-array/src/arrays/union/vtable/validate.rs +++ b/vortex-array/src/arrays/union/vtable/validate.rs @@ -6,7 +6,7 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure_eq; use crate::ArrayRef; -use crate::arrays::union::array::CHILDREN_OFFSET; +use crate::arrays::union::UnionSlots; use crate::arrays::union::union_type_ids_dtype; use crate::dtype::DType; @@ -23,8 +23,8 @@ pub(super) fn validate_union_components( variant_arrays.len(), variants.len(), "UnionArray has {} slots but expected {}", - CHILDREN_OFFSET + variant_arrays.len(), - CHILDREN_OFFSET + variants.len() + UnionSlots::CHILDREN_OFFSET + variant_arrays.len(), + UnionSlots::CHILDREN_OFFSET + variants.len() ); let expected_union_type_ids_dtype = union_type_ids_dtype(*nullability); diff --git a/vortex-array/src/arrays/union/vtable/validity.rs b/vortex-array/src/arrays/union/vtable/validity.rs index 80e043534e5..1da64995b14 100644 --- a/vortex-array/src/arrays/union/vtable/validity.rs +++ b/vortex-array/src/arrays/union/vtable/validity.rs @@ -6,7 +6,7 @@ use vortex_error::VortexResult; use crate::array::ArrayView; use crate::array::ValidityVTable; use crate::arrays::Union; -use crate::arrays::union::UnionArrayExt; +use crate::arrays::union::UnionArraySlotsExt; use crate::validity::Validity; impl ValidityVTable for Union { diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 0447aef517a..ffcc91fdebd 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -25,6 +25,7 @@ use vortex_array::arrays::masked::MaskedArraySlotsExt; use vortex_array::arrays::scalar_fn::AnyScalarFn; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::arrays::union::UnionArrayExt; +use vortex_array::arrays::union::UnionArraySlotsExt; use vortex_array::arrays::variant::VariantArrayExt; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; From d9cacd3b73102e84d9520e6d6ac3ed5f5aac2027 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 27 Jul 2026 11:31:19 +0100 Subject: [PATCH 2/3] fix: repair build and rustdoc breakage on the variadic slots branch The merge with develop left `vortex-compressor` importing a `VariantArrayExt` that does not exist, which broke the workspace build and cascaded into nearly every CI job (Python, Java, C++, generated-files, all Rust test matrices). - Import `VariantArraySlotsExt` in the cascade compressor. - Update `vortex-duckdb`'s chunked exporter for the `chunk_offsets` -> `chunk_offset_values` rename; `chunk_offsets` now names the generated slot accessor returning `&ArrayRef`. - Drop the intra-doc link to `::vortex_array::SlotSlice` from the `#[array_slots]` docs. `vortex-array-macros` does not depend on `vortex-array`, so the link failed `RUSTDOCFLAGS="-D warnings"`. - Remove a dead `#[allow(clippy::len_without_is_empty)]` on `SlotSlice::len`, which does define `is_empty`. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Joe Isaacs --- vortex-array-macros/src/lib.rs | 6 ++---- vortex-array/src/array/mod.rs | 1 - vortex-compressor/src/compressor/cascade.rs | 2 +- vortex-duckdb/src/exporter/chunked.rs | 2 +- 4 files changed, 4 insertions(+), 7 deletions(-) diff --git a/vortex-array-macros/src/lib.rs b/vortex-array-macros/src/lib.rs index 1522cc75ed9..f8087644f76 100644 --- a/vortex-array-macros/src/lib.rs +++ b/vortex-array-macros/src/lib.rs @@ -132,8 +132,8 @@ use syn::spanned::Spanned; /// } /// ``` /// -/// The view field and ext trait accessor for the tail are a [`SlotSlice`], a -/// borrowed run of required slots supporting `len()`, `get()`, `iter()`, and +/// The view field and ext trait accessor for the tail are a `vortex_array::SlotSlice`, +/// a borrowed run of required slots supporting `len()`, `get()`, `iter()`, and /// indexing: /// /// ```ignore @@ -148,8 +148,6 @@ use syn::spanned::Spanned; /// fn slots_view(&self) -> ChunkedSlotsView<'_> { ... } /// } /// ``` -/// -/// [`SlotSlice`]: ::vortex_array::SlotSlice #[proc_macro_attribute] pub fn array_slots(attr: TokenStream, item: TokenStream) -> TokenStream { let encoding = parse_macro_input!(attr as Path); diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index a65c8ecf708..8b9b2812bfa 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -72,7 +72,6 @@ impl<'a> SlotSlice<'a> { } /// The number of slots in the run. - #[allow(clippy::len_without_is_empty)] pub fn len(&self) -> usize { self.slots.len() } diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 8a3e80b3227..eb81e81241e 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -27,7 +27,7 @@ use vortex_array::arrays::scalar_fn::AnyScalarFn; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::arrays::union::UnionArrayExt; use vortex_array::arrays::union::UnionArraySlotsExt; -use vortex_array::arrays::variant::VariantArrayExt; +use vortex_array::arrays::variant::VariantArraySlotsExt; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; diff --git a/vortex-duckdb/src/exporter/chunked.rs b/vortex-duckdb/src/exporter/chunked.rs index 6e6b6642030..0c13b77aa3e 100644 --- a/vortex-duckdb/src/exporter/chunked.rs +++ b/vortex-duckdb/src/exporter/chunked.rs @@ -29,7 +29,7 @@ pub(crate) fn new_exporter_with_flatten( return canonical::new_exporter(array.into_array(), cache, ctx); } - let chunk_offsets = array.chunk_offsets().to_vec(); + let chunk_offsets = array.chunk_offset_values().to_vec(); let chunks = array .chunks() .iter() From 9cc62b870b91bc79eb5953895e165ebdcec486e0 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 27 Jul 2026 11:30:19 +0000 Subject: [PATCH 3/3] Avoid intermediate Vec allocations in slot construction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Union and Struct construction helpers routed borrowed children through `Slots { tail: slice.to_vec() }.into_slots()`, which allocated a throwaway `Vec`, cloned every child into it, then moved each element again into the `SmallVec` — an extra allocation and N moves versus the one-pass build they had before the macro migration. - make_union_parts / make_struct_slots now clone the borrowed children straight into an `ArraySlots` sized with `with_capacity`, in a single pass. `into_slots` is kept only where the caller genuinely owns the `Vec` (Chunked::new_unchecked), where it moves with zero clones. - UnionArray::into_data_parts collects children directly into the `Arc<[ArrayRef]>` via `iter_children().cloned().collect()` instead of `children().to_vec().into()` (drops an intermediate `Vec`). - Export ChunkedSlots/ChunkedSlotsView and StructSlots/StructSlotsView so the generated zero-copy view/accessors are public API (parity with UnionSlots). Read paths are unchanged: the generated tail accessor returns a `SlotSlice` borrow with no allocation. The remaining owned-conversion `to_vec`/`collect` sites are the minimal clone required at a borrow boundary for construction, matching pre-migration behavior. Verification: - cargo test -p vortex-array --lib -- arrays::union arrays::struct_ arrays::chunked arrays::masked patched::array::tests (192 passed) - cargo clippy --all-targets -p vortex-array - cargo check -p vortex-compressor -p vortex-file -p vortex-layout -p vortex-row -p vortex-arrow - cargo +nightly fmt --all Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_017BF3K2eZTEnCsqMsN3wGtc --- vortex-array/src/arrays/chunked/mod.rs | 2 ++ vortex-array/src/arrays/struct_/array.rs | 11 ++++++----- vortex-array/src/arrays/struct_/mod.rs | 2 ++ vortex-array/src/arrays/union/array.rs | 14 ++++++++------ 4 files changed, 18 insertions(+), 11 deletions(-) diff --git a/vortex-array/src/arrays/chunked/mod.rs b/vortex-array/src/arrays/chunked/mod.rs index 4bd6cb43b7e..3bd361d66a1 100644 --- a/vortex-array/src/arrays/chunked/mod.rs +++ b/vortex-array/src/arrays/chunked/mod.rs @@ -4,6 +4,8 @@ mod array; pub use array::ChunkedArrayExt; pub use array::ChunkedData; +pub use array::ChunkedSlots; +pub use array::ChunkedSlotsView; pub use vtable::ChunkedArray; pub(crate) mod compute; diff --git a/vortex-array/src/arrays/struct_/array.rs b/vortex-array/src/arrays/struct_/array.rs index 11a8212f7ac..ddea7efb543 100644 --- a/vortex-array/src/arrays/struct_/array.rs +++ b/vortex-array/src/arrays/struct_/array.rs @@ -178,11 +178,12 @@ pub(super) fn make_struct_slots( validity: &Validity, length: usize, ) -> ArraySlots { - StructSlots { - validity: validity_to_child(validity, length), - fields: fields.to_vec(), - } - .into_slots() + // `fields` is borrowed, so its arrays must be cloned regardless; clone them straight into the + // `SmallVec` rather than through an intermediate `Vec`. + let mut slots = ArraySlots::with_capacity(StructSlots::FIELDS_OFFSET + fields.len()); + slots.push(validity_to_child(validity, length)); + slots.extend(fields.iter().cloned().map(Some)); + slots } pub trait StructArrayExt: TypedArrayRef { diff --git a/vortex-array/src/arrays/struct_/mod.rs b/vortex-array/src/arrays/struct_/mod.rs index da49411c3bc..cdfc632083d 100644 --- a/vortex-array/src/arrays/struct_/mod.rs +++ b/vortex-array/src/arrays/struct_/mod.rs @@ -4,6 +4,8 @@ mod array; pub use array::StructArrayExt; pub use array::StructDataParts; +pub use array::StructSlots; +pub use array::StructSlotsView; pub use vtable::StructArray; pub(crate) mod compute; diff --git a/vortex-array/src/arrays/union/array.rs b/vortex-array/src/arrays/union/array.rs index 7df432214f8..6e876cfe2b5 100644 --- a/vortex-array/src/arrays/union/array.rs +++ b/vortex-array/src/arrays/union/array.rs @@ -9,6 +9,7 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_err; use crate::ArrayRef; +use crate::ArraySlots; use crate::IntoArray; use crate::array::Array; use crate::array::ArrayParts; @@ -38,11 +39,12 @@ pub(super) fn make_union_parts( ) -> ArrayParts { let len = type_ids.len(); let nullability = type_ids.dtype().nullability(); - let slots = UnionSlots { - type_ids, - children: children.to_vec(), - } - .into_slots(); + // Build the slot storage in a single pass. `children` is borrowed, so the child arrays must + // be cloned regardless; clone them straight into the `SmallVec` rather than through an + // intermediate `Vec`. + let mut slots = ArraySlots::with_capacity(UnionSlots::CHILDREN_OFFSET + children.len()); + slots.push(Some(type_ids)); + slots.extend(children.iter().cloned().map(Some)); ArrayParts::new( Union, @@ -167,7 +169,7 @@ impl Array { pub fn into_data_parts(self) -> UnionDataParts { let variants = self.variants().clone(); let type_ids = self.type_ids().clone(); - let children = self.children().to_vec().into(); + let children: Arc<[ArrayRef]> = self.iter_children().cloned().collect(); UnionDataParts { variants, type_ids,