From 711d73621fbc1ae60ce0eb34ffc7f1eb48ff78c1 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 21 Jul 2026 15:47:43 +0100 Subject: [PATCH 1/6] Improve take on Chunked(FSL) and FSL(Chunked) array alongside implementation of take using PiecewiseSequence Signed-off-by: Robert Kruszewski --- .../src/arrays/chunked/compute/take.rs | 673 +++++++++++++++++- 1 file changed, 658 insertions(+), 15 deletions(-) diff --git a/vortex-array/src/arrays/chunked/compute/take.rs b/vortex-array/src/arrays/chunked/compute/take.rs index 05fc1183563..43f0cab2fb9 100644 --- a/vortex-array/src/arrays/chunked/compute/take.rs +++ b/vortex-array/src/arrays/chunked/compute/take.rs @@ -1,27 +1,88 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use itertools::Itertools; +use num_traits::AsPrimitive; +use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_mask::Mask; use crate::ArrayRef; use crate::Canonical; +use crate::Columnar; use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; +use crate::arrays::ConstantArray; +use crate::arrays::FixedSizeList; +use crate::arrays::FixedSizeListArray; +use crate::arrays::PiecewiseSequence; +use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::dict::TakeExecute; +use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::piecewise_sequence::constant_unsigned_usize; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; +use crate::arrays::primitive::PrimitiveArrayExt; +use crate::builders::ArrayBuilder; +use crate::builders::builder_with_capacity; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::PType; use crate::executor::ExecutionCtx; +use crate::match_each_unsigned_integer_ptype; 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. +/// Flattens per-chunk take/filter results into a single array. +/// +/// Flat dtypes append directly into a canonical builder. Nested dtypes instead collect the chunks +/// and canonicalize them as a chunked array, which reuses the chunks' children zero-copy (e.g. +/// chunked FixedSizeLists canonicalize into one FixedSizeList over the chained elements) and lets +/// the follow-up take push down lazily instead of deep-copying every child. +enum ChunkFlattener { + Builder(Box), + Chunks(Vec), +} + +impl ChunkFlattener { + fn new(dtype: &DType, capacity: usize) -> Self { + if dtype.is_nested() { + Self::Chunks(Vec::new()) + } else { + Self::Builder(builder_with_capacity(dtype, capacity)) + } + } + + fn push(&mut self, chunk: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> { + match self { + Self::Builder(builder) => chunk.append_to_builder(builder.as_mut(), ctx), + Self::Chunks(chunks) => { + chunks.push(chunk); + Ok(()) + } + } + } + + fn finish(self, dtype: &DType, ctx: &mut ExecutionCtx) -> VortexResult { + match self { + Self::Builder(mut builder) => Ok(builder.finish()), + // SAFETY: every chunk is a filter or take of a chunk of a chunked array with `dtype`, + // which leaves the dtype unchanged. + Self::Chunks(chunks) => unsafe { ChunkedArray::new_unchecked(chunks, dtype.clone()) } + .into_array() + .execute::(ctx) + .map(IntoArray::into_array), + } + } +} + +// TODO(joe): we want to return a chunked array ideally. fn take_chunked( array: ArrayView<'_, Chunked>, indices: &ArrayRef, @@ -53,7 +114,7 @@ fn take_chunked( // - Scatter final_take[orig_pos] = dedup_idx for every pair let chunk_offsets = array.chunk_offset_values(); let nchunks = array.nchunks(); - let mut chunks = Vec::with_capacity(nchunks); + let mut flattener = ChunkFlattener::new(array.dtype(), pairs.len()); let mut final_take = BufferMut::::with_capacity(n); final_take.push_n(0u64, n); @@ -83,19 +144,13 @@ fn take_chunked( } let filter_mask = Mask::from_indices(chunk_len, local_indices); - chunks.push(array.chunk(chunk_idx).filter(filter_mask)?); + flattener.push(array.chunk(chunk_idx).filter(filter_mask)?, ctx)?; } 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(); + let flat = flattener.finish(array.dtype(), ctx)?; // 4. Single take to restore original order and expand duplicates. // Carry the original index validity so null indices produce null outputs. @@ -115,30 +170,621 @@ impl TakeExecute for Chunked { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + // A single chunk is logically identical to the chunked array, so delegate directly. + if array.nchunks() == 1 { + return array.chunk(0).take(indices.clone()).map(Some); + } + + if let Some(taken) = take_chunked_fsl(array, indices)? { + return Ok(Some(taken)); + } + + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_piecewise_chunked(array, piecewise_indices, ctx)? + { + return Ok(Some(taken)); + } + take_chunked(array, indices, ctx).map(Some) } } +/// Rewrites take over a chunked array of [`FixedSizeList`] chunks as take over a single +/// [`FixedSizeListArray`] whose elements child chains the chunks' elements. +/// +/// Each FSL chunk stores exactly `chunk.len() * list_size` elements starting at its first list, +/// so the chunks' elements children concatenate zero-copy into the elements of the combined +/// array, exactly as chunked canonicalization does. The FSL take implementation then gathers the +/// elements with a single `PiecewiseSequenceArray`, which [`take_piecewise_chunked`] resolves per +/// chunk without expanding one index per element. Chunks that are not FSL-encoded fall back to +/// the generic path since accessing their elements would require canonicalizing whole chunks. +fn take_chunked_fsl( + array: ArrayView<'_, Chunked>, + indices: &ArrayRef, +) -> VortexResult> { + let DType::FixedSizeList(element_dtype, list_size, _) = array.dtype() else { + return Ok(None); + }; + + let mut element_chunks = Vec::with_capacity(array.nchunks()); + for chunk in array.iter_chunks() { + let Some(fsl) = chunk.as_opt::() else { + return Ok(None); + }; + element_chunks.push(fsl.elements().clone()); + } + + let validity = array.array().validity()?; + // SAFETY: every chunk is a FixedSizeList with element dtype `element_dtype`. + let elements = + unsafe { ChunkedArray::new_unchecked(element_chunks, element_dtype.as_ref().clone()) } + .into_array(); + // SAFETY: each FSL chunk holds exactly `chunk.len() * list_size` elements, so the chained + // elements hold `array.len() * list_size` entries, and the chunked validity covers + // `array.len()` lists with the array's nullability. + let fsl = unsafe { + FixedSizeListArray::new_unchecked(elements, *list_size, validity, array.as_ref().len()) + }; + fsl.into_array().take(indices.clone()).map(Some) +} + +/// A per-chunk gather plan: chunk-local sub-piece runs to take from one chunk, in output order. +#[derive(Default)] +struct ChunkGather { + starts: Vec, + lengths: Vec, + total: usize, +} + +/// Take for [`PiecewiseSequence`] indices with unit multipliers. +/// +/// Each piece is a contiguous index run, so instead of expanding one index per element and +/// sorting them (as [`take_chunked`] does), pieces are split at chunk boundaries and gathered +/// from each chunk with a chunk-local `PiecewiseSequenceArray`. When the sub-pieces visit chunks +/// in non-decreasing order the gathered chunks concatenate directly into the result; otherwise +/// the gathered chunks are canonicalized once and a second piecewise take restores output order. +fn take_piecewise_chunked( + array: ArrayView<'_, Chunked>, + indices: ArrayView<'_, PiecewiseSequence>, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { + return Ok(None); + }; + + let output_len = indices.as_ref().len(); + let starts: Vec = match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + starts + .as_slice::() + .iter() + .map(|&start| start.as_()) + .collect() + }); + let lengths: Vec = match lengths { + Columnar::Constant(lengths) => vec![constant_unsigned_usize(&lengths); starts.len()], + Columnar::Canonical(lengths) => { + let lengths = lengths.into_primitive(); + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + lengths + .as_slice::() + .iter() + .map(|&length| length.as_()) + .collect() + }) + } + }; + + let array_len = array.as_ref().len(); + let chunk_offsets = array.chunk_offsets(); + let nchunks = array.nchunks(); + + let mut plans: Vec = Vec::new(); + plans.resize_with(nchunks, ChunkGather::default); + // Sub-pieces in output order: (chunk index, offset within that chunk's gathered output, run + // length). + let mut sub_pieces: Vec<(usize, usize, usize)> = Vec::with_capacity(starts.len()); + let mut monotonic = true; + let mut prev_chunk = 0usize; + let mut total_len = 0usize; + + for (&start, &length) in starts.iter().zip_eq(&lengths) { + if length == 0 { + continue; + } + let end = start + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + if end > array_len { + vortex_bail!(OutOfBounds: end - 1, 0, array_len); + } + total_len = total_len + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + + // Locate the chunk containing `start`; `<=` skips empty chunks sharing the same offset. + let mut chunk_idx = chunk_offsets.partition_point(|&offset| offset <= start) - 1; + let mut cursor = start; + let mut remaining = length; + while remaining > 0 { + while chunk_offsets[chunk_idx + 1] <= cursor { + chunk_idx += 1; + } + let run = remaining.min(chunk_offsets[chunk_idx + 1] - cursor); + monotonic &= chunk_idx >= prev_chunk; + prev_chunk = chunk_idx; + + let plan = &mut plans[chunk_idx]; + sub_pieces.push((chunk_idx, plan.total, run)); + plan.starts.push((cursor - chunk_offsets[chunk_idx]) as u64); + plan.lengths.push(run as u64); + plan.total += run; + + cursor += run; + remaining -= run; + } + } + + vortex_ensure!( + total_len == output_len, + "PiecewiseSequenceArray expanded length {total_len} does not match declared length {output_len}" + ); + + // Gather the sub-pieces of each touched chunk with one chunk-local piecewise take. When the + // sub-pieces visit chunks in order, the gathered chunks already concatenate into the result. + if monotonic { + let mut gathered = Vec::new(); + for (chunk_idx, plan) in plans.into_iter().enumerate() { + if plan.starts.is_empty() { + continue; + } + let chunk_indices = contiguous_runs(plan.starts, plan.lengths, plan.total); + gathered.push(array.chunk(chunk_idx).take(chunk_indices)?); + } + let result = if gathered.len() == 1 { + gathered.swap_remove(0) + } else { + // SAFETY: every gathered chunk is a take of a chunk with dtype `array.dtype()`, and + // the non-nullable piecewise indices leave the dtype unchanged. + unsafe { ChunkedArray::new_unchecked(gathered, array.dtype().clone()) }.into_array() + }; + return Ok(Some(result)); + } + + // Out-of-order sub-pieces: flatten each chunk's gather once through a [`ChunkFlattener`], + // then take the sub-piece runs from the flattened result in output order. + let mut bases = vec![0usize; nchunks]; + let mut running = 0usize; + let mut flattener = ChunkFlattener::new(array.dtype(), output_len); + for (chunk_idx, plan) in plans.into_iter().enumerate() { + if plan.starts.is_empty() { + continue; + } + bases[chunk_idx] = running; + running += plan.total; + let chunk_indices = contiguous_runs(plan.starts, plan.lengths, plan.total); + flattener.push(array.chunk(chunk_idx).take(chunk_indices)?, ctx)?; + } + let flat = flattener.finish(array.dtype(), ctx)?; + + let mut reorder_starts = Vec::with_capacity(sub_pieces.len()); + let mut reorder_lengths = Vec::with_capacity(sub_pieces.len()); + for &(chunk_idx, offset, run) in &sub_pieces { + reorder_starts.push((bases[chunk_idx] + offset) as u64); + reorder_lengths.push(run as u64); + } + flat.take(contiguous_runs(reorder_starts, reorder_lengths, output_len)) + .map(Some) +} + +/// Builds a `PiecewiseSequenceArray` of contiguous (unit multiplier) runs whose lengths sum to +/// `total`. +fn contiguous_runs(starts: Vec, lengths: Vec, total: usize) -> ArrayRef { + let count = starts.len(); + debug_assert_eq!(count, lengths.len()); + let starts = PrimitiveArray::new(Buffer::from(starts), Validity::NonNullable).into_array(); + let lengths = match lengths.first() { + Some(&first) if lengths.iter().all(|&length| length == first) => { + ConstantArray::new(first, count).into_array() + } + _ => PrimitiveArray::new(Buffer::from(lengths), Validity::NonNullable).into_array(), + }; + let multipliers = ConstantArray::new(1u64, count).into_array(); + // SAFETY: starts, lengths, and multipliers are non-nullable u64 arrays of equal length, and + // `total` is the sum of the lengths. + unsafe { PiecewiseSequenceArray::new_unchecked(starts, lengths, multipliers, total) } + .into_array() +} + #[cfg(test)] mod test { + use vortex_buffer::Buffer; use vortex_buffer::bitbuffer; use vortex_buffer::buffer; use vortex_error::VortexResult; + use crate::ArrayRef; + use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; + use crate::arrays::ConstantArray; + use crate::arrays::FixedSizeListArray; + use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::assert_arrays_eq; use crate::compute::conformance::take::test_take_conformance; + use crate::dtype::DType; use crate::dtype::FieldNames; use crate::dtype::Nullability; + use crate::dtype::PType; use crate::validity::Validity; + fn chunked_i32() -> VortexResult { + ChunkedArray::try_new( + vec![ + buffer![0i32, 1, 2, 3, 4].into_array(), + buffer![5i32, 6, 7, 8, 9].into_array(), + buffer![10i32, 11, 12, 13, 14].into_array(), + ], + DType::Primitive(PType::I32, Nullability::NonNullable), + ) + } + + fn contiguous_pieces(starts: &[u64], lengths: &[u64]) -> VortexResult { + let len = usize::try_from(lengths.iter().sum::())?; + Ok(PiecewiseSequenceArray::try_new( + starts.iter().copied().collect::>().into_array(), + lengths + .iter() + .copied() + .collect::>() + .into_array(), + ConstantArray::new(1u64, starts.len()).into_array(), + len, + )? + .into_array()) + } + + #[test] + fn test_take_piecewise_monotonic_spanning_chunks() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + // The second piece crosses the first chunk boundary, the third spans the last two chunks. + let indices = contiguous_pieces(&[1, 4, 9], &[3, 4, 6])?; + let result = arr.take(indices)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6, 7, 9, 10, 11, 12, 13, 14]), + &mut ctx + ); + Ok(()) + } + + #[test] + fn test_take_piecewise_interleaved() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + // Pieces visit chunks out of order, forcing the reorder take. + let indices = contiguous_pieces(&[12, 2, 7, 0], &[3, 2, 3, 1])?; + let result = arr.take(indices)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([12i32, 13, 14, 2, 3, 7, 8, 9, 0]), + &mut ctx + ); + Ok(()) + } + + #[test] + fn test_take_piecewise_whole_array() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + let indices = contiguous_pieces(&[0], &[15])?; + let result = arr.take(indices)?; + + assert_arrays_eq!(result, PrimitiveArray::from_iter(0i32..15), &mut ctx); + Ok(()) + } + + #[test] + fn test_take_piecewise_across_empty_chunk() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = ChunkedArray::try_new( + vec![ + buffer![0i32, 1, 2, 3, 4].into_array(), + PrimitiveArray::empty::(Nullability::NonNullable).into_array(), + buffer![5i32, 6, 7, 8, 9].into_array(), + ], + DType::Primitive(PType::I32, Nullability::NonNullable), + )?; + + let indices = contiguous_pieces(&[3], &[4])?; + let result = arr.take(indices)?; + + assert_arrays_eq!(result, PrimitiveArray::from_iter([3i32, 4, 5, 6]), &mut ctx); + Ok(()) + } + + #[test] + fn test_take_piecewise_zero_length_pieces() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + let indices = contiguous_pieces(&[5, 0, 2], &[0, 4, 0])?; + let result = arr.take(indices)?; + + assert_arrays_eq!(result, PrimitiveArray::from_iter([0i32, 1, 2, 3]), &mut ctx); + Ok(()) + } + + #[test] + fn test_take_piecewise_out_of_bounds() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + let indices = contiguous_pieces(&[12], &[5])?; + let result = arr + .take(indices) + .and_then(|taken| taken.execute::(&mut ctx)); + + assert!(result.is_err()); + Ok(()) + } + + #[test] + fn test_take_piecewise_non_unit_multiplier() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + // Multiplier 2 falls back to the generic path but must stay correct. + let indices = PiecewiseSequenceArray::try_new( + buffer![0u64, 1].into_array(), + buffer![5u64, 3].into_array(), + buffer![2u64, 4].into_array(), + 8, + )? + .into_array(); + let result = arr.take(indices)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([0i32, 2, 4, 6, 8, 1, 5, 9]), + &mut ctx + ); + Ok(()) + } + + #[test] + fn test_take_piecewise_nullable_values() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = ChunkedArray::try_new( + vec![ + PrimitiveArray::from_option_iter([Some(0i32), None, Some(2)]).into_array(), + PrimitiveArray::from_option_iter([None, Some(4i32), Some(5)]).into_array(), + ], + DType::Primitive(PType::I32, Nullability::Nullable), + )?; + + let indices = contiguous_pieces(&[4, 1], &[2, 3])?; + let result = arr.take(indices)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([Some(4i32), Some(5), None, Some(2), None]), + &mut ctx + ); + Ok(()) + } + + #[test] + fn test_take_fsl_over_chunked_elements() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // Chunk boundaries at 8 and 13 do not line up with the list size of 3. + let elements = ChunkedArray::try_new( + vec![ + PrimitiveArray::from_iter(0i32..8).into_array(), + PrimitiveArray::from_iter(8i32..13).into_array(), + PrimitiveArray::from_iter(13i32..18).into_array(), + ], + DType::Primitive(PType::I32, Nullability::NonNullable), + )? + .into_array(); + let fsl = FixedSizeListArray::try_new(elements, 3, Validity::NonNullable, 6)?; + + let indices = buffer![4u64, 0, 2, 5, 1, 4].into_array(); + let result = fsl.take(indices.clone())?; + let expected = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..18).into_array(), + 3, + Validity::NonNullable, + 6, + )? + .take(indices)?; + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_take_fsl_over_chunked_elements_conformance() -> VortexResult<()> { + let elements = ChunkedArray::try_new( + vec![ + PrimitiveArray::from_iter(0i32..8).into_array(), + PrimitiveArray::from_iter(8i32..13).into_array(), + PrimitiveArray::from_iter(13i32..18).into_array(), + ], + DType::Primitive(PType::I32, Nullability::NonNullable), + )? + .into_array(); + let fsl = FixedSizeListArray::try_new(elements, 3, Validity::NonNullable, 6)?; + test_take_conformance( + &fsl.into_array(), + &mut array_session().create_execution_ctx(), + ); + Ok(()) + } + + #[test] + fn test_take_chunked_fsl() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let c0 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..6).into_array(), + 2, + Validity::NonNullable, + 3, + )?; + let c1 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(6i32..10).into_array(), + 2, + Validity::NonNullable, + 2, + )?; + let dtype = c0.dtype().clone(); + let arr = ChunkedArray::try_new(vec![c0.into_array(), c1.into_array()], dtype)?; + + let indices = buffer![4u64, 0, 3, 3, 1].into_array(); + let result = arr.take(indices.clone())?; + let expected = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..10).into_array(), + 2, + Validity::NonNullable, + 5, + )? + .take(indices)?; + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_take_chunked_fsl_nullable() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let c0 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..6).into_array(), + 2, + Validity::Array(bitbuffer![1 0 1].into_array()), + 3, + )?; + let c1 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(6i32..10).into_array(), + 2, + Validity::AllValid, + 2, + )?; + let dtype = c0.dtype().clone(); + let arr = ChunkedArray::try_new(vec![c0.into_array(), c1.into_array()], dtype)?; + + let indices = + PrimitiveArray::from_option_iter([Some(4u64), None, Some(0), Some(1)]).into_array(); + let result = arr.take(indices.clone())?; + let expected = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..10).into_array(), + 2, + Validity::Array(bitbuffer![1 0 1 1 1].into_array()), + 5, + )? + .take(indices)?; + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_take_chunked_fsl_non_fsl_chunk_falls_back() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let c0 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..6).into_array(), + 2, + Validity::NonNullable, + 3, + )?; + let c1 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(6i32..10).into_array(), + 2, + Validity::NonNullable, + 2, + )?; + let dtype = c0.dtype().clone(); + // Wrap one chunk in a nested chunked array so it is not FSL-encoded. + let nested = ChunkedArray::try_new(vec![c1.into_array()], dtype.clone())?; + let arr = ChunkedArray::try_new(vec![c0.into_array(), nested.into_array()], dtype)?; + + let indices = buffer![4u64, 0, 3, 1].into_array(); + let result = arr.take(indices.clone())?; + let expected = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..10).into_array(), + 2, + Validity::NonNullable, + 5, + )? + .take(indices)?; + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_take_chunked_struct_nested_flatten() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let s0 = StructArray::try_new( + ["a"].into(), + vec![buffer![0i32, 1, 2].into_array()], + 3, + Validity::NonNullable, + )?; + let s1 = StructArray::try_new( + ["a"].into(), + vec![buffer![3i32, 4].into_array()], + 2, + Validity::NonNullable, + )?; + let dtype = s0.dtype().clone(); + let arr = ChunkedArray::try_new(vec![s0.into_array(), s1.into_array()], dtype)?; + + let result = arr.take(buffer![4u64, 0, 2, 4].into_array())?; + let expected = StructArray::try_new( + ["a"].into(), + vec![buffer![4i32, 0, 2, 4].into_array()], + 4, + Validity::NonNullable, + )?; + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_take_chunked_fsl_conformance() -> VortexResult<()> { + let c0 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..6).into_array(), + 2, + Validity::NonNullable, + 3, + )?; + let c1 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(6i32..10).into_array(), + 2, + Validity::NonNullable, + 2, + )?; + let dtype = c0.dtype().clone(); + let arr = ChunkedArray::try_new(vec![c0.into_array(), c1.into_array()], dtype)?; + test_take_conformance( + &arr.into_array(), + &mut array_session().create_execution_ctx(), + ); + Ok(()) + } + #[test] fn test_take() { let mut ctx = array_session().create_execution_ctx(); @@ -289,10 +935,7 @@ mod test { indices.swap(i, j); } - let indices_arr = PrimitiveArray::new( - vortex_buffer::Buffer::from(indices.clone()), - Validity::NonNullable, - ); + let indices_arr = PrimitiveArray::new(Buffer::from(indices.clone()), Validity::NonNullable); let result = arr.take(indices_arr.into_array())?; // Verify every element. From 0b4bf83bceafb68fbddea8b944e435b047f08650 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 21 Jul 2026 17:22:55 +0100 Subject: [PATCH 2/6] simpler Signed-off-by: Robert Kruszewski --- vortex-array/src/arrays/chunked/compute/take.rs | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/vortex-array/src/arrays/chunked/compute/take.rs b/vortex-array/src/arrays/chunked/compute/take.rs index 43f0cab2fb9..83524ffe451 100644 --- a/vortex-array/src/arrays/chunked/compute/take.rs +++ b/vortex-array/src/arrays/chunked/compute/take.rs @@ -115,8 +115,7 @@ fn take_chunked( let chunk_offsets = array.chunk_offset_values(); let nchunks = array.nchunks(); let mut flattener = ChunkFlattener::new(array.dtype(), pairs.len()); - let mut final_take = BufferMut::::with_capacity(n); - final_take.push_n(0u64, n); + let mut final_take = BufferMut::::zeroed(n); let mut cursor = 0usize; let mut dedup_idx = 0u64; @@ -154,13 +153,7 @@ fn take_chunked( // 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(), - ); + let take_validity = Validity::from_mask(indices_mask, indices.dtype().nullability()); flat.take(PrimitiveArray::new(final_take.freeze(), take_validity).into_array()) } From 031b39f97dda93828feeda82f52bb188f22fbf64 Mon Sep 17 00:00:00 2001 From: Dimitar Dimitrov Date: Mon, 27 Jul 2026 15:58:55 +0100 Subject: [PATCH 3/6] bench: add compressed FSL take coverage Signed-off-by: Dimitar Dimitrov --- vortex-btrblocks/Cargo.toml | 5 ++ vortex-btrblocks/benches/take_fsl.rs | 102 +++++++++++++++++++++++++++ 2 files changed, 107 insertions(+) create mode 100644 vortex-btrblocks/benches/take_fsl.rs diff --git a/vortex-btrblocks/Cargo.toml b/vortex-btrblocks/Cargo.toml index 255d2fbd075..72706a1fe1c 100644 --- a/vortex-btrblocks/Cargo.toml +++ b/vortex-btrblocks/Cargo.toml @@ -64,3 +64,8 @@ test = false name = "compress_listview" harness = false test = false + +[[bench]] +name = "take_fsl" +harness = false +test = false diff --git a/vortex-btrblocks/benches/take_fsl.rs b/vortex-btrblocks/benches/take_fsl.rs new file mode 100644 index 00000000000..eb3d612b0ca --- /dev/null +++ b/vortex-btrblocks/benches/take_fsl.rs @@ -0,0 +1,102 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks take from a chunked array of compressed fixed-size list chunks. + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use divan::counter::BytesCount; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex_array::ArrayRef; +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::FixedSizeList; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; +use vortex_array::validity::Validity; +use vortex_btrblocks::BtrBlocksCompressor; +use vortex_buffer::Buffer; +use vortex_session::VortexSession; + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(array_session); + +const NUM_LISTS: usize = 500; +const NUM_CHUNKS: usize = 8; +const NUM_INDICES: &[usize] = &[64]; +const LIST_SIZES: &[usize] = &[8, 16, 32]; + +fn create_random_indices(num_indices: usize) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(42); + let indices: Buffer = (0..num_indices) + .map(|_| rng.random_range(0..NUM_LISTS) as u64) + .collect(); + indices.into_array() +} + +fn create_chunked_compressed_fsl(list_size: usize) -> ArrayRef { + let compressor = BtrBlocksCompressor::default(); + let lists_per_chunk = NUM_LISTS.div_ceil(NUM_CHUNKS); + let mut chunks = Vec::new(); + let mut start = 0usize; + let mut ctx = SESSION.create_execution_ctx(); + + while start < NUM_LISTS { + let end = (start + lists_per_chunk).min(NUM_LISTS); + let elements: Buffer = + ((start * list_size) as i64..(end * list_size) as i64).collect(); + let fsl = FixedSizeListArray::new( + elements.into_array(), + u32::try_from(list_size).unwrap(), + Validity::NonNullable, + end - start, + ) + .into_array(); + let compressed = compressor.compress(&fsl, &mut ctx).unwrap(); + + let compressed_fsl = compressed.as_::(); + assert!( + !compressed_fsl.elements().is::(), + "expected compressed FSL elements" + ); + chunks.push(compressed); + start = end; + } + + let dtype = chunks[0].dtype().clone(); + ChunkedArray::try_new(chunks, dtype).unwrap().into_array() +} + +fn bench_take(bencher: Bencher, num_indices: usize) { + let array = create_chunked_compressed_fsl(LIST_SIZE); + let indices = create_random_indices(num_indices); + + bencher + .counter(BytesCount::of_many::(num_indices * LIST_SIZE)) + .with_inputs(|| (&array, &indices, SESSION.create_execution_ctx())) + .bench_refs(|(array, indices, ctx)| { + array + .take((*indices).clone()) + .unwrap() + .execute::(ctx) + .unwrap() + }); +} + +#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] +fn take_chunked_compressed_fsl(bencher: Bencher, num_indices: usize) { + bench_take::(bencher, num_indices); +} From c81ba0163ba82f357f8d4a72c5969c37023fc440 Mon Sep 17 00:00:00 2001 From: Dimitar Dimitrov Date: Tue, 28 Jul 2026 16:41:38 +0100 Subject: [PATCH 4/6] Execute chunked FSL before take (#9012) Depends on #8881. Use the existing chunked FixedSizeList canonicalization path before applying the specialized take. This removes the physical FSL encoding requirement and lets logically FSL-typed chunks use the optimization. --------- Signed-off-by: Dimitar Dimitrov Signed-off-by: Robert Kruszewski Co-authored-by: Robert Kruszewski Co-authored-by: Claude Fable 5 --- .../src/arrays/chunked/compute/take.rs | 48 ++++++------------- 1 file changed, 15 insertions(+), 33 deletions(-) diff --git a/vortex-array/src/arrays/chunked/compute/take.rs b/vortex-array/src/arrays/chunked/compute/take.rs index 83524ffe451..26e0fcfadc5 100644 --- a/vortex-array/src/arrays/chunked/compute/take.rs +++ b/vortex-array/src/arrays/chunked/compute/take.rs @@ -19,14 +19,12 @@ use crate::array::ArrayView; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; -use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::PiecewiseSequence; use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::dict::TakeExecute; -use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::piecewise_sequence::constant_unsigned_usize; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; @@ -168,7 +166,7 @@ impl TakeExecute for Chunked { return array.chunk(0).take(indices.clone()).map(Some); } - if let Some(taken) = take_chunked_fsl(array, indices)? { + if let Some(taken) = take_chunked_fsl(array, indices, ctx)? { return Ok(Some(taken)); } @@ -182,42 +180,23 @@ impl TakeExecute for Chunked { } } -/// Rewrites take over a chunked array of [`FixedSizeList`] chunks as take over a single -/// [`FixedSizeListArray`] whose elements child chains the chunks' elements. +/// Rewrites take over a chunked fixed-size list array as take over a single +/// [`FixedSizeListArray`]. /// -/// Each FSL chunk stores exactly `chunk.len() * list_size` elements starting at its first list, -/// so the chunks' elements children concatenate zero-copy into the elements of the combined -/// array, exactly as chunked canonicalization does. The FSL take implementation then gathers the -/// elements with a single `PiecewiseSequenceArray`, which [`take_piecewise_chunked`] resolves per -/// chunk without expanding one index per element. Chunks that are not FSL-encoded fall back to -/// the generic path since accessing their elements would require canonicalizing whole chunks. +/// Chunked FSL canonicalization executes each logical FSL chunk into its canonical representation, +/// then concatenates the chunks' elements children zero-copy. The FSL take implementation gathers +/// those elements with a single `PiecewiseSequenceArray`, which [`take_piecewise_chunked`] resolves +/// per chunk without expanding one index per element. fn take_chunked_fsl( array: ArrayView<'_, Chunked>, indices: &ArrayRef, + ctx: &mut ExecutionCtx, ) -> VortexResult> { - let DType::FixedSizeList(element_dtype, list_size, _) = array.dtype() else { + let DType::FixedSizeList(..) = array.dtype() else { return Ok(None); }; - let mut element_chunks = Vec::with_capacity(array.nchunks()); - for chunk in array.iter_chunks() { - let Some(fsl) = chunk.as_opt::() else { - return Ok(None); - }; - element_chunks.push(fsl.elements().clone()); - } - - let validity = array.array().validity()?; - // SAFETY: every chunk is a FixedSizeList with element dtype `element_dtype`. - let elements = - unsafe { ChunkedArray::new_unchecked(element_chunks, element_dtype.as_ref().clone()) } - .into_array(); - // SAFETY: each FSL chunk holds exactly `chunk.len() * list_size` elements, so the chained - // elements hold `array.len() * list_size` entries, and the chunked validity covers - // `array.len()` lists with the array's nullability. - let fsl = unsafe { - FixedSizeListArray::new_unchecked(elements, *list_size, validity, array.as_ref().len()) - }; + let fsl = array.as_ref().clone().execute::(ctx)?; fsl.into_array().take(indices.clone()).map(Some) } @@ -394,6 +373,7 @@ mod test { use vortex_buffer::bitbuffer; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_error::vortex_bail; use crate::ArrayRef; use crate::Canonical; @@ -692,7 +672,7 @@ mod test { } #[test] - fn test_take_chunked_fsl_non_fsl_chunk_falls_back() -> VortexResult<()> { + fn test_take_chunked_fsl_executes_non_fsl_chunk() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let c0 = FixedSizeListArray::try_new( PrimitiveArray::from_iter(0i32..6).into_array(), @@ -712,7 +692,9 @@ mod test { let arr = ChunkedArray::try_new(vec![c0.into_array(), nested.into_array()], dtype)?; let indices = buffer![4u64, 0, 3, 1].into_array(); - let result = arr.take(indices.clone())?; + let Some(result) = super::take_chunked_fsl(arr.as_view(), &indices, &mut ctx)? else { + vortex_bail!("fixed-size list optimization did not execute"); + }; let expected = FixedSizeListArray::try_new( PrimitiveArray::from_iter(0i32..10).into_array(), 2, From f375f7feadc378d05d8a1afbd1cc36d6d00098cc Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 29 Jul 2026 00:04:22 +0100 Subject: [PATCH 5/6] less Signed-off-by: Robert Kruszewski --- .../src/arrays/chunked/compute/take.rs | 137 ++++++++++++++++-- 1 file changed, 121 insertions(+), 16 deletions(-) diff --git a/vortex-array/src/arrays/chunked/compute/take.rs b/vortex-array/src/arrays/chunked/compute/take.rs index 26e0fcfadc5..35952527f6c 100644 --- a/vortex-array/src/arrays/chunked/compute/take.rs +++ b/vortex-array/src/arrays/chunked/compute/take.rs @@ -19,12 +19,14 @@ use crate::array::ArrayView; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; +use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::PiecewiseSequence; use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::dict::TakeExecute; +use crate::arrays::fixed_size_list::FixedSizeListDataParts; use crate::arrays::piecewise_sequence::constant_unsigned_usize; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; @@ -166,16 +168,18 @@ impl TakeExecute for Chunked { return array.chunk(0).take(indices.clone()).map(Some); } - if let Some(taken) = take_chunked_fsl(array, indices, ctx)? { - return Ok(Some(taken)); - } - + // Piecewise indices go first: the piecewise take only touches chunks the pieces + // intersect, while the FSL rewrite canonicalizes every chunk up front. if let Some(piecewise_indices) = indices.as_opt::() && let Some(taken) = take_piecewise_chunked(array, piecewise_indices, ctx)? { return Ok(Some(taken)); } + if let Some(taken) = take_chunked_fsl(array, indices, ctx)? { + return Ok(Some(taken)); + } + take_chunked(array, indices, ctx).map(Some) } } @@ -183,21 +187,44 @@ impl TakeExecute for Chunked { /// Rewrites take over a chunked fixed-size list array as take over a single /// [`FixedSizeListArray`]. /// -/// Chunked FSL canonicalization executes each logical FSL chunk into its canonical representation, -/// then concatenates the chunks' elements children zero-copy. The FSL take implementation gathers -/// those elements with a single `PiecewiseSequenceArray`, which [`take_piecewise_chunked`] resolves -/// per chunk without expanding one index per element. +/// Chunked FSL canonicalization concatenates the chunks' elements children zero-copy. The FSL +/// take implementation then gathers those elements with a single `PiecewiseSequenceArray`, which +/// [`take_piecewise_chunked`] resolves per chunk without expanding one index per element. +/// +/// Applies only when every chunk is FSL-encoded, so canonicalization reuses the chunks' elements +/// children zero-copy. Converting any other encoding (constant, dict, ...) would materialize +/// whole chunks the indices may never touch, so those fall back to the generic filter-based path, +/// which only decodes selected rows. fn take_chunked_fsl( array: ArrayView<'_, Chunked>, indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let DType::FixedSizeList(..) = array.dtype() else { + let DType::FixedSizeList(_, list_size, nullability) = array.dtype() else { return Ok(None); }; + if !array.iter_chunks().all(|chunk| chunk.is::()) { + return Ok(None); + } + let len = array.as_ref().len(); let fsl = array.as_ref().clone().execute::(ctx)?; - fsl.into_array().take(indices.clone()).map(Some) + let FixedSizeListDataParts { + elements, validity, .. + } = fsl.into_data_parts(); + // The canonicalized validity chains the chunks' validities as a chunked bool array; letting + // the FSL take gather from it would recurse into a second chunked take over the same indices. + // Execute it into a flat mask once so the validity take is a plain bitmap gather. + let validity = match validity { + validity @ Validity::Array(_) => { + Validity::from_mask(validity.execute_mask(len, ctx)?, *nullability) + } + validity => validity, + }; + FixedSizeListArray::try_new(elements, *list_size, validity, len)? + .into_array() + .take(indices.clone()) + .map(Some) } /// A per-chunk gather plan: chunk-local sub-piece runs to take from one chunk, in output order. @@ -247,7 +274,7 @@ fn take_piecewise_chunked( }; let array_len = array.as_ref().len(); - let chunk_offsets = array.chunk_offsets(); + let chunk_offsets = array.chunk_offset_values(); let nchunks = array.nchunks(); let mut plans: Vec = Vec::new(); @@ -373,7 +400,6 @@ mod test { use vortex_buffer::bitbuffer; use vortex_buffer::buffer; use vortex_error::VortexResult; - use vortex_error::vortex_bail; use crate::ArrayRef; use crate::Canonical; @@ -394,6 +420,7 @@ mod test { use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::scalar::Scalar; use crate::validity::Validity; fn chunked_i32() -> VortexResult { @@ -672,7 +699,7 @@ mod test { } #[test] - fn test_take_chunked_fsl_executes_non_fsl_chunk() -> VortexResult<()> { + fn test_take_chunked_fsl_non_fsl_chunk_falls_back() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let c0 = FixedSizeListArray::try_new( PrimitiveArray::from_iter(0i32..6).into_array(), @@ -692,9 +719,12 @@ mod test { let arr = ChunkedArray::try_new(vec![c0.into_array(), nested.into_array()], dtype)?; let indices = buffer![4u64, 0, 3, 1].into_array(); - let Some(result) = super::take_chunked_fsl(arr.as_view(), &indices, &mut ctx)? else { - vortex_bail!("fixed-size list optimization did not execute"); - }; + assert!( + super::take_chunked_fsl(arr.as_view(), &indices, &mut ctx)?.is_none(), + "expected fallback to the generic take path" + ); + + let result = arr.take(indices.clone())?; let expected = FixedSizeListArray::try_new( PrimitiveArray::from_iter(0i32..10).into_array(), 2, @@ -707,6 +737,81 @@ mod test { Ok(()) } + #[test] + fn test_take_chunked_fsl_constant_chunk_falls_back() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let c0 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..6).into_array(), + 2, + Validity::AllValid, + 3, + )?; + let dtype = c0.dtype().clone(); + // A constant chunk is only logically FSL; converting it materializes the whole chunk. + let c1 = ConstantArray::new(Scalar::null(dtype.clone()), 2).into_array(); + let arr = ChunkedArray::try_new(vec![c0.into_array(), c1], dtype)?; + + let indices = buffer![4u64, 0, 3, 1].into_array(); + assert!( + super::take_chunked_fsl(arr.as_view(), &indices, &mut ctx)?.is_none(), + "expected fallback to the generic take path" + ); + + let result = arr.take(indices.clone())?; + let expected = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..10).into_array(), + 2, + Validity::Array(bitbuffer![1 1 1 0 0].into_array()), + 5, + )? + .take(indices)?; + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_take_piecewise_chunked_fsl() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let c0 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..6).into_array(), + 2, + Validity::NonNullable, + 3, + )?; + let c1 = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(6i32..10).into_array(), + 2, + Validity::NonNullable, + 2, + )?; + let dtype = c0.dtype().clone(); + let arr = ChunkedArray::try_new(vec![c0.into_array(), c1.into_array()], dtype)?; + let reference = FixedSizeListArray::try_new( + PrimitiveArray::from_iter(0i32..10).into_array(), + 2, + Validity::NonNullable, + 5, + )?; + + // A monotonic run crossing the chunk boundary. + let indices = contiguous_pieces(&[1], &[3])?; + assert_arrays_eq!( + arr.take(indices.clone())?, + reference.take(indices)?, + &mut ctx + ); + + // Out-of-order pieces forcing the reorder take. + let indices = contiguous_pieces(&[3, 0], &[2, 2])?; + assert_arrays_eq!( + arr.take(indices.clone())?, + reference.take(indices)?, + &mut ctx + ); + Ok(()) + } + #[test] fn test_take_chunked_struct_nested_flatten() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); From 11b7ae62719221d5c1a64661c3bb69b616fc3cbf Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 29 Jul 2026 17:22:28 +0100 Subject: [PATCH 6/6] more Signed-off-by: Robert Kruszewski --- .../src/arrays/chunked/compute/take.rs | 280 ++++++++++++------ .../arrays/fixed_size_list/compute/take.rs | 109 +++++++ .../src/arrays/fixed_size_list/tests/take.rs | 84 ++++++ vortex-btrblocks/benches/take_fsl.rs | 31 +- 4 files changed, 415 insertions(+), 89 deletions(-) diff --git a/vortex-array/src/arrays/chunked/compute/take.rs b/vortex-array/src/arrays/chunked/compute/take.rs index 35952527f6c..3ec3691a97d 100644 --- a/vortex-array/src/arrays/chunked/compute/take.rs +++ b/vortex-array/src/arrays/chunked/compute/take.rs @@ -19,14 +19,12 @@ use crate::array::ArrayView; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; -use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::PiecewiseSequence; use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::dict::TakeExecute; -use crate::arrays::fixed_size_list::FixedSizeListDataParts; use crate::arrays::piecewise_sequence::constant_unsigned_usize; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; @@ -39,12 +37,9 @@ use crate::executor::ExecutionCtx; use crate::match_each_unsigned_integer_ptype; use crate::validity::Validity; -/// Flattens per-chunk take/filter results into a single array. -/// -/// Flat dtypes append directly into a canonical builder. Nested dtypes instead collect the chunks -/// and canonicalize them as a chunked array, which reuses the chunks' children zero-copy (e.g. -/// chunked FixedSizeLists canonicalize into one FixedSizeList over the chained elements) and lets -/// the follow-up take push down lazily instead of deep-copying every child. +/// Flattens per-chunk take/filter results into a single array. Flat dtypes append directly into +/// a canonical builder; nested dtypes collect the chunks and canonicalize them as a chunked +/// array, which reuses the chunks' children zero-copy. enum ChunkFlattener { Builder(Box), Chunks(Vec), @@ -99,8 +94,12 @@ 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. + if !indices.as_ref().dtype().is_nullable() && indices_values.is_sorted_by(|a, b| a < b) { + return take_chunked_sorted(array, indices_values); + } + + // Sort (value, orig_pos) pairs so indices for the same chunk are contiguous. Null indices + // are skipped: their final_take slots stay 0 and are masked null below. let mut pairs: Vec<(u64, usize)> = indices_values .iter() .enumerate() @@ -109,9 +108,8 @@ fn take_chunked( .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 + // Walk sorted pairs against chunk boundaries: build per-chunk filter masks, dedup repeated + // values inline, and scatter final_take[orig_pos] = dedup_idx. let chunk_offsets = array.chunk_offset_values(); let nchunks = array.nchunks(); let mut flattener = ChunkFlattener::new(array.dtype(), pairs.len()); @@ -121,6 +119,9 @@ fn take_chunked( let mut dedup_idx = 0u64; for chunk_idx in 0..nchunks { + if cursor == pairs.len() { + break; + } let chunk_start = chunk_offsets[chunk_idx]; let chunk_end = chunk_offsets[chunk_idx + 1]; let chunk_len = chunk_end - chunk_start; @@ -130,7 +131,7 @@ fn take_chunked( let chunk_pairs = &pairs[cursor..range_end]; if !chunk_pairs.is_empty() { - let mut local_indices: Vec = Vec::new(); + let mut local_indices: Vec = Vec::with_capacity(chunk_pairs.len()); for (i, &(val, orig_pos)) in chunk_pairs.iter().enumerate() { if cursor + i > 0 && val != pairs[cursor + i - 1].0 { dedup_idx += 1; @@ -148,35 +149,78 @@ fn take_chunked( cursor = range_end; } + if cursor != pairs.len() { + vortex_bail!(OutOfBounds: usize::try_from(pairs[cursor].0)?, 0, array.as_ref().len()); + } let flat = flattener.finish(array.dtype(), ctx)?; - // 4. Single take to restore original order and expand duplicates. - // Carry the original index validity so null indices produce null outputs. + // A single take restores original order and expands duplicates. Carrying the original index + // validity makes 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 for strictly increasing, non-nullable indices: every row is selected at most once and in +/// order, so the per-chunk filters concatenate directly into the result with no dedup or reorder +/// take. +fn take_chunked_sorted( + array: ArrayView<'_, Chunked>, + indices_values: &[u64], +) -> VortexResult { + let chunk_offsets = array.chunk_offset_values(); + let mut chunks = Vec::new(); + let mut cursor = 0usize; + + for chunk_idx in 0..array.nchunks() { + if cursor == indices_values.len() { + break; + } + let chunk_start = chunk_offsets[chunk_idx]; + let chunk_end = chunk_offsets[chunk_idx + 1]; + let chunk_end_u64 = u64::try_from(chunk_end)?; + + let range_end = cursor + indices_values[cursor..].partition_point(|&v| v < chunk_end_u64); + if range_end > cursor { + let mut local_indices = Vec::with_capacity(range_end - cursor); + for &val in &indices_values[cursor..range_end] { + local_indices.push(usize::try_from(val)? - chunk_start); + } + let filter_mask = Mask::from_indices(chunk_end - chunk_start, local_indices); + chunks.push(array.chunk(chunk_idx).filter(filter_mask)?); + } + cursor = range_end; + } + if cursor != indices_values.len() { + vortex_bail!( + OutOfBounds: usize::try_from(indices_values[cursor])?, 0, array.as_ref().len() + ); + } + + if chunks.len() == 1 { + return Ok(chunks.swap_remove(0)); + } + // SAFETY: every chunk is a filter of a chunk of `array`, which leaves the dtype unchanged. + Ok(unsafe { ChunkedArray::new_unchecked(chunks, array.dtype().clone()) }.into_array()) +} + impl TakeExecute for Chunked { fn take( array: ArrayView<'_, Chunked>, indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - // A single chunk is logically identical to the chunked array, so delegate directly. if array.nchunks() == 1 { return array.chunk(0).take(indices.clone()).map(Some); } - // Piecewise indices go first: the piecewise take only touches chunks the pieces - // intersect, while the FSL rewrite canonicalizes every chunk up front. - if let Some(piecewise_indices) = indices.as_opt::() - && let Some(taken) = take_piecewise_chunked(array, piecewise_indices, ctx)? - { + if let Some(taken) = take_chunked_fsl(array, indices, ctx)? { return Ok(Some(taken)); } - if let Some(taken) = take_chunked_fsl(array, indices, ctx)? { + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_piecewise_chunked(array, piecewise_indices, ctx)? + { return Ok(Some(taken)); } @@ -185,46 +229,20 @@ impl TakeExecute for Chunked { } /// Rewrites take over a chunked fixed-size list array as take over a single -/// [`FixedSizeListArray`]. -/// -/// Chunked FSL canonicalization concatenates the chunks' elements children zero-copy. The FSL -/// take implementation then gathers those elements with a single `PiecewiseSequenceArray`, which -/// [`take_piecewise_chunked`] resolves per chunk without expanding one index per element. -/// -/// Applies only when every chunk is FSL-encoded, so canonicalization reuses the chunks' elements -/// children zero-copy. Converting any other encoding (constant, dict, ...) would materialize -/// whole chunks the indices may never touch, so those fall back to the generic filter-based path, -/// which only decodes selected rows. +/// [`FixedSizeListArray`]. The swizzle is structural: FSL-encoded chunks execute to themselves, +/// so their elements chain zero-copy, and the FSL take then gathers per-chunk element runs +/// instead of taking each chunk. fn take_chunked_fsl( array: ArrayView<'_, Chunked>, indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let DType::FixedSizeList(_, list_size, nullability) = array.dtype() else { - return Ok(None); - }; - if !array.iter_chunks().all(|chunk| chunk.is::()) { + if !matches!(array.dtype(), DType::FixedSizeList(..)) { return Ok(None); } - let len = array.as_ref().len(); let fsl = array.as_ref().clone().execute::(ctx)?; - let FixedSizeListDataParts { - elements, validity, .. - } = fsl.into_data_parts(); - // The canonicalized validity chains the chunks' validities as a chunked bool array; letting - // the FSL take gather from it would recurse into a second chunked take over the same indices. - // Execute it into a flat mask once so the validity take is a plain bitmap gather. - let validity = match validity { - validity @ Validity::Array(_) => { - Validity::from_mask(validity.execute_mask(len, ctx)?, *nullability) - } - validity => validity, - }; - FixedSizeListArray::try_new(elements, *list_size, validity, len)? - .into_array() - .take(indices.clone()) - .map(Some) + fsl.into_array().take(indices.clone()).map(Some) } /// A per-chunk gather plan: chunk-local sub-piece runs to take from one chunk, in output order. @@ -237,11 +255,11 @@ struct ChunkGather { /// Take for [`PiecewiseSequence`] indices with unit multipliers. /// -/// Each piece is a contiguous index run, so instead of expanding one index per element and -/// sorting them (as [`take_chunked`] does), pieces are split at chunk boundaries and gathered -/// from each chunk with a chunk-local `PiecewiseSequenceArray`. When the sub-pieces visit chunks -/// in non-decreasing order the gathered chunks concatenate directly into the result; otherwise -/// the gathered chunks are canonicalized once and a second piecewise take restores output order. +/// Each piece is a contiguous index run, so instead of expanding one index per element, pieces +/// are split at chunk boundaries and gathered from each chunk with a chunk-local +/// `PiecewiseSequenceArray`. Sub-pieces that visit chunks in non-decreasing order concatenate +/// directly into the result; otherwise the gathered chunks are flattened once and a second +/// piecewise take restores output order. fn take_piecewise_chunked( array: ArrayView<'_, Chunked>, indices: ArrayView<'_, PiecewiseSequence>, @@ -313,9 +331,24 @@ fn take_piecewise_chunked( prev_chunk = chunk_idx; let plan = &mut plans[chunk_idx]; - sub_pieces.push((chunk_idx, plan.total, run)); - plan.starts.push((cursor - chunk_offsets[chunk_idx]) as u64); - plan.lengths.push(run as u64); + // Consecutive sub-pieces in the same chunk stay adjacent in its gathered output, so + // they merge into one reorder run. + match sub_pieces.last_mut() { + Some((last_chunk, _, last_run)) if *last_chunk == chunk_idx => *last_run += run, + _ => sub_pieces.push((chunk_idx, plan.total, run)), + } + // A run that continues exactly where the chunk's previous run ended extends it, so + // contiguous spans gather as one run. + let local_start = (cursor - chunk_offsets[chunk_idx]) as u64; + match (plan.starts.last(), plan.lengths.last_mut()) { + (Some(&prev_start), Some(prev_len)) if prev_start + *prev_len == local_start => { + *prev_len += run as u64; + } + _ => { + plan.starts.push(local_start); + plan.lengths.push(run as u64); + } + } plan.total += run; cursor += run; @@ -328,16 +361,14 @@ fn take_piecewise_chunked( "PiecewiseSequenceArray expanded length {total_len} does not match declared length {output_len}" ); - // Gather the sub-pieces of each touched chunk with one chunk-local piecewise take. When the - // sub-pieces visit chunks in order, the gathered chunks already concatenate into the result. + // Chunks visited in order: the per-chunk gathers already concatenate into the result. if monotonic { let mut gathered = Vec::new(); for (chunk_idx, plan) in plans.into_iter().enumerate() { if plan.starts.is_empty() { continue; } - let chunk_indices = contiguous_runs(plan.starts, plan.lengths, plan.total); - gathered.push(array.chunk(chunk_idx).take(chunk_indices)?); + gathered.push(gather_chunk(array.chunk(chunk_idx), plan)?); } let result = if gathered.len() == 1 { gathered.swap_remove(0) @@ -349,8 +380,8 @@ fn take_piecewise_chunked( return Ok(Some(result)); } - // Out-of-order sub-pieces: flatten each chunk's gather once through a [`ChunkFlattener`], - // then take the sub-piece runs from the flattened result in output order. + // Out-of-order sub-pieces: flatten the per-chunk gathers, then take the sub-piece runs from + // the flattened result in output order. let mut bases = vec![0usize; nchunks]; let mut running = 0usize; let mut flattener = ChunkFlattener::new(array.dtype(), output_len); @@ -360,8 +391,7 @@ fn take_piecewise_chunked( } bases[chunk_idx] = running; running += plan.total; - let chunk_indices = contiguous_runs(plan.starts, plan.lengths, plan.total); - flattener.push(array.chunk(chunk_idx).take(chunk_indices)?, ctx)?; + flattener.push(gather_chunk(array.chunk(chunk_idx), plan)?, ctx)?; } let flat = flattener.finish(array.dtype(), ctx)?; @@ -375,6 +405,16 @@ fn take_piecewise_chunked( .map(Some) } +/// Gathers a chunk's planned runs: a single run is a zero-copy slice, multiple runs become a +/// chunk-local piecewise take. +fn gather_chunk(chunk: &ArrayRef, plan: ChunkGather) -> VortexResult { + if let [start] = plan.starts.as_slice() { + let start = usize::try_from(*start)?; + return chunk.slice(start..start + plan.total); + } + chunk.take(contiguous_runs(plan.starts, plan.lengths, plan.total)) +} + /// Builds a `PiecewiseSequenceArray` of contiguous (unit multiplier) runs whose lengths sum to /// `total`. fn contiguous_runs(starts: Vec, lengths: Vec, total: usize) -> ArrayRef { @@ -395,7 +435,7 @@ fn contiguous_runs(starts: Vec, lengths: Vec, total: usize) -> ArrayRe } #[cfg(test)] -mod test { +mod tests { use vortex_buffer::Buffer; use vortex_buffer::bitbuffer; use vortex_buffer::buffer; @@ -514,6 +554,37 @@ mod test { Ok(()) } + #[test] + fn test_take_piecewise_adjacent_pieces_merge() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + // The first two pieces are contiguous within one chunk, so they merge into a single run + // that gathers as a zero-copy slice. + let indices = contiguous_pieces(&[0, 2, 7], &[2, 2, 2])?; + let result = arr.take(indices)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([0i32, 1, 2, 3, 7, 8]), + &mut ctx + ); + Ok(()) + } + + #[test] + fn test_take_piecewise_same_chunk_out_of_order() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + // Non-contiguous pieces in the same chunk must stay separate runs in piece order. + let indices = contiguous_pieces(&[2, 0], &[2, 2])?; + let result = arr.take(indices)?; + + assert_arrays_eq!(result, PrimitiveArray::from_iter([2i32, 3, 0, 1]), &mut ctx); + Ok(()) + } + #[test] fn test_take_piecewise_zero_length_pieces() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); @@ -699,7 +770,7 @@ mod test { } #[test] - fn test_take_chunked_fsl_non_fsl_chunk_falls_back() -> VortexResult<()> { + fn test_take_chunked_fsl_non_fsl_chunk() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let c0 = FixedSizeListArray::try_new( PrimitiveArray::from_iter(0i32..6).into_array(), @@ -714,16 +785,12 @@ mod test { 2, )?; let dtype = c0.dtype().clone(); - // Wrap one chunk in a nested chunked array so it is not FSL-encoded. + // Wrap one chunk in a nested chunked array so it is not FSL-encoded; the dtype-based + // swizzle must still handle it. let nested = ChunkedArray::try_new(vec![c1.into_array()], dtype.clone())?; let arr = ChunkedArray::try_new(vec![c0.into_array(), nested.into_array()], dtype)?; let indices = buffer![4u64, 0, 3, 1].into_array(); - assert!( - super::take_chunked_fsl(arr.as_view(), &indices, &mut ctx)?.is_none(), - "expected fallback to the generic take path" - ); - let result = arr.take(indices.clone())?; let expected = FixedSizeListArray::try_new( PrimitiveArray::from_iter(0i32..10).into_array(), @@ -738,7 +805,7 @@ mod test { } #[test] - fn test_take_chunked_fsl_constant_chunk_falls_back() -> VortexResult<()> { + fn test_take_chunked_fsl_constant_chunk() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let c0 = FixedSizeListArray::try_new( PrimitiveArray::from_iter(0i32..6).into_array(), @@ -747,16 +814,11 @@ mod test { 3, )?; let dtype = c0.dtype().clone(); - // A constant chunk is only logically FSL; converting it materializes the whole chunk. + // A constant chunk is only logically FSL; the swizzle canonicalizes it in place. let c1 = ConstantArray::new(Scalar::null(dtype.clone()), 2).into_array(); let arr = ChunkedArray::try_new(vec![c0.into_array(), c1], dtype)?; let indices = buffer![4u64, 0, 3, 1].into_array(); - assert!( - super::take_chunked_fsl(arr.as_view(), &indices, &mut ctx)?.is_none(), - "expected fallback to the generic take path" - ); - let result = arr.take(indices.clone())?; let expected = FixedSizeListArray::try_new( PrimitiveArray::from_iter(0i32..10).into_array(), @@ -1063,6 +1125,54 @@ mod test { Ok(()) } + #[test] + fn test_take_sorted_indices() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + // Strictly increasing indices spanning chunks hit the sorted fast path. + let indices = buffer![1u64, 4, 5, 9, 10, 14].into_array(); + let result = arr.take(indices)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([1i32, 4, 5, 9, 10, 14]), + &mut ctx + ); + Ok(()) + } + + #[test] + fn test_take_out_of_bounds() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + // Sorted indices hit the fast path, unsorted the generic path; both must fail. + for indices in [buffer![0u64, 100], buffer![100u64, 0]] { + let result = arr + .take(indices.into_array()) + .and_then(|taken| taken.execute::(&mut ctx)); + assert!(result.is_err(), "expected out-of-bounds error"); + } + Ok(()) + } + + #[test] + fn test_take_all_null_indices() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let arr = chunked_i32()?; + + let indices = PrimitiveArray::from_option_iter([None::, None]).into_array(); + let result = arr.take(indices)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([None::, None]), + &mut ctx + ); + Ok(()) + } + #[test] fn test_take_chunked_conformance() { let a = buffer![1i32, 2, 3].into_array(); diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index 4e41ffd7ad1..a95cde9f4f2 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools; +use num_traits::AsPrimitive; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; @@ -10,23 +11,28 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_err; use crate::ArrayRef; +use crate::Columnar; use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::ConstantArray; use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; +use crate::arrays::PiecewiseSequence; use crate::arrays::PiecewiseSequenceArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use crate::arrays::piecewise_sequence::constant_unsigned_usize; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; use crate::builders::builder_with_capacity; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::executor::ExecutionCtx; use crate::match_each_integer_ptype; +use crate::match_each_unsigned_integer_ptype; use crate::optimizer::ArrayOptimizer; use crate::validity::Validity; @@ -105,12 +111,115 @@ fn take_non_empty_fsl( return take_non_empty_degenerate_fsl(array, indices, ctx); } + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_piecewise_fsl(array, piecewise_indices, ctx)? + { + return Ok(taken); + } + let indices_array = indices.clone().execute::(ctx)?; match_each_integer_ptype!(indices_array.ptype(), |I| { take_non_empty_non_degenerate_fsl::(array, indices, indices_array.as_view(), ctx) }) } +/// Take for [`PiecewiseSequence`] indices with unit multipliers: a piece of `length` consecutive +/// lists gathers `length * list_size` consecutive elements, so each piece scales by `list_size` +/// instead of expanding into one element run per list. +fn take_piecewise_fsl( + array: ArrayView<'_, FixedSizeList>, + indices: ArrayView<'_, PiecewiseSequence>, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { + return Ok(None); + }; + + let list_size = array.list_size() as usize; + let array_len = array.as_ref().len(); + let new_len = indices.as_ref().len(); + let elements_len = new_len.checked_mul(list_size).ok_or_else(|| { + vortex_err!( + "FixedSizeList take output length overflow: {new_len} lists of size {list_size}" + ) + })?; + + let starts: Vec = match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + starts + .as_slice::() + .iter() + .map(|&start| start.as_()) + .collect() + }); + let lengths: Vec = match lengths { + Columnar::Constant(lengths) => vec![constant_unsigned_usize(&lengths); starts.len()], + Columnar::Canonical(lengths) => { + let lengths = lengths.into_primitive(); + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + lengths + .as_slice::() + .iter() + .map(|&length| length.as_()) + .collect() + }) + } + }; + + let mut element_starts = BufferMut::::with_capacity(starts.len()); + let mut element_lengths = BufferMut::::with_capacity(starts.len()); + let mut total_len = 0usize; + for (&start, &length) in starts.iter().zip_eq(&lengths) { + if length == 0 { + continue; + } + let end = start + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + if end > array_len { + vortex_bail!(OutOfBounds: end - 1, 0, array_len); + } + total_len += length; + element_starts.push(u64::try_from(start * list_size)?); + element_lengths.push(u64::try_from(length * list_size)?); + } + vortex_ensure!( + total_len == new_len, + "PiecewiseSequenceArray expanded length {total_len} does not match declared length {new_len}" + ); + + let new_elements = + if let ([start], [length]) = (element_starts.as_slice(), element_lengths.as_slice()) { + let start = usize::try_from(*start)?; + array + .elements() + .slice(start..start + usize::try_from(*length)?)? + } else { + let run_count = element_starts.len(); + let starts = + PrimitiveArray::new(element_starts.freeze(), Validity::NonNullable).into_array(); + let lengths = + PrimitiveArray::new(element_lengths.freeze(), Validity::NonNullable).into_array(); + let multipliers = ConstantArray::new(1u64, run_count).into_array(); + // SAFETY: starts, lengths, and multipliers are non-nullable u64 arrays of equal length, + // and the piece lengths were validated to sum to `new_len`, so the element runs sum to + // `elements_len`. + let element_indices = unsafe { + PiecewiseSequenceArray::new_unchecked(starts, lengths, multipliers, elements_len) + }; + array.elements().take(element_indices.into_array())? + }; + let new_validity = array.validity()?.take(indices.as_ref())?; + + // SAFETY: `new_elements` has `new_len * list_size` elements, and `Validity::take` produces + // validity for `new_len`. + unsafe { + FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) + } + .into_array() + .optimize_ctx(ctx.session()) + .map(Some) +} + fn take_non_empty_degenerate_fsl( array: ArrayView<'_, FixedSizeList>, indices: &ArrayRef, diff --git a/vortex-array/src/arrays/fixed_size_list/tests/take.rs b/vortex-array/src/arrays/fixed_size_list/tests/take.rs index 93262a74b37..7f80e2ec7ed 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -2,7 +2,9 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use rstest::rstest; +use vortex_buffer::Buffer; use vortex_buffer::buffer; +use vortex_error::VortexResult; use super::common::create_basic_fsl; use super::common::create_empty_fsl; @@ -10,10 +12,13 @@ use super::common::create_large_fsl; use super::common::create_nullable_fsl; use super::common::create_single_element_fsl; use crate::ArrayRef; +use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; +use crate::arrays::ConstantArray; use crate::arrays::FixedSizeListArray; +use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; @@ -169,6 +174,85 @@ fn test_take_fsl_null_index_ignores_out_of_bounds_payload() { assert_arrays_eq!(expected, result, &mut ctx); } +fn contiguous_pieces(starts: &[u64], lengths: &[u64]) -> VortexResult { + let len = usize::try_from(lengths.iter().sum::())?; + Ok(PiecewiseSequenceArray::try_new( + starts.iter().copied().collect::>().into_array(), + lengths + .iter() + .copied() + .collect::>() + .into_array(), + ConstantArray::new(1u64, starts.len()).into_array(), + len, + )? + .into_array()) +} + +fn nullable_fsl_i32() -> FixedSizeListArray { + FixedSizeListArray::new( + PrimitiveArray::from_iter(0i32..12).into_array(), + 3, + Validity::from_iter([true, false, true, true]), + 4, + ) +} + +// Piecewise indices with unit multipliers scale to element runs instead of expanding one run per +// list; results must match the equivalent expanded primitive indices. +#[rstest] +#[case::multi_piece(&[2, 0], &[2, 1], &[2u64, 3, 0])] +#[case::single_piece_slices(&[1], &[3], &[1u64, 2, 3])] +#[case::zero_length_piece(&[3, 0, 1], &[0, 1, 2], &[0u64, 1, 2])] +fn test_take_piecewise( + #[case] starts: &[u64], + #[case] lengths: &[u64], + #[case] expanded: &[u64], +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let fsl = nullable_fsl_i32(); + + let result = fsl.take(contiguous_pieces(starts, lengths)?)?; + let expected = fsl.take(Buffer::from(expanded.to_vec()).into_array())?; + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_take_piecewise_out_of_bounds() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let fsl = nullable_fsl_i32(); + + let indices = contiguous_pieces(&[2], &[3])?; + let result = fsl + .take(indices) + .and_then(|taken| taken.execute::(&mut ctx)); + + assert!(result.is_err(), "expected out-of-bounds error"); + Ok(()) +} + +#[test] +fn test_take_piecewise_non_unit_multiplier() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let fsl = nullable_fsl_i32(); + + // Multiplier 2 falls back to the generic path but must stay correct. + let indices = PiecewiseSequenceArray::try_new( + buffer![0u64].into_array(), + buffer![2u64].into_array(), + buffer![2u64].into_array(), + 2, + )? + .into_array(); + let result = fsl.take(indices)?; + let expected = fsl.take(buffer![0u64, 2].into_array())?; + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) +} + // Element index overflow: with u8 indices and list_size=16, data_idx=16 produces element index // 16*16=256 which overflows u8. The take kernel must widen the element index type. #[rstest] diff --git a/vortex-btrblocks/benches/take_fsl.rs b/vortex-btrblocks/benches/take_fsl.rs index eb3d612b0ca..6e6ed8ec5fe 100644 --- a/vortex-btrblocks/benches/take_fsl.rs +++ b/vortex-btrblocks/benches/take_fsl.rs @@ -18,10 +18,12 @@ use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::array_session; use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeList; use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::PiecewiseSequenceArray; use vortex_array::arrays::Primitive; -use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; +use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressor; use vortex_buffer::Buffer; @@ -80,9 +82,22 @@ fn create_chunked_compressed_fsl(list_size: usize) -> ArrayRef { ChunkedArray::try_new(chunks, dtype).unwrap().into_array() } -fn bench_take(bencher: Bencher, num_indices: usize) { +fn create_piecewise_indices(num_indices: usize) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(42); + let num_pieces = 8; + let piece_len = num_indices / num_pieces; + let starts: Buffer = (0..num_pieces) + .map(|_| rng.random_range(0..NUM_LISTS - piece_len) as u64) + .collect(); + let lengths = ConstantArray::new(piece_len as u64, num_pieces).into_array(); + let multipliers = ConstantArray::new(1u64, num_pieces).into_array(); + PiecewiseSequenceArray::try_new(starts.into_array(), lengths, multipliers, num_indices) + .unwrap() + .into_array() +} + +fn bench_take(bencher: Bencher, num_indices: usize, indices: ArrayRef) { let array = create_chunked_compressed_fsl(LIST_SIZE); - let indices = create_random_indices(num_indices); bencher .counter(BytesCount::of_many::(num_indices * LIST_SIZE)) @@ -98,5 +113,13 @@ fn bench_take(bencher: Bencher, num_indices: usize) { #[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] fn take_chunked_compressed_fsl(bencher: Bencher, num_indices: usize) { - bench_take::(bencher, num_indices); + bench_take::(bencher, num_indices, create_random_indices(num_indices)); +} + +#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] +fn take_chunked_compressed_fsl_piecewise( + bencher: Bencher, + num_indices: usize, +) { + bench_take::(bencher, num_indices, create_piecewise_indices(num_indices)); }