From 33ecc68921e814d08b1ff38cdd113cc420738075 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 25 Jul 2026 18:49:28 +0000 Subject: [PATCH 01/22] refactor(array): canonicalize chunked nested types through the builder `pack_struct_chunks`, `swizzle_fixed_size_list_chunks` and `swizzle_list_chunks` existed because `append_to_builder` used to decode a builder's children: without them, canonicalizing a `ChunkedArray` would concatenate every chunk's children instead of reusing them. Their doc comments describe what the builders now do on their own, so all three are dead weight - the generic builder path produces the same swizzled array, with each chunk's children kept as chunks of the combined child. Only `Variant` still needs a hand-written pack, because there is no variant builder. This drops the one canonicalization path that allocated its buffers through the session allocator: `swizzle_list_chunks` allocated its offsets and sizes with `ctx.allocator()`, whereas `builder_with_capacity_in` still ignores the allocator it is handed, so chunked primitives, structs and FSLs already went around it. `list_canonicalize_uses_memory_session_allocator` guarded that one path and goes with it; restoring the property means teaching the builders to allocate through a `HostAllocatorRef`, not keeping a bespoke list swizzle alive. Signed-off-by: Claude Signed-off-by: Robert Kruszewski Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Robert Kruszewski --- .../src/arrays/chunked/vtable/canonical.rs | 281 ++++-------------- 1 file changed, 54 insertions(+), 227 deletions(-) diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index a521fd0e6cc..86df7fd3850 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools as _; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; @@ -14,22 +13,11 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; -use crate::arrays::FixedSizeListArray; -use crate::arrays::ListViewArray; -use crate::arrays::PrimitiveArray; -use crate::arrays::StructArray; use crate::arrays::VariantArray; use crate::arrays::chunked::ChunkedArrayExt; -use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; -use crate::arrays::listview::ListViewArraySlotsExt; -use crate::arrays::listview::ListViewRebuildMode; use crate::arrays::variant::VariantArraySlotsExt; use crate::builders::builder_with_capacity_in; -use crate::builtins::ArrayBuiltins; use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::dtype::PType; -use crate::validity::Validity; pub(super) fn _canonicalize( array: ArrayView<'_, Chunked>, @@ -46,28 +34,11 @@ pub(super) fn _canonicalize( return array.chunk(0).clone().execute::(ctx); } - let owned_chunks: Vec = array.iter_chunks().cloned().collect(); Ok(match array.dtype() { - DType::Struct(..) => { - let struct_array = pack_struct_chunks(owned_chunks, ctx)?; - Canonical::Struct(struct_array) + DType::Variant(_) => { + let owned_chunks: Vec = array.iter_chunks().cloned().collect(); + Canonical::Variant(pack_variant_chunks(owned_chunks, ctx)?) } - DType::List(elem_dtype, _) => Canonical::List(swizzle_list_chunks( - &owned_chunks, - array.array().validity()?, - elem_dtype, - ctx, - )?), - DType::FixedSizeList(elem_dtype, list_size, _) => { - Canonical::FixedSizeList(swizzle_fixed_size_list_chunks( - &owned_chunks, - array.array().validity()?, - elem_dtype, - *list_size, - ctx, - )?) - } - DType::Variant(_) => Canonical::Variant(pack_variant_chunks(owned_chunks, ctx)?), _ => { let mut builder = builder_with_capacity_in(array.dtype(), array.len(), ctx.allocator()); array.array().append_to_builder(builder.as_mut(), ctx)?; @@ -76,17 +47,6 @@ pub(super) fn _canonicalize( }) } -/// Packs many [`StructArray`]s to instead be a single [`StructArray`], where the [`DynArrayData`](crate::array::DynArrayData) for each -/// field is a [`ChunkedArray`]. -/// -/// The caller guarantees there are at least 2 chunks. -fn pack_struct_chunks(chunks: Vec, ctx: &mut ExecutionCtx) -> VortexResult { - chunks - .into_iter() - .map(|c| c.execute::(ctx)) - .process_results(|iter| StructArray::try_concat(iter))? -} - /// Packs many [`VariantArray`]s into one [`VariantArray`] with chunked children. /// /// The caller guarantees there are at least 2 chunks. @@ -145,149 +105,14 @@ fn pack_variant_chunks( VariantArray::try_new(core_storage, shredded) } -/// Packs [`ListViewArray`]s together into a chunked `ListViewArray`. -/// -/// We use the existing arrays (chunks) to form a chunked array of `elements` (the child array). -/// -/// The caller guarantees there are at least 2 chunks. -fn swizzle_list_chunks( - chunks: &[ArrayRef], - validity: Validity, - elem_dtype: &DType, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let len: usize = chunks.iter().map(|c| c.len()).sum(); - - assert_eq!( - chunks[0] - .dtype() - .as_list_element_opt() - .vortex_expect("DType was somehow not a list") - .as_ref(), - elem_dtype - ); - - // Since each list array in `chunks` has offsets local to each array, we can reuse the existing - // array's child `elements` as the chunks and recompute offsets. - - let mut list_elements_chunks = Vec::with_capacity(chunks.len()); - let mut num_elements = 0; - - // TODO(connor)[ListView]: We could potentially choose a smaller type here, but that would make - // this much more complicated. - // We (somewhat arbitrarily) choose `u64` for our offsets and sizes here. These can always be - // narrowed later by the compressor. - let allocator = ctx.allocator(); - let mut offsets = allocator.zeroed::(len); - let mut sizes = allocator.zeroed::(len); - let offsets_out = offsets.as_mut_slice(); - let sizes_slice_out = sizes.as_mut_slice(); - let mut next_list = 0usize; - - for chunk in chunks { - let chunk_array = chunk.clone().execute::(ctx)?; - // By rebuilding as zero-copy to `List` and trimming all elements (to prevent gaps), we make - // the final output `ListView` also zero-copyable to `List`. - let chunk_array = chunk_array.rebuild(ListViewRebuildMode::MakeExact, ctx)?; - - // Add the `elements` of the current array as a new chunk. - list_elements_chunks.push(chunk_array.elements().clone()); - - // Cast offsets and sizes to `u64`. - let offsets_arr = chunk_array - .offsets() - .clone() - .cast(DType::Primitive(PType::U64, Nullability::NonNullable)) - .vortex_expect("Must be able to fit array offsets in u64") - .execute::(ctx)?; - - let sizes_arr = chunk_array - .sizes() - .clone() - .cast(DType::Primitive(PType::U64, Nullability::NonNullable)) - .vortex_expect("Must be able to fit array offsets in u64") - .execute::(ctx)?; - - let offsets_slice = offsets_arr.as_slice::(); - let sizes_slice = sizes_arr.as_slice::(); - - // Append offsets and sizes, adjusting offsets to point into the combined array. - for (&offset, &size) in offsets_slice.iter().zip(sizes_slice.iter()) { - offsets_out[next_list] = offset + num_elements; - sizes_slice_out[next_list] = size; - next_list += 1; - } - - num_elements += chunk_array.elements().len() as u64; - } - debug_assert_eq!(next_list, len); - - // SAFETY: elements are sliced from valid `ListViewArray`s (from `to_listview()`). - let chunked_elements = - unsafe { ChunkedArray::new_unchecked(list_elements_chunks, elem_dtype.clone()) } - .into_array(); - - let offsets = PrimitiveArray::new(offsets.freeze(), Validity::NonNullable).into_array(); - let sizes = PrimitiveArray::new(sizes.freeze(), Validity::NonNullable).into_array(); - - // SAFETY: - // - `offsets` and `sizes` are non-nullable u64 arrays of the same length - // - Each `offset[i] + size[i]` list view is within bounds of elements array because it came - // from valid chunks - // - Validity came from the outer chunked array so it must have the same length - // - Since we made sure that all chunks were zero-copyable to a list above, we know that the - // final concatenated output is also zero-copyable to a list. - Ok(unsafe { - ListViewArray::new_unchecked(chunked_elements, offsets, sizes, validity) - .with_zero_copy_to_list(true) - }) -} - -/// Packs [`FixedSizeListArray`]s together into a single [`FixedSizeListArray`] whose `elements` -/// child is a [`ChunkedArray`]. -/// -/// Every chunk shares the same `list_size`, and each chunk's `elements` child is exactly -/// `list_size * chunk.len()` long and starts at the first list, so we can reuse the chunks' -/// `elements` children directly as the chunks of a combined `elements` array without copying. -/// -/// The caller guarantees there are at least 2 chunks. -fn swizzle_fixed_size_list_chunks( - chunks: &[ArrayRef], - validity: Validity, - elem_dtype: &DType, - list_size: u32, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let len: usize = chunks.iter().map(|c| c.len()).sum(); - - let mut element_chunks = Vec::with_capacity(chunks.len()); - for chunk in chunks { - let chunk_array = chunk.clone().execute::(ctx)?; - // A canonical `FixedSizeListArray` keeps its `elements` child trimmed to exactly - // `list_size * chunk.len()` starting at the first list, so the children concatenate - // cleanly into the combined `elements` array. - element_chunks.push(chunk_array.elements().clone()); - } - - let chunked_elements = ChunkedArray::try_new(element_chunks, elem_dtype.clone())?.into_array(); - - FixedSizeListArray::try_new(chunked_elements, list_size, validity, len) -} - #[cfg(test)] mod tests { - use std::alloc::Layout; - use std::ptr::NonNull; use std::sync::Arc; use std::sync::LazyLock; - use std::sync::atomic::AtomicUsize; - use std::sync::atomic::Ordering; - use allocator_api2::alloc::AllocError; - use allocator_api2::alloc::Allocator; - use allocator_api2::alloc::Global; - use vortex_buffer::BufferAllocatorRef; + use rstest::rstest; use vortex_buffer::buffer; + use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; @@ -297,15 +122,22 @@ mod tests { use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::arrays::Chunked; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; + use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::ListArray; + use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; + use crate::arrays::Struct; use crate::arrays::StructArray; use crate::arrays::VarBinViewArray; use crate::arrays::VariantArray; + use crate::arrays::chunked::ChunkedArrayExt; + use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; + use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::struct_::StructArrayExt; use crate::arrays::variant::VariantArraySlotsExt; use crate::assert_arrays_eq; @@ -314,31 +146,12 @@ mod tests { use crate::dtype::DType::Variant as VariantDType; use crate::dtype::Nullability::NonNullable; use crate::dtype::PType::I32; - use crate::memory::MemorySessionExt; use crate::scalar::Scalar; use crate::validity::Validity; /// A shared session for these chunked-array tests, used to create execution contexts. static SESSION: LazyLock = LazyLock::new(crate::array_session); - #[derive(Debug)] - struct CountingAllocator { - allocations: Arc, - } - - // SAFETY: this forwards memory operations to Global and only counts allocations. - unsafe impl Allocator for CountingAllocator { - fn allocate(&self, layout: Layout) -> Result, AllocError> { - self.allocations.fetch_add(1, Ordering::Relaxed); - Global.allocate(layout) - } - - unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { - // SAFETY: ptr and layout came from Global. - unsafe { Global.deallocate(ptr, layout) } - } - } - fn variant_scalar(value: i32) -> Scalar { Scalar::variant(Scalar::primitive(value, NonNullable)) } @@ -655,39 +468,53 @@ mod tests { Ok(()) } - #[test] - fn list_canonicalize_uses_memory_session_allocator() { - let allocations = Arc::new(AtomicUsize::new(0)); - let session = - crate::array_session().with_allocator(BufferAllocatorRef::new(CountingAllocator { - allocations: Arc::clone(&allocations), - })); - let mut ctx = session.create_execution_ctx(); - - let l1 = ListArray::try_new( - buffer![1, 2, 3, 4].into_array(), - buffer![0, 3].into_array(), - Validity::NonNullable, - ) - .unwrap(); - let l2 = ListArray::try_new( - buffer![5, 6].into_array(), + /// Canonicalizing a `ChunkedArray` reuses each chunk's children instead of concatenating + /// them: a nested builder keeps a child chunked on exactly the boundaries it was appended on, + /// however short the appended chunk is. + #[rstest] + #[case::struct_( + StructArray::try_from_iter([("a", buffer![1i32, 2])]) + .vortex_expect("struct array") + .into_array(), + |array: &ArrayRef| array.as_::().unmasked_field(0).clone() + )] + #[case::fixed_size_list( + FixedSizeListArray::new(buffer![1i32, 2].into_array(), 2, Validity::NonNullable, 1) + .into_array(), + |array: &ArrayRef| array.as_::().elements().clone() + )] + #[case::list( + ListArray::try_new( + buffer![1i32, 2].into_array(), buffer![0, 2].into_array(), Validity::NonNullable, ) - .unwrap(); + .vortex_expect("list array") + .into_array(), + |array: &ArrayRef| array.as_::().elements().clone() + )] + fn canonicalize_reuses_chunk_children( + #[case] chunk: ArrayRef, + #[case] child_of: fn(&ArrayRef) -> ArrayRef, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); - let chunked_list = ChunkedArray::try_new( - vec![l1.into_array(), l2.into_array()], - List(Arc::new(Primitive(I32, NonNullable)), NonNullable), - ) - .unwrap() - .into_array(); + let chunked = + ChunkedArray::try_new(vec![chunk.clone(), chunk.clone()], chunk.dtype().clone())? + .into_array(); + let canonical = chunked.execute::(&mut ctx)?.into_array(); - drop(chunked_list.execute::(&mut ctx).unwrap()); - assert!( - allocations.load(Ordering::Relaxed) >= 2, - "expected offset+size allocations through MemorySession" + let child = child_of(&canonical); + assert_eq!( + child.as_::().nchunks(), + 2, + "each chunk's child should have become a chunk of the combined child", ); + + let expected = + ChunkedArray::try_new(vec![chunk.clone(), chunk], canonical.dtype().clone())?; + assert_arrays_eq!(&canonical, &expected, &mut ctx); + + Ok(()) } } From 72d6fc514149d226b8bad2b1f27ea0494cd19b58 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 21 Aug 2026 11:42:07 +0100 Subject: [PATCH 02/22] Better error when unpacking a decimal scalar Signed-off-by: Robert Kruszewski --- vortex-array/src/builders/child.rs | 56 +++++++++++++++----- vortex-array/src/builders/fixed_size_list.rs | 16 +++--- vortex-array/src/builders/list.rs | 16 +++--- vortex-array/src/builders/listview.rs | 16 +++--- vortex-array/src/builders/map.rs | 1 - vortex-array/src/builders/mod.rs | 1 - 6 files changed, 63 insertions(+), 43 deletions(-) diff --git a/vortex-array/src/builders/child.rs b/vortex-array/src/builders/child.rs index a9f5b46fc84..9f9270ebf80 100644 --- a/vortex-array/src/builders/child.rs +++ b/vortex-array/src/builders/child.rs @@ -11,6 +11,7 @@ use crate::IntoArray; use crate::arrays::ChunkedArray; use crate::builders::ArrayBuilder; use crate::builders::builder_with_capacity_in; +use crate::canonical::Canonical; use crate::dtype::DType; use crate::scalar::Scalar; @@ -37,8 +38,15 @@ pub struct ChildBuilder { /// The summed length of `chunks`. chunks_len: usize, - /// Builder holding the scalars appended after the last chunk. - pending: Box, + /// Builder holding the scalars appended after the last chunk, materialized by the first append + /// that is not a whole array. + pending: Option>, + + /// The capacity the scalar builder is materialized with, grown by + /// [`reserve_exact`](Self::reserve_exact). + pending_capacity: usize, + + allocator: BufferAllocatorRef, } impl ChildBuilder { @@ -48,13 +56,15 @@ impl ChildBuilder { dtype: dtype.clone(), chunks: Vec::new(), chunks_len: 0, - pending: builder_with_capacity_in(dtype, capacity, allocator), + pending: None, + pending_capacity: capacity, + allocator: allocator.clone(), } } /// The number of values appended so far. pub fn len(&self) -> usize { - self.chunks_len + self.pending.len() + self.chunks_len + self.pending.as_ref().map_or(0, |pending| pending.len()) } /// Appends every value of `array` to the child as a chunk of its own, keeping its encoding. @@ -95,40 +105,50 @@ impl ChildBuilder { /// Appends a single [`Scalar`] to the child. pub fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> { - self.pending.append_scalar(scalar) + self.pending().append_scalar(scalar) } /// Appends `n` "zero" values to the child. /// /// See [`ArrayBuilder::append_zeros`]. pub fn append_zeros(&mut self, n: usize) { - self.pending.append_zeros(n) + self.pending().append_zeros(n) } /// Appends `n` null values to the child. /// /// See [`ArrayBuilder::append_nulls`]. pub fn append_nulls(&mut self, n: usize) { - self.pending.append_nulls(n) + self.pending().append_nulls(n) } /// Appends `n` default values to the child. /// /// See [`ArrayBuilder::append_defaults`]. pub fn append_defaults(&mut self, n: usize) { - self.pending.append_defaults(n) + self.pending().append_defaults(n) } /// Allocates space for `additional` more values in the scalar builder. + /// + /// While the scalar builder is still unmaterialized this only records the request, so that + /// reserving on a child that goes on to receive nothing but arrays allocates nothing. pub fn reserve_exact(&mut self, additional: usize) { - self.pending.reserve_exact(additional) + match self.pending.as_mut() { + Some(pending) => pending.reserve_exact(additional), + None => self.pending_capacity += additional, + } } /// Finishes the child, combining the accumulated chunks into a [`ChunkedArray`] when there is /// more than one of them. pub fn finish(&mut self) -> ArrayRef { if self.chunks.is_empty() { - return self.pending.finish(); + return match self.pending.as_mut() { + Some(pending) => pending.finish(), + // Nothing was ever appended, so there is no builder to ask for an empty array. + None => Canonical::empty(&self.dtype).into_array(), + }; } self.flush_pending(); @@ -142,13 +162,23 @@ impl ChildBuilder { unsafe { ChunkedArray::new_unchecked(chunks, self.dtype.clone()) }.into_array() } + /// The scalar builder, materialized on first use. + fn pending(&mut self) -> &mut dyn ArrayBuilder { + self.pending + .get_or_insert_with(|| builder_with_capacity_in(&self.dtype, self.pending_capacity, &self.allocator)) + .as_mut() + } + /// Moves whatever the scalar builder holds into `chunks`, keeping the chunks in logical order. fn flush_pending(&mut self) { - if self.pending.is_empty() { + let Some(pending) = self.pending.as_mut() else { + return; + }; + if pending.is_empty() { return; } - self.chunks_len += self.pending.len(); - let pending = self.pending.finish(); + self.chunks_len += pending.len(); + let pending = pending.finish(); self.chunks.push(pending); } } diff --git a/vortex-array/src/builders/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index e9a66b62d83..f02ac1dfccd 100644 --- a/vortex-array/src/builders/fixed_size_list.rs +++ b/vortex-array/src/builders/fixed_size_list.rs @@ -246,14 +246,14 @@ impl FixedSizeListBuilder { "elements length must be equal to the array length times the list size" ); - // TODO(connor): Use `new_unchecked` here. - FixedSizeListArray::try_new( - self.elements_builder.finish(), - self.list_size(), - self.nulls.finish_with_nullability(self.dtype.nullability()), - final_len, - ) - .vortex_expect("tried to create an invalid `FixedSizeListArray` from a builder") + unsafe { + FixedSizeListArray::new_unchecked( + self.elements_builder.finish(), + self.list_size(), + self.nulls.finish_with_nullability(self.dtype.nullability()), + final_len, + ) + } } /// The [`DType`] of the inner elements. Note that this is **not** the same as the [`DType`] of diff --git a/vortex-array/src/builders/list.rs b/vortex-array/src/builders/list.rs index 5e97f122cba..b891a3a6629 100644 --- a/vortex-array/src/builders/list.rs +++ b/vortex-array/src/builders/list.rs @@ -71,17 +71,16 @@ impl ListBuilder { Self::with_capacity_in( value_dtype, nullability, - // We arbitrarily choose 2 times the number of list scalars for the capacity of the - // elements builder since we cannot know this ahead of time. - DEFAULT_BUILDER_CAPACITY * 2, DEFAULT_BUILDER_CAPACITY, allocator, ) } - /// Create a new [`ListArray`] builder with a with the given `capacity`, as well as an initial - /// capacity for the `elements` builder (since we cannot know that ahead of time solely based on - /// the outer array `capacity`). + /// Create a new [`ListArray`] builder with the given `capacity`. + /// + /// The `elements` builder is left to size itself: how many elements the lists hold is not + /// something the outer `capacity` says, and an appended array becomes a chunk of the elements + /// rather than a copy into them. /// /// # Notes /// @@ -90,13 +89,11 @@ impl ListBuilder { pub fn with_capacity( value_dtype: Arc, nullability: Nullability, - elements_capacity: usize, capacity: usize, ) -> Self { Self::with_capacity_in( value_dtype, nullability, - elements_capacity, capacity, BufferAllocatorRef::static_ref(), ) @@ -106,12 +103,11 @@ impl ListBuilder { pub fn with_capacity_in( value_dtype: Arc, nullability: Nullability, - elements_capacity: usize, capacity: usize, allocator: &BufferAllocatorRef, ) -> Self { let elements_builder = - ChildBuilder::with_capacity(value_dtype.as_ref(), elements_capacity, allocator); + ChildBuilder::with_capacity(value_dtype.as_ref(), 2 * capacity, allocator); let mut offsets_builder = PrimitiveBuilder::::with_capacity_in(NonNullable, capacity + 1, allocator); diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index a0c326b7f37..d1ecde46961 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -95,17 +95,16 @@ impl ListViewBuilder { Self::with_capacity_in( element_dtype, nullability, - // We arbitrarily choose 2 times the number of list scalars for the capacity of the - // elements builder since we cannot know this ahead of time. - DEFAULT_BUILDER_CAPACITY * 2, DEFAULT_BUILDER_CAPACITY, allocator, ) } - /// Create a new [`ListViewArray`] builder with a with the given `capacity`, as well as an - /// initial capacity for the `elements` builder (since we cannot know that ahead of time solely - /// based on the outer array `capacity`). + /// Create a new [`ListViewArray`] builder with the given `capacity`. + /// + /// The `elements` builder is left to size itself: how many elements the lists hold is not + /// something the outer `capacity` says, and an appended array becomes a chunk of the elements + /// rather than a copy into them. /// /// # Panics /// @@ -114,13 +113,11 @@ impl ListViewBuilder { pub fn with_capacity( element_dtype: Arc, nullability: Nullability, - elements_capacity: usize, capacity: usize, ) -> Self { Self::with_capacity_in( element_dtype, nullability, - elements_capacity, capacity, BufferAllocatorRef::static_ref(), ) @@ -134,12 +131,11 @@ impl ListViewBuilder { pub fn with_capacity_in( element_dtype: Arc, nullability: Nullability, - elements_capacity: usize, capacity: usize, allocator: &BufferAllocatorRef, ) -> Self { let elements_builder = - ChildBuilder::with_capacity(&element_dtype, elements_capacity, allocator); + ChildBuilder::with_capacity(&element_dtype, 2 * capacity, allocator); let offsets_builder = PrimitiveBuilder::::with_capacity_in(Nullability::NonNullable, capacity, allocator); diff --git a/vortex-array/src/builders/map.rs b/vortex-array/src/builders/map.rs index 4169000d96f..a49a8a3531a 100644 --- a/vortex-array/src/builders/map.rs +++ b/vortex-array/src/builders/map.rs @@ -75,7 +75,6 @@ impl MapBuilder { let entries_builder = ListViewBuilder::with_capacity_in( Arc::new(map_dtype.entries_dtype()), nullability, - capacity.saturating_mul(2), capacity, allocator, ); diff --git a/vortex-array/src/builders/mod.rs b/vortex-array/src/builders/mod.rs index d5744ad3df1..08009444221 100644 --- a/vortex-array/src/builders/mod.rs +++ b/vortex-array/src/builders/mod.rs @@ -435,7 +435,6 @@ pub fn builder_with_capacity_in( DType::List(dtype, n) => Box::new(ListViewBuilder::::with_capacity_in( Arc::clone(dtype), *n, - 2 * capacity, // Arbitrarily choose 2 times the `offsets` capacity here. capacity, allocator, )), From 266df2ca7efe3ecc97cea071331de20ec216f648 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 21 Aug 2026 12:10:33 +0100 Subject: [PATCH 03/22] less Signed-off-by: Robert Kruszewski --- .../bitpacking/array/bitpack_decompress.rs | 3 -- encodings/sparse/src/canonical.rs | 4 -- .../benches/listview_builder_extend.rs | 2 - vortex-array/src/arrays/arbitrary.rs | 1 - .../src/arrays/constant/vtable/mod.rs | 1 - .../arrays/fixed_size_list/tests/nested.rs | 2 - vortex-array/src/arrays/list/test_harness.rs | 2 - vortex-array/src/arrays/list/tests.rs | 1 - vortex-array/src/builders/child.rs | 39 ++++++++++++++++++- vortex-array/src/builders/list.rs | 21 ---------- vortex-array/src/builders/listview.rs | 17 -------- vortex-array/src/builders/tests.rs | 1 - 12 files changed, 38 insertions(+), 56 deletions(-) diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs index 0684aea5e6a..b569ffd23e8 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs @@ -302,7 +302,6 @@ mod tests { Arc::clone(&element_dtype), Nullability::NonNullable, 0, - 0, ctx.allocator(), ); list.clone() @@ -314,7 +313,6 @@ mod tests { Arc::clone(&element_dtype), Nullability::NonNullable, 0, - 0, ctx.allocator(), ); list.clone() @@ -333,7 +331,6 @@ mod tests { element_dtype, Nullability::NonNullable, 0, - 0, vortex_buffer::BufferAllocatorRef::static_ref(), ); listview diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index 6232e7176be..0b417b61010 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -402,7 +402,6 @@ fn execute_sparse_lists( fill_value, values_dtype, len, - total_canonical_values, nullability, ctx, ) @@ -410,14 +409,12 @@ fn execute_sparse_lists( })) } -#[expect(clippy::too_many_arguments)] fn execute_sparse_lists_inner( patch_indices: &[I], patch_values: ListViewArray, fill_value: &Scalar, values_dtype: Arc, len: usize, - total_canonical_values: usize, nullability: Nullability, ctx: &mut ExecutionCtx, ) -> ArrayRef { @@ -426,7 +423,6 @@ fn execute_sparse_lists_inner( let mut builder = ListViewBuilder::::with_capacity_in( values_dtype, nullability, - total_canonical_values, len, ctx.allocator(), ); diff --git a/vortex-array/benches/listview_builder_extend.rs b/vortex-array/benches/listview_builder_extend.rs index 92e0c95e54b..2e8d1da5367 100644 --- a/vortex-array/benches/listview_builder_extend.rs +++ b/vortex-array/benches/listview_builder_extend.rs @@ -75,7 +75,6 @@ fn extend_from_array_zctl(bencher: Bencher, (num_lists, list_size): (usize, usiz let mut builder = ListViewBuilder::::with_capacity_in( Arc::new(DType::Primitive(I32, NonNullable)), NonNullable, - num_lists * list_size, num_lists, ctx.allocator(), ); @@ -99,7 +98,6 @@ fn extend_from_array_non_zctl_overlapping( let mut builder = ListViewBuilder::::with_capacity_in( Arc::new(DType::Primitive(I32, NonNullable)), Nullable, - num_lists * list_size, num_lists, ctx.allocator(), ); diff --git a/vortex-array/src/arrays/arbitrary.rs b/vortex-array/src/arrays/arbitrary.rs index 810157eaf41..f350d0a561f 100644 --- a/vortex-array/src/arrays/arbitrary.rs +++ b/vortex-array/src/arrays/arbitrary.rs @@ -315,7 +315,6 @@ fn random_list_with_offset_type( let mut builder = ListViewBuilder::::with_capacity_in( Arc::clone(elem_dtype), null, - array_length, 10, vortex_buffer::BufferAllocatorRef::static_ref(), ); diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index c0796a0e7d0..2709d43a768 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -627,7 +627,6 @@ mod tests { element_dtype, Nullability::NonNullable, 0, - 0, vortex_buffer::BufferAllocatorRef::static_ref(), ); array.append_to_builder(&mut builder, &mut ctx)?; diff --git a/vortex-array/src/arrays/fixed_size_list/tests/nested.rs b/vortex-array/src/arrays/fixed_size_list/tests/nested.rs index b754229ff90..f2d35be3883 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/nested.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/nested.rs @@ -411,7 +411,6 @@ fn test_fsl_of_list() { let mut list_builder = ListBuilder::::with_capacity_in( Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), Nullability::NonNullable, - 12, 6, vortex_buffer::BufferAllocatorRef::static_ref(), ); @@ -493,7 +492,6 @@ fn test_fsl_of_nullable_list() { let mut list_builder = ListBuilder::::with_capacity_in( Arc::new(DType::Primitive(PType::U16, Nullability::NonNullable)), Nullability::Nullable, - 8, 4, vortex_buffer::BufferAllocatorRef::static_ref(), ); diff --git a/vortex-array/src/arrays/list/test_harness.rs b/vortex-array/src/arrays/list/test_harness.rs index 573c4d35a3c..4130fb32b93 100644 --- a/vortex-array/src/arrays/list/test_harness.rs +++ b/vortex-array/src/arrays/list/test_harness.rs @@ -30,7 +30,6 @@ impl ListArray { let mut builder = ListBuilder::::with_capacity_in( Arc::clone(&dtype), crate::dtype::Nullability::NonNullable, - 2 * iter.size_hint().0, iter.size_hint().0, BufferAllocatorRef::static_ref(), ); @@ -58,7 +57,6 @@ impl ListArray { let mut builder = ListBuilder::::with_capacity_in( Arc::clone(&dtype), crate::dtype::Nullability::Nullable, - 2 * iter.size_hint().0, iter.size_hint().0, BufferAllocatorRef::static_ref(), ); diff --git a/vortex-array/src/arrays/list/tests.rs b/vortex-array/src/arrays/list/tests.rs index 93daf3dc4a4..4ca4dd707c8 100644 --- a/vortex-array/src/arrays/list/tests.rs +++ b/vortex-array/src/arrays/list/tests.rs @@ -389,7 +389,6 @@ fn test_offset_to_0() { let mut builder = ListBuilder::::with_capacity_in( Arc::new(I32.into()), Nullability::NonNullable, - 10, 5, vortex_buffer::BufferAllocatorRef::static_ref(), ); diff --git a/vortex-array/src/builders/child.rs b/vortex-array/src/builders/child.rs index 9f9270ebf80..9d000b04a58 100644 --- a/vortex-array/src/builders/child.rs +++ b/vortex-array/src/builders/child.rs @@ -43,7 +43,7 @@ pub struct ChildBuilder { pending: Option>, /// The capacity the scalar builder is materialized with, grown by - /// [`reserve_exact`](Self::reserve_exact). + /// [`reserve_exact`](Self::reserve_exact) until it is. pending_capacity: usize, allocator: BufferAllocatorRef, @@ -51,6 +51,9 @@ pub struct ChildBuilder { impl ChildBuilder { /// Creates a child builder with the provided allocator and capacity. + /// + /// The scalar builder is allocated on first use. Children that only receive whole arrays + /// never allocate a scalar builder. pub fn with_capacity(dtype: &DType, capacity: usize, allocator: &BufferAllocatorRef) -> Self { Self { dtype: dtype.clone(), @@ -398,4 +401,38 @@ mod tests { Ok(()) } + + /// Reserving is recorded while the scalar builder is unmaterialized and forwarded once it + /// exists; either way the scalars appended after it are the ones that come back. + #[test] + fn test_reserving_before_and_after_the_first_scalar() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); + + builder.reserve_exact(2); + builder.append_scalar(&1i32.into())?; + builder.reserve_exact(2); + builder.append_scalar(&2i32.into())?; + + let expected = PrimitiveArray::new(buffer![1i32, 2], NonNullable.into()).into_array(); + assert_arrays_eq!(&builder.finish(), &expected, &mut ctx); + + Ok(()) + } + + /// A child that is only ever reserved never materializes a scalar builder, and still finishes + /// as an empty array of its own dtype. + #[test] + fn test_reserving_alone_finishes_empty() -> VortexResult<()> { + let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); + + builder.reserve_exact(CHUNK_LEN); + + assert_eq!(builder.len(), 0); + let child = builder.finish(); + assert!(child.is_empty()); + assert!(child.is::()); + + Ok(()) + } } diff --git a/vortex-array/src/builders/list.rs b/vortex-array/src/builders/list.rs index b891a3a6629..7412aed5198 100644 --- a/vortex-array/src/builders/list.rs +++ b/vortex-array/src/builders/list.rs @@ -78,10 +78,6 @@ impl ListBuilder { /// Create a new [`ListArray`] builder with the given `capacity`. /// - /// The `elements` builder is left to size itself: how many elements the lists hold is not - /// something the outer `capacity` says, and an appended array becomes a chunk of the elements - /// rather than a copy into them. - /// /// # Notes /// /// The number of offsets is one more than the length (# of list scalars) in the array. @@ -456,7 +452,6 @@ mod tests { Arc::new(I32.into()), NonNullable, 0, - 0, BufferAllocatorRef::static_ref(), ); @@ -471,7 +466,6 @@ mod tests { Arc::clone(&dtype), NonNullable, 0, - 0, BufferAllocatorRef::static_ref(), ); @@ -514,7 +508,6 @@ mod tests { Arc::clone(&dtype), NonNullable, 0, - 0, BufferAllocatorRef::static_ref(), ); @@ -532,7 +525,6 @@ mod tests { Arc::clone(&dtype), Nullable, 0, - 0, BufferAllocatorRef::static_ref(), ); @@ -587,7 +579,6 @@ mod tests { let mut builder = ListBuilder::::with_capacity_in( Arc::new(I32.into()), Nullable, - 18, 9, BufferAllocatorRef::static_ref(), ); @@ -672,7 +663,6 @@ mod tests { let mut lv_u64_u32 = ListViewBuilder::::with_capacity_in( elem_dtype(), Nullable, - 8, 4, BufferAllocatorRef::static_ref(), ); @@ -682,7 +672,6 @@ mod tests { let mut lv_i64_i32 = ListViewBuilder::::with_capacity_in( elem_dtype(), Nullable, - 8, 4, BufferAllocatorRef::static_ref(), ); @@ -692,7 +681,6 @@ mod tests { let mut lv_u32_u32 = ListViewBuilder::::with_capacity_in( elem_dtype(), Nullable, - 8, 4, BufferAllocatorRef::static_ref(), ); @@ -704,7 +692,6 @@ mod tests { let mut list_builder = ListBuilder::::with_capacity_in( elem_dtype(), Nullable, - 8, 4, BufferAllocatorRef::static_ref(), ); @@ -714,7 +701,6 @@ mod tests { let mut list_builder_i32 = ListBuilder::::with_capacity_in( elem_dtype(), Nullable, - 8, 4, BufferAllocatorRef::static_ref(), ); @@ -744,7 +730,6 @@ mod tests { Arc::clone(&dtype), Nullable, 0, - 0, BufferAllocatorRef::static_ref(), ); builder.append_list_array(source.as_view(), &mut ctx)?; @@ -756,7 +741,6 @@ mod tests { dtype, Nullable, 0, - 0, BufferAllocatorRef::static_ref(), ); builder.append_listview_array(source_listview.as_view(), &mut ctx)?; @@ -807,7 +791,6 @@ mod tests { dtype, Nullable, 0, - 0, BufferAllocatorRef::static_ref(), ); builder.append_listview_array(overlapping.as_view(), &mut ctx)?; @@ -893,7 +876,6 @@ mod tests { let mut builder = ListBuilder::::with_capacity_in( Arc::clone(&dtype), Nullable, - 20, 10, BufferAllocatorRef::static_ref(), ); @@ -970,7 +952,6 @@ mod tests { let mut builder = ListBuilder::::with_capacity_in( dtype, NonNullable, - 20, 10, BufferAllocatorRef::static_ref(), ); @@ -985,7 +966,6 @@ mod tests { let mut builder = ListBuilder::::with_capacity_in( Arc::clone(&dtype), NonNullable, - 20, 10, BufferAllocatorRef::static_ref(), ); @@ -1040,7 +1020,6 @@ mod tests { let mut builder = ListBuilder::::with_capacity_in( dtype, NonNullable, - 20, 10, BufferAllocatorRef::static_ref(), ); diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index d1ecde46961..f83dccd7176 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -102,10 +102,6 @@ impl ListViewBuilder { /// Create a new [`ListViewArray`] builder with the given `capacity`. /// - /// The `elements` builder is left to size itself: how many elements the lists hold is not - /// something the outer `capacity` says, and an appended array becomes a chunk of the elements - /// rather than a copy into them. - /// /// # Panics /// /// Panics if the size type `S` cannot fit within the offset type `O`. @@ -696,7 +692,6 @@ mod tests { Arc::new(I32.into()), NonNullable, 0, - 0, BufferAllocatorRef::static_ref(), ); @@ -712,7 +707,6 @@ mod tests { Arc::clone(&dtype), Nullable, 0, - 0, BufferAllocatorRef::static_ref(), ); @@ -782,7 +776,6 @@ mod tests { Arc::clone(&dtype), NonNullable, 0, - 0, BufferAllocatorRef::static_ref(), ); @@ -831,7 +824,6 @@ mod tests { Arc::clone(&dtype2), NonNullable, 0, - 0, BufferAllocatorRef::static_ref(), ); @@ -864,7 +856,6 @@ mod tests { Arc::clone(&dtype), Nullable, 0, - 0, BufferAllocatorRef::static_ref(), ); @@ -930,7 +921,6 @@ mod tests { Arc::clone(&dtype), Nullable, 0, - 0, BufferAllocatorRef::static_ref(), ); @@ -1012,7 +1002,6 @@ mod tests { Arc::clone(&dtype), Nullable, 0, - 0, BufferAllocatorRef::static_ref(), ); builder.append_list_array(source.as_view(), &mut ctx)?; @@ -1050,7 +1039,6 @@ mod tests { element_dtype, NonNullable, 0, - 0, BufferAllocatorRef::static_ref(), ); constant.append_to_builder(&mut builder, &mut ctx)?; @@ -1088,7 +1076,6 @@ mod tests { Arc::clone(&dtype), NonNullable, 0, - 0, BufferAllocatorRef::static_ref(), ); builder.append_listview_array(middle.as_view(), &mut ctx)?; @@ -1138,7 +1125,6 @@ mod tests { Arc::clone(&dtype), Nullable, 0, - 0, BufferAllocatorRef::static_ref(), ); builder @@ -1178,7 +1164,6 @@ mod tests { Arc::clone(&dtype), NonNullable, 0, - 0, BufferAllocatorRef::static_ref(), ); @@ -1204,7 +1189,6 @@ mod tests { let mut builder = ListViewBuilder::::with_capacity_in( Arc::clone(&dtype), NonNullable, - 20, 10, BufferAllocatorRef::static_ref(), ); @@ -1266,7 +1250,6 @@ mod tests { let mut builder = ListViewBuilder::::with_capacity_in( dtype, NonNullable, - 20, 10, BufferAllocatorRef::static_ref(), ); diff --git a/vortex-array/src/builders/tests.rs b/vortex-array/src/builders/tests.rs index 828cde17020..31e245455dc 100644 --- a/vortex-array/src/builders/tests.rs +++ b/vortex-array/src/builders/tests.rs @@ -1264,7 +1264,6 @@ fn test_list_builder_offsets_are_rebased_across_element_chunks() -> VortexResult element_dtype, Nullability::NonNullable, 0, - 0, BufferAllocatorRef::static_ref(), ); for value in 0..3i32 { From 6e1eb3f3e775c15136b46990e409c38d75a0b5f6 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 21 Aug 2026 12:30:02 +0100 Subject: [PATCH 04/22] builders Signed-off-by: Robert Kruszewski --- encodings/sparse/src/canonical.rs | 8 +-- .../src/arrays/chunked/vtable/canonical.rs | 63 +++++++++++++++++++ vortex-array/src/arrays/extension/array.rs | 12 +++- vortex-array/src/arrays/varbin/builder.rs | 4 +- vortex-array/src/builders/bool.rs | 13 ++-- vortex-array/src/builders/child.rs | 10 ++- vortex-array/src/builders/decimal.rs | 13 +++- vortex-array/src/builders/dict/bytes.rs | 3 +- vortex-array/src/builders/dict/primitive.rs | 9 ++- vortex-array/src/builders/extension.rs | 4 +- vortex-array/src/builders/fixed_size_list.rs | 3 +- vortex-array/src/builders/list.rs | 22 ++++--- vortex-array/src/builders/listview.rs | 6 +- vortex-array/src/builders/map.rs | 12 ++-- vortex-array/src/builders/primitive.rs | 13 +++- vortex-array/src/builders/struct_.rs | 22 +++---- 16 files changed, 167 insertions(+), 50 deletions(-) diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index 0b417b61010..ffb4281d40e 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -420,12 +420,8 @@ fn execute_sparse_lists_inner( ) -> ArrayRef { // Create the builder with appropriate types. It is easy to just use the same type for both // `offsets` and `sizes` since we have no other constraints. - let mut builder = ListViewBuilder::::with_capacity_in( - values_dtype, - nullability, - len, - ctx.allocator(), - ); + let mut builder = + ListViewBuilder::::with_capacity_in(values_dtype, nullability, len, ctx.allocator()); // The fill's elements become an array once, up front. Every gap then appends that same array, // so the fill's elements are stored once for the whole result however many gaps reference them. let fill_elements = list_scalar_elements_array(fill_value.as_list(), ctx.allocator()); diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index 86df7fd3850..68f471763ae 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -107,10 +107,18 @@ fn pack_variant_chunks( #[cfg(test)] mod tests { + use std::alloc::Layout; + use std::ptr::NonNull; use std::sync::Arc; use std::sync::LazyLock; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use allocator_api2::alloc::AllocError; + use allocator_api2::alloc::Allocator; + use allocator_api2::alloc::Global; use rstest::rstest; + use vortex_buffer::BufferAllocatorRef; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -146,12 +154,31 @@ mod tests { use crate::dtype::DType::Variant as VariantDType; use crate::dtype::Nullability::NonNullable; use crate::dtype::PType::I32; + use crate::memory::MemorySessionExt; use crate::scalar::Scalar; use crate::validity::Validity; /// A shared session for these chunked-array tests, used to create execution contexts. static SESSION: LazyLock = LazyLock::new(crate::array_session); + #[derive(Debug)] + struct CountingAllocator { + allocations: Arc, + } + + // SAFETY: this forwards memory operations to Global and only counts allocations. + unsafe impl Allocator for CountingAllocator { + fn allocate(&self, layout: Layout) -> Result, AllocError> { + self.allocations.fetch_add(1, Ordering::Relaxed); + Global.allocate(layout) + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + // SAFETY: ptr and layout came from Global. + unsafe { Global.deallocate(ptr, layout) } + } + } + fn variant_scalar(value: i32) -> Scalar { Scalar::variant(Scalar::primitive(value, NonNullable)) } @@ -517,4 +544,40 @@ mod tests { Ok(()) } + + #[test] + fn list_canonicalize_uses_memory_session_allocator() { + let allocations = Arc::new(AtomicUsize::new(0)); + let session = + crate::array_session().with_allocator(BufferAllocatorRef::new(CountingAllocator { + allocations: Arc::clone(&allocations), + })); + let mut ctx = session.create_execution_ctx(); + + let l1 = ListArray::try_new( + buffer![1, 2, 3, 4].into_array(), + buffer![0, 3].into_array(), + Validity::NonNullable, + ) + .unwrap(); + let l2 = ListArray::try_new( + buffer![5, 6].into_array(), + buffer![0, 2].into_array(), + Validity::NonNullable, + ) + .unwrap(); + + let chunked_list = ChunkedArray::try_new( + vec![l1.into_array(), l2.into_array()], + List(Arc::new(Primitive(I32, NonNullable)), NonNullable), + ) + .unwrap() + .into_array(); + + drop(chunked_list.execute::(&mut ctx).unwrap()); + assert!( + allocations.load(Ordering::Relaxed) >= 2, + "expected offset+size allocations through MemorySession" + ); + } } diff --git a/vortex-array/src/arrays/extension/array.rs b/vortex-array/src/arrays/extension/array.rs index 144e2d82622..475c94b2567 100644 --- a/vortex-array/src/arrays/extension/array.rs +++ b/vortex-array/src/arrays/extension/array.rs @@ -57,6 +57,16 @@ impl Array { "Tried to create an `ExtensionArray` with an incompatible storage array" ); + // SAFETY: the storage array's dtype was just checked against the extension dtype. + Ok(unsafe { Self::new_unchecked(ext_dtype, storage_array) }) + } + + /// Creates a new `ExtensionArray` without validating its storage array. + /// + /// # Safety + /// + /// The caller must ensure that `storage_array` has the extension dtype's storage dtype. + pub unsafe fn new_unchecked(ext_dtype: ExtDTypeRef, storage_array: ArrayRef) -> Self { let dtype = DType::Extension(ext_dtype); let len = storage_array.len(); @@ -67,7 +77,7 @@ impl Array { .into_slots(), ); - Ok(unsafe { Array::from_parts_unchecked(parts) }) + unsafe { Array::from_parts_unchecked(parts) } } /// Creates a new [`ExtensionArray`](crate::arrays::ExtensionArray) from a vtable, metadata, and diff --git a/vortex-array/src/arrays/varbin/builder.rs b/vortex-array/src/arrays/varbin/builder.rs index 39e694dd164..411448d6071 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -415,7 +415,9 @@ impl VarBinBuilder { let offsets = self.offsets.take(); self.offsets.push(O::zero()); - let offsets = PrimitiveArray::new(offsets.freeze(), Validity::NonNullable); + // SAFETY: a non-nullable validity carries no length to match against the offsets. + let offsets = + unsafe { PrimitiveArray::new_unchecked(offsets.freeze(), Validity::NonNullable) }; let data = self.data.take(); let nulls = self.validity.take().freeze(); diff --git a/vortex-array/src/builders/bool.rs b/vortex-array/src/builders/bool.rs index b5371d3a146..ffe7e138853 100644 --- a/vortex-array/src/builders/bool.rs +++ b/vortex-array/src/builders/bool.rs @@ -79,11 +79,14 @@ impl BoolBuilder { "Null count and value count should match when calling BoolBuilder::finish." ); - let inner = self.inner.take().freeze(); - BoolArray::new( - inner, - self.nulls.finish_with_nullability(self.dtype.nullability()), - ) + // SAFETY: the assert above establishes the only invariant, that the validity covers + // exactly the bits appended. + unsafe { + BoolArray::new_unchecked( + self.inner.take().freeze(), + self.nulls.finish_with_nullability(self.dtype.nullability()), + ) + } } pub(crate) fn append_bool_array( diff --git a/vortex-array/src/builders/child.rs b/vortex-array/src/builders/child.rs index 9d000b04a58..609a39cc3d2 100644 --- a/vortex-array/src/builders/child.rs +++ b/vortex-array/src/builders/child.rs @@ -168,7 +168,9 @@ impl ChildBuilder { /// The scalar builder, materialized on first use. fn pending(&mut self) -> &mut dyn ArrayBuilder { self.pending - .get_or_insert_with(|| builder_with_capacity_in(&self.dtype, self.pending_capacity, &self.allocator)) + .get_or_insert_with(|| { + builder_with_capacity_in(&self.dtype, self.pending_capacity, &self.allocator) + }) .as_mut() } @@ -407,7 +409,8 @@ mod tests { #[test] fn test_reserving_before_and_after_the_first_scalar() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); + let mut builder = + ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); builder.reserve_exact(2); builder.append_scalar(&1i32.into())?; @@ -424,7 +427,8 @@ mod tests { /// as an empty array of its own dtype. #[test] fn test_reserving_alone_finishes_empty() -> VortexResult<()> { - let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); + let mut builder = + ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); builder.reserve_exact(CHUNK_LEN); diff --git a/vortex-array/src/builders/decimal.rs b/vortex-array/src/builders/decimal.rs index 6d848b82aba..e953f34ee65 100644 --- a/vortex-array/src/builders/decimal.rs +++ b/vortex-array/src/builders/decimal.rs @@ -184,7 +184,18 @@ impl DecimalBuilder { let decimal_dtype = *self.decimal_dtype(); delegate_fn!(self.values.take(), |T, values| { - DecimalArray::new::(values.freeze(), decimal_dtype, validity) + assert!( + validity.maybe_len().is_none_or(|len| len == values.len()), + "validity of length {:?} does not cover the {} values appended", + validity.maybe_len(), + values.len() + ); + + // SAFETY: the values are this builder's own `Buffer`, so they are aligned and a + // whole number of elements, and the assert above pairs an array-backed validity with + // exactly those values. Whether they fit the precision is the appender's business, as + // it is on the checked path, which does not verify it either. + unsafe { DecimalArray::new_unchecked::(values.freeze(), decimal_dtype, validity) } }) } diff --git a/vortex-array/src/builders/dict/bytes.rs b/vortex-array/src/builders/dict/bytes.rs index 4e822967df1..197a01745b3 100644 --- a/vortex-array/src/builders/dict/bytes.rs +++ b/vortex-array/src/builders/dict/bytes.rs @@ -242,7 +242,8 @@ impl BytesDictBuilder { // Restore lookup dictionary back into the struct self.lookup = Some(local_lookup); - Ok(PrimitiveArray::new(codes, Validity::NonNullable)) + // SAFETY: a non-nullable validity carries no length to match against the codes. + Ok(unsafe { PrimitiveArray::new_unchecked(codes.freeze(), Validity::NonNullable) }) } fn encode_varbin( diff --git a/vortex-array/src/builders/dict/primitive.rs b/vortex-array/src/builders/dict/primitive.rs index ba2e86c6ee3..db94daee3e7 100644 --- a/vortex-array/src/builders/dict/primitive.rs +++ b/vortex-array/src/builders/dict/primitive.rs @@ -193,7 +193,8 @@ where } } - Ok(PrimitiveArray::new(codes, Validity::NonNullable)) + // SAFETY: a non-nullable validity carries no length to match against the codes. + Ok(unsafe { PrimitiveArray::new_unchecked(codes.freeze(), Validity::NonNullable) }) } fn reset(&mut self) -> ArrayRef { @@ -208,7 +209,11 @@ where BitBufferMut::empty_in(self.allocator.clone()), ) .freeze(); - PrimitiveArray::new(values, Validity::from_bit_buffer(nulls, self.nullability)).into_array() + let validity = Validity::from_bit_buffer(nulls, self.nullability); + + // SAFETY: every dictionary value is pushed alongside its validity bit, so an array-backed + // validity covers exactly the values. + unsafe { PrimitiveArray::new_unchecked(values.freeze(), validity) }.into_array() } fn codes_ptype(&self) -> PType { diff --git a/vortex-array/src/builders/extension.rs b/vortex-array/src/builders/extension.rs index 544b119a0a6..3f808680aab 100644 --- a/vortex-array/src/builders/extension.rs +++ b/vortex-array/src/builders/extension.rs @@ -75,7 +75,9 @@ impl ExtensionBuilder { /// Finishes the builder directly into a [`ExtensionArray`]. pub fn finish_into_extension(&mut self) -> ExtensionArray { let storage = self.storage.finish(); - ExtensionArray::new(self.ext_dtype(), storage) + // SAFETY: the storage child builder was created with this extension dtype's storage dtype + // and rejects an appended array of any other dtype. + unsafe { ExtensionArray::new_unchecked(self.ext_dtype(), storage) } } /// The [`ExtDType`] of this builder. diff --git a/vortex-array/src/builders/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index f02ac1dfccd..fec5669087c 100644 --- a/vortex-array/src/builders/fixed_size_list.rs +++ b/vortex-array/src/builders/fixed_size_list.rs @@ -5,7 +5,6 @@ use std::any::Any; use std::sync::Arc; use vortex_buffer::BufferAllocatorRef; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -246,6 +245,8 @@ impl FixedSizeListBuilder { "elements length must be equal to the array length times the list size" ); + // SAFETY: the assert above checks the element count, and every append records one + // validity entry per list. The child builder preserves the element dtype. unsafe { FixedSizeListArray::new_unchecked( self.elements_builder.finish(), diff --git a/vortex-array/src/builders/list.rs b/vortex-array/src/builders/list.rs index 7412aed5198..1f5b41621da 100644 --- a/vortex-array/src/builders/list.rs +++ b/vortex-array/src/builders/list.rs @@ -85,6 +85,7 @@ impl ListBuilder { pub fn with_capacity( value_dtype: Arc, nullability: Nullability, + _elements_capacity: usize, capacity: usize, ) -> Self { Self::with_capacity_in( @@ -95,7 +96,7 @@ impl ListBuilder { ) } - /// Creates a list builder with the given capacities using `allocator`. + /// Creates a list builder with the given capacity using `allocator`. pub fn with_capacity_in( value_dtype: Arc, nullability: Nullability, @@ -181,12 +182,19 @@ impl ListBuilder { "offsets length must be one more than nulls length." ); - ListArray::try_new( - self.elements_builder.finish(), - self.offsets_builder.finish(), - self.nulls.finish_with_nullability(self.dtype.nullability()), - ) - .vortex_expect("Buffer, offsets, and validity must have same length.") + // SAFETY: + // - The offsets come from a non-nullable integer builder seeded with a leading zero, so + // they are a non-nullable integer array of at least one non-negative value. + // - Every append records `elements_builder.len()` after appending its elements, so the + // offsets increase monotonically and the last one is exactly the elements length. + // - The assert above pairs one validity entry with each list. + unsafe { + ListArray::new_unchecked( + self.elements_builder.finish(), + self.offsets_builder.finish(), + self.nulls.finish_with_nullability(self.dtype.nullability()), + ) + } } /// The [`DType`] of the inner elements. Note that this is **not** the same as the [`DType`] of diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index f83dccd7176..3ff92e105cd 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -109,6 +109,7 @@ impl ListViewBuilder { pub fn with_capacity( element_dtype: Arc, nullability: Nullability, + _elements_capacity: usize, capacity: usize, ) -> Self { Self::with_capacity_in( @@ -119,7 +120,7 @@ impl ListViewBuilder { ) } - /// Creates a list-view builder with the given capacities using `allocator`. + /// Creates a list-view builder with the given capacity using `allocator`. /// /// # Panics /// @@ -130,8 +131,7 @@ impl ListViewBuilder { capacity: usize, allocator: &BufferAllocatorRef, ) -> Self { - let elements_builder = - ChildBuilder::with_capacity(&element_dtype, 2 * capacity, allocator); + let elements_builder = ChildBuilder::with_capacity(&element_dtype, 2 * capacity, allocator); let offsets_builder = PrimitiveBuilder::::with_capacity_in(Nullability::NonNullable, capacity, allocator); diff --git a/vortex-array/src/builders/map.rs b/vortex-array/src/builders/map.rs index a49a8a3531a..05b8264cdb4 100644 --- a/vortex-array/src/builders/map.rs +++ b/vortex-array/src/builders/map.rs @@ -111,10 +111,14 @@ impl MapBuilder { /// Finishes the builder directly into a [`MapArray`]. pub fn finish_into_map(&mut self) -> MapArray { - MapArray::new( - self.map_dtype.clone(), - self.entries_builder.finish_into_listview(), - ) + // SAFETY: the entries builder was constructed from `map_dtype.entries_dtype()` with the + // map's nullability, so its lists hold exactly that entry struct. + unsafe { + MapArray::new_unchecked( + self.map_dtype.clone(), + self.entries_builder.finish_into_listview(), + ) + } } /// Appends the values of a [`Map`]-encoded `array` to this builder. diff --git a/vortex-array/src/builders/primitive.rs b/vortex-array/src/builders/primitive.rs index 5f43e059a34..42369ee56ad 100644 --- a/vortex-array/src/builders/primitive.rs +++ b/vortex-array/src/builders/primitive.rs @@ -134,12 +134,21 @@ impl PrimitiveBuilder { /// Finishes the builder directly into a [`PrimitiveArray`]. pub fn finish_into_primitive(&mut self) -> PrimitiveArray { + let values = self.values.take().freeze(); let validity = self .nulls .finish_with_nullability(self.dtype().nullability()); - let values = self.values.take().freeze(); - PrimitiveArray::new(values, validity) + assert!( + validity.maybe_len().is_none_or(|len| len == values.len()), + "validity of length {:?} does not cover the {} values appended", + validity.maybe_len(), + values.len() + ); + + // SAFETY: the assert above establishes the only invariant, that an array-backed validity + // covers exactly the values appended. The other validities carry no length at all. + unsafe { PrimitiveArray::new_unchecked(values, validity) } } /// Extends the primitive array with an iterator. diff --git a/vortex-array/src/builders/struct_.rs b/vortex-array/src/builders/struct_.rs index 21f4361ee36..9eaa95a9198 100644 --- a/vortex-array/src/builders/struct_.rs +++ b/vortex-array/src/builders/struct_.rs @@ -5,7 +5,6 @@ use std::any::Any; use itertools::Itertools; use vortex_buffer::BufferAllocatorRef; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -125,21 +124,20 @@ impl StructBuilder { .map(|builder| builder.finish()) .collect::>(); - if fields.len() > 1 { - let expected_length = fields[0].len(); - for (index, field) in fields[1..].iter().enumerate() { - assert_eq!( - field.len(), - expected_length, - "Field {index} does not have expected length {expected_length}" - ); - } + for (index, field) in fields.iter().enumerate() { + assert_eq!( + field.len(), + len, + "Field {index} does not have expected length {len}" + ); } let validity = self.nulls.finish_with_nullability(self.dtype.nullability()); - StructArray::try_new_with_dtype(fields, self.struct_fields().clone(), len, validity) - .vortex_expect("Fields must all have same length.") + // SAFETY: each field is finished by the child builder this builder made for that field of + // the struct dtype, so the fields match the dtype in order, and the assert above pairs + // every one of them with the `len` rows the validity covers. + unsafe { StructArray::new_unchecked(fields, self.struct_fields().clone(), len, validity) } } /// The [`StructFields`] of this struct builder. From 235a32cc58437e591c7ec4e352a6139646a08106 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 21 Aug 2026 12:34:35 +0100 Subject: [PATCH 05/22] DCO Remediation Commit for Robert Kruszewski I, Robert Kruszewski , hereby add my Signed-off-by to this commit: 5c8401d8cf065b99697d54bbec7084089d14a505 Signed-off-by: Robert Kruszewski From 09a414178adf801738674f6df6775506c5adbfe9 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 16 Sep 2026 03:47:31 +0000 Subject: [PATCH 06/22] perf(array): append runs of empty lists in bulk `ListBuilder` and `ListViewBuilder` wrote one offset - and for list views one size - per row when appending empty or null lists, pushing the same value `n` times into a `PrimitiveBuilder`. `append_n_values` writes the run in one go. This is not a rare path: a sparse list array with a null or empty fill calls `append_nulls(count)` with the gap between patches, which can be millions of rows. `append_nulls(1_000_000)` on a `ListViewBuilder` goes from 22.2ms to 979us. Also let a builder be told how many arrays are about to be appended. `Chunked::append_to_builder` hands over one array per chunk, and a nested builder keeps each as a chunk of its child, so the chunk lists grew one push at a time. `reserve_chunks` lets them grow once. In isolation this is worth 36% of `ChildBuilder::append_array` at 100k chunks; end to end it measures as noise, since canonicalization is dominated by the per-chunk work around it. Signed-off-by: Claude Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01CZTu9XUrdRZZ91rx852DGf --- vortex-array/src/arrays/chunked/vtable/mod.rs | 3 + vortex-array/src/builders/child.rs | 7 ++ vortex-array/src/builders/extension.rs | 4 ++ vortex-array/src/builders/fixed_size_list.rs | 5 ++ vortex-array/src/builders/list.rs | 63 ++++++++++++++---- vortex-array/src/builders/listview.rs | 65 +++++++++++++++---- vortex-array/src/builders/map.rs | 4 ++ vortex-array/src/builders/mod.rs | 11 ++++ vortex-array/src/builders/struct_.rs | 7 ++ vortex-array/src/builders/validity.rs | 8 +++ 10 files changed, 153 insertions(+), 24 deletions(-) diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index 6f0462c6c29..19dd7cab5fb 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -239,6 +239,9 @@ impl VTable for Chunked { builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { + // The builder is about to receive one array per chunk, so let any chunk lists it keeps + // grow once rather than on the way through. + builder.reserve_chunks(array.nchunks()); for chunk in array.iter_chunks() { chunk.append_to_builder(builder, ctx)?; } diff --git a/vortex-array/src/builders/child.rs b/vortex-array/src/builders/child.rs index 609a39cc3d2..a88c5c562d6 100644 --- a/vortex-array/src/builders/child.rs +++ b/vortex-array/src/builders/child.rs @@ -143,6 +143,13 @@ impl ChildBuilder { } } + /// Reserves room for `additional` more chunks. + /// + /// See [`ArrayBuilder::reserve_chunks`]. + pub fn reserve_chunks(&mut self, additional: usize) { + self.chunks.reserve(additional); + } + /// Finishes the child, combining the accumulated chunks into a [`ChunkedArray`] when there is /// more than one of them. pub fn finish(&mut self) -> ArrayRef { diff --git a/vortex-array/src/builders/extension.rs b/vortex-array/src/builders/extension.rs index 3f808680aab..c2f4be2df5a 100644 --- a/vortex-array/src/builders/extension.rs +++ b/vortex-array/src/builders/extension.rs @@ -128,6 +128,10 @@ impl ArrayBuilder for ExtensionBuilder { self.append_value(scalar.as_extension()) } + fn reserve_chunks(&mut self, additional: usize) { + self.storage.reserve_chunks(additional) + } + fn reserve_exact(&mut self, capacity: usize) { self.storage.reserve_exact(capacity) } diff --git a/vortex-array/src/builders/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index fec5669087c..f9a42425ce7 100644 --- a/vortex-array/src/builders/fixed_size_list.rs +++ b/vortex-array/src/builders/fixed_size_list.rs @@ -339,6 +339,11 @@ impl ArrayBuilder for FixedSizeListBuilder { /// This will increase the capacity if extending with this `array` would go past the original /// capacity. + fn reserve_chunks(&mut self, additional: usize) { + self.elements_builder.reserve_chunks(additional); + self.nulls.reserve_runs(additional); + } + fn reserve_exact(&mut self, additional: usize) { self.elements_builder .reserve_exact(additional * self.list_size() as usize); diff --git a/vortex-array/src/builders/list.rs b/vortex-array/src/builders/list.rs index 1f5b41621da..72c9b859519 100644 --- a/vortex-array/src/builders/list.rs +++ b/vortex-array/src/builders/list.rs @@ -371,23 +371,21 @@ impl ArrayBuilder for ListBuilder { fn append_zeros(&mut self, n: usize) { let curr_len = self.elements_builder.len(); - for _ in 0..n { - self.offsets_builder.append_value( - O::from_usize(curr_len).vortex_expect("Failed to convert from usize to "), - ) - } + self.offsets_builder.append_n_values( + O::from_usize(curr_len).vortex_expect("Failed to convert from usize to "), + n, + ); self.nulls.append_n_non_nulls(n); } unsafe fn append_nulls_unchecked(&mut self, n: usize) { let curr_len = self.elements_builder.len(); - for _ in 0..n { - // A list with a null element is can be a list with a zero-span offset and a validity - // bit set - self.offsets_builder.append_value( - O::from_usize(curr_len).vortex_expect("Failed to convert from usize to "), - ) - } + // A list with a null element can be a list with a zero-span offset and a cleared validity + // bit. + self.offsets_builder.append_n_values( + O::from_usize(curr_len).vortex_expect("Failed to convert from usize to "), + n, + ); self.nulls.append_n_nulls(n); } @@ -402,6 +400,11 @@ impl ArrayBuilder for ListBuilder { self.append_value(scalar.as_list()) } + fn reserve_chunks(&mut self, additional: usize) { + self.elements_builder.reserve_chunks(additional); + self.nulls.reserve_runs(additional); + } + fn reserve_exact(&mut self, additional: usize) { self.elements_builder.reserve_exact(additional); self.offsets_builder.reserve_exact(additional); @@ -427,6 +430,7 @@ mod tests { use Nullability::NonNullable; use Nullability::Nullable; + use rstest::rstest; use vortex_buffer::BufferAllocatorRef; use vortex_buffer::buffer; use vortex_error::VortexExpect; @@ -454,6 +458,41 @@ mod tests { use crate::scalar::Scalar; use crate::validity::Validity; + /// The `ListBuilder` twin of the `ListViewBuilder` bulk empty-list case: `n` empty lists must + /// all land at the current end of `elements`, with validity from the method that was called. + #[rstest] + #[case::zeros(false)] + #[case::nulls(true)] + fn test_bulk_empty_lists(#[case] null: bool) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype: Arc = Arc::new(I32.into()); + let mut builder = ListBuilder::::with_capacity_in( + Arc::clone(&dtype), + Nullable, + 0, + BufferAllocatorRef::static_ref(), + ); + + builder.append_array_as_list(&buffer![1i32, 2].into_array(), &mut ctx)?; + + const N: usize = 5; + if null { + builder.append_nulls(N); + } else { + builder.append_zeros(N); + } + + let list = builder.finish_into_list(); + assert_eq!(list.len(), N + 1); + let listview = list.into_array().execute::(&mut ctx)?; + for i in 1..=N { + assert_eq!(listview.size_at(i), 0); + assert_eq!(listview.validity()?.execute_is_valid(i, &mut ctx)?, !null,); + } + + Ok(()) + } + #[test] fn test_empty() { let mut builder = ListBuilder::::with_capacity_in( diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index 3ff92e105cd..6eb43b31d35 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -479,12 +479,11 @@ impl ArrayBuilder for ListViewBuil // Since we consider the "zero" element of a list an empty list, we simply update the // `offsets` and `sizes` metadata to add an empty list. - for _ in 0..n { - self.offsets_builder.append_value( - O::from_usize(curr_offset).vortex_expect("Failed to convert from usize to `O`"), - ); - self.sizes_builder.append_value(S::zero()); - } + self.offsets_builder.append_n_values( + O::from_usize(curr_offset).vortex_expect("Failed to convert from usize to `O`"), + n, + ); + self.sizes_builder.append_n_values(S::zero(), n); self.nulls.append_n_non_nulls(n); } @@ -497,12 +496,11 @@ impl ArrayBuilder for ListViewBuil let curr_offset = self.elements_builder.len(); // A null list can have any representation, but we choose to use the zero representation. - for _ in 0..n { - self.offsets_builder.append_value( - O::from_usize(curr_offset).vortex_expect("Failed to convert from usize to `O`"), - ); - self.sizes_builder.append_value(S::zero()); - } + self.offsets_builder.append_n_values( + O::from_usize(curr_offset).vortex_expect("Failed to convert from usize to `O`"), + n, + ); + self.sizes_builder.append_n_values(S::zero(), n); // This is the only difference from `append_zeros`. self.nulls.append_n_nulls(n); @@ -520,6 +518,11 @@ impl ArrayBuilder for ListViewBuil self.append_value(list_scalar) } + fn reserve_chunks(&mut self, additional: usize) { + self.elements_builder.reserve_chunks(additional); + self.nulls.reserve_runs(additional); + } + fn reserve_exact(&mut self, capacity: usize) { self.elements_builder.reserve_exact(capacity * 2); self.offsets_builder.reserve_exact(capacity); @@ -662,6 +665,7 @@ fn extend_converted_sizes( mod tests { use std::sync::Arc; + use rstest::rstest; use vortex_buffer::BufferAllocatorRef; use vortex_buffer::buffer; use vortex_error::VortexExpect; @@ -848,6 +852,43 @@ mod tests { } } + /// Empty lists are appended in bulk rather than one offset/size pair at a time, so a large + /// `n` has to land exactly as a per-row loop would: `n` empty lists, all sharing the current + /// end of `elements`, with validity taken from which method was called. + #[rstest] + #[case::zeros(false)] + #[case::nulls(true)] + fn test_bulk_empty_lists(#[case] null: bool) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype: Arc = Arc::new(I32.into()); + let mut builder = ListViewBuilder::::with_capacity_in( + Arc::clone(&dtype), + Nullable, + 0, + BufferAllocatorRef::static_ref(), + ); + + // Put something in `elements` first, so the empty lists have a non-zero offset to share. + builder.append_array_as_list(&buffer![1i32, 2].into_array(), &mut ctx)?; + + const N: usize = 5; + if null { + builder.append_nulls(N); + } else { + builder.append_zeros(N); + } + + let listview = builder.finish_into_listview(); + assert_eq!(listview.len(), N + 1); + for i in 1..=N { + assert_eq!(listview.size_at(i), 0); + assert_eq!(listview.offset_at(i), 2); + assert_eq!(listview.validity()?.execute_is_valid(i, &mut ctx)?, !null); + } + + Ok(()) + } + #[test] fn test_builder_trait_methods() { let mut ctx = array_session().create_execution_ctx(); diff --git a/vortex-array/src/builders/map.rs b/vortex-array/src/builders/map.rs index 05b8264cdb4..1b13e495451 100644 --- a/vortex-array/src/builders/map.rs +++ b/vortex-array/src/builders/map.rs @@ -173,6 +173,10 @@ impl ArrayBuilder for MapBuilder Date: Wed, 16 Sep 2026 04:14:20 +0000 Subject: [PATCH 07/22] refactor(array): append typed arrays to builders through ArrayView The vtable hands `append_to_builder` an `ArrayView<'_, V>` - two borrows, `Copy` - while the typed builder methods took `&Array`. Bridging the two meant `array.into_owned()`, which clones the `ArrayRef`, so every appended array paid an atomic refcount pair for a builder that only reads from it. The builders read through `TypedArrayRef`, which `ArrayView` implements, so taking the view directly needs no other change. `BuffersWithOffsets:: from_array` and `VarBinViewBuilder::push_view` follow their caller. `buffer_utilizations` and the two helpers it uses were inherent to `VarBinViewArray` and so unreachable from a view. They read nothing an `ArrayView` cannot supply, so they move to a blanket `VarBinViewCompactExt` over `TypedArrayRef`; the compaction entry points stay inherent. Signed-off-by: Claude Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01CZTu9XUrdRZZ91rx852DGf --- vortex-array/src/arrays/bool/vtable/mod.rs | 2 +- vortex-array/src/arrays/decimal/vtable/mod.rs | 2 +- .../src/arrays/extension/vtable/mod.rs | 2 +- .../src/arrays/fixed_size_list/vtable/mod.rs | 2 +- .../src/arrays/primitive/vtable/mod.rs | 2 +- vortex-array/src/arrays/struct_/vtable/mod.rs | 2 +- vortex-array/src/arrays/varbinview/compact.rs | 74 ++++++++++--------- .../src/arrays/varbinview/vtable/mod.rs | 2 +- vortex-array/src/builders/bool.rs | 6 +- vortex-array/src/builders/decimal.rs | 4 +- vortex-array/src/builders/extension.rs | 4 +- vortex-array/src/builders/fixed_size_list.rs | 4 +- vortex-array/src/builders/primitive.rs | 4 +- vortex-array/src/builders/struct_.rs | 4 +- vortex-array/src/builders/varbinview.rs | 9 ++- 15 files changed, 73 insertions(+), 50 deletions(-) diff --git a/vortex-array/src/arrays/bool/vtable/mod.rs b/vortex-array/src/arrays/bool/vtable/mod.rs index 393414924b3..38c298bad86 100644 --- a/vortex-array/src/arrays/bool/vtable/mod.rs +++ b/vortex-array/src/arrays/bool/vtable/mod.rs @@ -200,7 +200,7 @@ impl VTable for Bool { let Some(builder) = builder.as_any_mut().downcast_mut::() else { vortex_bail!("append_to_builder for Bool requires a BoolBuilder"); }; - builder.append_bool_array(&array.into_owned(), ctx) + builder.append_bool_array(array, ctx) } fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { diff --git a/vortex-array/src/arrays/decimal/vtable/mod.rs b/vortex-array/src/arrays/decimal/vtable/mod.rs index 5f2e087e983..3dd3769b197 100644 --- a/vortex-array/src/arrays/decimal/vtable/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/mod.rs @@ -195,7 +195,7 @@ impl VTable for Decimal { let Some(builder) = builder.as_any_mut().downcast_mut::() else { vortex_bail!("append_to_builder for Decimal requires a DecimalBuilder"); }; - builder.append_decimal_array(&array.into_owned(), ctx) + builder.append_decimal_array(array, ctx) } fn reduce_parent( diff --git a/vortex-array/src/arrays/extension/vtable/mod.rs b/vortex-array/src/arrays/extension/vtable/mod.rs index 2937a2b98b9..b144b13b745 100644 --- a/vortex-array/src/arrays/extension/vtable/mod.rs +++ b/vortex-array/src/arrays/extension/vtable/mod.rs @@ -196,7 +196,7 @@ impl VTable for Extension { let Some(builder) = builder.as_any_mut().downcast_mut::() else { vortex_bail!("append_to_builder for Extension requires an ExtensionBuilder"); }; - builder.append_extension_array(&array.into_owned(), ctx) + builder.append_extension_array(array, ctx) } fn reduce(array: ArrayView<'_, Self>) -> VortexResult> { diff --git a/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs b/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs index 656a08b0f66..b1586016296 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs @@ -211,6 +211,6 @@ impl VTable for FixedSizeList { let Some(builder) = builder.as_any_mut().downcast_mut::() else { vortex_bail!("append_to_builder for FixedSizeList requires a FixedSizeListBuilder"); }; - builder.append_fixed_size_list_array(&array.into_owned(), ctx) + builder.append_fixed_size_list_array(array, ctx) } } diff --git a/vortex-array/src/arrays/primitive/vtable/mod.rs b/vortex-array/src/arrays/primitive/vtable/mod.rs index 45256a7a498..5472e8df4a6 100644 --- a/vortex-array/src/arrays/primitive/vtable/mod.rs +++ b/vortex-array/src/arrays/primitive/vtable/mod.rs @@ -188,7 +188,7 @@ impl VTable for Primitive { ) -> VortexResult<()> { match_each_native_ptype!(array.ptype(), |P| { if let Some(builder) = builder.as_any_mut().downcast_mut::>() { - return builder.append_primitive_array(&array.into_owned(), ctx); + return builder.append_primitive_array(array, ctx); } }); diff --git a/vortex-array/src/arrays/struct_/vtable/mod.rs b/vortex-array/src/arrays/struct_/vtable/mod.rs index c1639501288..4f924758234 100644 --- a/vortex-array/src/arrays/struct_/vtable/mod.rs +++ b/vortex-array/src/arrays/struct_/vtable/mod.rs @@ -207,7 +207,7 @@ impl VTable for Struct { let Some(builder) = builder.as_any_mut().downcast_mut::() else { vortex_bail!("append_to_builder for Struct requires a StructBuilder"); }; - builder.append_struct_array(&array.into_owned(), ctx) + builder.append_struct_array(array, ctx) } fn reduce_parent( diff --git a/vortex-array/src/arrays/varbinview/compact.rs b/vortex-array/src/arrays/varbinview/compact.rs index 38549b9cf85..35be476db52 100644 --- a/vortex-array/src/arrays/varbinview/compact.rs +++ b/vortex-array/src/arrays/varbinview/compact.rs @@ -11,6 +11,8 @@ use vortex_error::VortexResult; use vortex_mask::Mask; use crate::ExecutionCtx; +use crate::array::TypedArrayRef; +use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::arrays::varbinview::Ref; use crate::builders::VarBinViewBuilder; @@ -63,6 +65,41 @@ impl VarBinViewArray { Ok((bytes_referenced as f64 / buffer_total_bytes as f64) < DEFAULT_COMPACTION_THRESHOLD) } + /// Returns a compacted copy of the input array using selective buffer compaction. + /// + /// This method analyzes each buffer's utilization and applies one of three strategies: + /// - **KeepFull** (zero-copy): Well-utilized buffers are kept unchanged + /// - **Slice** (zero-copy): Buffers with contiguous ranges of used data are sliced to that range + /// - **Rewrite**: Poorly-utilized buffers have their data copied to new compact buffers + /// + /// By preserving or slicing well-utilized buffers, compaction becomes zero-copy in many cases. + /// + /// # Arguments + /// + /// * `buffer_utilization_threshold` - Threshold in range [0, 1]. Buffers with utilization + /// below this value will be compacted. Use 0.0 for no compaction, 1.0 for aggressive + /// compaction of any buffer with wasted space. + pub fn compact_with_threshold( + &self, + buffer_utilization_threshold: f64, // [0, 1] + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let mut builder = VarBinViewBuilder::with_compaction_in( + self.dtype().clone(), + self.len(), + buffer_utilization_threshold, + ctx.allocator().clone(), + ); + builder.append_varbinview_array(self.as_view(), ctx)?; + Ok(builder.finish_into_varbinview()) + } +} + +/// Read-only compaction statistics over a `VarBinView` array. +/// +/// These live on a trait rather than on [`VarBinViewArray`] so that a borrowed +/// [`ArrayView`] can ask for them without cloning the array behind it. +pub(crate) trait VarBinViewCompactExt: TypedArrayRef { /// Iterates over all valid, non-inlined views, calling the provided /// closure for each one. #[allow(clippy::inline_always)] @@ -103,10 +140,7 @@ impl VarBinViewArray { Ok(total) } - pub(crate) fn buffer_utilizations( - &self, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { + fn buffer_utilizations(&self, ctx: &mut ExecutionCtx) -> VortexResult> { let mut utilizations: Vec = self .data_buffers() .iter() @@ -122,37 +156,10 @@ impl VarBinViewArray { Ok(utilizations) } - - /// Returns a compacted copy of the input array using selective buffer compaction. - /// - /// This method analyzes each buffer's utilization and applies one of three strategies: - /// - **KeepFull** (zero-copy): Well-utilized buffers are kept unchanged - /// - **Slice** (zero-copy): Buffers with contiguous ranges of used data are sliced to that range - /// - **Rewrite**: Poorly-utilized buffers have their data copied to new compact buffers - /// - /// By preserving or slicing well-utilized buffers, compaction becomes zero-copy in many cases. - /// - /// # Arguments - /// - /// * `buffer_utilization_threshold` - Threshold in range [0, 1]. Buffers with utilization - /// below this value will be compacted. Use 0.0 for no compaction, 1.0 for aggressive - /// compaction of any buffer with wasted space. - pub fn compact_with_threshold( - &self, - buffer_utilization_threshold: f64, // [0, 1] - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let mut builder = VarBinViewBuilder::with_compaction_in( - self.dtype().clone(), - self.len(), - buffer_utilization_threshold, - ctx.allocator().clone(), - ); - builder.append_varbinview_array(self, ctx)?; - Ok(builder.finish_into_varbinview()) - } } +impl> VarBinViewCompactExt for T {} + pub(crate) struct BufferUtilization { len: u32, used: u32, @@ -204,6 +211,7 @@ mod tests { use rstest::rstest; use vortex_buffer::buffer; + use super::VarBinViewCompactExt; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; diff --git a/vortex-array/src/arrays/varbinview/vtable/mod.rs b/vortex-array/src/arrays/varbinview/vtable/mod.rs index 0d2963cd6b2..fafa2ad26d0 100644 --- a/vortex-array/src/arrays/varbinview/vtable/mod.rs +++ b/vortex-array/src/arrays/varbinview/vtable/mod.rs @@ -249,7 +249,7 @@ impl VTable for VarBinView { ctx: &mut ExecutionCtx, ) -> VortexResult<()> { if let Some(builder) = builder.as_any_mut().downcast_mut::() { - return builder.append_varbinview_array(&array.into_owned(), ctx); + return builder.append_varbinview_array(array, ctx); } if let Some(result) = match_each_varbin_builder!(builder, |builder| builder.append_varbinview(array, ctx)) diff --git a/vortex-array/src/builders/bool.rs b/vortex-array/src/builders/bool.rs index ffe7e138853..1d62bca7a8b 100644 --- a/vortex-array/src/builders/bool.rs +++ b/vortex-array/src/builders/bool.rs @@ -11,6 +11,8 @@ use vortex_error::vortex_ensure; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; use crate::builders::ArrayBuilder; @@ -91,12 +93,12 @@ impl BoolBuilder { pub(crate) fn append_bool_array( &mut self, - array: &BoolArray, + array: ArrayView<'_, Bool>, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { self.inner.append_buffer(&array.to_bit_buffer()); self.nulls - .append_validity_mask(&BoolArrayExt::validity(array).execute_mask(array.len(), ctx)?); + .append_validity_mask(&BoolArrayExt::validity(&array).execute_mask(array.len(), ctx)?); Ok(()) } } diff --git a/vortex-array/src/builders/decimal.rs b/vortex-array/src/builders/decimal.rs index e953f34ee65..1fe345ddd53 100644 --- a/vortex-array/src/builders/decimal.rs +++ b/vortex-array/src/builders/decimal.rs @@ -14,6 +14,8 @@ use vortex_error::vortex_panic; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Decimal; use crate::arrays::DecimalArray; use crate::builders::ArrayBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; @@ -159,7 +161,7 @@ impl DecimalBuilder { /// storage type to the builder's type as needed. pub(crate) fn append_decimal_array( &mut self, - array: &DecimalArray, + array: ArrayView<'_, Decimal>, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { match_each_decimal_value_type!(array.values_type(), |D| { diff --git a/vortex-array/src/builders/extension.rs b/vortex-array/src/builders/extension.rs index c2f4be2df5a..b9337b426a4 100644 --- a/vortex-array/src/builders/extension.rs +++ b/vortex-array/src/builders/extension.rs @@ -10,6 +10,8 @@ use vortex_error::vortex_ensure; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Extension; use crate::arrays::ExtensionArray; use crate::arrays::extension::ExtensionArrayExt; use crate::builders::ArrayBuilder; @@ -66,7 +68,7 @@ impl ExtensionBuilder { /// storage array to the underlying storage builder. pub(crate) fn append_extension_array( &mut self, - array: &ExtensionArray, + array: ArrayView<'_, Extension>, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { self.storage.append_array(array.storage_array(), ctx) diff --git a/vortex-array/src/builders/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index f9a42425ce7..bb7400a2f26 100644 --- a/vortex-array/src/builders/fixed_size_list.rs +++ b/vortex-array/src/builders/fixed_size_list.rs @@ -13,7 +13,9 @@ use vortex_error::vortex_panic; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; +use crate::array::ArrayView; use crate::arrays::ChunkedArray; +use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::builders::ArrayBuilder; @@ -191,7 +193,7 @@ impl FixedSizeListBuilder { /// elements builder. pub(crate) fn append_fixed_size_list_array( &mut self, - array: &FixedSizeListArray, + array: ArrayView<'_, FixedSizeList>, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { if array.is_empty() { diff --git a/vortex-array/src/builders/primitive.rs b/vortex-array/src/builders/primitive.rs index 42369ee56ad..92af8ba187c 100644 --- a/vortex-array/src/builders/primitive.rs +++ b/vortex-array/src/builders/primitive.rs @@ -14,6 +14,8 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::builders::ArrayBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; @@ -159,7 +161,7 @@ impl PrimitiveBuilder { pub(crate) fn append_primitive_array( &mut self, - array: &PrimitiveArray, + array: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { debug_assert_eq!( diff --git a/vortex-array/src/builders/struct_.rs b/vortex-array/src/builders/struct_.rs index 6cdb76b0044..f3f67cc45b3 100644 --- a/vortex-array/src/builders/struct_.rs +++ b/vortex-array/src/builders/struct_.rs @@ -13,6 +13,8 @@ use vortex_error::vortex_panic; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Struct; use crate::arrays::StructArray; use crate::arrays::struct_::StructArrayExt; use crate::builders::ArrayBuilder; @@ -153,7 +155,7 @@ impl StructBuilder { /// field's builder. pub(crate) fn append_struct_array( &mut self, - array: &StructArray, + array: ArrayView<'_, Struct>, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { for (field, builder) in array diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index e27336fe998..8dfb9ad7362 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -25,12 +25,15 @@ use vortex_utils::aliases::hash_map::HashMap; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::arrays::varbinview::VarBinViewArrayExt; use crate::arrays::varbinview::build_views::BinaryView; use crate::arrays::varbinview::build_views::MAX_BUFFER_LEN; use crate::arrays::varbinview::build_views::extend_views; use crate::arrays::varbinview::compact::BufferUtilization; +use crate::arrays::varbinview::compact::VarBinViewCompactExt; use crate::builders::ArrayBuilder; use crate::builders::LazyBitBufferBuilder; use crate::canonical::Canonical; @@ -755,7 +758,7 @@ impl VarBinViewBuilder { pub(crate) fn append_varbinview_array( &mut self, - array: &VarBinViewArray, + array: ArrayView<'_, VarBinView>, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { self.flush_in_progress(); @@ -881,7 +884,7 @@ impl VarBinViewBuilder { &mut self, view: BinaryView, adjustment: &RewritingViewAdjustment, - array: &VarBinViewArray, + array: ArrayView<'_, VarBinView>, idx: usize, ) -> BinaryView { if view.is_inlined() { @@ -1164,7 +1167,7 @@ enum BuffersWithOffsets { impl BuffersWithOffsets { pub fn from_array( - array: &VarBinViewArray, + array: ArrayView<'_, VarBinView>, compaction_threshold: f64, ctx: &mut ExecutionCtx, ) -> Self { From 627a6fafb600be2dbc2fa05092397b761949038a Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 16 Sep 2026 13:36:31 +0000 Subject: [PATCH 08/22] perf(array): stop paying per-chunk builder overhead Canonicalizing a chunked array through the builder now costs roughly what the bespoke swizzle it replaced did, instead of ~11% more instructions. Three things were being paid once per chunk that only need paying once per appended array: - `ValidityBuilder` recorded a run per chunk for validities that are uniform, so a 32-chunk non-nullable array left 32 runs for `Validity::concat` to collapse and for `finish_with_nullability` to walk. Adjacent uniform runs of the same kind now extend rather than push. - `append_to_builder` re-derived and re-compared the builder's dtype, and re-checked the builder's length, for every chunk. `Chunked` now appends its chunks through an `append_to_builder_unchecked` that skips both: the enclosing call already checked the chunked array's dtype, which every chunk shares, and its length, which the chunks' lengths sum to. - `ChildBuilder::append_array` compared dtypes on every appended chunk. `ChildBuilder` is `pub(crate)`, and each of its callers is a nested builder passing down a child of an array whose dtype was checked against that builder on the way in, so the child's dtype follows from the parent's. The comparison becomes a debug assertion, which is what the invariant is. Measured with callgrind on a 32-chunk x 1000-list x 1024-element `FixedSizeList` canonicalization, against the pre-builder swizzle: swizzle 12,610,357 Ir builder, before 14,022,180 Ir +11.2% builder, after 12,791,687 Ir +1.4% Signed-off-by: Claude Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01CZTu9XUrdRZZ91rx852DGf --- vortex-array/src/array/erased.rs | 21 ++++++++++++++++ vortex-array/src/array/mod.rs | 25 +++++++++++++++++-- vortex-array/src/arrays/chunked/vtable/mod.rs | 5 +++- vortex-array/src/builders/child.rs | 20 ++++++++++----- vortex-array/src/builders/validity.rs | 22 ++++++++++++++++ 5 files changed, 84 insertions(+), 9 deletions(-) diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index 0608eecb5dc..f0ec60d56cc 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -421,6 +421,27 @@ impl ArrayRef { self.0.data.append_to_builder(self, builder, ctx) } + /// [`append_to_builder`](Self::append_to_builder) for a caller that is itself inside an + /// [`append_to_builder`](Self::append_to_builder), and so is already covered by its dtype check + /// and its length post-condition. + /// + /// Appending the chunks of a [`ChunkedArray`](crate::arrays::ChunkedArray) is the case this + /// exists for. The chunked array's own dtype was checked on the way in and every chunk shares + /// it, and if the chunks together grow the builder by the wrong number of values, the enclosing + /// check on the chunked array says so — both are worth paying once rather than once per chunk. + pub(crate) fn append_to_builder_unchecked( + &self, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + debug_assert_eq!( + builder.dtype(), + self.dtype(), + "append_to_builder_unchecked called with a mismatched builder", + ); + self.0.data.append_to_builder_unchecked(self, builder, ctx) + } + /// Returns the statistics of the array. pub fn statistics(&self) -> StatsSetRef<'_> { self.0.stats.to_ref(self) diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index c3fb88a8c62..b7470b161ba 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -137,6 +137,18 @@ pub(crate) trait DynArrayData: 'static + private::Sealed + Send + Sync + Debug { ctx: &mut ExecutionCtx, ) -> VortexResult<()>; + /// [`append_to_builder`](Self::append_to_builder) without its dtype check or its length + /// post-condition. + /// + /// The caller must have established that `builder.dtype() == this.dtype()`, and is responsible + /// for checking that the builder grew by the number of values it was handed. + fn append_to_builder_unchecked( + &self, + this: &ArrayRef, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()>; + // --- Visitor methods (formerly in ArrayVisitor) --- /// Returns the buffers of the array. @@ -308,8 +320,7 @@ impl DynArrayData for ArrayData { } let len = builder.len(); - let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; - V::append_to_builder(view, builder, ctx)?; + self.append_to_builder_unchecked(this, builder, ctx)?; assert_eq!( len + this.len(), @@ -320,6 +331,16 @@ impl DynArrayData for ArrayData { Ok(()) } + fn append_to_builder_unchecked( + &self, + this: &ArrayRef, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; + V::append_to_builder(view, builder, ctx) + } + fn buffers(&self, this: &ArrayRef) -> Vec { let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; (0..V::nbuffers(view)) diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index 19dd7cab5fb..e597e92e068 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -242,8 +242,11 @@ impl VTable for Chunked { // The builder is about to receive one array per chunk, so let any chunk lists it keeps // grow once rather than on the way through. builder.reserve_chunks(array.nchunks()); + // The dtype check and the length post-condition around this call already cover the chunks: + // every chunk shares this array's dtype, and their lengths sum to its own. Paying for both + // once beats paying per chunk. for chunk in array.iter_chunks() { - chunk.append_to_builder(builder, ctx)?; + chunk.append_to_builder_unchecked(builder, ctx)?; } Ok(()) } diff --git a/vortex-array/src/builders/child.rs b/vortex-array/src/builders/child.rs index a88c5c562d6..4b552d7acd6 100644 --- a/vortex-array/src/builders/child.rs +++ b/vortex-array/src/builders/child.rs @@ -3,7 +3,6 @@ use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use crate::ArrayRef; use crate::ExecutionCtx; @@ -86,9 +85,16 @@ impl ChildBuilder { /// Nothing is decoded here, so `_ctx` goes unused; it stays in the signature so that the /// nested builders forwarding their [`ExecutionCtx`] here do not have to explain why they /// don't. + /// + /// `array` must have the child's dtype. This is a crate-internal invariant, checked under + /// `debug_assertions` only: every caller is a nested builder handing over a child of an array + /// whose own dtype was checked against that builder on the way in, so the child's dtype follows + /// from the parent's and re-deriving it per appended array would cost a dtype comparison per + /// chunk on the hot path. pub fn append_array(&mut self, array: &ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult<()> { - vortex_ensure!( - array.dtype() == &self.dtype, + debug_assert_eq!( + array.dtype(), + &self.dtype, "Cannot append an array of dtype {} to a child builder of dtype {}", array.dtype(), self.dtype, @@ -350,18 +356,20 @@ mod tests { Ok(()) } - /// The dtype check has to run before the empty check, so that a mismatched array is rejected + /// The debug assertion has to run before the empty check, so that a mismatched array is caught /// whether or not it would have become a chunk. + #[cfg(debug_assertions)] #[rstest] #[case::empty(0)] #[case::non_empty(CHUNK_LEN)] - fn test_appending_a_mismatched_dtype_is_rejected(#[case] len: usize) { + #[should_panic(expected = "Cannot append an array of dtype")] + fn test_appending_a_mismatched_dtype_is_caught(#[case] len: usize) { let mut ctx = array_session().create_execution_ctx(); let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); let wrong_dtype = ConstantArray::new(1i64, len).into_array(); - assert!(builder.append_array(&wrong_dtype, &mut ctx).is_err()); + drop(builder.append_array(&wrong_dtype, &mut ctx)); } /// Everything the scalar builder can produce has to be flushed ahead of the next chunk. diff --git a/vortex-array/src/builders/validity.rs b/vortex-array/src/builders/validity.rs index 4ea56a1a40e..678048c6a13 100644 --- a/vortex-array/src/builders/validity.rs +++ b/vortex-array/src/builders/validity.rs @@ -77,6 +77,19 @@ impl ValidityBuilder { self.flush_pending(); self.runs_len += len; + + // Two adjacent runs that are uniform and agree describe one longer run. Extending beats + // pushing: a chunked array of non-nullable chunks would otherwise record a run per chunk + // for a validity that `Validity::concat` collapses at the end anyway, and every run is + // walked again by `finish_with_nullability`. + if let Some((last, last_len)) = self.runs.last_mut() + && Self::is_uniform(last) + && std::mem::discriminant(last) == std::mem::discriminant(&validity) + { + *last_len += len; + return; + } + self.runs.push((validity, len)); } @@ -125,6 +138,15 @@ impl ValidityBuilder { validity } + /// Whether a validity describes every value it covers the same way, so that two adjacent runs + /// of it can be recorded as one. + fn is_uniform(validity: &Validity) -> bool { + matches!( + validity, + Validity::NonNullable | Validity::AllValid | Validity::AllInvalid + ) + } + /// Moves whatever the null buffer holds into `runs`, keeping the runs in logical order. fn flush_pending(&mut self) { let len = self.pending.len(); From 77cd2594ad68f8a5f123fd36157d92171cf15cb1 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 16 Sep 2026 16:00:36 -0400 Subject: [PATCH 09/22] more Signed-off-by: Robert Kruszewski --- vortex-array/src/arrays/chunked/array.rs | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/vortex-array/src/arrays/chunked/array.rs b/vortex-array/src/arrays/chunked/array.rs index 807d526c17b..234dd23e9cf 100644 --- a/vortex-array/src/arrays/chunked/array.rs +++ b/vortex-array/src/arrays/chunked/array.rs @@ -79,20 +79,18 @@ pub trait ChunkedArrayExt: TypedArrayRef { .vortex_expect("validated chunk slot") } - fn iter_chunks<'a>(&'a self) -> Box + 'a> { - Box::new( - self.as_ref().slots()[ChunkedSlots::CHUNKS_OFFSET..] - .iter() - .map(|slot| slot.as_ref().vortex_expect("validated chunk slot")), - ) + fn iter_chunks(&self) -> impl Iterator { + self.as_ref().slots()[ChunkedSlots::CHUNKS_OFFSET..] + .iter() + .map(|slot| slot.as_ref().vortex_expect("validated chunk slot")) } fn chunks(&self) -> Vec { self.iter_chunks().cloned().collect() } - fn non_empty_chunks<'a>(&'a self) -> Box + 'a> { - Box::new(self.iter_chunks().filter(|chunk| !chunk.is_empty())) + fn non_empty_chunks(&self) -> impl Iterator { + self.iter_chunks().filter(|chunk| !chunk.is_empty()) } /// Returns the cached chunk boundary offsets. From f9f3d9e437268b5fc320ebd7348eb07cef75dc9a Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 16 Sep 2026 16:07:55 -0400 Subject: [PATCH 10/22] fix Signed-off-by: Robert Kruszewski --- vortex-array/src/builders/child.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/vortex-array/src/builders/child.rs b/vortex-array/src/builders/child.rs index 4b552d7acd6..0c0df16f777 100644 --- a/vortex-array/src/builders/child.rs +++ b/vortex-array/src/builders/child.rs @@ -203,7 +203,6 @@ impl ChildBuilder { #[cfg(test)] mod tests { - use rstest::rstest; use vortex_buffer::BufferAllocatorRef; use vortex_buffer::buffer; use vortex_error::VortexResult; From 82911e592fd144d5ff6f275ce5dc60cf55ededb2 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 16 Sep 2026 16:11:25 -0400 Subject: [PATCH 11/22] fixes Signed-off-by: Robert Kruszewski --- vortex-array/src/builders/child.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-array/src/builders/child.rs b/vortex-array/src/builders/child.rs index 0c0df16f777..e06617f6ef9 100644 --- a/vortex-array/src/builders/child.rs +++ b/vortex-array/src/builders/child.rs @@ -203,6 +203,7 @@ impl ChildBuilder { #[cfg(test)] mod tests { + use rstest::rstest; use vortex_buffer::BufferAllocatorRef; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -357,7 +358,6 @@ mod tests { /// The debug assertion has to run before the empty check, so that a mismatched array is caught /// whether or not it would have become a chunk. - #[cfg(debug_assertions)] #[rstest] #[case::empty(0)] #[case::non_empty(CHUNK_LEN)] From 9bac9bb581a535a949a980ce8ac77176152be0d5 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 16 Sep 2026 21:10:32 +0000 Subject: [PATCH 12/22] docs(array): qualify the ArrayView intra-doc link `ArrayView` is not in scope in `varbinview::compact`, so rustdoc could not resolve the bare link and the workspace doc build failed under `-D warnings`. Link through `crate::array::ArrayView`, matching how the rest of the crate writes out-of-scope doc links. Signed-off-by: Robert Kruszewski Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01CZTu9XUrdRZZ91rx852DGf --- vortex-array/src/arrays/varbinview/compact.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-array/src/arrays/varbinview/compact.rs b/vortex-array/src/arrays/varbinview/compact.rs index 35be476db52..9d22afeaddb 100644 --- a/vortex-array/src/arrays/varbinview/compact.rs +++ b/vortex-array/src/arrays/varbinview/compact.rs @@ -98,7 +98,7 @@ impl VarBinViewArray { /// Read-only compaction statistics over a `VarBinView` array. /// /// These live on a trait rather than on [`VarBinViewArray`] so that a borrowed -/// [`ArrayView`] can ask for them without cloning the array behind it. +/// [`ArrayView`](crate::array::ArrayView) can ask for them without cloning the array behind it. pub(crate) trait VarBinViewCompactExt: TypedArrayRef { /// Iterates over all valid, non-inlined views, calling the provided /// closure for each one. From 4580bc7e60ccde7f68592e6edb96bb4256582f25 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 16 Sep 2026 19:54:35 -0400 Subject: [PATCH 13/22] better Signed-off-by: Robert Kruszewski --- .../src/arrays/chunked/vtable/canonical.rs | 315 +----------------- vortex-array/src/arrays/chunked/vtable/mod.rs | 5 +- 2 files changed, 20 insertions(+), 300 deletions(-) diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index 68f471763ae..9795ad7f52d 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -16,47 +16,40 @@ use crate::arrays::ChunkedArray; use crate::arrays::VariantArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::variant::VariantArraySlotsExt; -use crate::builders::builder_with_capacity_in; -use crate::dtype::DType; pub(super) fn _canonicalize( array: ArrayView<'_, Chunked>, ctx: &mut ExecutionCtx, ) -> VortexResult { + vortex_ensure!( + array.dtype().is_variant(), + "only variant needs recursive swizzling" + ); + if array.nchunks() == 0 { - if matches!(array.dtype(), DType::Variant(_)) { - return VariantArray::try_new(array.array().clone().into_array(), None) - .map(Canonical::Variant); - } - return Ok(Canonical::empty(array.dtype())); + return VariantArray::try_new(array.array().clone().into_array(), None) + .map(Canonical::Variant); } if array.nchunks() == 1 { return array.chunk(0).clone().execute::(ctx); } - Ok(match array.dtype() { - DType::Variant(_) => { - let owned_chunks: Vec = array.iter_chunks().cloned().collect(); - Canonical::Variant(pack_variant_chunks(owned_chunks, ctx)?) - } - _ => { - let mut builder = builder_with_capacity_in(array.dtype(), array.len(), ctx.allocator()); - array.array().append_to_builder(builder.as_mut(), ctx)?; - builder.finish_into_canonical(ctx) - } - }) + Ok(Canonical::Variant(pack_variant_chunks( + array.iter_chunks(), + ctx, + )?)) } /// Packs many [`VariantArray`]s into one [`VariantArray`] with chunked children. /// /// The caller guarantees there are at least 2 chunks. -fn pack_variant_chunks( - chunks: Vec, +fn pack_variant_chunks<'a>( + chunks: impl Iterator, ctx: &mut ExecutionCtx, ) -> VortexResult { let variant_chunks: Vec = chunks .into_iter() - .map(|chunk| chunk.execute::(ctx)) + .map(|chunk| chunk.clone().execute::(ctx)) .try_collect()?; let outer_dtype = variant_chunks[0].dtype().clone(); @@ -107,83 +100,37 @@ fn pack_variant_chunks( #[cfg(test)] mod tests { - use std::alloc::Layout; - use std::ptr::NonNull; - use std::sync::Arc; use std::sync::LazyLock; - use std::sync::atomic::AtomicUsize; - use std::sync::atomic::Ordering; - - use allocator_api2::alloc::AllocError; - use allocator_api2::alloc::Allocator; - use allocator_api2::alloc::Global; - use rstest::rstest; - use vortex_buffer::BufferAllocatorRef; - use vortex_buffer::buffer; - use vortex_error::VortexExpect; + use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_session::VortexSession; - use crate::ArrayRef; use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::arrays::Chunked; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; - use crate::arrays::FixedSizeList; - use crate::arrays::FixedSizeListArray; - use crate::arrays::ListArray; - use crate::arrays::ListView; - use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; - use crate::arrays::Struct; - use crate::arrays::StructArray; - use crate::arrays::VarBinViewArray; use crate::arrays::VariantArray; - use crate::arrays::chunked::ChunkedArrayExt; - use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; - use crate::arrays::listview::ListViewArraySlotsExt; - use crate::arrays::struct_::StructArrayExt; use crate::arrays::variant::VariantArraySlotsExt; use crate::assert_arrays_eq; - use crate::dtype::DType::List; use crate::dtype::DType::Primitive; use crate::dtype::DType::Variant as VariantDType; use crate::dtype::Nullability::NonNullable; use crate::dtype::PType::I32; - use crate::memory::MemorySessionExt; use crate::scalar::Scalar; - use crate::validity::Validity; + use crate::{ArrayRef, array_session}; /// A shared session for these chunked-array tests, used to create execution contexts. - static SESSION: LazyLock = LazyLock::new(crate::array_session); - - #[derive(Debug)] - struct CountingAllocator { - allocations: Arc, - } - - // SAFETY: this forwards memory operations to Global and only counts allocations. - unsafe impl Allocator for CountingAllocator { - fn allocate(&self, layout: Layout) -> Result, AllocError> { - self.allocations.fetch_add(1, Ordering::Relaxed); - Global.allocate(layout) - } - - unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { - // SAFETY: ptr and layout came from Global. - unsafe { Global.deallocate(ptr, layout) } - } - } + static SESSION: LazyLock = LazyLock::new(array_session); fn variant_scalar(value: i32) -> Scalar { Scalar::variant(Scalar::primitive(value, NonNullable)) } - fn variant_core(values: impl IntoIterator) -> VortexResult { + fn variant_core(values: impl IntoIterator) -> VortexResult { Ok(ChunkedArray::try_new( values .into_iter() @@ -354,230 +301,4 @@ mod tests { assert!(variant.shredded().is_some()); assert_variant_values(&variant, &[10, 20]) } - - #[test] - pub fn pack_nested_structs() { - let mut ctx = SESSION.create_execution_ctx(); - let struct_array = StructArray::try_new( - ["a"].into(), - vec![VarBinViewArray::from_iter_str(["foo", "bar", "baz", "quak"]).into_array()], - 4, - Validity::NonNullable, - ) - .unwrap(); - let dtype = struct_array.dtype().clone(); - let chunked = ChunkedArray::try_new( - vec![ - ChunkedArray::try_new(vec![struct_array.clone().into_array()], dtype.clone()) - .unwrap() - .into_array(), - ], - dtype, - ) - .unwrap() - .into_array(); - let canonical_struct = chunked.execute::(&mut ctx).unwrap(); - let canonical_varbin = canonical_struct - .unmasked_field(0) - .clone() - .execute::(&mut ctx) - .unwrap(); - let original_varbin = struct_array - .unmasked_field(0) - .clone() - .execute::(&mut ctx) - .unwrap(); - let orig_mask = original_varbin - .validity() - .unwrap() - .execute_mask(original_varbin.len(), &mut ctx) - .unwrap(); - let orig_values = (0..original_varbin.len()) - .map(|i| { - orig_mask - .value(i) - .then(|| original_varbin.bytes_at(i).to_vec()) - }) - .collect::>(); - let canon_mask = canonical_varbin - .validity() - .unwrap() - .execute_mask(canonical_varbin.len(), &mut ctx) - .unwrap(); - let canon_values = (0..canonical_varbin.len()) - .map(|i| { - canon_mask - .value(i) - .then(|| canonical_varbin.bytes_at(i).to_vec()) - }) - .collect::>(); - assert_eq!(orig_values, canon_values); - } - - #[test] - pub fn pack_nested_lists() { - let mut ctx = SESSION.create_execution_ctx(); - let l1 = ListArray::try_new( - buffer![1, 2, 3, 4].into_array(), - buffer![0, 3].into_array(), - Validity::NonNullable, - ) - .unwrap(); - - let l2 = ListArray::try_new( - buffer![5, 6].into_array(), - buffer![0, 2].into_array(), - Validity::NonNullable, - ) - .unwrap(); - - let chunked_list = ChunkedArray::try_new( - vec![l1.clone().into_array(), l2.clone().into_array()], - List(Arc::new(Primitive(I32, NonNullable)), NonNullable), - ); - - let canon_values = chunked_list - .unwrap() - .as_array() - .clone() - .execute::(&mut ctx) - .unwrap(); - - assert_eq!( - l1.execute_scalar(0, &mut ctx).unwrap(), - canon_values.execute_scalar(0, &mut ctx).unwrap() - ); - assert_eq!( - l2.execute_scalar(0, &mut ctx).unwrap(), - canon_values.execute_scalar(1, &mut ctx).unwrap() - ); - } - - #[test] - fn pack_fixed_size_lists() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let f1 = FixedSizeListArray::try_new( - buffer![1, 2, 3, 4, 5, 6].into_array(), - 2, - Validity::NonNullable, - 3, - )?; - let f2 = FixedSizeListArray::try_new( - buffer![7, 8, 9, 10].into_array(), - 2, - Validity::NonNullable, - 2, - )?; - let dtype = f1.dtype().clone(); - - let chunked = - ChunkedArray::try_new(vec![f1.into_array(), f2.into_array()], dtype)?.into_array(); - - let canonical = chunked.clone().execute::(&mut ctx)?; - let fsl = match canonical { - Canonical::FixedSizeList(fsl) => fsl, - other => vortex_bail!("expected FixedSizeList canonical array, got {other:?}"), - }; - - assert_eq!(fsl.len(), 5); - let expected = FixedSizeListArray::try_new( - buffer![1, 2, 3, 4, 5, 6, 7, 8, 9, 10].into_array(), - 2, - Validity::NonNullable, - 5, - )?; - for idx in 0..5 { - assert_eq!( - chunked.execute_scalar(idx, &mut ctx)?, - expected.execute_scalar(idx, &mut ctx)?, - ); - } - Ok(()) - } - - /// Canonicalizing a `ChunkedArray` reuses each chunk's children instead of concatenating - /// them: a nested builder keeps a child chunked on exactly the boundaries it was appended on, - /// however short the appended chunk is. - #[rstest] - #[case::struct_( - StructArray::try_from_iter([("a", buffer![1i32, 2])]) - .vortex_expect("struct array") - .into_array(), - |array: &ArrayRef| array.as_::().unmasked_field(0).clone() - )] - #[case::fixed_size_list( - FixedSizeListArray::new(buffer![1i32, 2].into_array(), 2, Validity::NonNullable, 1) - .into_array(), - |array: &ArrayRef| array.as_::().elements().clone() - )] - #[case::list( - ListArray::try_new( - buffer![1i32, 2].into_array(), - buffer![0, 2].into_array(), - Validity::NonNullable, - ) - .vortex_expect("list array") - .into_array(), - |array: &ArrayRef| array.as_::().elements().clone() - )] - fn canonicalize_reuses_chunk_children( - #[case] chunk: ArrayRef, - #[case] child_of: fn(&ArrayRef) -> ArrayRef, - ) -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - - let chunked = - ChunkedArray::try_new(vec![chunk.clone(), chunk.clone()], chunk.dtype().clone())? - .into_array(); - let canonical = chunked.execute::(&mut ctx)?.into_array(); - - let child = child_of(&canonical); - assert_eq!( - child.as_::().nchunks(), - 2, - "each chunk's child should have become a chunk of the combined child", - ); - - let expected = - ChunkedArray::try_new(vec![chunk.clone(), chunk], canonical.dtype().clone())?; - assert_arrays_eq!(&canonical, &expected, &mut ctx); - - Ok(()) - } - - #[test] - fn list_canonicalize_uses_memory_session_allocator() { - let allocations = Arc::new(AtomicUsize::new(0)); - let session = - crate::array_session().with_allocator(BufferAllocatorRef::new(CountingAllocator { - allocations: Arc::clone(&allocations), - })); - let mut ctx = session.create_execution_ctx(); - - let l1 = ListArray::try_new( - buffer![1, 2, 3, 4].into_array(), - buffer![0, 3].into_array(), - Validity::NonNullable, - ) - .unwrap(); - let l2 = ListArray::try_new( - buffer![5, 6].into_array(), - buffer![0, 2].into_array(), - Validity::NonNullable, - ) - .unwrap(); - - let chunked_list = ChunkedArray::try_new( - vec![l1.into_array(), l2.into_array()], - List(Arc::new(Primitive(I32, NonNullable)), NonNullable), - ) - .unwrap() - .into_array(); - - drop(chunked_list.execute::(&mut ctx).unwrap()); - assert!( - allocations.load(Ordering::Relaxed) >= 2, - "expected offset+size allocations through MemorySession" - ); - } } diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index e597e92e068..268bbab3bfe 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -266,9 +266,8 @@ impl VTable for Chunked { and every sparse child along identical chunk boundaries" ) } - // Struct, List, FixedSizeList, and Variant need child swizzling that the builder path - // cannot express. - DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) | DType::Variant(..) => { + // Variant need child swizzling that the builder path cannot express. + DType::Variant(..) => { // TODO(joe)[#7674]: iterative execution here too Ok(ExecutionResult::done(_canonicalize(array.as_view(), ctx)?)) } From 537e2742ee2e79c34027fa8e63321db72cf7ab4b Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 16 Sep 2026 19:58:28 -0400 Subject: [PATCH 14/22] more Signed-off-by: Robert Kruszewski --- vortex-array/src/arrays/chunked/vtable/canonical.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index 9795ad7f52d..b419fd4d77f 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -107,9 +107,11 @@ mod tests { use vortex_error::vortex_err; use vortex_session::VortexSession; + use crate::ArrayRef; use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; @@ -121,7 +123,6 @@ mod tests { use crate::dtype::Nullability::NonNullable; use crate::dtype::PType::I32; use crate::scalar::Scalar; - use crate::{ArrayRef, array_session}; /// A shared session for these chunked-array tests, used to create execution contexts. static SESSION: LazyLock = LazyLock::new(array_session); @@ -130,7 +131,7 @@ mod tests { Scalar::variant(Scalar::primitive(value, NonNullable)) } - fn variant_core(values: impl IntoIterator) -> VortexResult { + fn variant_core(values: impl IntoIterator) -> VortexResult { Ok(ChunkedArray::try_new( values .into_iter() From 6009ab5c0029fdf49d118a9e52fe70fe8e3d1add Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 17 Sep 2026 00:03:08 +0000 Subject: [PATCH 15/22] refactor(array): drop append_to_builder_unchecked Chunked appends its chunks through the checked `append_to_builder` again, paying the builder dtype comparison and the length post-condition per chunk rather than once for the whole chunked array. Signed-off-by: Robert Kruszewski Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01CZTu9XUrdRZZ91rx852DGf --- vortex-array/src/array/erased.rs | 21 ---------------- vortex-array/src/array/mod.rs | 25 ++----------------- vortex-array/src/arrays/chunked/vtable/mod.rs | 5 +--- 3 files changed, 3 insertions(+), 48 deletions(-) diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index f0ec60d56cc..0608eecb5dc 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -421,27 +421,6 @@ impl ArrayRef { self.0.data.append_to_builder(self, builder, ctx) } - /// [`append_to_builder`](Self::append_to_builder) for a caller that is itself inside an - /// [`append_to_builder`](Self::append_to_builder), and so is already covered by its dtype check - /// and its length post-condition. - /// - /// Appending the chunks of a [`ChunkedArray`](crate::arrays::ChunkedArray) is the case this - /// exists for. The chunked array's own dtype was checked on the way in and every chunk shares - /// it, and if the chunks together grow the builder by the wrong number of values, the enclosing - /// check on the chunked array says so — both are worth paying once rather than once per chunk. - pub(crate) fn append_to_builder_unchecked( - &self, - builder: &mut dyn ArrayBuilder, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - debug_assert_eq!( - builder.dtype(), - self.dtype(), - "append_to_builder_unchecked called with a mismatched builder", - ); - self.0.data.append_to_builder_unchecked(self, builder, ctx) - } - /// Returns the statistics of the array. pub fn statistics(&self) -> StatsSetRef<'_> { self.0.stats.to_ref(self) diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index b7470b161ba..c3fb88a8c62 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -137,18 +137,6 @@ pub(crate) trait DynArrayData: 'static + private::Sealed + Send + Sync + Debug { ctx: &mut ExecutionCtx, ) -> VortexResult<()>; - /// [`append_to_builder`](Self::append_to_builder) without its dtype check or its length - /// post-condition. - /// - /// The caller must have established that `builder.dtype() == this.dtype()`, and is responsible - /// for checking that the builder grew by the number of values it was handed. - fn append_to_builder_unchecked( - &self, - this: &ArrayRef, - builder: &mut dyn ArrayBuilder, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()>; - // --- Visitor methods (formerly in ArrayVisitor) --- /// Returns the buffers of the array. @@ -320,7 +308,8 @@ impl DynArrayData for ArrayData { } let len = builder.len(); - self.append_to_builder_unchecked(this, builder, ctx)?; + let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; + V::append_to_builder(view, builder, ctx)?; assert_eq!( len + this.len(), @@ -331,16 +320,6 @@ impl DynArrayData for ArrayData { Ok(()) } - fn append_to_builder_unchecked( - &self, - this: &ArrayRef, - builder: &mut dyn ArrayBuilder, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; - V::append_to_builder(view, builder, ctx) - } - fn buffers(&self, this: &ArrayRef) -> Vec { let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; (0..V::nbuffers(view)) diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index 268bbab3bfe..4a551946e63 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -242,11 +242,8 @@ impl VTable for Chunked { // The builder is about to receive one array per chunk, so let any chunk lists it keeps // grow once rather than on the way through. builder.reserve_chunks(array.nchunks()); - // The dtype check and the length post-condition around this call already cover the chunks: - // every chunk shares this array's dtype, and their lengths sum to its own. Paying for both - // once beats paying per chunk. for chunk in array.iter_chunks() { - chunk.append_to_builder_unchecked(builder, ctx)?; + chunk.append_to_builder(builder, ctx)?; } Ok(()) } From 8c6dde37324f1c15d608ba7378e382ebd71b36d8 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 17 Sep 2026 00:54:42 +0000 Subject: [PATCH 16/22] perf(array): stop repeating per-chunk work in the executor loop Canonicalizing an N-chunk array through `execute_until` runs N+2 iterations, and four things were being paid on every one of them that only need paying once. Measured on `chunked_fsl_canonicalize`, median times drop 30-42%, with the gain growing in the chunk count: chunks before after 2 1.770 us 1.218 us -31% 8 3.214 us 1.996 us -38% 32 8.535 us 4.979 us -42% The timings are flat in list size, which is the tell: this is per-chunk bookkeeping, not data movement. - `AnyCanonical::matches` asked twelve encodings `array.is::()` in turn, each a virtual `as_any` call plus a `TypeId` comparison. It now compares the array's encoding id against the twelve canonical ids, interned once. - The loop then ran that check twice per iteration. `execute::` enters with `M = AnyCanonical`, so while the stack is empty the target predicate *is* the canonical stop condition. A `Matcher::IS_ANY_CANONICAL` associated constant lets the root case answer both with one scan, resolved at monomorphization rather than by comparing function pointers. - The builder was never told how many arrays were coming. `reserve_chunks` only ran in `Chunked::append_to_builder`, which `execute_until` bypasses now that `Chunked::execute` drives the builder through `AppendChild`, so every nested builder's chunk list grew incrementally: a 128-chunk array of 8-field structs took 40 reallocations, now none. - `Chunked` built a real empty array of its own dtype to carry the final `Done`, which `finalize_done` discards in favour of the builder. For a nested dtype that is a recursive construction of empty children, so `ExecutionResult::done_into_builder` now names that case and carries a null placeholder instead. That, plus the reserves, takes a 128-chunk 8-field struct from 99 allocations to 81. Also defers the executor's `expected_dtype` clone to debug builds, where `finalize_done` is the only thing that reads it back. Signed-off-by: Robert Kruszewski Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01CZTu9XUrdRZZ91rx852DGf --- vortex-array/src/arrays/chunked/vtable/mod.rs | 7 +- vortex-array/src/canonical.rs | 41 ++++++++---- vortex-array/src/executor.rs | 64 +++++++++++++++---- vortex-array/src/matcher.rs | 8 +++ vortex-array/src/test_harness/trace/tests.rs | 2 +- 5 files changed, 94 insertions(+), 28 deletions(-) diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index 4a551946e63..5fa22096220 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -276,10 +276,15 @@ impl VTable for Chunked { array.with_next_builder_slot(slot_idx + 1), slot_idx, )) - } else { + } else if slot_idx == ChunkedSlots::CHUNKS_OFFSET { + // No chunks at all, so nothing was ever appended and the executor holds no + // builder to finish. This array really is empty. Ok(ExecutionResult::done( Canonical::empty(array.dtype()).into_array(), )) + } else { + // Every chunk has gone to the builder, which is what the executor will finish. + Ok(ExecutionResult::done_into_builder()) } } } diff --git a/vortex-array/src/canonical.rs b/vortex-array/src/canonical.rs index ede547d95d2..b75ab9d75dd 100644 --- a/vortex-array/src/canonical.rs +++ b/vortex-array/src/canonical.rs @@ -4,6 +4,7 @@ //! Encodings that enable zero-copy sharing of data with Arrow. use std::sync::Arc; +use std::sync::LazyLock; use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; @@ -17,8 +18,10 @@ use crate::ArraySlots; use crate::Executable; use crate::ExecutionCtx; use crate::IntoArray; +use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::child_to_validity; +use crate::array::vtable::VTable as _; use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::arrays::Decimal; @@ -1230,25 +1233,39 @@ impl CanonicalView<'_> { } } +/// The encoding ids of the canonical encodings, interned once. +/// +/// [`AnyCanonical::matches`] is the executor's per-iteration predicate, so it compares these +/// rather than asking each encoding in turn: `array.is::()` is a virtual `as_any` call plus a +/// `TypeId` comparison, and twelve of those per iteration dwarf a scan of twelve interned ids. +static CANONICAL_IDS: LazyLock<[ArrayId; 12]> = LazyLock::new(|| { + [ + Null.id(), + Bool.id(), + Primitive.id(), + Decimal.id(), + Struct.id(), + Union.id(), + ListView.id(), + Map.id(), + FixedSizeList.id(), + VarBinView.id(), + Variant.id(), + Extension.id(), + ] +}); + /// A matcher for any canonical array type. pub struct AnyCanonical; impl Matcher for AnyCanonical { type Match<'a> = CanonicalView<'a>; + const IS_ANY_CANONICAL: bool = true; + #[inline] fn matches(array: &ArrayRef) -> bool { - array.is::() - || array.is::() - || array.is::() - || array.is::() - || array.is::() - || array.is::() - || array.is::() - || array.is::() - || array.is::() - || array.is::() - || array.is::() - || array.is::() + let id = array.encoding_id(); + CANONICAL_IDS.contains(&id) } #[inline] diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 18afd45e16c..4e7b9c36b90 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -186,12 +186,24 @@ impl ArrayRef { current_builder.is_some(), )); - let is_done = stack - .last() - .map_or(M::matches as DonePredicate, |frame| frame.done); - - let done_target = is_done(¤t_array); - let done_canonical = AnyCanonical::matches(¤t_array); + let (done_target, done_canonical) = match stack.last() { + // At the root the target matcher decides. `execute::` arrives with + // `M = AnyCanonical`, which is also the loop's universal stop condition, so one + // scan answers both rather than scanning the encoding twice per iteration. + None => { + let done_target = M::matches(¤t_array); + let done_canonical = if M::IS_ANY_CANONICAL { + done_target + } else { + AnyCanonical::matches(¤t_array) + }; + (done_target, done_canonical) + } + Some(frame) => ( + (frame.done)(¤t_array), + AnyCanonical::matches(¤t_array), + ), + }; trace_op!(record_execute_until_done_check(done_target, done_canonical)); if done_target || done_canonical { @@ -267,7 +279,10 @@ impl ArrayRef { } let expected_len = current_array.len(); - let expected_dtype = current_array.dtype().clone(); + // `finalize_done` only reads the dtype back under debug assertions, and only the + // `Done` arm reads it at all. Cloning a nested dtype is a pair of refcount ops, so + // skip it on the `AppendChild` and `ExecuteSlot` iterations that discard it. + let expected_dtype = cfg!(debug_assertions).then(|| current_array.dtype().clone()); let stats = current_array.statistics().to_array_stats(); let encoding_id = current_array.encoding_id(); trace_op!(record_execute_encoding(¤t_array)); @@ -292,11 +307,15 @@ impl ArrayRef { ExecutionStep::AppendChild(i) => { if current_builder.is_none() { trace_op!(record_builder_start(&array)); - current_builder = Some(builder_with_capacity_in( - array.dtype(), - array.len(), - ctx.allocator(), - )); + let mut builder = + builder_with_capacity_in(array.dtype(), array.len(), ctx.allocator()); + // The parent hands its children over one `AppendChild` at a time, so tell + // the builder how many are coming. A nested builder keeps each appended + // array as a chunk of its child, and without this every one of those chunk + // lists grows incrementally: a 128-chunk array of 8-field structs pays 40 + // reallocations for what is eight reserves. + builder.reserve_chunks(array.nchildren()); + current_builder = Some(builder); } let (parent, child) = unsafe { array.take_slot_unchecked(i) }?; @@ -591,7 +610,7 @@ fn finalize_done( result: ArrayRef, mut builder: Option>, expected_len: usize, - expected_dtype: DType, + expected_dtype: Option, stats: ArrayStats, encoding_id: ArrayId, ) -> VortexResult<(ArrayRef, Option>)> { @@ -601,7 +620,7 @@ fn finalize_done( result }; - if cfg!(debug_assertions) { + if let Some(expected_dtype) = expected_dtype { vortex_ensure!( output.len() == expected_len, "Result length mismatch for {:?}", @@ -781,6 +800,23 @@ impl ExecutionResult { } } + /// Signal that execution is complete and the result is in the executor's active builder. + /// + /// The executor finishes that builder and discards the array carried here, so an encoding + /// that has already handed every value over via [`ExecutionStep::AppendChild`] should use + /// this rather than building a real empty array of its own dtype only for it to be dropped - + /// for a nested dtype that is a recursive construction of empty children. + /// + /// Only valid once at least one [`ExecutionStep::AppendChild`] has been returned, which is + /// what guarantees the builder exists. [`finalize_done`] debug-asserts the resulting dtype, + /// so a mistake here shows up as a dtype mismatch rather than silently wrong data. + pub fn done_into_builder() -> Self { + Self { + array: Canonical::empty(&DType::Null).into_array(), + step: ExecutionStep::Done, + } + } + /// Request execution of slot at `slot_idx` until it matches the given [`Matcher`]. /// /// The provided array is the (possibly modified) parent that still needs its slot executed. diff --git a/vortex-array/src/matcher.rs b/vortex-array/src/matcher.rs index 4bda03f420e..5b4eea023ea 100644 --- a/vortex-array/src/matcher.rs +++ b/vortex-array/src/matcher.rs @@ -7,6 +7,14 @@ use crate::ArrayRef; pub trait Matcher { type Match<'a>; + /// Whether this matcher is [`AnyCanonical`](crate::AnyCanonical), which the executor checks + /// every iteration as its universal stop condition on top of the target matcher. + /// + /// Executing to canonical form - by far the common case - makes those two the same predicate, + /// and this lets the executor answer both with one scan instead of repeating it per iteration. + /// Only `AnyCanonical` may override it. + const IS_ANY_CANONICAL: bool = false; + /// Check if the given array matches this matcher type #[inline] fn matches(array: &ArrayRef) -> bool { diff --git a/vortex-array/src/test_harness/trace/tests.rs b/vortex-array/src/test_harness/trace/tests.rs index bfe99c3857b..76bf952da9e 100644 --- a/vortex-array/src/test_harness/trace/tests.rs +++ b/vortex-array/src/test_harness/trace/tests.rs @@ -588,7 +588,7 @@ fn trace_execution_chunked_append_child_flow() -> VortexResult<()> { AppendChild slot=3 parent=vortex.chunked(i32, len=5) child=vortex.primitive(i32, len=2) builder append child=vortex.primitive(i32, len=2) iter 3 current=vortex.chunked(i32, len=5) builder_active=true - Done array=vortex.primitive(i32, len=0) + Done array=vortex.null(null, len=0) builder finish output=vortex.primitive(i32, len=5) iter 4 current=vortex.primitive(i32, len=5) builder_active=false return output=vortex.primitive(i32, len=5) From 8d89e72e08fea626abb1408e8f29aded6e615d52 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 17 Sep 2026 01:08:31 +0000 Subject: [PATCH 17/22] refactor(array): keep the canonical-target shortcut out of the Matcher trait `Matcher::IS_ANY_CANONICAL` put an executor implementation detail on a public trait, where every implementor saw it and any of them could set it wrongly - a `true` on a matcher that is not `AnyCanonical` would make the executor stop early, which is a correctness bug rather than a performance one. The executor now asks the question itself, comparing `TypeId::of::()` against `AnyCanonical` once per `execute_until` call. It benchmarks the same as the constant did, and the trait goes back to what it was. Also closes a hole in the id-based `AnyCanonical::matches`: an encoding id does not identify a `VTable` type on its own, because `ForeignArray`, `ScalarFn` and the Python vtable each return a per-instance `self.id` and nothing stops one being registered under a canonical encoding's id. The id scan is now only a fast reject, with a hit still confirmed by the downcast in `try_match`. Without that, `matches` could answer yes where `try_match` answers `None`, and `Canonical::execute` turns that disagreement into a panic. Signed-off-by: Robert Kruszewski Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01CZTu9XUrdRZZ91rx852DGf --- vortex-array/src/canonical.rs | 11 +++++++---- vortex-array/src/executor.rs | 13 +++++++++++-- vortex-array/src/matcher.rs | 8 -------- 3 files changed, 18 insertions(+), 14 deletions(-) diff --git a/vortex-array/src/canonical.rs b/vortex-array/src/canonical.rs index b75ab9d75dd..8f0f9c9e19b 100644 --- a/vortex-array/src/canonical.rs +++ b/vortex-array/src/canonical.rs @@ -1260,12 +1260,15 @@ pub struct AnyCanonical; impl Matcher for AnyCanonical { type Match<'a> = CanonicalView<'a>; - const IS_ANY_CANONICAL: bool = true; - #[inline] fn matches(array: &ArrayRef) -> bool { - let id = array.encoding_id(); - CANONICAL_IDS.contains(&id) + // The id scan only rejects. An encoding id does not identify a `VTable` type on its own - + // `ForeignArray`, `ScalarFn` and the Python vtable all return a per-instance `self.id`, and + // nothing stops one of those being registered under a canonical encoding's id - so a hit + // still has to go through the downcast that `try_match` performs. Otherwise `matches` could + // answer yes where `try_match` answers `None`, and `Canonical::execute` turns that + // disagreement into a panic. + CANONICAL_IDS.contains(&array.encoding_id()) && Self::try_match(array).is_some() } #[inline] diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 4e7b9c36b90..2ee2f0b388b 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -12,6 +12,7 @@ //! See for the full execution //! narrative, diagrams, and walkthroughs. +use std::any::TypeId; use std::env::VarError; use std::fmt; use std::fmt::Display; @@ -166,7 +167,15 @@ impl ArrayRef { /// parent rewrite would observe inconsistent state and could discard accumulated builder /// data. #[allow(clippy::cognitive_complexity)] - pub fn execute_until(self, ctx: &mut ExecutionCtx) -> VortexResult { + pub fn execute_until( + self, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // Every iteration checks the target matcher and then `AnyCanonical`, the loop's universal + // "no further progress possible" stop condition. `execute::` passes + // `AnyCanonical` as the target too, which makes those the same call, and this folds to a + // constant per monomorphization so that case scans the encoding once rather than twice. + let target_is_any_canonical = TypeId::of::() == TypeId::of::(); let mut current_array = self; let mut current_builder: Option> = None; let mut stack: Vec = Vec::new(); @@ -192,7 +201,7 @@ impl ArrayRef { // scan answers both rather than scanning the encoding twice per iteration. None => { let done_target = M::matches(¤t_array); - let done_canonical = if M::IS_ANY_CANONICAL { + let done_canonical = if target_is_any_canonical { done_target } else { AnyCanonical::matches(¤t_array) diff --git a/vortex-array/src/matcher.rs b/vortex-array/src/matcher.rs index 5b4eea023ea..4bda03f420e 100644 --- a/vortex-array/src/matcher.rs +++ b/vortex-array/src/matcher.rs @@ -7,14 +7,6 @@ use crate::ArrayRef; pub trait Matcher { type Match<'a>; - /// Whether this matcher is [`AnyCanonical`](crate::AnyCanonical), which the executor checks - /// every iteration as its universal stop condition on top of the target matcher. - /// - /// Executing to canonical form - by far the common case - makes those two the same predicate, - /// and this lets the executor answer both with one scan instead of repeating it per iteration. - /// Only `AnyCanonical` may override it. - const IS_ANY_CANONICAL: bool = false; - /// Check if the given array matches this matcher type #[inline] fn matches(array: &ArrayRef) -> bool { From f8097f5f08737356e1c8f81009b225b15dc8624e Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 17 Sep 2026 02:11:15 +0000 Subject: [PATCH 18/22] perf(array): select the canonical encoding by id instead of walking the list `AnyCanonical::try_match` offered the array to each canonical encoding in turn, so a `FixedSizeList` array paid nine `as_opt` calls - each a virtual `as_any` plus a `TypeId` comparison - to reach its own arm, and an `Extension` array all twelve. Counting them over a chunked canonicalize showed 18 downcasts per execution against a single canonical array. Comparing the interned encoding id picks the one encoding worth asking, leaving one downcast to confirm it. That confirmation is what the previous commit paid for with a second walk of the list. It is still needed, because an encoding id does not identify a `VTable` type on its own, but it now costs one downcast rather than nine. `matches` keeps its own expansion rather than delegating to `try_match`: returning `bool` keeps the rejection path clear of the `CanonicalView` that the `Option` would carry, and rejecting with a single `||` reduction over the ids keeps it branchless. Both were worth about 4% on their own. Measured on `chunked_fsl_canonicalize` at list size 1024, medians of three runs, against the pre-optimization baseline re-measured in the same conditions: chunks before after 2 1.796 us 1.243 us -31% 8 3.219 us 2.069 us -36% 32 8.612 us 5.107 us -41% Both halves of the matcher are generated from one list of `field => vtable => variant` triples, so they cannot disagree about which encodings are canonical, and each arm's id comes from the same vtable that arm downcasts to. The pairing of a vtable with its `CanonicalView` variant is enforced by the compiler, since the view borrows `ArrayView<'_, V>`. A test covers the remaining gap: that every canonical encoding appears in the list at all, checked through `matches` and `try_match` alike. Signed-off-by: Robert Kruszewski Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01CZTu9XUrdRZZ91rx852DGf --- vortex-array/src/canonical.rs | 238 +++++++++++++++++++++++++--------- 1 file changed, 176 insertions(+), 62 deletions(-) diff --git a/vortex-array/src/canonical.rs b/vortex-array/src/canonical.rs index 8f0f9c9e19b..490d08cab9f 100644 --- a/vortex-array/src/canonical.rs +++ b/vortex-array/src/canonical.rs @@ -1233,82 +1233,111 @@ impl CanonicalView<'_> { } } -/// The encoding ids of the canonical encodings, interned once. -/// -/// [`AnyCanonical::matches`] is the executor's per-iteration predicate, so it compares these -/// rather than asking each encoding in turn: `array.is::()` is a virtual `as_any` call plus a -/// `TypeId` comparison, and twelve of those per iteration dwarf a scan of twelve interned ids. -static CANONICAL_IDS: LazyLock<[ArrayId; 12]> = LazyLock::new(|| { - [ - Null.id(), - Bool.id(), - Primitive.id(), - Decimal.id(), - Struct.id(), - Union.id(), - ListView.id(), - Map.id(), - FixedSizeList.id(), - VarBinView.id(), - Variant.id(), - Extension.id(), - ] -}); - /// A matcher for any canonical array type. pub struct AnyCanonical; -impl Matcher for AnyCanonical { - type Match<'a> = CanonicalView<'a>; - - #[inline] - fn matches(array: &ArrayRef) -> bool { - // The id scan only rejects. An encoding id does not identify a `VTable` type on its own - - // `ForeignArray`, `ScalarFn` and the Python vtable all return a per-instance `self.id`, and - // nothing stops one of those being registered under a canonical encoding's id - so a hit - // still has to go through the downcast that `try_match` performs. Otherwise `matches` could - // answer yes where `try_match` answers `None`, and `Canonical::execute` turns that - // disagreement into a panic. - CANONICAL_IDS.contains(&array.encoding_id()) && Self::try_match(array).is_some() - } - #[inline] - fn try_match(array: &ArrayRef) -> Option> { - if let Some(a) = array.as_opt::() { - Some(CanonicalView::Null(a)) - } else if let Some(a) = array.as_opt::() { - Some(CanonicalView::Bool(a)) - } else if let Some(a) = array.as_opt::() { - Some(CanonicalView::Primitive(a)) - } else if let Some(a) = array.as_opt::() { - Some(CanonicalView::Decimal(a)) - } else if let Some(a) = array.as_opt::() { - Some(CanonicalView::Struct(a)) - } else if let Some(a) = array.as_opt::() { - Some(CanonicalView::Union(a)) - } else if let Some(a) = array.as_opt::() { - Some(CanonicalView::List(a)) - } else if let Some(a) = array.as_opt::() { - Some(CanonicalView::Map(a)) - } else if let Some(a) = array.as_opt::() { - Some(CanonicalView::FixedSizeList(a)) - } else if let Some(a) = array.as_opt::() { - Some(CanonicalView::VarBinView(a)) - } else if let Some(a) = array.as_opt::() { - Some(CanonicalView::Variant(a)) - } else { - array.as_opt::().map(CanonicalView::Extension) +/// The canonical encodings, as `field => vtable => CanonicalView variant` triples. +/// +/// [`canonical_matcher`] expands this into the id cache and both halves of the [`Matcher`] impl, +/// so the two cannot disagree about which encodings are canonical. +macro_rules! with_canonical_encodings { + ($mac:ident) => { + $mac! { + null => Null => Null, + bool_ => Bool => Bool, + primitive => Primitive => Primitive, + decimal => Decimal => Decimal, + struct_ => Struct => Struct, + union_ => Union => Union, + list => ListView => List, + map => Map => Map, + fixed_size_list => FixedSizeList => FixedSizeList, + varbinview => VarBinView => VarBinView, + variant => Variant => Variant, + extension => Extension => Extension, } - } + }; } +/// Defines the canonical encoding id cache and [`AnyCanonical`]'s [`Matcher`] impl. +macro_rules! canonical_matcher { + ($($field:ident => $vtable:ident => $variant:ident,)+) => { + /// The encoding ids of the canonical encodings, interned once. + /// + /// `VTable::id` takes `&self` and reads through a `CachedId`, so these are resolved once + /// rather than on every match. + struct CanonicalIds { + $($field: ArrayId,)+ + } + + static CANONICAL_IDS: LazyLock = LazyLock::new(|| CanonicalIds { + $($field: $vtable.id(),)+ + }); + + impl Matcher for AnyCanonical { + type Match<'a> = CanonicalView<'a>; + + #[inline] + fn matches(array: &ArrayRef) -> bool { + // Comparing the interned encoding id picks the one encoding worth asking, instead + // of offering the array to each in turn: `array.is::()` is a virtual `as_any` + // call plus a `TypeId` comparison, and this is the executor's per-iteration + // predicate. + // + // The id only selects, it does not decide. An encoding id does not identify a + // `VTable` type on its own - `ForeignArray`, `ScalarFn` and the Python vtable each + // return a per-instance `self.id`, and nothing stops one being registered under a + // canonical encoding's id - so the selected encoding still confirms by downcast. + // Without that, `matches` could answer yes where `try_match` answers `None`, and + // `Canonical::execute` turns that disagreement into a panic. + // + // This deliberately does not delegate to `try_match`: returning `bool` keeps the + // rejection path, which is the common one, clear of the `CanonicalView` the + // `Option` would carry. + let ids = &*CANONICAL_IDS; + let id = array.encoding_id(); + + // Reject as one `||` reduction over the ids rather than as a dozen separate + // branches, so the common non-canonical case stays branchless. Only an id that + // survives it walks the chain below to find its encoding. + if !($(id == ids.$field ||)+ false) { + return false; + } + + $(if id == ids.$field { + return array.is::<$vtable>(); + })+ + false + } + + #[inline] + fn try_match(array: &ArrayRef) -> Option> { + // Selected by id and confirmed by downcast, as in `matches` above. Walking the + // encodings in order instead cost a `FixedSizeList` array nine `as_opt` calls to + // reach its own arm, and an `Extension` array all twelve. + let ids = &*CANONICAL_IDS; + let id = array.encoding_id(); + $(if id == ids.$field { + return array.as_opt::<$vtable>().map(CanonicalView::$variant); + })+ + None + } + } + }; +} + +with_canonical_encodings!(canonical_matcher); + #[cfg(test)] mod test { + use std::sync::Arc; use std::sync::LazyLock; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_session::VortexSession; + use crate::AnyCanonical; use crate::ArrayRef; use crate::Canonical; use crate::CanonicalValidity; @@ -1323,12 +1352,97 @@ mod test { use crate::arrays::struct_::StructArrayExt; use crate::arrays::variant::VariantArraySlotsExt; use crate::canonical::StructArray; + use crate::dtype::DType; + use crate::dtype::DecimalDType; + use crate::dtype::MapDType; use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::dtype::StructFields; + use crate::dtype::UnionVariants; + use crate::extension::datetime::Time; + use crate::extension::datetime::TimeUnit; + use crate::matcher::Matcher; use crate::scalar::Scalar; /// A shared session for these canonical tests, used to create execution contexts. static SESSION: LazyLock = LazyLock::new(crate::array_session); + /// One empty array per canonical encoding, covering every arm the matcher generates. + fn one_array_per_canonical_encoding() -> VortexResult> { + let i32_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let dtypes = [ + DType::Null, + DType::Bool(Nullability::NonNullable), + i32_dtype.clone(), + DType::Decimal(DecimalDType::new(10, 2), Nullability::NonNullable), + DType::Utf8(Nullability::NonNullable), + DType::List(Arc::new(i32_dtype.clone()), Nullability::NonNullable), + DType::Map( + MapDType::try_new(i32_dtype.clone(), i32_dtype.clone(), false)?, + Nullability::NonNullable, + ), + DType::FixedSizeList(Arc::new(i32_dtype.clone()), 2, Nullability::NonNullable), + DType::Struct( + StructFields::new(["a"].into(), vec![i32_dtype.clone()]), + Nullability::NonNullable, + ), + DType::Union( + UnionVariants::try_new(["a"].into(), vec![i32_dtype], vec![0])?, + Nullability::NonNullable, + ), + DType::Extension(Time::new(TimeUnit::Seconds, Nullability::NonNullable).erased()), + ]; + + let mut arrays: Vec = dtypes + .iter() + .map(|dtype| Canonical::empty(dtype).into_array()) + .collect(); + // `Canonical::empty` rejects `DType::Variant`, so build that one directly. + arrays.push(VariantArray::try_new(variant_core_storage(0), None)?.into_array()); + + Ok(arrays) + } + + /// Every canonical encoding must reach its own arm of both halves of [`AnyCanonical`]. + /// + /// Each arm is selected by comparing the interned encoding id and then confirmed by a + /// downcast, so an arm that pairs an id with the wrong encoding answers no here rather than + /// viewing the array as the wrong variant. Pairing an id with the wrong `CanonicalView` + /// variant does not need covering: the view borrows `ArrayView<'_, V>`, so the compiler + /// rejects it. + /// + /// Checking `matches` as well as `try_match` is what keeps the two in step: they are separate + /// expansions, and the executor stops on `matches` while `Canonical::execute` then unwraps + /// `try_match`, so a disagreement between them is a panic. + /// + /// A loop rather than `rstest` cases because several of these dtypes are built fallibly. + #[test] + fn every_canonical_encoding_matches_any_canonical() -> VortexResult<()> { + let arrays = one_array_per_canonical_encoding()?; + assert_eq!( + arrays.len(), + 12, + "expected one array per canonical encoding" + ); + + for array in arrays { + assert!( + AnyCanonical::matches(&array), + "{} array of dtype {} did not match AnyCanonical", + array.encoding_id(), + array.dtype(), + ); + assert!( + AnyCanonical::try_match(&array).is_some(), + "{} array of dtype {} did not view as AnyCanonical", + array.encoding_id(), + array.dtype(), + ); + } + + Ok(()) + } + fn variant_core_storage(len: usize) -> ArrayRef { ConstantArray::new( Scalar::variant(Scalar::primitive(1i32, Nullability::NonNullable)), From 6b651cb194526cb7e92604eae71c57b504ad8daf Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 17 Sep 2026 02:34:07 +0000 Subject: [PATCH 19/22] perf(array): drop the TypeId downcast from Array::data_mut `data_mut` reached its own encoding data through `as_any_mut().downcast_mut()`, spending a virtual call and a `TypeId` comparison to learn what `Array` already guarantees. The shared path beside it, `downcast_inner`, does not: it asserts the invariant in debug and casts. This makes the two consistent. The executor pays this once per chunk, through `Chunked::execute` -> `with_next_builder_slot`, which threads the builder slot counter through the array itself. Callgrind over a 32-chunk FSL canonicalize attributed 1,977 instructions per execution to that function, a third of them in `core::any`. It is now 1,417, and the whole canonicalize drops from 41,071 instructions to 40,314. `as_any_mut` had no other callers, so `DynArrayData` loses it. Wall clock on `chunked_fsl_canonicalize` at list size 1024, medians of three runs, against the pre-optimization baseline measured in the same conditions: chunks before after 2 1.796 us 1.227 us -32% 8 3.219 us 1.989 us -38% 32 8.612 us 4.999 us -42% Signed-off-by: Robert Kruszewski Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01CZTu9XUrdRZZ91rx852DGf --- vortex-array/src/array/mod.rs | 7 ------- vortex-array/src/array/typed.rs | 11 +++++++++-- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index c3fb88a8c62..69c16c9b3a6 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -121,9 +121,6 @@ pub(crate) trait DynArrayData: 'static + private::Sealed + Send + Sync + Debug { /// Returns the array as a reference to a generic [`Any`] trait object. fn as_any(&self) -> &dyn Any; - /// Returns the array as a mutable reference to a generic [`Any`] trait object. - fn as_any_mut(&mut self) -> &mut dyn Any; - /// Returns the [`Validity`] of the array. fn validity(&self, this: &ArrayRef) -> VortexResult; @@ -271,10 +268,6 @@ impl DynArrayData for ArrayData { self } - fn as_any_mut(&mut self) -> &mut dyn Any { - self - } - fn validity(&self, this: &ArrayRef) -> VortexResult { if this.dtype().is_nullable() { let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; diff --git a/vortex-array/src/array/typed.rs b/vortex-array/src/array/typed.rs index a9def0318bf..4f5c806a749 100644 --- a/vortex-array/src/array/typed.rs +++ b/vortex-array/src/array/typed.rs @@ -306,8 +306,15 @@ impl Array { /// Returns `None` when this handle is not the unique owner of the backing allocation. pub fn data_mut(&mut self) -> Option<&mut V::TypedArrayData> { let store = self.inner.inner_mut()?; - let array_inner = store.data.as_any_mut().downcast_mut::>(); - Some(&mut array_inner?.data) + // NOTE(ngates): use downcast_mut_unchecked when it becomes stable + debug_assert!(store.data.as_any().is::>()); + // SAFETY: `Array` guarantees the inner is `ArrayData`, the same invariant the shared + // path in `downcast_inner` relies on. Reaching it through `Any` instead spends a virtual + // `as_any_mut` and a `TypeId` comparison on a type the compiler already knows, which the + // executor pays once per chunk through `with_next_builder_slot`. + let array_inner = + unsafe { &mut *std::ptr::from_mut(&mut store.data).cast::>() }; + Some(&mut array_inner.data) } /// Returns the full typed array construction parts if this handle owns the allocation. From ce0bd9cef488e64a474929fbe49e42e6f51187d9 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 16 Sep 2026 22:31:12 -0400 Subject: [PATCH 20/22] maybe Signed-off-by: Robert Kruszewski --- vortex-array/src/arrays/chunked/array.rs | 19 ++++++++++--------- vortex-array/src/builders/validity.rs | 16 +--------------- 2 files changed, 11 insertions(+), 24 deletions(-) diff --git a/vortex-array/src/arrays/chunked/array.rs b/vortex-array/src/arrays/chunked/array.rs index 234dd23e9cf..668ef5d7c38 100644 --- a/vortex-array/src/arrays/chunked/array.rs +++ b/vortex-array/src/arrays/chunked/array.rs @@ -10,7 +10,7 @@ use std::fmt::Display; use std::fmt::Formatter; use futures::stream; -use vortex_buffer::BufferMut; +use vortex_buffer::{Buffer, BufferMut}; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -138,13 +138,14 @@ impl ChunkedData { } 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) } - } - PrimitiveArray::new(chunk_offsets_buf.freeze(), Validity::NonNullable).into_array() + let chunk_offsets_buf = + Buffer::from_trusted_len_iter(chunk_offsets.iter().copied().map(|offset| { + u64::try_from(offset) + .vortex_expect("chunk offset must fit in u64 for serialization") + })); + + unsafe { PrimitiveArray::new_unchecked(chunk_offsets_buf, Validity::NonNullable) } + .into_array() } /// Validates the components that would be used to create a `ChunkedArray`. @@ -194,7 +195,7 @@ impl Array { return self; } // This is the slow path that will be hit at most once per execution since the second one - // *MUST* have execlusive access due to this copy. + // *MUST* have exclusive access due to this copy. let stats = self.statistics().to_owned(); let mut data = self.data().clone(); data.next_builder_slot = next_builder_slot; diff --git a/vortex-array/src/builders/validity.rs b/vortex-array/src/builders/validity.rs index 678048c6a13..a64e48ba898 100644 --- a/vortex-array/src/builders/validity.rs +++ b/vortex-array/src/builders/validity.rs @@ -78,12 +78,8 @@ impl ValidityBuilder { self.flush_pending(); self.runs_len += len; - // Two adjacent runs that are uniform and agree describe one longer run. Extending beats - // pushing: a chunked array of non-nullable chunks would otherwise record a run per chunk - // for a validity that `Validity::concat` collapses at the end anyway, and every run is - // walked again by `finish_with_nullability`. if let Some((last, last_len)) = self.runs.last_mut() - && Self::is_uniform(last) + && !matches!(last, Validity::Array(_)) && std::mem::discriminant(last) == std::mem::discriminant(&validity) { *last_len += len; @@ -137,16 +133,6 @@ impl ValidityBuilder { } validity } - - /// Whether a validity describes every value it covers the same way, so that two adjacent runs - /// of it can be recorded as one. - fn is_uniform(validity: &Validity) -> bool { - matches!( - validity, - Validity::NonNullable | Validity::AllValid | Validity::AllInvalid - ) - } - /// Moves whatever the null buffer holds into `runs`, keeping the runs in logical order. fn flush_pending(&mut self) { let len = self.pending.len(); From b22210342b8539cd79a2feb4adec51732e5c25e0 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 17 Sep 2026 00:40:25 -0400 Subject: [PATCH 21/22] keep swizzling Signed-off-by: Robert Kruszewski --- vortex-array/Cargo.toml | 4 + .../benches/chunked_nested_canonicalize.rs | 77 +++++++ vortex-array/src/array/typed.rs | 16 +- vortex-array/src/arrays/chunked/array.rs | 26 +-- vortex-array/src/arrays/chunked/tests.rs | 173 ++++++++++++++++ .../src/arrays/chunked/vtable/canonical.rs | 196 ++++++++++++++++++ vortex-array/src/arrays/chunked/vtable/mod.rs | 16 +- vortex-array/src/executor.rs | 16 +- vortex-array/src/test_harness/trace/tests.rs | 2 +- 9 files changed, 491 insertions(+), 35 deletions(-) create mode 100644 vortex-array/benches/chunked_nested_canonicalize.rs diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index d0a1afcf917..0681d66b621 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -171,6 +171,10 @@ harness = false name = "chunked_fsl_canonicalize" harness = false +[[bench]] +name = "chunked_nested_canonicalize" +harness = false + [[bench]] name = "scalar_at_struct" harness = false diff --git a/vortex-array/benches/chunked_nested_canonicalize.rs b/vortex-array/benches/chunked_nested_canonicalize.rs new file mode 100644 index 00000000000..f5d32912a4c --- /dev/null +++ b/vortex-array/benches/chunked_nested_canonicalize.rs @@ -0,0 +1,77 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_session::VortexSession; + +static SESSION: LazyLock = LazyLock::new(array_session); +const ROWS: usize = 1_024; +const CHUNKS: &[usize] = &[2, 32]; + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +fn canonicalize(bencher: Bencher, chunk: ArrayRef, nchunks: usize) { + let array = ChunkedArray::try_new( + std::iter::repeat_n(chunk.clone(), nchunks), + chunk.dtype().clone(), + ) + .unwrap() + .into_array(); + bencher + .with_inputs(|| (&array, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| array.clone().execute::(ctx).unwrap()); +} + +#[divan::bench(args = CHUNKS, consts = [1, 8])] +fn structs(bencher: Bencher, nchunks: usize) { + let field = PrimitiveArray::from_iter(0..ROWS as u64).into_array(); + let chunk = StructArray::try_from_iter( + (0..FIELDS).map(|idx| (format!("field_{idx}"), field.clone())), + ) + .unwrap(); + canonicalize(bencher, chunk.into_array(), nchunks); +} + +#[divan::bench(args = CHUNKS)] +fn lists(bencher: Bencher, nchunks: usize) { + let elements = PrimitiveArray::from_iter(0..4 * ROWS as u64).into_array(); + let offsets: Buffer = (0..=ROWS).map(|idx| 4 * idx as u64).collect(); + let chunk = ListArray::try_new(elements, offsets.into_array(), Validity::NonNullable).unwrap(); + canonicalize(bencher, chunk.into_array(), nchunks); +} + +#[divan::bench(args = CHUNKS, consts = [false, true])] +fn listviews(bencher: Bencher, nchunks: usize) { + let step = if OVERLAPPING { 1 } else { 4 }; + let elements = PrimitiveArray::from_iter(0..(step * ROWS + 4) as u64).into_array(); + let offsets: Buffer = (0..ROWS).map(|idx| (step * idx) as u64).collect(); + let sizes: Buffer = std::iter::repeat_n(4, ROWS).collect(); + let chunk = ListViewArray::new( + elements, + offsets.into_array(), + sizes.into_array(), + Validity::NonNullable, + ); + // SAFETY: with step 4, every view ends exactly where the next begins. + let chunk = unsafe { chunk.with_zero_copy_to_list(!OVERLAPPING) }; + canonicalize(bencher, chunk.into_array(), nchunks); +} diff --git a/vortex-array/src/array/typed.rs b/vortex-array/src/array/typed.rs index 4f5c806a749..975055d2633 100644 --- a/vortex-array/src/array/typed.rs +++ b/vortex-array/src/array/typed.rs @@ -42,6 +42,7 @@ use crate::validity::Validity; /// `ArrayRef` stores `Arc>` — a single 16-byte fat pointer. /// Metadata is accessed via `self.0.*` (a normal struct field read through the Arc), /// while encoding-specific methods go through `self.0.data` (vtable dispatch). +#[derive(Clone)] pub(crate) struct ArrayInner { pub(crate) len: usize, pub(crate) encoding_id: ArrayId, @@ -311,12 +312,25 @@ impl Array { // SAFETY: `Array` guarantees the inner is `ArrayData`, the same invariant the shared // path in `downcast_inner` relies on. Reaching it through `Any` instead spends a virtual // `as_any_mut` and a `TypeId` comparison on a type the compiler already knows, which the - // executor pays once per chunk through `with_next_builder_slot`. + // executor pays once per chunk through `with_next_child_slot`. let array_inner = unsafe { &mut *std::ptr::from_mut(&mut store.data).cast::>() }; Some(&mut array_inner.data) } + /// Mutates encoding data, copying shared array metadata and slots while retaining statistics. + /// + /// The update must preserve the logical values and array invariants. + pub(crate) fn with_data_mut(self, update: impl FnOnce(&mut V::TypedArrayData)) -> Self { + // SAFETY: Array guarantees the inner is ArrayData. + let mut inner = unsafe { self.inner.downcast_inner_unchecked::() }; + update(&mut Arc::make_mut(&mut inner).data.data); + Self { + inner: ArrayRef::from_inner(inner), + _phantom: PhantomData, + } + } + /// Returns the full typed array construction parts if this handle owns the allocation. pub fn try_into_parts(self) -> Result, Self> { let Self { inner, _phantom } = self; diff --git a/vortex-array/src/arrays/chunked/array.rs b/vortex-array/src/arrays/chunked/array.rs index 668ef5d7c38..33c700819b0 100644 --- a/vortex-array/src/arrays/chunked/array.rs +++ b/vortex-array/src/arrays/chunked/array.rs @@ -49,8 +49,8 @@ pub struct ChunkedSlots { #[derive(Clone, Debug)] pub struct ChunkedData { pub(super) chunk_offsets: Vec, - /// This is used to find the next child to execute when in executing into a builder. - pub(super) next_builder_slot: usize, + /// The next child to visit while executing slots or appending into a builder. + pub(super) next_child_slot: usize, } impl Display for ChunkedData { @@ -133,7 +133,7 @@ impl ChunkedData { pub(super) fn new(chunk_offsets: Vec) -> Self { Self { chunk_offsets, - next_builder_slot: ChunkedSlots::CHUNKS_OFFSET, + next_child_slot: ChunkedSlots::CHUNKS_OFFSET, } } @@ -189,24 +189,8 @@ impl Array { Ok(ArrayParts::new(Chunked, dtype, len, ChunkedData::new(chunk_offsets)).with_slots(slots)) } - pub(super) fn with_next_builder_slot(mut self, next_builder_slot: usize) -> Self { - if let Some(data) = self.data_mut() { - data.next_builder_slot = next_builder_slot; - return self; - } - // This is the slow path that will be hit at most once per execution since the second one - // *MUST* have exclusive access due to this copy. - let stats = self.statistics().to_owned(); - let mut data = self.data().clone(); - data.next_builder_slot = next_builder_slot; - // SAFETY: we only modified next_builder_slot which doesn't affect array invariants. - unsafe { - Array::from_parts_unchecked( - ArrayParts::new(Chunked, self.dtype().clone(), self.len(), data) - .with_slots(self.slots().iter().cloned().collect::()), - ) - } - .with_stats_set(stats) + pub(super) fn with_next_child_slot(self, next_child_slot: usize) -> Self { + self.with_data_mut(|data| data.next_child_slot = next_child_slot) } /// Constructs a new `ChunkedArray`. diff --git a/vortex-array/src/arrays/chunked/tests.rs b/vortex-array/src/arrays/chunked/tests.rs index e0755108e9c..41eb7e13850 100644 --- a/vortex-array/src/arrays/chunked/tests.rs +++ b/vortex-array/src/arrays/chunked/tests.rs @@ -9,12 +9,14 @@ use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; +use crate::ArrayRef; use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; +use crate::arrays::FixedSizeListArray; use crate::arrays::ListArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; @@ -22,6 +24,7 @@ use crate::arrays::StructArray; use crate::arrays::VarBinViewArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::dict_test::gen_dict_primitive_chunks; +use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; use crate::builders::builder_with_capacity_in; @@ -167,6 +170,176 @@ fn execute_path_repeated_shared_chunked_dict_execution() { assert_arrays_eq!(second, expected, &mut ctx); } +#[rstest::rstest] +#[case(false, 2)] +#[case(true, 2)] +#[case(false, 0)] +#[case(true, 0)] +fn execute_path_repeated_shared_chunked_fsl( + #[case] nested: bool, + #[case] list_size: u32, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let elements = PrimitiveArray::from_iter(0..(2 * list_size)).into_array(); + let chunk = FixedSizeListArray::new(elements, list_size, Validity::from_iter([true, false]), 2) + .into_array(); + let child = if nested { + ChunkedArray::try_new(vec![chunk.clone()], chunk.dtype().clone())?.into_array() + } else { + chunk.clone() + }; + let array = + ChunkedArray::try_new(vec![child.clone(), child], chunk.dtype().clone())?.into_array(); + let expected = FixedSizeListArray::new( + PrimitiveArray::from_iter((0..(2 * list_size)).chain(0..(2 * list_size))).into_array(), + list_size, + Validity::from_iter([true, false, true, false]), + 4, + ); + + for _ in 0..2 { + let output = array.clone().execute::(&mut ctx)?.into_array(); + assert_arrays_eq!(output, expected.clone(), &mut ctx); + assert_arrays_eq!( + array.as_::().chunk(0).clone(), + chunk.clone(), + &mut ctx + ); + assert_arrays_eq!( + array.as_::().chunk(1).clone(), + chunk.clone(), + &mut ctx + ); + } + Ok(()) +} + +#[rstest::rstest] +#[case(false, 0)] +#[case(true, 0)] +#[case(false, 2)] +#[case(true, 2)] +fn swizzle_shared_structs(#[case] nested: bool, #[case] nfields: usize) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let field = buffer![10u64, 20].into_array(); + let names = (0..nfields).map(|idx| format!("field_{idx}")).collect(); + let chunk = StructArray::try_new( + names, + std::iter::repeat_n(field.clone(), nfields), + 2, + Validity::from_iter([true, false]), + )?; + let expected = StructArray::try_new( + chunk.names().clone(), + std::iter::repeat_n(buffer![10u64, 20, 10, 20].into_array(), nfields), + 4, + Validity::from_iter([true, false, true, false]), + )?; + let child = if nested { + ChunkedArray::try_new(vec![chunk.clone().into_array()], chunk.dtype().clone())?.into_array() + } else { + chunk.clone().into_array() + }; + let array = + ChunkedArray::try_new(vec![child.clone(), child], chunk.dtype().clone())?.into_array(); + + for _ in 0..2 { + let output = array.clone().execute::(&mut ctx)?; + if !nested { + for idx in 0..nfields { + assert!(ArrayRef::ptr_eq( + output.unmasked_field(idx).as_::().chunk(0), + &field, + )); + } + } + assert_arrays_eq!(output, expected.clone(), &mut ctx); + } + Ok(()) +} + +#[rstest::rstest] +#[case(false, false, false)] +#[case(true, false, false)] +#[case(false, true, false)] +#[case(true, true, false)] +#[case(false, true, true)] +#[case(true, true, true)] +fn swizzle_shared_lists( + #[case] nested: bool, + #[case] views: bool, + #[case] overlapping: bool, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let validity = Validity::from_iter([true, false, true]); + let chunk = if overlapping { + ListViewArray::new( + buffer![99i32, 10, 11, 12, 88].into_array(), + buffer![2i64, 2, 1].into_array(), + buffer![2i16, 1, 2].into_array(), + validity, + ) + .into_array() + } else if views { + let view = ListViewArray::new( + buffer![99i32, 11, 12, 42, 10, 11, 88].into_array(), + buffer![1u32, 3, 4].into_array(), + buffer![2u16, 1, 2].into_array(), + validity, + ); + // SAFETY: the views are contiguous and non-overlapping, with unreferenced ends. + unsafe { view.with_zero_copy_to_list(true) }.into_array() + } else { + ListArray::try_new( + buffer![99i32, 11, 12, 42, 10, 11, 88].into_array(), + buffer![1i32, 3, 4, 6].into_array(), + validity, + )? + .into_array() + }; + let child = if nested { + ChunkedArray::try_new(vec![chunk.clone()], chunk.dtype().clone())?.into_array() + } else { + chunk.clone() + }; + let array = + ChunkedArray::try_new(vec![child.clone(), child], chunk.dtype().clone())?.into_array(); + let expected = ListArray::try_new( + buffer![11i32, 12, 0, 10, 11, 11, 12, 0, 10, 11].into_array(), + buffer![0u64, 2, 3, 5, 7, 8, 10].into_array(), + Validity::from_iter([true, false, true, true, false, true]), + )?; + + for _ in 0..2 { + let output = array.clone().execute::(&mut ctx)?; + assert_eq!(output.elements().len(), if overlapping { 6 } else { 10 }); + assert_eq!(output.is_zero_copy_to_list(), !overlapping); + assert_arrays_eq!(output, expected.clone(), &mut ctx); + } + Ok(()) +} + +#[test] +fn swizzle_empty_list_rows() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let chunk = ListArray::try_new( + buffer![1u64, 2, 3].into_array(), + buffer![2u32, 2, 2].into_array(), + Validity::from_iter([true, false]), + )? + .into_array(); + let array = ChunkedArray::try_new(vec![chunk.clone(), chunk.clone()], chunk.dtype().clone())?; + let output = array.into_array().execute::(&mut ctx)?; + let expected = ListArray::try_new( + Buffer::::empty().into_array(), + buffer![0u32, 0, 0, 0, 0].into_array(), + Validity::from_iter([true, false, true, false]), + )?; + assert!(output.elements().is_empty()); + assert_arrays_eq!(output, expected, &mut ctx); + Ok(()) +} + #[test] fn execute_path_nested_chunked_dict_of_dict_into_canonical() { let mut ctx = SESSION.create_execution_ctx(); diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index b419fd4d77f..5a6e93d5c0a 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools as _; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; @@ -9,13 +10,208 @@ use vortex_error::vortex_err; use crate::ArrayRef; use crate::Canonical; use crate::ExecutionCtx; +use crate::ExecutionResult; use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; +use crate::arrays::FixedSizeList; +use crate::arrays::FixedSizeListArray; +use crate::arrays::List; +use crate::arrays::ListView; +use crate::arrays::ListViewArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::Struct; +use crate::arrays::StructArray; use crate::arrays::VariantArray; use crate::arrays::chunked::ChunkedArrayExt; +use crate::arrays::chunked::array::ChunkedSlots; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use crate::arrays::list::ListArraySlotsExt; +use crate::arrays::listview::ListViewArraySlotsExt; +use crate::arrays::struct_::StructArrayExt; use crate::arrays::variant::VariantArraySlotsExt; +use crate::dtype::DType; +use crate::match_each_integer_ptype; +use crate::matcher::Matcher; +use crate::validity::Validity; + +pub(super) struct ListChunks; + +impl Matcher for ListChunks { + type Match<'a> = (); + + fn try_match(array: &ArrayRef) -> Option<()> { + (array.is::() || array.is::()).then_some(()) + } +} + +/// Executes chunks into the required encoding one at a time, then transposes their children. +pub(super) fn swizzle( + array: ChunkedArray, + finish: impl FnOnce(ChunkedArray) -> VortexResult, +) -> VortexResult { + let next_chunk = array + .next_child_slot + .saturating_sub(ChunkedSlots::CHUNKS_OFFSET); + let needs_execution = array + .iter_chunks() + .enumerate() + .skip(next_chunk) + .find(|(_, chunk)| !V::matches(chunk)) + .map(|(idx, _)| idx + ChunkedSlots::CHUNKS_OFFSET); + if let Some(slot) = needs_execution { + return Ok(ExecutionResult::execute_slot::( + array.with_next_child_slot(slot + 1), + slot, + )); + } + + Ok(ExecutionResult::done(finish(array)?)) +} + +pub(super) fn swizzle_fixed_size_list(array: ChunkedArray) -> VortexResult { + let DType::FixedSizeList(element_dtype, list_size, _) = array.dtype() else { + unreachable!("called only for a fixed-size list dtype") + }; + // Canonical FSL children are trimmed to list_size * len; their elements concatenate directly. + let element_chunks: Vec<_> = array + .iter_chunks() + .map(|chunk| chunk.as_::().elements().clone()) + .collect(); + let validity = array.validity()?; + // SAFETY: all chunks share the parent's FSL dtype, so their elements share element_dtype + // and their lengths sum to list_size * array.len(). The parent supplies the combined validity. + Ok(unsafe { + let elements = + ChunkedArray::new_unchecked(element_chunks, element_dtype.as_ref().clone()) + .into_array(); + FixedSizeListArray::new_unchecked(elements, *list_size, validity, array.len()).into_array() + }) +} + +pub(super) fn swizzle_struct(array: ChunkedArray) -> VortexResult { + let struct_fields = array.dtype().as_struct_fields(); + let chunks: Vec<_> = array + .iter_chunks() + .map(|chunk| chunk.as_::()) + .collect(); + let fields = struct_fields.fields().enumerate().map(|(idx, dtype)| { + let children = chunks.iter().map(|chunk| chunk.unmasked_field(idx).clone()); + // SAFETY: each chunk has the parent's struct dtype, which fixes each field's dtype. + unsafe { ChunkedArray::new_unchecked(children, dtype) }.into_array() + }); + let validity = array.validity()?; + // SAFETY: each field concatenates the same chunk lengths, summing to array.len(). Their + // dtypes come from the parent and its validity applies to exactly those rows, including nulls. + Ok(unsafe { + StructArray::new_unchecked(fields, struct_fields.clone(), array.len(), validity) + .into_array() + }) +} + +pub(super) fn swizzle_list(array: ChunkedArray, ctx: &mut ExecutionCtx) -> VortexResult { + let DType::List(element_dtype, _) = array.dtype() else { + unreachable!("called only for a list dtype") + }; + let mut elements = Vec::with_capacity(array.nchunks()); + let mut offsets = ctx.allocator().zeroed::(array.len()); + let mut sizes = ctx.allocator().zeroed::(array.len()); + let mut element_base = 0usize; + let mut row = 0; + let mut zero_copy_to_list = true; + + for chunk in array.iter_chunks() { + if chunk.is_empty() { + continue; + } + let offsets_out = &mut offsets.as_mut_slice()[row..row + chunk.len()]; + let sizes_out = &mut sizes.as_mut_slice()[row..row + chunk.len()]; + let (child, start, end, chunk_zero_copy_to_list) = if let Some(list) = chunk.as_opt::() { + let chunk_offsets = list.offsets().clone().execute::(ctx)?; + match_each_integer_ptype!(chunk_offsets.ptype(), |O| { + for ((out, size), pair) in offsets_out + .iter_mut() + .zip(sizes_out.iter_mut()) + .zip(chunk_offsets.as_slice::().windows(2)) + { + let start = u64::try_from(pair[0]) + .vortex_expect("validated list offset is nonnegative"); + let end = u64::try_from(pair[1]) + .vortex_expect("validated list offset is nonnegative"); + *out = start; + *size = end - start; + } + }); + let last = chunk.len() - 1; + ( + list.elements().clone(), + offsets_out[0], + offsets_out[last] + sizes_out[last], + true, + ) + } else { + let list = chunk.as_::(); + let chunk_offsets = list.offsets().clone().execute::(ctx)?; + let chunk_sizes = list.sizes().clone().execute::(ctx)?; + match_each_integer_ptype!(chunk_offsets.ptype(), |O| { + for (out, &offset) in offsets_out.iter_mut().zip(chunk_offsets.as_slice::()) { + *out = u64::try_from(offset) + .vortex_expect("validated list offset is nonnegative"); + } + }); + match_each_integer_ptype!(chunk_sizes.ptype(), |S| { + for (out, &size) in sizes_out.iter_mut().zip(chunk_sizes.as_slice::()) { + *out = u64::try_from(size).vortex_expect("validated list size is nonnegative"); + } + }); + // Exact views bound their window with the first and last row. Other layouts need + // both extrema so that overlaps and interior gaps remain unchanged. + let (start, end) = if list.is_zero_copy_to_list() { + let last = chunk.len() - 1; + (offsets_out[0], offsets_out[last] + sizes_out[last]) + } else { + offsets_out.iter().zip(sizes_out.iter()).fold( + (u64::MAX, 0), + |(start, end), (&offset, &size)| (start.min(offset), end.max(offset + size)), + ) + }; + ( + list.elements().clone(), + start, + end, + list.is_zero_copy_to_list(), + ) + }; + let start = usize::try_from(start).vortex_expect("offset is bounded by elements.len()"); + let end = usize::try_from(end).vortex_expect("view end is bounded by elements.len()"); + let next_base = element_base + .checked_add(end - start) + .ok_or_else(|| vortex_err!("combined list elements length overflow"))?; + for offset in offsets_out { + *offset = (*offset - start as u64) + element_base as u64; + } + elements.push(child.slice(start..end)?); + zero_copy_to_list &= chunk_zero_copy_to_list; + element_base = next_base; + row += chunk.len(); + } + let validity = array.validity()?; + // SAFETY: element dtypes come from the common list dtype. Offsets/sizes are non-nullable u64 + // arrays with one entry per row. Trimming and rebasing preserve each validated view's bounds, + // and the checked combined length prevents overflow. Validity comes from the parent. If every + // chunk is zero-copyable, trimming makes adjacent chunks meet without gaps or overlaps. + Ok(unsafe { + let elements = + ChunkedArray::new_unchecked(elements, element_dtype.as_ref().clone()).into_array(); + let offsets = + PrimitiveArray::new_unchecked(offsets.freeze(), Validity::NonNullable).into_array(); + let sizes = PrimitiveArray::new_unchecked(sizes.freeze(), Validity::NonNullable).into_array(); + ListViewArray::new_unchecked(elements, offsets, sizes, validity) + .with_zero_copy_to_list(zero_copy_to_list) + .into_array() + }) +} pub(super) fn _canonicalize( array: ArrayView<'_, Chunked>, diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index 5fa22096220..41f10191b59 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -29,12 +29,19 @@ use crate::array::ArrayParts; use crate::array::ArrayView; use crate::array::VTable; use crate::array::with_empty_buffers; +use crate::arrays::FixedSizeList; use crate::arrays::PrimitiveArray; +use crate::arrays::Struct; use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::chunked::ChunkedData; use crate::arrays::chunked::array::ChunkedSlots; use crate::arrays::chunked::compute::rules::PARENT_RULES; +use crate::arrays::chunked::vtable::canonical::ListChunks; use crate::arrays::chunked::vtable::canonical::_canonicalize; +use crate::arrays::chunked::vtable::canonical::swizzle; +use crate::arrays::chunked::vtable::canonical::swizzle_fixed_size_list; +use crate::arrays::chunked::vtable::canonical::swizzle_list; +use crate::arrays::chunked::vtable::canonical::swizzle_struct; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::dtype::DType; @@ -263,6 +270,9 @@ impl VTable for Chunked { and every sparse child along identical chunk boundaries" ) } + DType::FixedSizeList(..) => swizzle::(array, swizzle_fixed_size_list), + DType::Struct(..) => swizzle::(array, swizzle_struct), + DType::List(..) => swizzle::(array, |array| swizzle_list(array, ctx)), // Variant need child swizzling that the builder path cannot express. DType::Variant(..) => { // TODO(joe)[#7674]: iterative execution here too @@ -270,10 +280,10 @@ impl VTable for Chunked { } // For all other types, use the builder path via AppendChild. _ => { - let slot_idx = array.next_builder_slot.max(ChunkedSlots::CHUNKS_OFFSET); + let slot_idx = array.next_child_slot.max(ChunkedSlots::CHUNKS_OFFSET); if slot_idx < array.slots().len() { Ok(ExecutionResult::append_child( - array.with_next_builder_slot(slot_idx + 1), + array.with_next_child_slot(slot_idx + 1), slot_idx, )) } else if slot_idx == ChunkedSlots::CHUNKS_OFFSET { @@ -284,7 +294,7 @@ impl VTable for Chunked { )) } else { // Every chunk has gone to the builder, which is what the executor will finish. - Ok(ExecutionResult::done_into_builder()) + Ok(ExecutionResult::done_into_builder(array)) } } } diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 2ee2f0b388b..394e04e8232 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -811,17 +811,15 @@ impl ExecutionResult { /// Signal that execution is complete and the result is in the executor's active builder. /// - /// The executor finishes that builder and discards the array carried here, so an encoding - /// that has already handed every value over via [`ExecutionStep::AppendChild`] should use - /// this rather than building a real empty array of its own dtype only for it to be dropped - - /// for a nested dtype that is a recursive construction of empty children. + /// Pass the consumed parent array: the executor discards it after finishing the builder, + /// avoiding an allocation for a placeholder result. /// - /// Only valid once at least one [`ExecutionStep::AppendChild`] has been returned, which is - /// what guarantees the builder exists. [`finalize_done`] debug-asserts the resulting dtype, - /// so a mistake here shows up as a dtype mismatch rather than silently wrong data. - pub fn done_into_builder() -> Self { + /// Only valid once at least one [`ExecutionStep::AppendChild`] has been returned, which + /// guarantees that the executor has an active builder. The parent may have empty child slots + /// because its children have already been appended; it must not escape the executor. + pub fn done_into_builder(array: impl IntoArray) -> Self { Self { - array: Canonical::empty(&DType::Null).into_array(), + array: array.into_array(), step: ExecutionStep::Done, } } diff --git a/vortex-array/src/test_harness/trace/tests.rs b/vortex-array/src/test_harness/trace/tests.rs index 76bf952da9e..4e2b098cd7e 100644 --- a/vortex-array/src/test_harness/trace/tests.rs +++ b/vortex-array/src/test_harness/trace/tests.rs @@ -588,7 +588,7 @@ fn trace_execution_chunked_append_child_flow() -> VortexResult<()> { AppendChild slot=3 parent=vortex.chunked(i32, len=5) child=vortex.primitive(i32, len=2) builder append child=vortex.primitive(i32, len=2) iter 3 current=vortex.chunked(i32, len=5) builder_active=true - Done array=vortex.null(null, len=0) + Done array=vortex.chunked(i32, len=5) builder finish output=vortex.primitive(i32, len=5) iter 4 current=vortex.primitive(i32, len=5) builder_active=false return output=vortex.primitive(i32, len=5) From e0929fb3851b4fad402957e6fd988f47e5f723e8 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 2 Oct 2026 10:35:43 +0100 Subject: [PATCH 22/22] less Signed-off-by: Robert Kruszewski --- .../benches/chunked_nested_canonicalize.rs | 7 +- vortex-array/src/arrays/chunked/array.rs | 2 +- .../src/arrays/chunked/vtable/canonical.rs | 157 +++++++++++------- vortex-array/src/arrays/chunked/vtable/mod.rs | 2 +- 4 files changed, 100 insertions(+), 68 deletions(-) diff --git a/vortex-array/benches/chunked_nested_canonicalize.rs b/vortex-array/benches/chunked_nested_canonicalize.rs index f5d32912a4c..f8d02c07a7a 100644 --- a/vortex-array/benches/chunked_nested_canonicalize.rs +++ b/vortex-array/benches/chunked_nested_canonicalize.rs @@ -44,10 +44,9 @@ fn canonicalize(bencher: Bencher, chunk: ArrayRef, nchunks: usize) { #[divan::bench(args = CHUNKS, consts = [1, 8])] fn structs(bencher: Bencher, nchunks: usize) { let field = PrimitiveArray::from_iter(0..ROWS as u64).into_array(); - let chunk = StructArray::try_from_iter( - (0..FIELDS).map(|idx| (format!("field_{idx}"), field.clone())), - ) - .unwrap(); + let chunk = + StructArray::try_from_iter((0..FIELDS).map(|idx| (format!("field_{idx}"), field.clone()))) + .unwrap(); canonicalize(bencher, chunk.into_array(), nchunks); } diff --git a/vortex-array/src/arrays/chunked/array.rs b/vortex-array/src/arrays/chunked/array.rs index 33c700819b0..6e226d39a8d 100644 --- a/vortex-array/src/arrays/chunked/array.rs +++ b/vortex-array/src/arrays/chunked/array.rs @@ -10,7 +10,7 @@ use std::fmt::Display; use std::fmt::Formatter; use futures::stream; -use vortex_buffer::{Buffer, BufferMut}; +use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index 5a6e93d5c0a..17c7b9ef69a 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools as _; +use num_traits::AsPrimitive; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -83,9 +84,8 @@ pub(super) fn swizzle_fixed_size_list(array: ChunkedArray) -> VortexResult Vorte } let offsets_out = &mut offsets.as_mut_slice()[row..row + chunk.len()]; let sizes_out = &mut sizes.as_mut_slice()[row..row + chunk.len()]; - let (child, start, end, chunk_zero_copy_to_list) = if let Some(list) = chunk.as_opt::() { - let chunk_offsets = list.offsets().clone().execute::(ctx)?; - match_each_integer_ptype!(chunk_offsets.ptype(), |O| { - for ((out, size), pair) in offsets_out - .iter_mut() - .zip(sizes_out.iter_mut()) - .zip(chunk_offsets.as_slice::().windows(2)) - { - let start = u64::try_from(pair[0]) - .vortex_expect("validated list offset is nonnegative"); - let end = u64::try_from(pair[1]) - .vortex_expect("validated list offset is nonnegative"); - *out = start; - *size = end - start; - } - }); - let last = chunk.len() - 1; - ( - list.elements().clone(), - offsets_out[0], - offsets_out[last] + sizes_out[last], - true, - ) + let window = if let Some(list) = chunk.as_opt::() { + list_window(list, offsets_out, sizes_out, ctx)? } else { - let list = chunk.as_::(); - let chunk_offsets = list.offsets().clone().execute::(ctx)?; - let chunk_sizes = list.sizes().clone().execute::(ctx)?; - match_each_integer_ptype!(chunk_offsets.ptype(), |O| { - for (out, &offset) in offsets_out.iter_mut().zip(chunk_offsets.as_slice::()) { - *out = u64::try_from(offset) - .vortex_expect("validated list offset is nonnegative"); - } - }); - match_each_integer_ptype!(chunk_sizes.ptype(), |S| { - for (out, &size) in sizes_out.iter_mut().zip(chunk_sizes.as_slice::()) { - *out = u64::try_from(size).vortex_expect("validated list size is nonnegative"); - } - }); - // Exact views bound their window with the first and last row. Other layouts need - // both extrema so that overlaps and interior gaps remain unchanged. - let (start, end) = if list.is_zero_copy_to_list() { - let last = chunk.len() - 1; - (offsets_out[0], offsets_out[last] + sizes_out[last]) - } else { - offsets_out.iter().zip(sizes_out.iter()).fold( - (u64::MAX, 0), - |(start, end), (&offset, &size)| (start.min(offset), end.max(offset + size)), - ) - }; - ( - list.elements().clone(), - start, - end, - list.is_zero_copy_to_list(), - ) + list_view_window(chunk.as_::(), offsets_out, sizes_out, ctx)? }; - let start = usize::try_from(start).vortex_expect("offset is bounded by elements.len()"); - let end = usize::try_from(end).vortex_expect("view end is bounded by elements.len()"); + + let start = + usize::try_from(window.start).vortex_expect("offset is bounded by elements.len()"); + let end = + usize::try_from(window.end).vortex_expect("view end is bounded by elements.len()"); let next_base = element_base .checked_add(end - start) .ok_or_else(|| vortex_err!("combined list elements length overflow"))?; for offset in offsets_out { *offset = (*offset - start as u64) + element_base as u64; } - elements.push(child.slice(start..end)?); - zero_copy_to_list &= chunk_zero_copy_to_list; + elements.push(window.elements.slice(start..end)?); + zero_copy_to_list &= window.zero_copy_to_list; element_base = next_base; row += chunk.len(); } @@ -206,13 +158,94 @@ pub(super) fn swizzle_list(array: ChunkedArray, ctx: &mut ExecutionCtx) -> Vorte ChunkedArray::new_unchecked(elements, element_dtype.as_ref().clone()).into_array(); let offsets = PrimitiveArray::new_unchecked(offsets.freeze(), Validity::NonNullable).into_array(); - let sizes = PrimitiveArray::new_unchecked(sizes.freeze(), Validity::NonNullable).into_array(); + let sizes = + PrimitiveArray::new_unchecked(sizes.freeze(), Validity::NonNullable).into_array(); ListViewArray::new_unchecked(elements, offsets, sizes, validity) .with_zero_copy_to_list(zero_copy_to_list) .into_array() }) } +/// The range of child elements that one list chunk references. +struct ElementWindow { + elements: ArrayRef, + start: u64, + end: u64, + zero_copy_to_list: bool, +} + +/// Writes the `u64` offsets and sizes of a [`List`] chunk and returns its element window. +fn list_window( + list: ArrayView<'_, List>, + offsets_out: &mut [u64], + sizes_out: &mut [u64], + ctx: &mut ExecutionCtx, +) -> VortexResult { + let chunk_offsets = list.offsets().clone().execute::(ctx)?; + match_each_integer_ptype!(chunk_offsets.ptype(), |O| { + for ((out, size), pair) in offsets_out + .iter_mut() + .zip(sizes_out.iter_mut()) + .zip(chunk_offsets.as_slice::().windows(2)) + { + let start: u64 = pair[0].as_(); + let end: u64 = pair[1].as_(); + *out = start; + *size = end - start; + } + }); + + let last = offsets_out.len() - 1; + Ok(ElementWindow { + elements: list.elements().clone(), + start: offsets_out[0], + end: offsets_out[last] + sizes_out[last], + zero_copy_to_list: true, + }) +} + +/// Writes the `u64` offsets and sizes of a [`ListView`] chunk and returns its element window. +fn list_view_window( + list: ArrayView<'_, ListView>, + offsets_out: &mut [u64], + sizes_out: &mut [u64], + ctx: &mut ExecutionCtx, +) -> VortexResult { + let chunk_offsets = list.offsets().clone().execute::(ctx)?; + let chunk_sizes = list.sizes().clone().execute::(ctx)?; + match_each_integer_ptype!(chunk_offsets.ptype(), |O| { + for (out, &offset) in offsets_out.iter_mut().zip(chunk_offsets.as_slice::()) { + *out = offset.as_(); + } + }); + match_each_integer_ptype!(chunk_sizes.ptype(), |S| { + for (out, &size) in sizes_out.iter_mut().zip(chunk_sizes.as_slice::()) { + *out = size.as_(); + } + }); + + // Exact views bound their window with the first and last row. Other layouts need + // both extrema so that overlaps and interior gaps remain unchanged. + let (start, end) = if list.is_zero_copy_to_list() { + let last = offsets_out.len() - 1; + (offsets_out[0], offsets_out[last] + sizes_out[last]) + } else { + offsets_out + .iter() + .zip(sizes_out.iter()) + .fold((u64::MAX, 0), |(start, end), (&offset, &size)| { + (start.min(offset), end.max(offset + size)) + }) + }; + + Ok(ElementWindow { + elements: list.elements().clone(), + start, + end, + zero_copy_to_list: list.is_zero_copy_to_list(), + }) +} + pub(super) fn _canonicalize( array: ArrayView<'_, Chunked>, ctx: &mut ExecutionCtx, diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index 41f10191b59..afbdfb02d8c 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -36,8 +36,8 @@ use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::chunked::ChunkedData; use crate::arrays::chunked::array::ChunkedSlots; use crate::arrays::chunked::compute::rules::PARENT_RULES; -use crate::arrays::chunked::vtable::canonical::ListChunks; use crate::arrays::chunked::vtable::canonical::_canonicalize; +use crate::arrays::chunked::vtable::canonical::ListChunks; use crate::arrays::chunked::vtable::canonical::swizzle; use crate::arrays::chunked::vtable::canonical::swizzle_fixed_size_list; use crate::arrays::chunked::vtable::canonical::swizzle_list;