From 531d3c911042da1108be88e6c29bd56f5b31d9fa Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 7 Jul 2026 02:20:10 +0100 Subject: [PATCH 1/3] Improve execution of Take(Chunked) for nested types Signed-off-by: Robert Kruszewski --- vortex-array/Cargo.toml | 4 + vortex-array/benches/take_chunked.rs | 213 +++++++++++ .../src/arrays/chunked/compute/take.rs | 348 +++++++++++++++--- vortex-array/src/arrays/chunked/vtable/mod.rs | 47 ++- vortex-array/src/arrays/dict/execute.rs | 94 ++--- vortex-array/src/arrays/dict/vtable/mod.rs | 5 +- 6 files changed, 590 insertions(+), 121 deletions(-) create mode 100644 vortex-array/benches/take_chunked.rs diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 4d8ef4ea2bb..5d3825972ad 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -228,6 +228,10 @@ harness = false name = "take_primitive" harness = false +[[bench]] +name = "take_chunked" +harness = false + [[bench]] name = "take_struct" harness = false diff --git a/vortex-array/benches/take_chunked.rs b/vortex-array/benches/take_chunked.rs new file mode 100644 index 00000000000..e695c05dcb8 --- /dev/null +++ b/vortex-array/benches/take_chunked.rs @@ -0,0 +1,213 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks for `take` on [`ChunkedArray`], materialized through the builder path. +//! +//! `take` is lazy (a `Dict` over the source), so the cost lands where the result is +//! materialized. These benches mirror the common consumer path of appending the taken +//! array into an [`ArrayBuilder`], and compare against `execute::`. +//! +//! Parameterized over: +//! - Number of indices to take +//! - Shuffled vs sorted indices (sortedness dominates `take_chunked`'s index grouping cost) +//! - Flat primitive chunks vs fixed-size-list shapes, where row indices are expanded into +//! `list_size` element indices before reaching the chunked elements child + +#![expect(clippy::cast_possible_truncation)] +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use half::f16; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::IntoArray; +use vortex_array::RecursiveCanonical; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::builders::builder_with_capacity; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_session::VortexSession; + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(array_session); + +/// Chunk length for the flat primitive benches. +const CHUNK_LEN: usize = 65_536; + +/// Number of chunks for the flat primitive benches. +const NCHUNKS: usize = 64; + +/// Number of indices to take in the flat primitive benches. +const NUM_INDICES: &[usize] = &[10_000, 100_000, 1_000_000]; + +/// Elements per list in the fixed-size-list benches. +const LIST_SIZE: usize = 128; + +/// Rows per chunk in the fixed-size-list benches. +const FSL_CHUNK_ROWS: usize = 2_048; + +/// Number of chunks in the fixed-size-list benches. +const FSL_NCHUNKS: usize = 32; + +/// Number of row indices to take in the fixed-size-list benches. +const FSL_NUM_INDICES: &[usize] = &[4_096, 32_768]; + +fn make_chunked_f32(chunk_len: usize, nchunks: usize) -> ArrayRef { + (0..nchunks) + .map(|c| { + PrimitiveArray::from_iter((0..chunk_len).map(|i| (c * chunk_len + i) as f32)) + .into_array() + }) + .collect::() + .into_array() +} + +fn f16_elements(num_rows: usize) -> Buffer { + (0..num_rows * LIST_SIZE) + .map(|i| f16::from_f32((i % 1024) as f32)) + .collect() +} + +/// `Chunked(FixedSizeList(f16))`: the shape produced by reading a fixed-size-list column +/// as row-group chunks. Take runs against the chunked array at row granularity. +fn make_chunked_of_fsl() -> ArrayRef { + (0..FSL_NCHUNKS) + .map(|_| { + FixedSizeListArray::new( + f16_elements(FSL_CHUNK_ROWS).into_array(), + LIST_SIZE as u32, + Validity::NonNullable, + FSL_CHUNK_ROWS, + ) + .into_array() + }) + .collect::() + .into_array() +} + +/// `FixedSizeList(Chunked(f16))`: the shape produced by canonicalizing +/// `Chunked(FixedSizeList)`, which swizzles chunking down into the elements child. Take +/// expands row indices into `LIST_SIZE` element indices before hitting the chunked child. +fn make_fsl_of_chunked() -> ArrayRef { + let elements = (0..FSL_NCHUNKS) + .map(|_| f16_elements(FSL_CHUNK_ROWS).into_array()) + .collect::() + .into_array(); + FixedSizeListArray::new( + elements, + LIST_SIZE as u32, + Validity::NonNullable, + FSL_NCHUNKS * FSL_CHUNK_ROWS, + ) + .into_array() +} + +fn shuffled_indices(num_indices: usize, max_index: usize) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(42); + (0..num_indices) + .map(|_| rng.random_range(0..max_index) as u64) + .collect::>() + .into_array() +} + +fn sorted_indices(num_indices: usize, max_index: usize) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(42); + let mut indices: Vec = (0..num_indices) + .map(|_| rng.random_range(0..max_index) as u64) + .collect(); + indices.sort_unstable(); + Buffer::from(indices).into_array() +} + +fn take_to_builder(array: &ArrayRef, indices: &ArrayRef, session: &VortexSession) -> ArrayRef { + let mut ctx = session.create_execution_ctx(); + let taken = array.take(indices.clone()).unwrap(); + let mut builder = builder_with_capacity(taken.dtype(), taken.len()); + taken.append_to_builder(builder.as_mut(), &mut ctx).unwrap(); + builder.finish() +} + +#[divan::bench(args = NUM_INDICES)] +fn take_chunked_f32_to_builder_shuffled(bencher: Bencher, num_indices: usize) { + let array = make_chunked_f32(CHUNK_LEN, NCHUNKS); + let indices = shuffled_indices(num_indices, CHUNK_LEN * NCHUNKS); + + bencher + .with_inputs(|| (&array, &indices)) + .bench_refs(|(array, indices)| take_to_builder(array, indices, &SESSION)); +} + +#[divan::bench(args = NUM_INDICES)] +fn take_chunked_f32_to_builder_sorted(bencher: Bencher, num_indices: usize) { + let array = make_chunked_f32(CHUNK_LEN, NCHUNKS); + let indices = sorted_indices(num_indices, CHUNK_LEN * NCHUNKS); + + bencher + .with_inputs(|| (&array, &indices)) + .bench_refs(|(array, indices)| take_to_builder(array, indices, &SESSION)); +} + +#[divan::bench(args = NUM_INDICES)] +fn take_chunked_f32_canonical_shuffled(bencher: Bencher, num_indices: usize) { + let array = make_chunked_f32(CHUNK_LEN, NCHUNKS); + let indices = shuffled_indices(num_indices, CHUNK_LEN * NCHUNKS); + + bencher + .with_inputs(|| (&array, &indices, SESSION.create_execution_ctx())) + .bench_refs(|(array, indices, ctx)| { + array + .take((*indices).clone()) + .unwrap() + .execute::(ctx) + .unwrap() + }); +} + +#[divan::bench(args = FSL_NUM_INDICES)] +fn take_chunked_of_fsl_to_builder_shuffled(bencher: Bencher, num_indices: usize) { + let array = make_chunked_of_fsl(); + let indices = shuffled_indices(num_indices, FSL_NCHUNKS * FSL_CHUNK_ROWS); + + bencher + .with_inputs(|| (&array, &indices)) + .bench_refs(|(array, indices)| take_to_builder(array, indices, &SESSION)); +} + +#[divan::bench(args = FSL_NUM_INDICES)] +fn take_fsl_of_chunked_to_builder_shuffled(bencher: Bencher, num_indices: usize) { + let array = make_fsl_of_chunked(); + let indices = shuffled_indices(num_indices, FSL_NCHUNKS * FSL_CHUNK_ROWS); + + bencher + .with_inputs(|| (&array, &indices)) + .bench_refs(|(array, indices)| take_to_builder(array, indices, &SESSION)); +} + +#[divan::bench(args = FSL_NUM_INDICES)] +fn take_fsl_of_chunked_canonical_shuffled(bencher: Bencher, num_indices: usize) { + let array = make_fsl_of_chunked(); + let indices = shuffled_indices(num_indices, FSL_NCHUNKS * FSL_CHUNK_ROWS); + + bencher + .with_inputs(|| (&array, &indices, SESSION.create_execution_ctx())) + .bench_refs(|(array, indices, ctx)| { + array + .take((*indices).clone()) + .unwrap() + .execute::(ctx) + .unwrap() + }); +} diff --git a/vortex-array/src/arrays/chunked/compute/take.rs b/vortex-array/src/arrays/chunked/compute/take.rs index 1b867981236..1f58175518c 100644 --- a/vortex-array/src/arrays/chunked/compute/take.rs +++ b/vortex-array/src/arrays/chunked/compute/take.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_buffer::BitBufferMut; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -14,14 +15,19 @@ use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::dict::TakeExecute; +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::executor::ExecutionCtx; use crate::validity::Validity; -// TODO(joe): this is pretty unoptimized but better than before. We want canonical using a builder -// we also want to return a chunked array ideally. +/// A bucket is "dense" in its chunk when building a chunk-length bitmap is cheaper than +/// comparison-sorting the bucket. Bitmap cost is O(chunk_len / 64) words; sort cost is +/// O(k log k) element moves, so the crossover sits around one hit per 64 rows. +const DENSE_BUCKET_DIVISOR: usize = 64; + fn take_chunked( array: ArrayView<'_, Chunked>, indices: &ArrayRef, @@ -38,75 +44,185 @@ fn take_chunked( let indices_values = indices.as_slice::(); let n = indices_values.len(); - // 1. Sort (value, orig_pos) pairs so indices for the same chunk are contiguous. - // Skip null indices — their final_take slots stay 0 and are masked null by validity. - let mut pairs: Vec<(u64, usize)> = indices_values - .iter() - .enumerate() - .filter(|&(i, _)| indices_mask.value(i)) - .map(|(i, &v)| (v, i)) - .collect(); - pairs.sort_unstable(); - - // 2. Fused pass: walk sorted pairs against chunk boundaries. - // - Dedup inline → build per-chunk filter masks - // - Scatter final_take[orig_pos] = dedup_idx for every pair + // Fast path: strictly increasing non-nullable indices contain no duplicates and preserve + // order, so the per-chunk filters concatenate directly into the result with no reorder. + if indices.dtype().nullability() == Nullability::NonNullable + && indices_values.is_sorted_by(|a, b| a < b) + { + return take_chunked_sorted_unique(array, indices_values); + } + let chunk_offsets = array.chunk_offsets(); let nchunks = array.nchunks(); + + // 1. Resolve each valid index to its chunk and count per-chunk occupancy. Grouping by + // chunk only needs a stable O(n) bucket partition, not a comparison sort of all + // indices. Null indices are skipped — their final_take slots stay 0 and are masked + // null by validity. + let mut chunk_ids = vec![0u32; n]; + let mut counts = vec![0usize; nchunks]; + if nchunks > 0 { + // Uniform chunk sizing (every chunk but the last has the same length) resolves with + // a divide; otherwise fall back to a binary search over the chunk offsets. + let chunk_len0 = chunk_offsets[1] - chunk_offsets[0]; + let uniform = chunk_len0 > 0 + && chunk_offsets[..nchunks] + .windows(2) + .all(|w| w[1] - w[0] == chunk_len0); + for i in 0..n { + if !indices_mask.value(i) { + continue; + } + let v = usize::try_from(indices_values[i])?; + let chunk_id = if uniform { + (v / chunk_len0).min(nchunks - 1) + } else { + chunk_offsets.partition_point(|&off| off <= v) - 1 + }; + chunk_ids[i] = u32::try_from(chunk_id)?; + counts[chunk_id] += 1; + } + } + + // 2. Scatter (local index, original position) into per-chunk buckets, preserving the + // original relative order within each chunk. + let mut bucket_starts = vec![0usize; nchunks + 1]; + for chunk_id in 0..nchunks { + bucket_starts[chunk_id + 1] = bucket_starts[chunk_id] + counts[chunk_id]; + } + let total_valid = bucket_starts[nchunks]; + let mut bucket_local = vec![0usize; total_valid]; + let mut bucket_pos = vec![0usize; total_valid]; + let mut fill = bucket_starts.clone(); + for i in 0..n { + if !indices_mask.value(i) { + continue; + } + let chunk_id = chunk_ids[i] as usize; + let slot = fill[chunk_id]; + fill[chunk_id] += 1; + bucket_local[slot] = usize::try_from(indices_values[i])? - chunk_offsets[chunk_id]; + bucket_pos[slot] = i; + } + + // 3. Per chunk: build a dedup filter mask and scatter final_take[orig_pos] = dedup_idx. + // Dense buckets use an idempotent bitmap (no sort at all); sparse buckets sort only + // their own few entries. let mut chunks = Vec::with_capacity(nchunks); let mut final_take = BufferMut::::with_capacity(n); final_take.push_n(0u64, n); + let mut dedup_base = 0u64; + let mut rank_scratch: Vec = Vec::new(); - let mut cursor = 0usize; - let mut dedup_idx = 0u64; - - for chunk_idx in 0..nchunks { - let chunk_start = chunk_offsets[chunk_idx]; - let chunk_end = chunk_offsets[chunk_idx + 1]; - let chunk_len = chunk_end - chunk_start; - let chunk_end_u64 = u64::try_from(chunk_end)?; - - let range_end = cursor + pairs[cursor..].partition_point(|&(v, _)| v < chunk_end_u64); - let chunk_pairs = &pairs[cursor..range_end]; - - if !chunk_pairs.is_empty() { - let mut local_indices: Vec = Vec::new(); - for (i, &(val, orig_pos)) in chunk_pairs.iter().enumerate() { - if cursor + i > 0 && val != pairs[cursor + i - 1].0 { - dedup_idx += 1; - } - let local = usize::try_from(val)? - chunk_start; - if local_indices.last() != Some(&local) { - local_indices.push(local); + for chunk_id in 0..nchunks { + let bucket = bucket_starts[chunk_id]..bucket_starts[chunk_id + 1]; + if bucket.is_empty() { + continue; + } + let locals = &bucket_local[bucket.clone()]; + let positions = &bucket_pos[bucket]; + let chunk_len = chunk_offsets[chunk_id + 1] - chunk_offsets[chunk_id]; + + let filter_mask = if locals.len() >= chunk_len / DENSE_BUCKET_DIVISOR { + let mut bits = BitBufferMut::new_unset(chunk_len); + for &local in locals { + bits.set(local); + } + let bits = bits.freeze(); + if rank_scratch.len() < chunk_len { + rank_scratch.resize(chunk_len, 0); + } + let mut next_rank = dedup_base; + bits.for_each_set_index(|local| { + rank_scratch[local] = next_rank; + next_rank += 1; + }); + for (&local, &pos) in locals.iter().zip(positions) { + final_take[pos] = rank_scratch[local]; + } + dedup_base = next_rank; + Mask::from_buffer(bits) + } else { + let mut pairs: Vec<(usize, usize)> = locals + .iter() + .copied() + .zip(positions.iter().copied()) + .collect(); + pairs.sort_unstable(); + let mut unique: Vec = Vec::with_capacity(pairs.len()); + for &(local, pos) in &pairs { + if unique.last() != Some(&local) { + unique.push(local); } - final_take[orig_pos] = dedup_idx; + final_take[pos] = dedup_base + (unique.len() - 1) as u64; } + dedup_base += unique.len() as u64; + Mask::from_indices(chunk_len, unique) + }; - let filter_mask = Mask::from_indices(chunk_len, local_indices); - chunks.push(array.chunk(chunk_idx).filter(filter_mask)?); + chunks.push(array.chunk(chunk_id).filter(filter_mask)?); + } + + // 4. Flatten the filtered chunks through a builder. Unlike execute::, this + // produces truly flat leaves for nested dtypes (Struct/List/FSL), so the reorder take + // below stays a flat gather instead of cascading a fresh chunked take into every leaf. + let flat = if chunks.is_empty() { + // SAFETY: an empty chunk list trivially satisfies the dtype invariant. + unsafe { ChunkedArray::new_unchecked(chunks, array.dtype().clone()) } + .into_array() + .execute::(ctx)? + .into_array() + } else { + let total: usize = chunks.iter().map(|c| c.len()).sum(); + let mut builder = builder_with_capacity_in(ctx.allocator(), array.dtype(), total); + for chunk in &chunks { + chunk.append_to_builder(builder.as_mut(), ctx)?; } + builder.finish_into_canonical().into_array() + }; + + // 5. Single take to restore original order and expand duplicates. + // Carry the original index validity so null indices produce null outputs. + let take_validity = Validity::from_mask(indices_mask, indices.dtype().nullability()); + flat.take(PrimitiveArray::new(final_take.freeze(), take_validity).into_array()) +} + +/// Take with strictly increasing, non-nullable indices: filter each chunk with the sorted +/// in-range indices and return the filtered chunks directly — order is already correct and +/// there are no duplicates, so no reorder take is needed. +fn take_chunked_sorted_unique( + array: ArrayView<'_, Chunked>, + indices_values: &[u64], +) -> VortexResult { + let chunk_offsets = array.chunk_offsets(); + let nchunks = array.nchunks(); + let mut chunks = Vec::new(); + let mut cursor = 0usize; + for chunk_id in 0..nchunks { + let chunk_start = chunk_offsets[chunk_id]; + let chunk_end_u64 = u64::try_from(chunk_offsets[chunk_id + 1])?; + let range_end = cursor + indices_values[cursor..].partition_point(|&v| v < chunk_end_u64); + if range_end > cursor { + let chunk_len = chunk_offsets[chunk_id + 1] - chunk_start; + let locals = indices_values[cursor..range_end] + .iter() + .map(|&v| usize::try_from(v).map(|v| v - chunk_start)) + .collect::, _>>()?; + let filter_mask = Mask::from_indices(chunk_len, locals); + chunks.push(array.chunk(chunk_id).filter(filter_mask)?); + } cursor = range_end; } - // SAFETY: every chunk came from a filter on a chunk with the same base dtype, - // unioned with the index nullability. - let flat = unsafe { ChunkedArray::new_unchecked(chunks, array.dtype().clone()) } - .into_array() - // TODO(joe): can we relax this. - .execute::(ctx)? - .into_array(); + if chunks.is_empty() { + return Ok(Canonical::empty(array.dtype()).into_array()); + } - // 4. Single take to restore original order and expand duplicates. - // Carry the original index validity so null indices produce null outputs. - let take_validity = Validity::from_mask( - indices - .as_ref() - .validity()? - .execute_mask(indices.as_ref().len(), ctx)?, - indices.dtype().nullability(), - ); - flat.take(PrimitiveArray::new(final_take.freeze(), take_validity).into_array()) + // SAFETY: every chunk came from a filter on a chunk with the same dtype, and the index + // nullability is NonNullable so the result dtype is unchanged. The result stays lazy; + // the executor drives it further as needed. + Ok(unsafe { ChunkedArray::new_unchecked(chunks, array.dtype().clone()) }.into_array()) } impl TakeExecute for Chunked { @@ -340,6 +456,126 @@ mod test { Ok(()) } + #[test] + fn test_take_shuffled_with_duplicates_and_nulls() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let c0 = buffer![10i32, 20, 30].into_array(); + let c1 = buffer![40i32, 50, 60].into_array(); + let c2 = buffer![70i32, 80, 90].into_array(); + let arr = ChunkedArray::try_new( + vec![c0, c1, c2], + PrimitiveArray::empty::(Nullability::NonNullable) + .dtype() + .clone(), + )?; + + // Duplicates within and across chunks, nulls interleaved, order shuffled. + let indices = PrimitiveArray::from_option_iter([ + Some(8u64), + Some(2), + None, + Some(8), + Some(0), + Some(5), + None, + Some(2), + Some(4), + ]); + let result = arr.take(indices.into_array())?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([ + Some(90i32), + Some(30), + None, + Some(90), + Some(10), + Some(60), + None, + Some(30), + Some(50) + ]), + &mut ctx + ); + Ok(()) + } + + #[test] + fn test_take_sorted_unique_fast_path() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let c0 = buffer![0i32, 1, 2, 3].into_array(); + let c1 = buffer![4i32, 5, 6, 7].into_array(); + let c2 = buffer![8i32, 9, 10, 11].into_array(); + let arr = ChunkedArray::try_new( + vec![c0, c1, c2], + PrimitiveArray::empty::(Nullability::NonNullable) + .dtype() + .clone(), + )?; + + // Strictly increasing indices spanning some (not all) chunks. + let indices = buffer![1u64, 3, 4, 10, 11].into_array(); + let result = arr.take(indices)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([1i32, 3, 4, 10, 11]), + &mut ctx + ); + Ok(()) + } + + #[test] + fn test_take_sparse_over_large_chunks() -> VortexResult<()> { + // Few indices over large chunks exercises the sparse per-bucket sort path. + let mut ctx = array_session().create_execution_ctx(); + let chunk_len = 10_000i32; + let chunks: Vec<_> = (0..3) + .map(|c| { + let start = c * chunk_len; + PrimitiveArray::from_iter(start..start + chunk_len).into_array() + }) + .collect(); + let dtype = chunks[0].dtype().clone(); + let arr = ChunkedArray::try_new(chunks, dtype)?; + + let indices = buffer![29_999u64, 5, 29_999, 10_001, 5].into_array(); + let result = arr.take(indices)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([29_999i32, 5, 29_999, 10_001, 5]), + &mut ctx + ); + Ok(()) + } + + #[test] + fn test_take_non_uniform_chunks() -> VortexResult<()> { + // Uneven chunk lengths force the binary-search chunk resolution path. + let mut ctx = array_session().create_execution_ctx(); + let c0 = buffer![0i32, 1].into_array(); + let c1 = buffer![2i32, 3, 4, 5, 6].into_array(); + let c2 = buffer![7i32].into_array(); + let arr = ChunkedArray::try_new( + vec![c0, c1, c2], + PrimitiveArray::empty::(Nullability::NonNullable) + .dtype() + .clone(), + )?; + + let indices = buffer![7u64, 0, 3, 7, 1, 6].into_array(); + let result = arr.take(indices)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([7i32, 0, 3, 7, 1, 6]), + &mut ctx + ); + Ok(()) + } + #[test] fn test_take_chunked_conformance() { let a = buffer![1i32, 2, 3].into_array(); diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index 2eabd883742..92c3cc7947e 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -29,7 +29,10 @@ use crate::array::ArrayParts; use crate::array::ArrayView; use crate::array::VTable; use crate::array::with_empty_buffers; -use crate::arrays::PrimitiveArray; +use crate::arrays::{FixedSizeList, PrimitiveArray}; +use crate::arrays::ListView; +use crate::arrays::Struct; +use crate::arrays::Variant; use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::chunked::ChunkedData; use crate::arrays::chunked::array::CHUNK_OFFSETS_SLOT; @@ -41,7 +44,9 @@ use crate::builders::ArrayBuilder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::matcher::Matcher; use crate::serde::ArrayChildren; + mod canonical; mod operations; mod validity; @@ -251,12 +256,14 @@ impl VTable for Chunked { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { match array.dtype() { - // 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)?)) - } + // Struct, List, FixedSizeList, and Variant canonicalize by swizzling chunking + // down into their children zero-copy, which the value-copying builder path cannot + // express. Drive every chunk to its canonical encoding through the scheduler + // first, then swizzle the canonical chunks as pure metadata. + DType::Struct(..) => execute_swizzle::(array, ctx), + DType::List(..) => execute_swizzle::(array, ctx), + DType::FixedSizeList(..) => execute_swizzle::(array, ctx), + DType::Variant(..) => execute_swizzle::(array, ctx), // For all other types, use the builder path via AppendChild. _ => { let slot_idx = array.next_builder_slot.max(CHUNKS_OFFSET); @@ -290,3 +297,29 @@ impl VTable for Chunked { PARENT_RULES.evaluate(array, parent, child_idx) } } + +/// Iteratively request execution of each chunk to its canonical encoding `M` on the +/// scheduler's explicit stack, then swizzle the canonical chunks without copying values. +/// +/// The `next_builder_slot` cursor tracks progress across re-entries; chunks that already +/// match `M` are skipped without a scheduler round-trip. +fn execute_swizzle( + array: Array, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mut slot_idx = array.next_builder_slot.max(CHUNKS_OFFSET); + while slot_idx < array.slots().len() + && array.slots()[slot_idx] + .as_ref() + .is_some_and(|chunk| M::matches(chunk)) + { + slot_idx += 1; + } + if slot_idx < array.slots().len() { + return Ok(ExecutionResult::execute_slot::( + array.with_next_builder_slot(slot_idx + 1), + slot_idx, + )); + } + Ok(ExecutionResult::done(_canonicalize(array.as_view(), ctx)?)) +} diff --git a/vortex-array/src/arrays/dict/execute.rs b/vortex-array/src/arrays/dict/execute.rs index 47a9a051666..ae74f4cc987 100644 --- a/vortex-array/src/arrays/dict/execute.rs +++ b/vortex-array/src/arrays/dict/execute.rs @@ -6,10 +6,10 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; +use crate::ArrayView; use crate::Canonical; use crate::CanonicalView; use crate::ExecutionCtx; -use crate::IntoArray; use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::arrays::Decimal; @@ -20,6 +20,7 @@ use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::ListView; use crate::arrays::ListViewArray; +use crate::arrays::Null; use crate::arrays::NullArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; @@ -38,24 +39,23 @@ use crate::arrays::variant::VariantArrayExt; /// by looking up each code in the values array. pub(crate) fn take_canonical( values: CanonicalView, - codes: &PrimitiveArray, + codes: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> VortexResult { - let values = Canonical::from(values); Ok(match values { - Canonical::Null(a) => Canonical::Null(take_null(&a, codes)), - Canonical::Bool(a) => Canonical::Bool(take_bool(&a, codes, ctx)?), - Canonical::Primitive(a) => Canonical::Primitive(take_primitive(&a, codes, ctx)), - Canonical::Decimal(a) => Canonical::Decimal(take_decimal(&a, codes, ctx)), - Canonical::VarBinView(a) => Canonical::VarBinView(take_varbinview(&a, codes, ctx)), - Canonical::List(a) => Canonical::List(take_listview(&a, codes, ctx)), - Canonical::FixedSizeList(a) => { - Canonical::FixedSizeList(take_fixed_size_list(&a, codes, ctx)) + CanonicalView::Null(a) => Canonical::Null(take_null(a, codes)), + CanonicalView::Bool(a) => Canonical::Bool(take_bool(a, codes, ctx)?), + CanonicalView::Primitive(a) => Canonical::Primitive(take_primitive(a, codes, ctx)), + CanonicalView::Decimal(a) => Canonical::Decimal(take_decimal(a, codes, ctx)), + CanonicalView::VarBinView(a) => Canonical::VarBinView(take_varbinview(a, codes, ctx)), + CanonicalView::List(a) => Canonical::List(take_listview(a, codes, ctx)), + CanonicalView::FixedSizeList(a) => { + Canonical::FixedSizeList(take_fixed_size_list(a, codes, ctx)) } - Canonical::Struct(a) => Canonical::Struct(take_struct(&a, codes)), - Canonical::Extension(a) => Canonical::Extension(take_extension(&a, codes, ctx)), - Canonical::Variant(a) => { - let indices = codes.clone().into_array(); + CanonicalView::Struct(a) => Canonical::Struct(take_struct(a, codes)), + CanonicalView::Extension(a) => Canonical::Extension(take_extension(a, codes, ctx)), + CanonicalView::Variant(a) => { + let indices = codes.array(); let taken_core_storage = a.core_storage().take(indices.clone())?; let taken_shredded = a .shredded() @@ -67,32 +67,28 @@ pub(crate) fn take_canonical( } /// Take for NullArray is trivial - just create a new NullArray with the new length. -fn take_null(_array: &NullArray, codes: &PrimitiveArray) -> NullArray { +fn take_null(_array: ArrayView<'_, Null>, codes: ArrayView<'_, Primitive>) -> NullArray { NullArray::new(codes.len()) } // TODO(joe): use dict_bool_take fn take_bool( - array: &BoolArray, - codes: &PrimitiveArray, + array: ArrayView<'_, Bool>, + codes: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> VortexResult { - let codes_ref = codes.clone().into_array(); - let array = array.as_view(); - Ok(::take(array, &codes_ref, ctx)? + Ok(::take(array, codes.array(), ctx)? .vortex_expect("take bool should not return None") .as_::() .into_owned()) } fn take_primitive( - array: &PrimitiveArray, - codes: &PrimitiveArray, + array: ArrayView<'_, Primitive>, + codes: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> PrimitiveArray { - let codes_ref = codes.clone().into_array(); - let array = array.as_view(); - ::take(array, &codes_ref, ctx) + ::take(array, codes.array(), ctx) .vortex_expect("take primitive array") .vortex_expect("take primitive should not return None") .as_::() @@ -100,13 +96,11 @@ fn take_primitive( } fn take_decimal( - array: &DecimalArray, - codes: &PrimitiveArray, + array: ArrayView<'_, Decimal>, + codes: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> DecimalArray { - let codes_ref = codes.clone().into_array(); - let array = array.as_view(); - ::take(array, &codes_ref, ctx) + ::take(array, codes.array(), ctx) .vortex_expect("take decimal array") .vortex_expect("take decimal should not return None") .as_::() @@ -114,13 +108,11 @@ fn take_decimal( } fn take_varbinview( - array: &VarBinViewArray, - codes: &PrimitiveArray, + array: ArrayView<'_, VarBinView>, + codes: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> VarBinViewArray { - let codes_ref = codes.clone().into_array(); - let array = array.as_view(); - ::take(array, &codes_ref, ctx) + ::take(array, codes.array(), ctx) .vortex_expect("take varbinview array") .vortex_expect("take varbinview should not return None") .as_::() @@ -128,13 +120,11 @@ fn take_varbinview( } fn take_listview( - array: &ListViewArray, - codes: &PrimitiveArray, + array: ArrayView<'_, ListView>, + codes: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> ListViewArray { - let codes_ref = codes.clone().into_array(); - let array = array.as_view(); - ::take(array, &codes_ref, ctx) + ::take(array, codes.array(), ctx) .vortex_expect("take listview execute") .vortex_expect("ListView TakeExecute should not return None") .as_::() @@ -142,23 +132,19 @@ fn take_listview( } fn take_fixed_size_list( - array: &FixedSizeListArray, - codes: &PrimitiveArray, + array: ArrayView<'_, FixedSizeList>, + codes: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> FixedSizeListArray { - let codes_ref = codes.clone().into_array(); - let array = array.as_view(); - ::take(array, &codes_ref, ctx) + ::take(array, codes.array(), ctx) .vortex_expect("take fixed size list array") .vortex_expect("take fixed size list should not return None") .as_::() .into_owned() } -fn take_struct(array: &StructArray, codes: &PrimitiveArray) -> StructArray { - let codes_ref = codes.clone().into_array(); - let array = array.as_view(); - ::take(array, &codes_ref) +fn take_struct(array: ArrayView<'_, Struct>, codes: ArrayView<'_, Primitive>) -> StructArray { + ::take(array, codes.array()) .vortex_expect("take struct array") .vortex_expect("take struct should not return None") .as_::() @@ -166,13 +152,11 @@ fn take_struct(array: &StructArray, codes: &PrimitiveArray) -> StructArray { } fn take_extension( - array: &ExtensionArray, - codes: &PrimitiveArray, + array: ArrayView<'_, Extension>, + codes: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, ) -> ExtensionArray { - let codes_ref = codes.clone().into_array(); - let array = array.as_view(); - ::take(array, &codes_ref, ctx) + ::take(array, codes.array(), ctx) .vortex_expect("take extension storage") .vortex_expect("take extension should not return None") .as_::() diff --git a/vortex-array/src/arrays/dict/vtable/mod.rs b/vortex-array/src/arrays/dict/vtable/mod.rs index d9ced9eeffa..3524cb4cc92 100644 --- a/vortex-array/src/arrays/dict/vtable/mod.rs +++ b/vortex-array/src/arrays/dict/vtable/mod.rs @@ -205,7 +205,7 @@ impl VTable for Dict { Ok(ExecutionResult::done(take_canonical( values.as_::(), - &codes.downcast::(), + codes.as_::(), ctx, )?)) } @@ -222,8 +222,7 @@ impl VTable for Dict { ) && !codes.validity()?.definitely_all_null() { - let codes = codes.into_owned(); - let canonical = take_canonical(values, &codes, ctx)?.into_array(); + let canonical = take_canonical(values, codes, ctx)?.into_array(); canonical.append_to_builder(builder, ctx)?; return Ok(()); } From 5ae3c09988d4bc533d200c0b6df5f0d576366b3a Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 7 Jul 2026 16:28:28 +0100 Subject: [PATCH 2/3] more Signed-off-by: Robert Kruszewski --- .../src/arrays/chunked/vtable/canonical.rs | 99 +++++++++---------- vortex-array/src/arrays/chunked/vtable/mod.rs | 1 + vortex-array/src/arrays/listview/rebuild.rs | 2 +- 3 files changed, 51 insertions(+), 51 deletions(-) diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index 12c3bbcd9e4..e7bd8d4c0fd 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -1,13 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use itertools::Itertools as _; -use vortex_buffer::Buffer; +use std::mem; +use std::mem::MaybeUninit; + +use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; +use crate::AnyCanonical; use crate::ArrayRef; use crate::Canonical; use crate::ExecutionCtx; @@ -15,22 +18,24 @@ 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::ListView; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; +use crate::arrays::Struct; use crate::arrays::StructArray; +use crate::arrays::Variant; use crate::arrays::VariantArray; use crate::arrays::chunked::ChunkedArrayExt; -use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListDataParts; use crate::arrays::listview::ListViewArrayExt; use crate::arrays::listview::ListViewRebuildMode; use crate::arrays::variant::VariantArrayExt; -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( @@ -45,35 +50,32 @@ pub(super) fn _canonicalize( return Ok(Canonical::empty(array.dtype())); } if array.nchunks() == 1 { - return array.chunk(0).clone().execute::(ctx); + return Ok(Canonical::from(array.chunk(0).as_::())); } let owned_chunks: Vec = array.iter_chunks().cloned().collect(); Ok(match array.dtype() { DType::Struct(..) => { - let struct_array = pack_struct_chunks(owned_chunks, ctx)?; + let struct_array = pack_struct_chunks(owned_chunks)?; Canonical::Struct(struct_array) } DType::List(elem_dtype, _) => Canonical::List(swizzle_list_chunks( - &owned_chunks, + owned_chunks, array.array().validity()?, elem_dtype, ctx, )?), DType::FixedSizeList(elem_dtype, list_size, _) => { Canonical::FixedSizeList(swizzle_fixed_size_list_chunks( - &owned_chunks, + owned_chunks, array.array().validity()?, elem_dtype, *list_size, - ctx, )?) } - DType::Variant(_) => Canonical::Variant(pack_variant_chunks(owned_chunks, ctx)?), + DType::Variant(_) => Canonical::Variant(pack_variant_chunks(owned_chunks)?), _ => { - let mut builder = builder_with_capacity_in(ctx.allocator(), array.dtype(), array.len()); - array.array().append_to_builder(builder.as_mut(), ctx)?; - builder.finish_into_canonical(ctx) + unreachable!("All non nested types should be handled by execute loop") } }) } @@ -82,24 +84,18 @@ pub(super) fn _canonicalize( /// 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))? +fn pack_struct_chunks(chunks: Vec) -> VortexResult { + StructArray::try_concat(chunks.into_iter().map(|c| c.downcast::())) } /// 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, - ctx: &mut ExecutionCtx, -) -> VortexResult { +fn pack_variant_chunks(chunks: Vec) -> VortexResult { let variant_chunks: Vec = chunks .into_iter() - .map(|chunk| chunk.execute::(ctx)) - .try_collect()?; + .map(|chunk| chunk.downcast::()) + .collect(); let outer_dtype = variant_chunks[0].dtype().clone(); let core_chunks = variant_chunks @@ -151,7 +147,7 @@ fn pack_variant_chunks( /// /// The caller guarantees there are at least 2 chunks. fn swizzle_list_chunks( - chunks: &[ArrayRef], + chunks: Vec, validity: Validity, elem_dtype: &DType, ctx: &mut ExecutionCtx, @@ -177,15 +173,14 @@ fn swizzle_list_chunks( // 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 offsets = BufferMut::::with_capacity(len); + let mut sizes = BufferMut::::with_capacity(len); + let offsets_out = &mut offsets.spare_capacity_mut()[..len]; + let sizes_slice_out = &mut sizes.spare_capacity_mut()[..len]; let mut next_list = 0usize; for chunk in chunks { - let chunk_array = chunk.clone().execute::(ctx)?; + let chunk_array = chunk.downcast::(); // 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)?; @@ -211,13 +206,25 @@ fn swizzle_list_chunks( let offsets_slice = offsets_arr.as_slice::(); let sizes_slice = sizes_arr.as_slice::(); + let chunk_len = chunk_array.len(); + + // SAFETY: MaybeUninit has the same memory layout as the underlying type + unsafe { + mem::transmute::<&mut [MaybeUninit], &mut [u64]>( + &mut sizes_slice_out[next_list..][..chunk_len], + ) + } + .copy_from_slice(sizes_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; + for (off_out, &offset) in offsets_out[next_list..][..chunk_len] + .iter_mut() + .zip(offsets_slice) + { + off_out.write(offset + num_elements); } + next_list += chunk_array.len(); num_elements += chunk_array.elements().len() as u64; } debug_assert_eq!(next_list, len); @@ -227,16 +234,8 @@ fn swizzle_list_chunks( 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(); + 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 @@ -260,21 +259,21 @@ fn swizzle_list_chunks( /// /// The caller guarantees there are at least 2 chunks. fn swizzle_fixed_size_list_chunks( - chunks: &[ArrayRef], + chunks: Vec, 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)?; + let FixedSizeListDataParts { elements, .. } = + chunk.downcast::().into_data_parts(); // 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()); + element_chunks.push(elements); } let chunked_elements = ChunkedArray::try_new(element_chunks, elem_dtype.clone())?.into_array(); diff --git a/vortex-array/src/arrays/chunked/vtable/mod.rs b/vortex-array/src/arrays/chunked/vtable/mod.rs index 92c3cc7947e..436a60ccef6 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -256,6 +256,7 @@ impl VTable for Chunked { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { match array.dtype() { + // TODO(joe)[#7674]: iterative execution here too // Struct, List, FixedSizeList, and Variant canonicalize by swizzling chunking // down into their children zero-copy, which the value-copying builder path cannot // express. Drive every chunk to its canonical encoding through the scheduler diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 342c9eccf19..730ecc261fc 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -162,7 +162,7 @@ impl ListViewArray { let total: u64 = sizes_canonical .as_slice::() .iter() - .map(|s| (*s).as_() as u64) + .map(|s| s.as_() as u64) .sum(); if Self::should_use_take(total, self.len()) { self.rebuild_with_take::(ctx) From 415d6d39f16a6c18da8f6f0079653ffeeeaad520 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 8 Jul 2026 14:03:58 +0100 Subject: [PATCH 3/3] fixes Signed-off-by: Robert Kruszewski --- vortex-array/benches/take_chunked.rs | 23 ++++++ .../src/arrays/chunked/compute/take.rs | 2 +- .../src/arrays/chunked/vtable/canonical.rs | 74 +++---------------- vortex-array/src/arrays/chunked/vtable/mod.rs | 3 +- vortex-array/src/executor.rs | 7 +- 5 files changed, 41 insertions(+), 68 deletions(-) diff --git a/vortex-array/benches/take_chunked.rs b/vortex-array/benches/take_chunked.rs index e695c05dcb8..7751d6dac17 100644 --- a/vortex-array/benches/take_chunked.rs +++ b/vortex-array/benches/take_chunked.rs @@ -23,6 +23,7 @@ use half::f16; use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; +use rand::seq::index::sample; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; @@ -132,6 +133,18 @@ fn sorted_indices(num_indices: usize, max_index: usize) -> ArrayRef { Buffer::from(indices).into_array() } +/// Strictly increasing indices with no duplicates (samples without replacement), the shape +/// that hits the sorted-unique fast path in `take_chunked`. +fn sorted_unique_indices(num_indices: usize, max_index: usize) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(42); + let mut indices: Vec = sample(&mut rng, max_index, num_indices) + .into_iter() + .map(|i| i as u64) + .collect(); + indices.sort_unstable(); + Buffer::from(indices).into_array() +} + fn take_to_builder(array: &ArrayRef, indices: &ArrayRef, session: &VortexSession) -> ArrayRef { let mut ctx = session.create_execution_ctx(); let taken = array.take(indices.clone()).unwrap(); @@ -160,6 +173,16 @@ fn take_chunked_f32_to_builder_sorted(bencher: Bencher, num_indices: usize) { .bench_refs(|(array, indices)| take_to_builder(array, indices, &SESSION)); } +#[divan::bench(args = NUM_INDICES)] +fn take_chunked_f32_to_builder_sorted_unique(bencher: Bencher, num_indices: usize) { + let array = make_chunked_f32(CHUNK_LEN, NCHUNKS); + let indices = sorted_unique_indices(num_indices, CHUNK_LEN * NCHUNKS); + + bencher + .with_inputs(|| (&array, &indices)) + .bench_refs(|(array, indices)| take_to_builder(array, indices, &SESSION)); +} + #[divan::bench(args = NUM_INDICES)] fn take_chunked_f32_canonical_shuffled(bencher: Bencher, num_indices: usize) { let array = make_chunked_f32(CHUNK_LEN, NCHUNKS); diff --git a/vortex-array/src/arrays/chunked/compute/take.rs b/vortex-array/src/arrays/chunked/compute/take.rs index 1f58175518c..08166415ae7 100644 --- a/vortex-array/src/arrays/chunked/compute/take.rs +++ b/vortex-array/src/arrays/chunked/compute/take.rs @@ -178,7 +178,7 @@ fn take_chunked( for chunk in &chunks { chunk.append_to_builder(builder.as_mut(), ctx)?; } - builder.finish_into_canonical().into_array() + builder.finish_into_canonical(ctx).into_array() }; // 5. Single take to restore original order and expand duplicates. diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index e7bd8d4c0fd..889f7573beb 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -208,13 +208,10 @@ fn swizzle_list_chunks( let chunk_len = chunk_array.len(); - // SAFETY: MaybeUninit has the same memory layout as the underlying type - unsafe { - mem::transmute::<&mut [MaybeUninit], &mut [u64]>( - &mut sizes_slice_out[next_list..][..chunk_len], - ) - } - .copy_from_slice(sizes_slice); + // SAFETY: `&[u64]` and `&[MaybeUninit]` have identical layout, and viewing + // initialized values as possibly-uninitialized is always sound (unlike the reverse). + let sizes_uninit = unsafe { mem::transmute::<&[u64], &[MaybeUninit]>(sizes_slice) }; + sizes_slice_out[next_list..][..chunk_len].copy_from_slice(sizes_uninit); // Append offsets and sizes, adjusting offsets to point into the combined array. for (off_out, &offset) in offsets_out[next_list..][..chunk_len] @@ -229,6 +226,12 @@ fn swizzle_list_chunks( } debug_assert_eq!(next_list, len); + // SAFETY: the loop above initialized exactly `len` offsets and sizes. + unsafe { + offsets.set_len(len); + sizes.set_len(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()) } @@ -285,8 +288,6 @@ fn swizzle_fixed_size_list_chunks( mod tests { use std::sync::Arc; use std::sync::LazyLock; - use std::sync::atomic::AtomicUsize; - use std::sync::atomic::Ordering; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -315,32 +316,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)) } @@ -655,39 +636,4 @@ 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(), - 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 436a60ccef6..c98c3c4b840 100644 --- a/vortex-array/src/arrays/chunked/vtable/mod.rs +++ b/vortex-array/src/arrays/chunked/vtable/mod.rs @@ -29,8 +29,9 @@ use crate::array::ArrayParts; use crate::array::ArrayView; use crate::array::VTable; use crate::array::with_empty_buffers; -use crate::arrays::{FixedSizeList, PrimitiveArray}; +use crate::arrays::FixedSizeList; use crate::arrays::ListView; +use crate::arrays::PrimitiveArray; use crate::arrays::Struct; use crate::arrays::Variant; use crate::arrays::chunked::ChunkedArrayExt; diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index a49f98b4d55..490e8be72de 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -545,8 +545,11 @@ fn finalize_done( if cfg!(debug_assertions) { vortex_ensure!( output.len() == expected_len, - "Result length mismatch for {:?}", - encoding_id + "Result length mismatch for {:?}: output {} expected {} (output encoding {:?})", + encoding_id, + output.len(), + expected_len, + output.encoding_id() ); vortex_ensure!( output.dtype() == &expected_dtype,