From e0e8c9641e517f69736c49f9daefd934fd5ff39e Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 29 May 2026 16:17:01 +0000 Subject: [PATCH] feat(chunked): swizzle FixedSizeList chunks during canonicalization MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Canonicalizing a `ChunkedArray` of `FixedSizeList` chunks previously fell through to the generic builder path, copying every element into a fresh contiguous buffer. Because each chunk's `elements` child is already exactly `list_size * chunk.len()` long and starts at the first list, the children can be reused directly as the chunks of a combined `ChunkedArray` of `elements`, mirroring the existing `swizzle_list_chunks` approach for `List`. Route `DType::FixedSizeList` through `_canonicalize` and add `swizzle_fixed_size_list_chunks`, which builds the combined elements and wraps them in a single `FixedSizeListArray` in O(nchunks) without copying element data. Adds a `chunked_fsl_canonicalize` divan benchmark and a regression test. The benchmark shows the swizzle is effectively O(1) in list size (e.g. 1024 elements/list across 32 chunks drops from ~175 ms to ~3 µs). Signed-off-by: Joe Isaacs --- vortex-array/Cargo.toml | 4 + .../benches/chunked_fsl_canonicalize.rs | 67 +++++++++++++++ .../src/arrays/chunked/vtable/canonical.rs | 86 +++++++++++++++++++ vortex-array/src/arrays/chunked/vtable/mod.rs | 5 +- 4 files changed, 160 insertions(+), 2 deletions(-) create mode 100644 vortex-array/benches/chunked_fsl_canonicalize.rs diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 666a23c02c4..e103097f8b6 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -129,6 +129,10 @@ harness = false name = "chunk_array_builder" harness = false +[[bench]] +name = "chunked_fsl_canonicalize" +harness = false + [[bench]] name = "scalar_at_struct" harness = false diff --git a/vortex-array/benches/chunked_fsl_canonicalize.rs b/vortex-array/benches/chunked_fsl_canonicalize.rs new file mode 100644 index 00000000000..d267415cc4e --- /dev/null +++ b/vortex-array/benches/chunked_fsl_canonicalize.rs @@ -0,0 +1,67 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks canonicalizing a [`ChunkedArray`] of [`FixedSizeListArray`] chunks. +//! +//! Parameterized over: +//! - Number of chunks +//! - Fixed size list length (elements per list) + +#![expect(clippy::cast_possible_truncation)] +#![expect(clippy::unwrap_used)] + +use divan::Bencher; +use vortex_array::Canonical; +use vortex_array::IntoArray; +use vortex_array::LEGACY_SESSION; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; + +fn main() { + divan::main(); +} + +/// Number of lists in each chunk. +const LISTS_PER_CHUNK: usize = 1_000; + +/// Number of chunks in the source array. +const NUM_CHUNKS: &[usize] = &[2, 8, 32]; + +/// Fixed size list lengths (elements per list). +const LIST_SIZES: &[usize] = &[16, 256, 1024]; + +/// Creates a `FixedSizeListArray` with the given list size and number of lists. +fn create_fsl(list_size: usize, num_lists: usize) -> FixedSizeListArray { + let total_elements = list_size * num_lists; + let elements: Buffer = (0..total_elements as i64).collect(); + FixedSizeListArray::new( + elements.into_array(), + list_size as u32, + Validity::NonNullable, + num_lists, + ) +} + +/// Builds a `ChunkedArray` of `FixedSizeListArray` chunks. +fn create_chunked_fsl(list_size: usize, num_chunks: usize) -> ChunkedArray { + let chunk = create_fsl(list_size, LISTS_PER_CHUNK); + let dtype = chunk.dtype().clone(); + let chunks = (0..num_chunks) + .map(|_| chunk.clone().into_array()) + .collect(); + ChunkedArray::try_new(chunks, dtype).unwrap() +} + +#[divan::bench(args = NUM_CHUNKS, consts = LIST_SIZES)] +fn canonicalize(bencher: Bencher, num_chunks: usize) { + let chunked = create_chunked_fsl(LIST_SIZE, num_chunks).into_array(); + + bencher + .with_inputs(|| (&chunked, LEGACY_SESSION.create_execution_ctx())) + .bench_refs(|(array, execution_ctx)| { + array.clone().execute::(execution_ctx).unwrap() + }); +} diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index 9c74862c63d..6da34e3cb07 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -15,11 +15,13 @@ 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::FixedSizeListArrayExt; use crate::arrays::listview::ListViewArrayExt; use crate::arrays::listview::ListViewRebuildMode; use crate::arrays::variant::VariantArrayExt; @@ -58,6 +60,15 @@ pub(super) fn _canonicalize( 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()); @@ -240,6 +251,37 @@ fn swizzle_list_chunks( }) } +/// 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; @@ -263,6 +305,7 @@ mod tests { use crate::accessor::ArrayAccessor; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; + use crate::arrays::FixedSizeListArray; use crate::arrays::ListArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; @@ -551,6 +594,49 @@ mod tests { ); } + #[test] + fn pack_fixed_size_lists() -> VortexResult<()> { + 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 LEGACY_SESSION.create_execution_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 LEGACY_SESSION.create_execution_ctx())?, + expected.execute_scalar(idx, &mut LEGACY_SESSION.create_execution_ctx())?, + ); + } + Ok(()) + } + #[test] fn list_canonicalize_uses_memory_session_allocator() { let allocations = Arc::new(AtomicUsize::new(0)); diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index c4a886268a6..68679a4c3b5 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -240,8 +240,9 @@ impl VTable for Chunked { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { match array.dtype() { - // Struct, List, and Variant need child swizzling that the builder path cannot express. - DType::Struct(..) | DType::List(..) | DType::Variant(..) => { + // Struct, List, FixedSizeList, and Variant need child swizzling that the builder path + // cannot express. + DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) | DType::Variant(..) => { // TODO(joe)[#7674]: iterative execution here too Ok(ExecutionResult::done(_canonicalize(array.as_view(), ctx)?)) }