From 5f173bfe3cf479a3556b0d45cbfa43429ca0d00e Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 25 Jul 2026 18:49:28 +0000 Subject: [PATCH] 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) --- .../src/arrays/chunked/vtable/canonical.rs | 285 ++++-------------- 1 file changed, 54 insertions(+), 231 deletions(-) diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index 8cc8c7aa4a2..3df7b9d5ab8 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -2,8 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools as _; -use vortex_buffer::Buffer; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; @@ -15,23 +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::memory::HostAllocatorExt; -use crate::validity::Validity; pub(super) fn _canonicalize( array: ArrayView<'_, Chunked>, @@ -48,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(ctx.allocator(), array.dtype(), array.len()); array.array().append_to_builder(builder.as_mut(), ctx)?; @@ -78,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. @@ -147,151 +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.allocate_typed::(len)?; - let mut sizes = allocator.allocate_typed::(len)?; - let offsets_out: &mut [u64] = offsets.as_mut_slice_typed::()?; - let sizes_slice_out: &mut [u64] = sizes.as_mut_slice_typed::()?; - 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( - Buffer::::from_byte_buffer(offsets.freeze()), - Validity::NonNullable, - ) - .into_array(); - let sizes = PrimitiveArray::new( - Buffer::::from_byte_buffer(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::sync::Arc; use std::sync::LazyLock; - use std::sync::atomic::AtomicUsize; - use std::sync::atomic::Ordering; + 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; @@ -301,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; @@ -318,32 +146,12 @@ mod tests { use crate::dtype::DType::Variant as VariantDType; use crate::dtype::Nullability::NonNullable; use crate::dtype::PType::I32; - use crate::memory::DefaultHostAllocator; - use crate::memory::HostAllocator; - use crate::memory::MemorySessionExt; - use crate::memory::WritableHostBuffer; 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, - } - - impl HostAllocator for CountingAllocator { - fn allocate( - &self, - len: usize, - alignment: vortex_buffer::Alignment, - ) -> VortexResult { - self.allocations.fetch_add(1, Ordering::Relaxed); - DefaultHostAllocator.allocate(len, alignment) - } - } - fn variant_scalar(value: i32) -> Scalar { Scalar::variant(Scalar::primitive(value, NonNullable)) } @@ -660,38 +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(Arc::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(()) } }