From 98e82f4c69c74aba9c39318d788fad71d39cd9de Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 16 Jul 2026 13:29:50 -0400 Subject: [PATCH 01/37] Add PiecewiseSequential index array Signed-off-by: Daniel King --- vortex-array/src/arrays/mod.rs | 4 + .../src/arrays/piecewise_sequential/array.rs | 58 ++++ .../src/arrays/piecewise_sequential/mod.rs | 104 +++++++ .../src/arrays/piecewise_sequential/tests.rs | 129 +++++++++ .../src/arrays/piecewise_sequential/vtable.rs | 255 ++++++++++++++++++ vortex-array/src/session/mod.rs | 2 + 6 files changed, 552 insertions(+) create mode 100644 vortex-array/src/arrays/piecewise_sequential/array.rs create mode 100644 vortex-array/src/arrays/piecewise_sequential/mod.rs create mode 100644 vortex-array/src/arrays/piecewise_sequential/tests.rs create mode 100644 vortex-array/src/arrays/piecewise_sequential/vtable.rs diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index 15b6b5aa6be..610f89bfdb8 100644 --- a/vortex-array/src/arrays/mod.rs +++ b/vortex-array/src/arrays/mod.rs @@ -88,6 +88,10 @@ pub mod patched; pub use patched::Patched; pub use patched::PatchedArray; +pub mod piecewise_sequential; +pub use piecewise_sequential::PiecewiseSequential; +pub use piecewise_sequential::PiecewiseSequentialArray; + pub mod primitive; pub use primitive::Primitive; pub use primitive::PrimitiveArray; diff --git a/vortex-array/src/arrays/piecewise_sequential/array.rs b/vortex-array/src/arrays/piecewise_sequential/array.rs new file mode 100644 index 00000000000..21d4ae8c586 --- /dev/null +++ b/vortex-array/src/arrays/piecewise_sequential/array.rs @@ -0,0 +1,58 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use smallvec::smallvec; +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::array::Array; +use crate::array::ArrayParts; +use crate::array::EmptyArrayData; +use crate::array::TypedArrayRef; +use crate::array_slots; +use crate::arrays::PiecewiseSequential; +use crate::dtype::PType; + +#[array_slots(PiecewiseSequential)] +pub struct PiecewiseSequentialSlots { + /// The inclusive start index of each sequential piece. + pub starts: ArrayRef, + /// The length of each sequential piece. + pub lengths: ArrayRef, +} + +/// Extension methods for [`PiecewiseSequentialArray`]. +pub trait PiecewiseSequentialArrayExt: + TypedArrayRef + PiecewiseSequentialArraySlotsExt +{ +} +impl> PiecewiseSequentialArrayExt for T {} + +impl Array { + /// Constructs a new `PiecewiseSequentialArray` from start and length arrays. + /// + /// This validates only structural invariants: both children must be non-nullable unsigned + /// integer arrays with matching lengths, and the outer array length is the declared expanded + /// index length. Individual ranges are checked when the index array is executed or consumed by + /// a take implementation. + pub fn try_new(starts: ArrayRef, lengths: ArrayRef, len: usize) -> VortexResult { + Array::try_from_parts( + ArrayParts::new(PiecewiseSequential, PType::U64.into(), len, EmptyArrayData) + .with_slots(smallvec![Some(starts), Some(lengths)]), + ) + } + + /// Constructs a new `PiecewiseSequentialArray` without validation. + /// + /// # Safety + /// + /// The caller must guarantee the same structural invariants as [`Self::try_new`]. + pub unsafe fn new_unchecked(starts: ArrayRef, lengths: ArrayRef, len: usize) -> Self { + unsafe { + Array::from_parts_unchecked( + ArrayParts::new(PiecewiseSequential, PType::U64.into(), len, EmptyArrayData) + .with_slots(smallvec![Some(starts), Some(lengths)]), + ) + } + } +} diff --git a/vortex-array/src/arrays/piecewise_sequential/mod.rs b/vortex-array/src/arrays/piecewise_sequential/mod.rs new file mode 100644 index 00000000000..1f60b2abde5 --- /dev/null +++ b/vortex-array/src/arrays/piecewise_sequential/mod.rs @@ -0,0 +1,104 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Index encoding for concatenated sequential ranges. +//! +//! A `PiecewiseSequentialArray` represents the expanded index sequence +//! `starts[i]..starts[i] + lengths[i]` for each piece `i`. It is intended for take operations that +//! can gather contiguous runs without materializing one index per element. + +use itertools::Itertools; +use num_traits::AsPrimitive; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; + +use crate::ArrayRef; +use crate::arrays::PrimitiveArray; +use crate::dtype::UnsignedPType; + +pub mod array; +mod vtable; + +#[cfg(test)] +mod tests; + +pub use array::PiecewiseSequentialArrayExt; +pub use vtable::*; + +pub(crate) fn check_index_arrays(starts: &ArrayRef, lengths: &ArrayRef) -> VortexResult<()> { + check_index_array("starts", starts)?; + check_index_array("lengths", lengths)?; + vortex_ensure!( + starts.len() == lengths.len(), + "PiecewiseSequentialArray starts length {} does not match lengths length {}", + starts.len(), + lengths.len() + ); + Ok(()) +} + +fn check_index_array(name: &str, array: &ArrayRef) -> VortexResult<()> { + vortex_ensure!( + array.dtype().is_unsigned_int(), + "PiecewiseSequentialArray {name} must have unsigned integer dtype, got {}", + array.dtype() + ); + vortex_ensure!( + !array.dtype().is_nullable(), + "PiecewiseSequentialArray {name} must be non-nullable, got {}", + array.dtype() + ); + Ok(()) +} + +#[inline] +pub(crate) fn index_value_to_usize(value: T) -> usize { + value.as_() +} + +#[inline] +pub(crate) fn index_value_to_u64>(value: T) -> u64 { + value.as_() +} + +#[inline] +pub(crate) fn checked_range_end(start: u64, length: usize) -> VortexResult { + start + .checked_add(length as u64) + .ok_or_else(|| vortex_error::vortex_err!("PiecewiseSequentialArray range overflows u64")) +} + +pub(crate) fn materialize_ranges( + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + output_len: usize, +) -> VortexResult> +where + S: UnsignedPType + AsPrimitive, + L: UnsignedPType, +{ + let starts = starts.as_slice::(); + let lengths = lengths.as_slice::(); + let mut values = BufferMut::with_capacity(output_len); + let mut computed_len = 0usize; + + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = index_value_to_u64(start); + let length = index_value_to_usize(length); + checked_range_end(start, length)?; + computed_len = computed_len.checked_add(length).ok_or_else(|| { + vortex_error::vortex_err!("PiecewiseSequentialArray output length overflows usize") + })?; + + values.extend((0..length).map(|offset| start + offset as u64)); + } + + if computed_len != output_len { + vortex_bail!( + "PiecewiseSequentialArray expanded length {computed_len} does not match declared length {output_len}" + ); + } + Ok(values) +} diff --git a/vortex-array/src/arrays/piecewise_sequential/tests.rs b/vortex-array/src/arrays/piecewise_sequential/tests.rs new file mode 100644 index 00000000000..2cab62fd676 --- /dev/null +++ b/vortex-array/src/arrays/piecewise_sequential/tests.rs @@ -0,0 +1,129 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::ByteBufferMut; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_session::registry::ReadContext; + +use crate::ArrayContext; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::ConstantArray; +use crate::arrays::PiecewiseSequential; +use crate::arrays::PiecewiseSequentialArray; +use crate::arrays::PrimitiveArray; +use crate::assert_arrays_eq; +use crate::serde::SerializeOptions; +use crate::serde::SerializedArray; + +#[test] +fn materializes_piecewise_indices() -> VortexResult<()> { + let starts = buffer![3u64, 15, 21].into_array(); + let lengths = buffer![3u64, 3, 3].into_array(); + let array = PiecewiseSequentialArray::try_new(starts, lengths, 9)?.into_array(); + + let expected = PrimitiveArray::from_iter([3u64, 4, 5, 15, 16, 17, 21, 22, 23]).into_array(); + assert_arrays_eq!(array, expected, &mut array_session().create_execution_ctx()); + Ok(()) +} + +#[test] +fn materializes_repeated_and_empty_ranges() -> VortexResult<()> { + let starts = buffer![5u64, 2, 5].into_array(); + let lengths = buffer![2u64, 0, 2].into_array(); + let array = PiecewiseSequentialArray::try_new(starts, lengths, 4)?.into_array(); + + let expected = PrimitiveArray::from_iter([5u64, 6, 5, 6]).into_array(); + assert_arrays_eq!(array, expected, &mut array_session().create_execution_ctx()); + Ok(()) +} + +#[test] +fn supports_constant_lengths() -> VortexResult<()> { + let starts = buffer![0u64, 10, 20].into_array(); + let lengths = ConstantArray::new(2u64, 3).into_array(); + let array = PiecewiseSequentialArray::try_new(starts, lengths, 6)?.into_array(); + + let expected = PrimitiveArray::from_iter([0u64, 1, 10, 11, 20, 21]).into_array(); + assert_arrays_eq!(array, expected, &mut array_session().create_execution_ctx()); + Ok(()) +} + +#[test] +fn scalar_at_maps_into_piece() -> VortexResult<()> { + let starts = buffer![3u64, 15, 21].into_array(); + let lengths = buffer![3u64, 3, 3].into_array(); + let array = PiecewiseSequentialArray::try_new(starts, lengths, 9)?.into_array(); + let mut ctx = array_session().create_execution_ctx(); + + assert_eq!(array.execute_scalar(0, &mut ctx)?, 3u64.into()); + assert_eq!(array.execute_scalar(4, &mut ctx)?, 16u64.into()); + assert_eq!(array.execute_scalar(8, &mut ctx)?, 23u64.into()); + Ok(()) +} + +#[test] +fn constructor_defers_range_value_validation() -> VortexResult<()> { + let starts = buffer![u64::MAX].into_array(); + let lengths = buffer![2u64].into_array(); + let array = PiecewiseSequentialArray::try_new(starts, lengths, 2)?.into_array(); + + assert!( + array + .execute::(&mut array_session().create_execution_ctx()) + .is_err() + ); + Ok(()) +} + +#[test] +fn execution_checks_declared_length() -> VortexResult<()> { + let starts = buffer![0u64, 3].into_array(); + let lengths = buffer![2u64, 2].into_array(); + let array = PiecewiseSequentialArray::try_new(starts, lengths, 3)?.into_array(); + + assert!( + array + .execute::(&mut array_session().create_execution_ctx()) + .is_err() + ); + Ok(()) +} + +#[test] +fn serde_roundtrip_preserves_piecewise_indices() -> VortexResult<()> { + let array = PiecewiseSequentialArray::try_new( + buffer![3u32, 15, 21].into_array(), + buffer![2u16, 0, 2].into_array(), + 4, + )? + .into_array(); + let dtype = array.dtype().clone(); + let len = array.len(); + + let array_ctx = ArrayContext::empty(); + let serialized = array.serialize(&array_ctx, &array_session(), &SerializeOptions::default())?; + + let mut concat = ByteBufferMut::empty(); + for buffer in serialized { + concat.extend_from_slice(buffer.as_ref()); + } + + let parts = SerializedArray::try_from(concat.freeze())?; + let decoded = parts.decode( + &dtype, + len, + &ReadContext::new(array_ctx.to_ids()), + &array_session(), + )?; + + assert!(decoded.is::()); + assert_arrays_eq!( + decoded, + PrimitiveArray::from_iter([3u64, 4, 21, 22]).into_array(), + &mut array_session().create_execution_ctx() + ); + Ok(()) +} diff --git a/vortex-array/src/arrays/piecewise_sequential/vtable.rs b/vortex-array/src/arrays/piecewise_sequential/vtable.rs new file mode 100644 index 00000000000..e2809cb683e --- /dev/null +++ b/vortex-array/src/arrays/piecewise_sequential/vtable.rs @@ -0,0 +1,255 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use itertools::Itertools; +use prost::Message; +use smallvec::smallvec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayParts; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::ExecutionResult; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayId; +use crate::array::ArrayView; +use crate::array::EmptyArrayData; +use crate::array::OperationsVTable; +use crate::array::VTable; +use crate::array::ValidityVTable; +use crate::array::with_empty_buffers; +use crate::arrays::PrimitiveArray; +use crate::arrays::piecewise_sequential::array::PiecewiseSequentialArraySlotsExt; +use crate::arrays::piecewise_sequential::array::PiecewiseSequentialSlots; +use crate::arrays::piecewise_sequential::check_index_arrays; +use crate::arrays::piecewise_sequential::index_value_to_u64; +use crate::arrays::piecewise_sequential::index_value_to_usize; +use crate::arrays::piecewise_sequential::materialize_ranges; +use crate::arrays::primitive::PrimitiveArrayExt; +use crate::buffer::BufferHandle; +use crate::dtype::DType; +use crate::dtype::PType; +use crate::match_each_unsigned_integer_ptype; +use crate::scalar::Scalar; +use crate::serde::ArrayChildren; +use crate::validity::Validity; + +/// A [`PiecewiseSequential`]-encoded Vortex index array. +pub type PiecewiseSequentialArray = Array; + +#[derive(Clone, prost::Message)] +struct PiecewiseSequentialMetadata { + #[prost(uint64, tag = "1")] + num_pieces: u64, + #[prost(enumeration = "PType", tag = "2")] + starts_ptype: i32, + #[prost(enumeration = "PType", tag = "3")] + lengths_ptype: i32, +} + +#[derive(Clone, Debug)] +pub struct PiecewiseSequential; + +impl VTable for PiecewiseSequential { + type TypedArrayData = EmptyArrayData; + type OperationsVTable = Self; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.piecewise-sequential"); + *ID + } + + fn validate( + &self, + _data: &Self::TypedArrayData, + dtype: &DType, + _len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure!( + dtype == &DType::from(PType::U64), + "PiecewiseSequentialArray dtype must be u64, got {dtype}" + ); + vortex_ensure!( + slots.len() == PiecewiseSequentialSlots::NAMES.len(), + "PiecewiseSequentialArray requires {} slots, got {}", + PiecewiseSequentialSlots::NAMES.len(), + slots.len() + ); + let starts = slots[PiecewiseSequentialSlots::STARTS] + .as_ref() + .ok_or_else(|| { + vortex_error::vortex_err!("PiecewiseSequentialArray starts slot must be present") + })?; + let lengths = slots[PiecewiseSequentialSlots::LENGTHS] + .as_ref() + .ok_or_else(|| { + vortex_error::vortex_err!("PiecewiseSequentialArray lengths slot must be present") + })?; + check_index_arrays(starts, lengths) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, _idx: usize) -> BufferHandle { + vortex_panic!("PiecewiseSequentialArray has no buffers") + } + + fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option { + None + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + PiecewiseSequentialSlots::NAMES[idx].to_string() + } + + fn serialize( + array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(Some( + PiecewiseSequentialMetadata { + num_pieces: u64::try_from(array.starts().len()).map_err(|_| { + vortex_err!( + "PiecewiseSequentialArray piece count {} overflowed u64", + array.starts().len() + ) + })?, + starts_ptype: PType::try_from(array.starts().dtype())? as i32, + lengths_ptype: PType::try_from(array.lengths().dtype())? as i32, + } + .encode_to_vec(), + )) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + let metadata = PiecewiseSequentialMetadata::decode(metadata)?; + vortex_ensure!( + dtype == &DType::from(PType::U64), + "PiecewiseSequentialArray dtype must be u64, got {dtype}" + ); + vortex_ensure!( + buffers.is_empty(), + "PiecewiseSequentialArray expects no buffers, got {}", + buffers.len() + ); + vortex_ensure!( + children.len() == PiecewiseSequentialSlots::NAMES.len(), + "PiecewiseSequentialArray expects {} children, got {}", + PiecewiseSequentialSlots::NAMES.len(), + children.len() + ); + + let num_pieces = usize::try_from(metadata.num_pieces).map_err(|_| { + vortex_err!( + "PiecewiseSequentialArray piece count {} does not fit in usize", + metadata.num_pieces + ) + })?; + let starts = children.get( + PiecewiseSequentialSlots::STARTS, + &metadata.starts_ptype().into(), + num_pieces, + )?; + let lengths = children.get( + PiecewiseSequentialSlots::LENGTHS, + &metadata.lengths_ptype().into(), + num_pieces, + )?; + + Ok( + ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData) + .with_slots(smallvec![Some(starts), Some(lengths)]), + ) + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let starts = array.starts().clone().execute::(ctx)?; + let lengths = array.lengths().clone().execute::(ctx)?; + check_index_arrays(starts.as_ref(), lengths.as_ref())?; + + let values = match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + materialize_ranges::(&starts, &lengths, array.len())? + }) + }); + Ok(ExecutionResult::done( + PrimitiveArray::new(values.freeze(), Validity::NonNullable).into_array(), + )) + } +} + +impl OperationsVTable for PiecewiseSequential { + fn scalar_at( + array: ArrayView<'_, PiecewiseSequential>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let starts = array.starts().clone().execute::(ctx)?; + let lengths = array.lengths().clone().execute::(ctx)?; + check_index_arrays(starts.as_ref(), lengths.as_ref())?; + + let value = match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + scalar_at::(&starts, &lengths, index)? + }) + }); + Ok(value.into()) + } +} + +impl ValidityVTable for PiecewiseSequential { + fn validity(_array: ArrayView<'_, PiecewiseSequential>) -> VortexResult { + Ok(Validity::NonNullable) + } +} + +fn scalar_at( + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + index: usize, +) -> VortexResult +where + S: crate::dtype::UnsignedPType + num_traits::AsPrimitive, + L: crate::dtype::UnsignedPType, +{ + let mut remaining = index; + for (&start, &length) in starts + .as_slice::() + .iter() + .zip_eq(lengths.as_slice::()) + { + let length = index_value_to_usize(length); + if remaining < length { + return Ok(index_value_to_u64(start) + remaining as u64); + } + remaining -= length; + } + vortex_bail!("PiecewiseSequentialArray index {index} out of bounds") +} diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index 2cb8414cbfd..6f5c9f9e03f 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -25,6 +25,7 @@ use crate::arrays::List; use crate::arrays::ListView; use crate::arrays::Masked; use crate::arrays::Null; +use crate::arrays::PiecewiseSequential; use crate::arrays::Primitive; use crate::arrays::Struct; use crate::arrays::VarBin; @@ -81,6 +82,7 @@ impl Default for ArraySession { this.register(Dict); this.register(List); this.register(Masked); + this.register(PiecewiseSequential); this.register(VarBin); this From 40fe633615ff99b2b54d9dadc32ad886b9bc63da Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 16 Jul 2026 15:08:54 -0400 Subject: [PATCH 02/37] Rename PiecewiseSequential to PiecewiseSequence Signed-off-by: Daniel King --- vortex-array/src/arrays/mod.rs | 6 +- .../array.rs | 24 +++--- .../mod.rs | 16 ++-- .../tests.rs | 38 ++++++---- .../vtable.rs | 76 +++++++++---------- vortex-array/src/session/mod.rs | 4 +- 6 files changed, 85 insertions(+), 79 deletions(-) rename vortex-array/src/arrays/{piecewise_sequential => piecewise_sequence}/array.rs (66%) rename vortex-array/src/arrays/{piecewise_sequential => piecewise_sequence}/mod.rs (80%) rename vortex-array/src/arrays/{piecewise_sequential => piecewise_sequence}/tests.rs (74%) rename vortex-array/src/arrays/{piecewise_sequential => piecewise_sequence}/vtable.rs (72%) diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index 610f89bfdb8..001e20dee65 100644 --- a/vortex-array/src/arrays/mod.rs +++ b/vortex-array/src/arrays/mod.rs @@ -88,9 +88,9 @@ pub mod patched; pub use patched::Patched; pub use patched::PatchedArray; -pub mod piecewise_sequential; -pub use piecewise_sequential::PiecewiseSequential; -pub use piecewise_sequential::PiecewiseSequentialArray; +pub mod piecewise_sequence; +pub use piecewise_sequence::PiecewiseSequence; +pub use piecewise_sequence::PiecewiseSequenceArray; pub mod primitive; pub use primitive::Primitive; diff --git a/vortex-array/src/arrays/piecewise_sequential/array.rs b/vortex-array/src/arrays/piecewise_sequence/array.rs similarity index 66% rename from vortex-array/src/arrays/piecewise_sequential/array.rs rename to vortex-array/src/arrays/piecewise_sequence/array.rs index 21d4ae8c586..d41e6ff171f 100644 --- a/vortex-array/src/arrays/piecewise_sequential/array.rs +++ b/vortex-array/src/arrays/piecewise_sequence/array.rs @@ -10,26 +10,26 @@ use crate::array::ArrayParts; use crate::array::EmptyArrayData; use crate::array::TypedArrayRef; use crate::array_slots; -use crate::arrays::PiecewiseSequential; +use crate::arrays::PiecewiseSequence; use crate::dtype::PType; -#[array_slots(PiecewiseSequential)] -pub struct PiecewiseSequentialSlots { +#[array_slots(PiecewiseSequence)] +pub struct PiecewiseSequenceSlots { /// The inclusive start index of each sequential piece. pub starts: ArrayRef, /// The length of each sequential piece. pub lengths: ArrayRef, } -/// Extension methods for [`PiecewiseSequentialArray`]. -pub trait PiecewiseSequentialArrayExt: - TypedArrayRef + PiecewiseSequentialArraySlotsExt +/// Extension methods for [`PiecewiseSequenceArray`]. +pub trait PiecewiseSequenceArrayExt: + TypedArrayRef + PiecewiseSequenceArraySlotsExt { } -impl> PiecewiseSequentialArrayExt for T {} +impl> PiecewiseSequenceArrayExt for T {} -impl Array { - /// Constructs a new `PiecewiseSequentialArray` from start and length arrays. +impl Array { + /// Constructs a new `PiecewiseSequenceArray` from start and length arrays. /// /// This validates only structural invariants: both children must be non-nullable unsigned /// integer arrays with matching lengths, and the outer array length is the declared expanded @@ -37,12 +37,12 @@ impl Array { /// a take implementation. pub fn try_new(starts: ArrayRef, lengths: ArrayRef, len: usize) -> VortexResult { Array::try_from_parts( - ArrayParts::new(PiecewiseSequential, PType::U64.into(), len, EmptyArrayData) + ArrayParts::new(PiecewiseSequence, PType::U64.into(), len, EmptyArrayData) .with_slots(smallvec![Some(starts), Some(lengths)]), ) } - /// Constructs a new `PiecewiseSequentialArray` without validation. + /// Constructs a new `PiecewiseSequenceArray` without validation. /// /// # Safety /// @@ -50,7 +50,7 @@ impl Array { pub unsafe fn new_unchecked(starts: ArrayRef, lengths: ArrayRef, len: usize) -> Self { unsafe { Array::from_parts_unchecked( - ArrayParts::new(PiecewiseSequential, PType::U64.into(), len, EmptyArrayData) + ArrayParts::new(PiecewiseSequence, PType::U64.into(), len, EmptyArrayData) .with_slots(smallvec![Some(starts), Some(lengths)]), ) } diff --git a/vortex-array/src/arrays/piecewise_sequential/mod.rs b/vortex-array/src/arrays/piecewise_sequence/mod.rs similarity index 80% rename from vortex-array/src/arrays/piecewise_sequential/mod.rs rename to vortex-array/src/arrays/piecewise_sequence/mod.rs index 1f60b2abde5..974c395c0d1 100644 --- a/vortex-array/src/arrays/piecewise_sequential/mod.rs +++ b/vortex-array/src/arrays/piecewise_sequence/mod.rs @@ -3,7 +3,7 @@ //! Index encoding for concatenated sequential ranges. //! -//! A `PiecewiseSequentialArray` represents the expanded index sequence +//! A `PiecewiseSequenceArray` represents the expanded index sequence //! `starts[i]..starts[i] + lengths[i]` for each piece `i`. It is intended for take operations that //! can gather contiguous runs without materializing one index per element. @@ -24,7 +24,7 @@ mod vtable; #[cfg(test)] mod tests; -pub use array::PiecewiseSequentialArrayExt; +pub use array::PiecewiseSequenceArrayExt; pub use vtable::*; pub(crate) fn check_index_arrays(starts: &ArrayRef, lengths: &ArrayRef) -> VortexResult<()> { @@ -32,7 +32,7 @@ pub(crate) fn check_index_arrays(starts: &ArrayRef, lengths: &ArrayRef) -> Vorte check_index_array("lengths", lengths)?; vortex_ensure!( starts.len() == lengths.len(), - "PiecewiseSequentialArray starts length {} does not match lengths length {}", + "PiecewiseSequenceArray starts length {} does not match lengths length {}", starts.len(), lengths.len() ); @@ -42,12 +42,12 @@ pub(crate) fn check_index_arrays(starts: &ArrayRef, lengths: &ArrayRef) -> Vorte fn check_index_array(name: &str, array: &ArrayRef) -> VortexResult<()> { vortex_ensure!( array.dtype().is_unsigned_int(), - "PiecewiseSequentialArray {name} must have unsigned integer dtype, got {}", + "PiecewiseSequenceArray {name} must have unsigned integer dtype, got {}", array.dtype() ); vortex_ensure!( !array.dtype().is_nullable(), - "PiecewiseSequentialArray {name} must be non-nullable, got {}", + "PiecewiseSequenceArray {name} must be non-nullable, got {}", array.dtype() ); Ok(()) @@ -67,7 +67,7 @@ pub(crate) fn index_value_to_u64>(value: T) pub(crate) fn checked_range_end(start: u64, length: usize) -> VortexResult { start .checked_add(length as u64) - .ok_or_else(|| vortex_error::vortex_err!("PiecewiseSequentialArray range overflows u64")) + .ok_or_else(|| vortex_error::vortex_err!("PiecewiseSequenceArray range overflows u64")) } pub(crate) fn materialize_ranges( @@ -89,7 +89,7 @@ where let length = index_value_to_usize(length); checked_range_end(start, length)?; computed_len = computed_len.checked_add(length).ok_or_else(|| { - vortex_error::vortex_err!("PiecewiseSequentialArray output length overflows usize") + vortex_error::vortex_err!("PiecewiseSequenceArray output length overflows usize") })?; values.extend((0..length).map(|offset| start + offset as u64)); @@ -97,7 +97,7 @@ where if computed_len != output_len { vortex_bail!( - "PiecewiseSequentialArray expanded length {computed_len} does not match declared length {output_len}" + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" ); } Ok(values) diff --git a/vortex-array/src/arrays/piecewise_sequential/tests.rs b/vortex-array/src/arrays/piecewise_sequence/tests.rs similarity index 74% rename from vortex-array/src/arrays/piecewise_sequential/tests.rs rename to vortex-array/src/arrays/piecewise_sequence/tests.rs index 2cab62fd676..2f9992b4d05 100644 --- a/vortex-array/src/arrays/piecewise_sequential/tests.rs +++ b/vortex-array/src/arrays/piecewise_sequence/tests.rs @@ -11,8 +11,8 @@ use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::ConstantArray; -use crate::arrays::PiecewiseSequential; -use crate::arrays::PiecewiseSequentialArray; +use crate::arrays::PiecewiseSequence; +use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::serde::SerializeOptions; @@ -22,7 +22,7 @@ use crate::serde::SerializedArray; fn materializes_piecewise_indices() -> VortexResult<()> { let starts = buffer![3u64, 15, 21].into_array(); let lengths = buffer![3u64, 3, 3].into_array(); - let array = PiecewiseSequentialArray::try_new(starts, lengths, 9)?.into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, 9)?.into_array(); let expected = PrimitiveArray::from_iter([3u64, 4, 5, 15, 16, 17, 21, 22, 23]).into_array(); assert_arrays_eq!(array, expected, &mut array_session().create_execution_ctx()); @@ -33,7 +33,7 @@ fn materializes_piecewise_indices() -> VortexResult<()> { fn materializes_repeated_and_empty_ranges() -> VortexResult<()> { let starts = buffer![5u64, 2, 5].into_array(); let lengths = buffer![2u64, 0, 2].into_array(); - let array = PiecewiseSequentialArray::try_new(starts, lengths, 4)?.into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, 4)?.into_array(); let expected = PrimitiveArray::from_iter([5u64, 6, 5, 6]).into_array(); assert_arrays_eq!(array, expected, &mut array_session().create_execution_ctx()); @@ -44,7 +44,7 @@ fn materializes_repeated_and_empty_ranges() -> VortexResult<()> { fn supports_constant_lengths() -> VortexResult<()> { let starts = buffer![0u64, 10, 20].into_array(); let lengths = ConstantArray::new(2u64, 3).into_array(); - let array = PiecewiseSequentialArray::try_new(starts, lengths, 6)?.into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, 6)?.into_array(); let expected = PrimitiveArray::from_iter([0u64, 1, 10, 11, 20, 21]).into_array(); assert_arrays_eq!(array, expected, &mut array_session().create_execution_ctx()); @@ -55,7 +55,7 @@ fn supports_constant_lengths() -> VortexResult<()> { fn scalar_at_maps_into_piece() -> VortexResult<()> { let starts = buffer![3u64, 15, 21].into_array(); let lengths = buffer![3u64, 3, 3].into_array(); - let array = PiecewiseSequentialArray::try_new(starts, lengths, 9)?.into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, 9)?.into_array(); let mut ctx = array_session().create_execution_ctx(); assert_eq!(array.execute_scalar(0, &mut ctx)?, 3u64.into()); @@ -68,12 +68,15 @@ fn scalar_at_maps_into_piece() -> VortexResult<()> { fn constructor_defers_range_value_validation() -> VortexResult<()> { let starts = buffer![u64::MAX].into_array(); let lengths = buffer![2u64].into_array(); - let array = PiecewiseSequentialArray::try_new(starts, lengths, 2)?.into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, 2)?.into_array(); + let err = array + .execute::(&mut array_session().create_execution_ctx()) + .unwrap_err(); assert!( - array - .execute::(&mut array_session().create_execution_ctx()) - .is_err() + err.to_string() + .contains("PiecewiseSequenceArray range overflows u64"), + "{err}" ); Ok(()) } @@ -82,19 +85,22 @@ fn constructor_defers_range_value_validation() -> VortexResult<()> { fn execution_checks_declared_length() -> VortexResult<()> { let starts = buffer![0u64, 3].into_array(); let lengths = buffer![2u64, 2].into_array(); - let array = PiecewiseSequentialArray::try_new(starts, lengths, 3)?.into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, 3)?.into_array(); + let err = array + .execute::(&mut array_session().create_execution_ctx()) + .unwrap_err(); assert!( - array - .execute::(&mut array_session().create_execution_ctx()) - .is_err() + err.to_string() + .contains("PiecewiseSequenceArray expanded length 4 does not match declared length 3"), + "{err}" ); Ok(()) } #[test] fn serde_roundtrip_preserves_piecewise_indices() -> VortexResult<()> { - let array = PiecewiseSequentialArray::try_new( + let array = PiecewiseSequenceArray::try_new( buffer![3u32, 15, 21].into_array(), buffer![2u16, 0, 2].into_array(), 4, @@ -119,7 +125,7 @@ fn serde_roundtrip_preserves_piecewise_indices() -> VortexResult<()> { &array_session(), )?; - assert!(decoded.is::()); + assert!(decoded.is::()); assert_arrays_eq!( decoded, PrimitiveArray::from_iter([3u64, 4, 21, 22]).into_array(), diff --git a/vortex-array/src/arrays/piecewise_sequential/vtable.rs b/vortex-array/src/arrays/piecewise_sequence/vtable.rs similarity index 72% rename from vortex-array/src/arrays/piecewise_sequential/vtable.rs rename to vortex-array/src/arrays/piecewise_sequence/vtable.rs index e2809cb683e..53dba5c5879 100644 --- a/vortex-array/src/arrays/piecewise_sequential/vtable.rs +++ b/vortex-array/src/arrays/piecewise_sequence/vtable.rs @@ -26,12 +26,12 @@ use crate::array::VTable; use crate::array::ValidityVTable; use crate::array::with_empty_buffers; use crate::arrays::PrimitiveArray; -use crate::arrays::piecewise_sequential::array::PiecewiseSequentialArraySlotsExt; -use crate::arrays::piecewise_sequential::array::PiecewiseSequentialSlots; -use crate::arrays::piecewise_sequential::check_index_arrays; -use crate::arrays::piecewise_sequential::index_value_to_u64; -use crate::arrays::piecewise_sequential::index_value_to_usize; -use crate::arrays::piecewise_sequential::materialize_ranges; +use crate::arrays::piecewise_sequence::array::PiecewiseSequenceArraySlotsExt; +use crate::arrays::piecewise_sequence::array::PiecewiseSequenceSlots; +use crate::arrays::piecewise_sequence::check_index_arrays; +use crate::arrays::piecewise_sequence::index_value_to_u64; +use crate::arrays::piecewise_sequence::index_value_to_usize; +use crate::arrays::piecewise_sequence::materialize_ranges; use crate::arrays::primitive::PrimitiveArrayExt; use crate::buffer::BufferHandle; use crate::dtype::DType; @@ -41,11 +41,11 @@ use crate::scalar::Scalar; use crate::serde::ArrayChildren; use crate::validity::Validity; -/// A [`PiecewiseSequential`]-encoded Vortex index array. -pub type PiecewiseSequentialArray = Array; +/// A [`PiecewiseSequence`]-encoded Vortex index array. +pub type PiecewiseSequenceArray = Array; #[derive(Clone, prost::Message)] -struct PiecewiseSequentialMetadata { +struct PiecewiseSequenceMetadata { #[prost(uint64, tag = "1")] num_pieces: u64, #[prost(enumeration = "PType", tag = "2")] @@ -55,15 +55,15 @@ struct PiecewiseSequentialMetadata { } #[derive(Clone, Debug)] -pub struct PiecewiseSequential; +pub struct PiecewiseSequence; -impl VTable for PiecewiseSequential { +impl VTable for PiecewiseSequence { type TypedArrayData = EmptyArrayData; type OperationsVTable = Self; type ValidityVTable = Self; fn id(&self) -> ArrayId { - static ID: CachedId = CachedId::new("vortex.piecewise-sequential"); + static ID: CachedId = CachedId::new("vortex.piecewise-sequence"); *ID } @@ -76,23 +76,23 @@ impl VTable for PiecewiseSequential { ) -> VortexResult<()> { vortex_ensure!( dtype == &DType::from(PType::U64), - "PiecewiseSequentialArray dtype must be u64, got {dtype}" + "PiecewiseSequenceArray dtype must be u64, got {dtype}" ); vortex_ensure!( - slots.len() == PiecewiseSequentialSlots::NAMES.len(), - "PiecewiseSequentialArray requires {} slots, got {}", - PiecewiseSequentialSlots::NAMES.len(), + slots.len() == PiecewiseSequenceSlots::NAMES.len(), + "PiecewiseSequenceArray requires {} slots, got {}", + PiecewiseSequenceSlots::NAMES.len(), slots.len() ); - let starts = slots[PiecewiseSequentialSlots::STARTS] + let starts = slots[PiecewiseSequenceSlots::STARTS] .as_ref() .ok_or_else(|| { - vortex_error::vortex_err!("PiecewiseSequentialArray starts slot must be present") + vortex_error::vortex_err!("PiecewiseSequenceArray starts slot must be present") })?; - let lengths = slots[PiecewiseSequentialSlots::LENGTHS] + let lengths = slots[PiecewiseSequenceSlots::LENGTHS] .as_ref() .ok_or_else(|| { - vortex_error::vortex_err!("PiecewiseSequentialArray lengths slot must be present") + vortex_error::vortex_err!("PiecewiseSequenceArray lengths slot must be present") })?; check_index_arrays(starts, lengths) } @@ -102,7 +102,7 @@ impl VTable for PiecewiseSequential { } fn buffer(_array: ArrayView<'_, Self>, _idx: usize) -> BufferHandle { - vortex_panic!("PiecewiseSequentialArray has no buffers") + vortex_panic!("PiecewiseSequenceArray has no buffers") } fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option { @@ -118,7 +118,7 @@ impl VTable for PiecewiseSequential { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - PiecewiseSequentialSlots::NAMES[idx].to_string() + PiecewiseSequenceSlots::NAMES[idx].to_string() } fn serialize( @@ -126,10 +126,10 @@ impl VTable for PiecewiseSequential { _session: &VortexSession, ) -> VortexResult>> { Ok(Some( - PiecewiseSequentialMetadata { + PiecewiseSequenceMetadata { num_pieces: u64::try_from(array.starts().len()).map_err(|_| { vortex_err!( - "PiecewiseSequentialArray piece count {} overflowed u64", + "PiecewiseSequenceArray piece count {} overflowed u64", array.starts().len() ) })?, @@ -149,36 +149,36 @@ impl VTable for PiecewiseSequential { children: &dyn ArrayChildren, _session: &VortexSession, ) -> VortexResult> { - let metadata = PiecewiseSequentialMetadata::decode(metadata)?; + let metadata = PiecewiseSequenceMetadata::decode(metadata)?; vortex_ensure!( dtype == &DType::from(PType::U64), - "PiecewiseSequentialArray dtype must be u64, got {dtype}" + "PiecewiseSequenceArray dtype must be u64, got {dtype}" ); vortex_ensure!( buffers.is_empty(), - "PiecewiseSequentialArray expects no buffers, got {}", + "PiecewiseSequenceArray expects no buffers, got {}", buffers.len() ); vortex_ensure!( - children.len() == PiecewiseSequentialSlots::NAMES.len(), - "PiecewiseSequentialArray expects {} children, got {}", - PiecewiseSequentialSlots::NAMES.len(), + children.len() == PiecewiseSequenceSlots::NAMES.len(), + "PiecewiseSequenceArray expects {} children, got {}", + PiecewiseSequenceSlots::NAMES.len(), children.len() ); let num_pieces = usize::try_from(metadata.num_pieces).map_err(|_| { vortex_err!( - "PiecewiseSequentialArray piece count {} does not fit in usize", + "PiecewiseSequenceArray piece count {} does not fit in usize", metadata.num_pieces ) })?; let starts = children.get( - PiecewiseSequentialSlots::STARTS, + PiecewiseSequenceSlots::STARTS, &metadata.starts_ptype().into(), num_pieces, )?; let lengths = children.get( - PiecewiseSequentialSlots::LENGTHS, + PiecewiseSequenceSlots::LENGTHS, &metadata.lengths_ptype().into(), num_pieces, )?; @@ -205,9 +205,9 @@ impl VTable for PiecewiseSequential { } } -impl OperationsVTable for PiecewiseSequential { +impl OperationsVTable for PiecewiseSequence { fn scalar_at( - array: ArrayView<'_, PiecewiseSequential>, + array: ArrayView<'_, PiecewiseSequence>, index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { @@ -224,8 +224,8 @@ impl OperationsVTable for PiecewiseSequential { } } -impl ValidityVTable for PiecewiseSequential { - fn validity(_array: ArrayView<'_, PiecewiseSequential>) -> VortexResult { +impl ValidityVTable for PiecewiseSequence { + fn validity(_array: ArrayView<'_, PiecewiseSequence>) -> VortexResult { Ok(Validity::NonNullable) } } @@ -251,5 +251,5 @@ where } remaining -= length; } - vortex_bail!("PiecewiseSequentialArray index {index} out of bounds") + vortex_bail!("PiecewiseSequenceArray index {index} out of bounds") } diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index 6f5c9f9e03f..98a769ceea9 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -25,7 +25,7 @@ use crate::arrays::List; use crate::arrays::ListView; use crate::arrays::Masked; use crate::arrays::Null; -use crate::arrays::PiecewiseSequential; +use crate::arrays::PiecewiseSequence; use crate::arrays::Primitive; use crate::arrays::Struct; use crate::arrays::VarBin; @@ -82,7 +82,7 @@ impl Default for ArraySession { this.register(Dict); this.register(List); this.register(Masked); - this.register(PiecewiseSequential); + this.register(PiecewiseSequence); this.register(VarBin); this From 0e0fe567c57789dea32bff03b5493ed88fe681bf Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 16 Jul 2026 15:33:03 -0400 Subject: [PATCH 03/37] Add PiecewiseSequence multipliers Signed-off-by: Daniel King --- .../src/arrays/piecewise_sequence/array.rs | 24 ++++-- .../src/arrays/piecewise_sequence/mod.rs | 83 ++++++++++++------- .../src/arrays/piecewise_sequence/tests.rs | 35 ++++++-- .../src/arrays/piecewise_sequence/vtable.rs | 61 +++++++++----- 4 files changed, 137 insertions(+), 66 deletions(-) diff --git a/vortex-array/src/arrays/piecewise_sequence/array.rs b/vortex-array/src/arrays/piecewise_sequence/array.rs index d41e6ff171f..5c98cc15bb2 100644 --- a/vortex-array/src/arrays/piecewise_sequence/array.rs +++ b/vortex-array/src/arrays/piecewise_sequence/array.rs @@ -19,6 +19,8 @@ pub struct PiecewiseSequenceSlots { pub starts: ArrayRef, /// The length of each sequential piece. pub lengths: ArrayRef, + /// The distance between consecutive indices in each piece. + pub multipliers: ArrayRef, } /// Extension methods for [`PiecewiseSequenceArray`]. @@ -29,16 +31,21 @@ pub trait PiecewiseSequenceArrayExt: impl> PiecewiseSequenceArrayExt for T {} impl Array { - /// Constructs a new `PiecewiseSequenceArray` from start and length arrays. + /// Constructs a new `PiecewiseSequenceArray` from start, length, and multiplier arrays. /// - /// This validates only structural invariants: both children must be non-nullable unsigned + /// This validates only structural invariants: all children must be non-nullable unsigned /// integer arrays with matching lengths, and the outer array length is the declared expanded /// index length. Individual ranges are checked when the index array is executed or consumed by /// a take implementation. - pub fn try_new(starts: ArrayRef, lengths: ArrayRef, len: usize) -> VortexResult { + pub fn try_new( + starts: ArrayRef, + lengths: ArrayRef, + multipliers: ArrayRef, + len: usize, + ) -> VortexResult { Array::try_from_parts( ArrayParts::new(PiecewiseSequence, PType::U64.into(), len, EmptyArrayData) - .with_slots(smallvec![Some(starts), Some(lengths)]), + .with_slots(smallvec![Some(starts), Some(lengths), Some(multipliers)]), ) } @@ -47,11 +54,16 @@ impl Array { /// # Safety /// /// The caller must guarantee the same structural invariants as [`Self::try_new`]. - pub unsafe fn new_unchecked(starts: ArrayRef, lengths: ArrayRef, len: usize) -> Self { + pub unsafe fn new_unchecked( + starts: ArrayRef, + lengths: ArrayRef, + multipliers: ArrayRef, + len: usize, + ) -> Self { unsafe { Array::from_parts_unchecked( ArrayParts::new(PiecewiseSequence, PType::U64.into(), len, EmptyArrayData) - .with_slots(smallvec![Some(starts), Some(lengths)]), + .with_slots(smallvec![Some(starts), Some(lengths), Some(multipliers)]), ) } } diff --git a/vortex-array/src/arrays/piecewise_sequence/mod.rs b/vortex-array/src/arrays/piecewise_sequence/mod.rs index 974c395c0d1..b914c5e9651 100644 --- a/vortex-array/src/arrays/piecewise_sequence/mod.rs +++ b/vortex-array/src/arrays/piecewise_sequence/mod.rs @@ -4,19 +4,23 @@ //! Index encoding for concatenated sequential ranges. //! //! A `PiecewiseSequenceArray` represents the expanded index sequence -//! `starts[i]..starts[i] + lengths[i]` for each piece `i`. It is intended for take operations that -//! can gather contiguous runs without materializing one index per element. +//! `starts[i] + j * multipliers[i]` for `j` in `0..lengths[i]` for each piece `i`. It is +//! intended for take operations that can gather regular runs without materializing one index per +//! element. use itertools::Itertools; -use num_traits::AsPrimitive; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use crate::ArrayRef; +use crate::array::ArrayView; use crate::arrays::PrimitiveArray; +use crate::arrays::piecewise_sequence::array::PiecewiseSequenceArraySlotsExt; use crate::dtype::UnsignedPType; +use crate::executor::ExecutionCtx; pub mod array; mod vtable; @@ -27,18 +31,40 @@ mod tests; pub use array::PiecewiseSequenceArrayExt; pub use vtable::*; -pub(crate) fn check_index_arrays(starts: &ArrayRef, lengths: &ArrayRef) -> VortexResult<()> { +pub(crate) fn check_index_arrays( + starts: &ArrayRef, + lengths: &ArrayRef, + multipliers: &ArrayRef, +) -> VortexResult<()> { check_index_array("starts", starts)?; check_index_array("lengths", lengths)?; + check_index_array("multipliers", multipliers)?; vortex_ensure!( starts.len() == lengths.len(), "PiecewiseSequenceArray starts length {} does not match lengths length {}", starts.len(), lengths.len() ); + vortex_ensure!( + starts.len() == multipliers.len(), + "PiecewiseSequenceArray starts length {} does not match multipliers length {}", + starts.len(), + multipliers.len() + ); Ok(()) } +pub(crate) fn execute_index_arrays( + array: ArrayView<'_, PiecewiseSequence>, + ctx: &mut ExecutionCtx, +) -> VortexResult<(PrimitiveArray, PrimitiveArray, PrimitiveArray)> { + let starts = array.starts().clone().execute::(ctx)?; + let lengths = array.lengths().clone().execute::(ctx)?; + let multipliers = array.multipliers().clone().execute::(ctx)?; + check_index_arrays(starts.as_ref(), lengths.as_ref(), multipliers.as_ref())?; + Ok((starts, lengths, multipliers)) +} + fn check_index_array(name: &str, array: &ArrayRef) -> VortexResult<()> { vortex_ensure!( array.dtype().is_unsigned_int(), @@ -53,46 +79,41 @@ fn check_index_array(name: &str, array: &ArrayRef) -> VortexResult<()> { Ok(()) } -#[inline] -pub(crate) fn index_value_to_usize(value: T) -> usize { - value.as_() -} - -#[inline] -pub(crate) fn index_value_to_u64>(value: T) -> u64 { - value.as_() -} - -#[inline] -pub(crate) fn checked_range_end(start: u64, length: usize) -> VortexResult { - start - .checked_add(length as u64) - .ok_or_else(|| vortex_error::vortex_err!("PiecewiseSequenceArray range overflows u64")) -} - -pub(crate) fn materialize_ranges( +pub(crate) fn materialize_ranges( starts: &PrimitiveArray, lengths: &PrimitiveArray, + multipliers: &PrimitiveArray, output_len: usize, ) -> VortexResult> where - S: UnsignedPType + AsPrimitive, + S: UnsignedPType, L: UnsignedPType, + M: UnsignedPType, { let starts = starts.as_slice::(); let lengths = lengths.as_slice::(); + let multipliers = multipliers.as_slice::(); let mut values = BufferMut::with_capacity(output_len); let mut computed_len = 0usize; - for (&start, &length) in starts.iter().zip_eq(lengths) { - let start = index_value_to_u64(start); - let length = index_value_to_usize(length); - checked_range_end(start, length)?; - computed_len = computed_len.checked_add(length).ok_or_else(|| { - vortex_error::vortex_err!("PiecewiseSequenceArray output length overflows usize") - })?; + for ((&start, &length), &multiplier) in starts.iter().zip_eq(lengths).zip_eq(multipliers) { + let start: usize = start.as_(); + let length: usize = length.as_(); + let multiplier: usize = multiplier.as_(); + if length != 0 { + let last_offset = length - 1; + let last_delta = last_offset + .checked_mul(multiplier) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + start + .checked_add(last_delta) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + } + computed_len = computed_len + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; - values.extend((0..length).map(|offset| start + offset as u64)); + values.extend((0..length).map(|offset| (start + offset * multiplier) as u64)); } if computed_len != output_len { diff --git a/vortex-array/src/arrays/piecewise_sequence/tests.rs b/vortex-array/src/arrays/piecewise_sequence/tests.rs index 2f9992b4d05..ab201943e3d 100644 --- a/vortex-array/src/arrays/piecewise_sequence/tests.rs +++ b/vortex-array/src/arrays/piecewise_sequence/tests.rs @@ -22,7 +22,8 @@ use crate::serde::SerializedArray; fn materializes_piecewise_indices() -> VortexResult<()> { let starts = buffer![3u64, 15, 21].into_array(); let lengths = buffer![3u64, 3, 3].into_array(); - let array = PiecewiseSequenceArray::try_new(starts, lengths, 9)?.into_array(); + let multipliers = ConstantArray::new(1u64, 3).into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 9)?.into_array(); let expected = PrimitiveArray::from_iter([3u64, 4, 5, 15, 16, 17, 21, 22, 23]).into_array(); assert_arrays_eq!(array, expected, &mut array_session().create_execution_ctx()); @@ -33,18 +34,32 @@ fn materializes_piecewise_indices() -> VortexResult<()> { fn materializes_repeated_and_empty_ranges() -> VortexResult<()> { let starts = buffer![5u64, 2, 5].into_array(); let lengths = buffer![2u64, 0, 2].into_array(); - let array = PiecewiseSequenceArray::try_new(starts, lengths, 4)?.into_array(); + let multipliers = ConstantArray::new(1u64, 3).into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 4)?.into_array(); let expected = PrimitiveArray::from_iter([5u64, 6, 5, 6]).into_array(); assert_arrays_eq!(array, expected, &mut array_session().create_execution_ctx()); Ok(()) } +#[test] +fn materializes_multiplied_ranges() -> VortexResult<()> { + let starts = buffer![3u64, 15].into_array(); + let lengths = buffer![3u64, 2].into_array(); + let multipliers = buffer![2u64, 4].into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 5)?.into_array(); + + let expected = PrimitiveArray::from_iter([3u64, 5, 7, 15, 19]).into_array(); + assert_arrays_eq!(array, expected, &mut array_session().create_execution_ctx()); + Ok(()) +} + #[test] fn supports_constant_lengths() -> VortexResult<()> { let starts = buffer![0u64, 10, 20].into_array(); let lengths = ConstantArray::new(2u64, 3).into_array(); - let array = PiecewiseSequenceArray::try_new(starts, lengths, 6)?.into_array(); + let multipliers = ConstantArray::new(1u64, 3).into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 6)?.into_array(); let expected = PrimitiveArray::from_iter([0u64, 1, 10, 11, 20, 21]).into_array(); assert_arrays_eq!(array, expected, &mut array_session().create_execution_ctx()); @@ -55,7 +70,8 @@ fn supports_constant_lengths() -> VortexResult<()> { fn scalar_at_maps_into_piece() -> VortexResult<()> { let starts = buffer![3u64, 15, 21].into_array(); let lengths = buffer![3u64, 3, 3].into_array(); - let array = PiecewiseSequenceArray::try_new(starts, lengths, 9)?.into_array(); + let multipliers = buffer![1u64, 1, 1].into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 9)?.into_array(); let mut ctx = array_session().create_execution_ctx(); assert_eq!(array.execute_scalar(0, &mut ctx)?, 3u64.into()); @@ -68,14 +84,15 @@ fn scalar_at_maps_into_piece() -> VortexResult<()> { fn constructor_defers_range_value_validation() -> VortexResult<()> { let starts = buffer![u64::MAX].into_array(); let lengths = buffer![2u64].into_array(); - let array = PiecewiseSequenceArray::try_new(starts, lengths, 2)?.into_array(); + let multipliers = buffer![1u64].into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 2)?.into_array(); let err = array .execute::(&mut array_session().create_execution_ctx()) .unwrap_err(); assert!( err.to_string() - .contains("PiecewiseSequenceArray range overflows u64"), + .contains("PiecewiseSequenceArray range overflows usize"), "{err}" ); Ok(()) @@ -85,7 +102,8 @@ fn constructor_defers_range_value_validation() -> VortexResult<()> { fn execution_checks_declared_length() -> VortexResult<()> { let starts = buffer![0u64, 3].into_array(); let lengths = buffer![2u64, 2].into_array(); - let array = PiecewiseSequenceArray::try_new(starts, lengths, 3)?.into_array(); + let multipliers = ConstantArray::new(1u64, 2).into_array(); + let array = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 3)?.into_array(); let err = array .execute::(&mut array_session().create_execution_ctx()) @@ -103,6 +121,7 @@ fn serde_roundtrip_preserves_piecewise_indices() -> VortexResult<()> { let array = PiecewiseSequenceArray::try_new( buffer![3u32, 15, 21].into_array(), buffer![2u16, 0, 2].into_array(), + buffer![2u8, 1, 3].into_array(), 4, )? .into_array(); @@ -128,7 +147,7 @@ fn serde_roundtrip_preserves_piecewise_indices() -> VortexResult<()> { assert!(decoded.is::()); assert_arrays_eq!( decoded, - PrimitiveArray::from_iter([3u64, 4, 21, 22]).into_array(), + PrimitiveArray::from_iter([3u64, 5, 21, 24]).into_array(), &mut array_session().create_execution_ctx() ); Ok(()) diff --git a/vortex-array/src/arrays/piecewise_sequence/vtable.rs b/vortex-array/src/arrays/piecewise_sequence/vtable.rs index 53dba5c5879..e4da87a11ff 100644 --- a/vortex-array/src/arrays/piecewise_sequence/vtable.rs +++ b/vortex-array/src/arrays/piecewise_sequence/vtable.rs @@ -29,8 +29,7 @@ use crate::arrays::PrimitiveArray; use crate::arrays::piecewise_sequence::array::PiecewiseSequenceArraySlotsExt; use crate::arrays::piecewise_sequence::array::PiecewiseSequenceSlots; use crate::arrays::piecewise_sequence::check_index_arrays; -use crate::arrays::piecewise_sequence::index_value_to_u64; -use crate::arrays::piecewise_sequence::index_value_to_usize; +use crate::arrays::piecewise_sequence::execute_index_arrays; use crate::arrays::piecewise_sequence::materialize_ranges; use crate::arrays::primitive::PrimitiveArrayExt; use crate::buffer::BufferHandle; @@ -52,6 +51,8 @@ struct PiecewiseSequenceMetadata { starts_ptype: i32, #[prost(enumeration = "PType", tag = "3")] lengths_ptype: i32, + #[prost(enumeration = "PType", tag = "4")] + multipliers_ptype: i32, } #[derive(Clone, Debug)] @@ -86,15 +87,16 @@ impl VTable for PiecewiseSequence { ); let starts = slots[PiecewiseSequenceSlots::STARTS] .as_ref() - .ok_or_else(|| { - vortex_error::vortex_err!("PiecewiseSequenceArray starts slot must be present") - })?; + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray starts slot must be present"))?; let lengths = slots[PiecewiseSequenceSlots::LENGTHS] + .as_ref() + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray lengths slot must be present"))?; + let multipliers = slots[PiecewiseSequenceSlots::MULTIPLIERS] .as_ref() .ok_or_else(|| { - vortex_error::vortex_err!("PiecewiseSequenceArray lengths slot must be present") + vortex_err!("PiecewiseSequenceArray multipliers slot must be present") })?; - check_index_arrays(starts, lengths) + check_index_arrays(starts, lengths, multipliers) } fn nbuffers(_array: ArrayView<'_, Self>) -> usize { @@ -135,6 +137,7 @@ impl VTable for PiecewiseSequence { })?, starts_ptype: PType::try_from(array.starts().dtype())? as i32, lengths_ptype: PType::try_from(array.lengths().dtype())? as i32, + multipliers_ptype: PType::try_from(array.multipliers().dtype())? as i32, } .encode_to_vec(), )) @@ -182,21 +185,26 @@ impl VTable for PiecewiseSequence { &metadata.lengths_ptype().into(), num_pieces, )?; + let multipliers = children.get( + PiecewiseSequenceSlots::MULTIPLIERS, + &metadata.multipliers_ptype().into(), + num_pieces, + )?; Ok( ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData) - .with_slots(smallvec![Some(starts), Some(lengths)]), + .with_slots(smallvec![Some(starts), Some(lengths), Some(multipliers)]), ) } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - let starts = array.starts().clone().execute::(ctx)?; - let lengths = array.lengths().clone().execute::(ctx)?; - check_index_arrays(starts.as_ref(), lengths.as_ref())?; + let (starts, lengths, multipliers) = execute_index_arrays(array.as_view(), ctx)?; let values = match_each_unsigned_integer_ptype!(starts.ptype(), |S| { match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { - materialize_ranges::(&starts, &lengths, array.len())? + match_each_unsigned_integer_ptype!(multipliers.ptype(), |M| { + materialize_ranges::(&starts, &lengths, &multipliers, array.len())? + }) }) }); Ok(ExecutionResult::done( @@ -211,13 +219,13 @@ impl OperationsVTable for PiecewiseSequence { index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let starts = array.starts().clone().execute::(ctx)?; - let lengths = array.lengths().clone().execute::(ctx)?; - check_index_arrays(starts.as_ref(), lengths.as_ref())?; + let (starts, lengths, multipliers) = execute_index_arrays(array, ctx)?; let value = match_each_unsigned_integer_ptype!(starts.ptype(), |S| { match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { - scalar_at::(&starts, &lengths, index)? + match_each_unsigned_integer_ptype!(multipliers.ptype(), |M| { + scalar_at::(&starts, &lengths, &multipliers, index)? + }) }) }); Ok(value.into()) @@ -230,24 +238,35 @@ impl ValidityVTable for PiecewiseSequence { } } -fn scalar_at( +fn scalar_at( starts: &PrimitiveArray, lengths: &PrimitiveArray, + multipliers: &PrimitiveArray, index: usize, ) -> VortexResult where - S: crate::dtype::UnsignedPType + num_traits::AsPrimitive, + S: crate::dtype::UnsignedPType, L: crate::dtype::UnsignedPType, + M: crate::dtype::UnsignedPType, { let mut remaining = index; - for (&start, &length) in starts + for ((&start, &length), &multiplier) in starts .as_slice::() .iter() .zip_eq(lengths.as_slice::()) + .zip_eq(multipliers.as_slice::()) { - let length = index_value_to_usize(length); + let length: usize = length.as_(); if remaining < length { - return Ok(index_value_to_u64(start) + remaining as u64); + let start: usize = start.as_(); + let multiplier: usize = multiplier.as_(); + let offset = remaining + .checked_mul(multiplier) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + let value = start + .checked_add(offset) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + return Ok(value as u64); } remaining -= length; } From 16ca5dc10ca9dcf4aca924bb438d4cb91abb9f9b Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 16 Jul 2026 19:36:43 -0400 Subject: [PATCH 04/37] Fix PiecewiseSequence rustdoc link Signed-off-by: Daniel King --- vortex-array/src/arrays/piecewise_sequence/array.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-array/src/arrays/piecewise_sequence/array.rs b/vortex-array/src/arrays/piecewise_sequence/array.rs index 5c98cc15bb2..b51e8ff043e 100644 --- a/vortex-array/src/arrays/piecewise_sequence/array.rs +++ b/vortex-array/src/arrays/piecewise_sequence/array.rs @@ -23,7 +23,7 @@ pub struct PiecewiseSequenceSlots { pub multipliers: ArrayRef, } -/// Extension methods for [`PiecewiseSequenceArray`]. +/// Extension methods for [`crate::arrays::PiecewiseSequenceArray`]. pub trait PiecewiseSequenceArrayExt: TypedArrayRef + PiecewiseSequenceArraySlotsExt { From 52affcd4cc83108481158814a616bf247c5ca744 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 16 Jul 2026 20:13:35 -0400 Subject: [PATCH 05/37] Use imported UnsignedPType in PiecewiseSequence Signed-off-by: Daniel King --- vortex-array/src/arrays/piecewise_sequence/vtable.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/vortex-array/src/arrays/piecewise_sequence/vtable.rs b/vortex-array/src/arrays/piecewise_sequence/vtable.rs index e4da87a11ff..ac2eba61116 100644 --- a/vortex-array/src/arrays/piecewise_sequence/vtable.rs +++ b/vortex-array/src/arrays/piecewise_sequence/vtable.rs @@ -35,6 +35,7 @@ use crate::arrays::primitive::PrimitiveArrayExt; use crate::buffer::BufferHandle; use crate::dtype::DType; use crate::dtype::PType; +use crate::dtype::UnsignedPType; use crate::match_each_unsigned_integer_ptype; use crate::scalar::Scalar; use crate::serde::ArrayChildren; @@ -245,9 +246,9 @@ fn scalar_at( index: usize, ) -> VortexResult where - S: crate::dtype::UnsignedPType, - L: crate::dtype::UnsignedPType, - M: crate::dtype::UnsignedPType, + S: UnsignedPType, + L: UnsignedPType, + M: UnsignedPType, { let mut remaining = index; for ((&start, &length), &multiplier) in starts From 2b6fa3e469e651ddb7feb205605fdcee346da4f2 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 16 Jul 2026 20:17:30 -0400 Subject: [PATCH 06/37] Shorten PiecewiseSequence doc link Signed-off-by: Daniel King --- vortex-array/src/arrays/piecewise_sequence/array.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-array/src/arrays/piecewise_sequence/array.rs b/vortex-array/src/arrays/piecewise_sequence/array.rs index b51e8ff043e..1559cf7d3ad 100644 --- a/vortex-array/src/arrays/piecewise_sequence/array.rs +++ b/vortex-array/src/arrays/piecewise_sequence/array.rs @@ -23,7 +23,7 @@ pub struct PiecewiseSequenceSlots { pub multipliers: ArrayRef, } -/// Extension methods for [`crate::arrays::PiecewiseSequenceArray`]. +/// Extension methods for [`Array`] values using the [`PiecewiseSequence`] encoding. pub trait PiecewiseSequenceArrayExt: TypedArrayRef + PiecewiseSequenceArraySlotsExt { From b9653f43dd3644d0920e07ad603e7787b3b13f30 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 16 Jul 2026 13:42:45 -0400 Subject: [PATCH 07/37] Use PiecewiseSequential indices for FSL take Signed-off-by: Daniel King --- vortex-array/src/arrays/bool/compute/take.rs | 58 ++++ .../src/arrays/decimal/compute/take.rs | 67 ++++ .../arrays/fixed_size_list/compute/take.rs | 326 ++++++++++-------- .../src/arrays/piecewise_sequence/mod.rs | 58 ++++ .../src/arrays/piecewise_sequence/tests.rs | 116 +++++++ .../src/arrays/primitive/compute/take/mod.rs | 59 ++++ .../src/arrays/varbin/compute/take.rs | 178 ++++++++++ .../src/arrays/varbinview/compute/take.rs | 68 ++++ 8 files changed, 786 insertions(+), 144 deletions(-) diff --git a/vortex-array/src/arrays/bool/compute/take.rs b/vortex-array/src/arrays/bool/compute/take.rs index 8469c7e25c2..fe7955fbbae 100644 --- a/vortex-array/src/arrays/bool/compute/take.rs +++ b/vortex-array/src/arrays/bool/compute/take.rs @@ -4,6 +4,7 @@ use itertools::Itertools as _; use num_traits::AsPrimitive; use vortex_buffer::BitBuffer; +use vortex_buffer::BitBufferMut; use vortex_buffer::BitBufferView; use vortex_buffer::get_bit; use vortex_error::VortexResult; @@ -15,12 +16,16 @@ use crate::array::ArrayView; use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; +use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; +use crate::arrays::piecewise_sequence::validate_index_ranges; use crate::builtins::ArrayBuiltins; use crate::executor::ExecutionCtx; use crate::match_each_integer_ptype; +use crate::match_each_unsigned_integer_ptype; use crate::scalar::Scalar; impl TakeExecute for Bool { @@ -29,6 +34,12 @@ impl TakeExecute for Bool { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_piecewise_sequence(array, piecewise_indices, indices, ctx)? + { + return Ok(Some(taken)); + } + let indices_nulls_zeroed = match indices.validity()?.execute_mask(indices.len(), ctx)? { Mask::AllTrue(_) => indices.clone(), Mask::AllFalse(_) => { @@ -55,6 +66,31 @@ impl TakeExecute for Bool { } } +fn take_piecewise_sequence( + array: ArrayView<'_, Bool>, + indices: ArrayView<'_, PiecewiseSequence>, + indices_ref: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some((starts, lengths)) = execute_unit_multiplier_index_arrays(indices, ctx)? else { + return Ok(None); + }; + let buffer = match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + take_piecewise_bits( + &array.to_bit_buffer(), + starts.as_slice::(), + lengths.as_slice::(), + indices_ref.len(), + )? + }) + }); + + Ok(Some( + BoolArray::new(buffer, array.validity()?.take(indices_ref)?).into_array(), + )) +} + fn take_valid_indices>(bools: BitBufferView<'_>, indices: &[I]) -> BitBuffer { // For boolean arrays that roughly fit into a single page (at least, on Linux), it's worth // the overhead to convert to a Vec. @@ -82,6 +118,28 @@ fn take_bool_impl>(bools: BitBufferView<'_>, indices: &[I] }) } +fn take_piecewise_bits( + source: &BitBuffer, + starts: &[S], + lengths: &[L], + output_len: usize, +) -> VortexResult +where + S: crate::dtype::UnsignedPType, + L: crate::dtype::UnsignedPType, +{ + validate_index_ranges(source.len(), starts, lengths, output_len)?; + + let mut values = BitBufferMut::with_capacity(output_len); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = start.as_(); + let length = length.as_(); + values.append_buffer(&source.slice(start..start + length)); + } + + Ok(values.freeze()) +} + #[cfg(test)] mod test { use rstest::rstest; diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index ce786f8e4ab..213d6db39bf 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -1,7 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use itertools::Itertools as _; use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; use vortex_error::VortexResult; use crate::ArrayRef; @@ -9,13 +11,17 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Decimal; use crate::arrays::DecimalArray; +use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; +use crate::arrays::piecewise_sequence::validate_index_ranges; use crate::dtype::IntegerPType; use crate::dtype::NativeDecimalType; use crate::executor::ExecutionCtx; use crate::match_each_decimal_value_type; use crate::match_each_integer_ptype; +use crate::match_each_unsigned_integer_ptype; impl TakeExecute for Decimal { fn take( @@ -23,6 +29,12 @@ impl TakeExecute for Decimal { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_piecewise_sequence(array, piecewise_indices, indices, ctx)? + { + return Ok(Some(taken)); + } + let indices = indices.clone().execute::(ctx)?; let validity = array.validity()?.take(&indices.clone().into_array())?; @@ -42,10 +54,65 @@ impl TakeExecute for Decimal { } } +fn take_piecewise_sequence( + array: ArrayView<'_, Decimal>, + indices: ArrayView<'_, PiecewiseSequence>, + indices_ref: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some((starts, lengths)) = execute_unit_multiplier_index_arrays(indices, ctx)? else { + return Ok(None); + }; + match_each_decimal_value_type!(array.values_type(), |D| { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + let values = take_piecewise_to_buffer::( + starts.as_slice::(), + lengths.as_slice::(), + array.buffer::().as_slice(), + indices_ref.len(), + )?; + let validity = array.validity()?.take(indices_ref)?; + + // SAFETY: contiguous gather preserves the decimal dtype and value representation. + Ok( + Some(unsafe { + DecimalArray::new_unchecked(values, array.decimal_dtype(), validity) + } + .into_array()), + ) + }) + }) + }) +} + fn take_to_buffer(indices: &[I], values: &[T]) -> Buffer { indices.iter().map(|idx| values[idx.as_()]).collect() } +fn take_piecewise_to_buffer( + starts: &[S], + lengths: &[L], + values: &[T], + output_len: usize, +) -> VortexResult> +where + S: crate::dtype::UnsignedPType, + L: crate::dtype::UnsignedPType, + T: NativeDecimalType, +{ + validate_index_ranges(values.len(), starts, lengths, output_len)?; + + let mut result = BufferMut::::with_capacity(output_len); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = start.as_(); + let length = length.as_(); + result.extend_from_slice(&values[start..start + length]); + } + + Ok(result.freeze()) +} + #[cfg(test)] mod tests { use rstest::rstest; 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 193730f62c4..b2533d132b9 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -1,221 +1,259 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_buffer::BitBufferMut; +use itertools::Itertools; +use vortex_buffer::Buffer; use vortex_buffer::BufferMut; -use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_panic; -use vortex_mask::Mask; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use crate::ArrayRef; use crate::IntoArray; use crate::array::ArrayView; -use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; +use crate::arrays::PiecewiseSequenceArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; -use crate::arrays::bool::BoolArrayExt; use crate::arrays::dict::TakeExecute; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::primitive::PrimitiveArrayExt; +use crate::builders::builder_with_capacity; +use crate::dtype::DType; use crate::dtype::IntegerPType; +use crate::dtype::Nullability; use crate::executor::ExecutionCtx; -use crate::match_each_unsigned_integer_ptype; -use crate::match_smallest_offset_type; +use crate::match_each_integer_ptype; use crate::validity::Validity; /// Take implementation for [`FixedSizeListArray`]. /// -/// Unlike `ListView`, `FixedSizeListArray` must rebuild the elements array because it requires -/// that elements start at offset 0 and be perfectly packed without gaps. We expand list indices -/// into element indices and push them down to the child elements array. +/// `FixedSizeListArray` must rebuild its elements array because selected lists need to become +/// packed from offset 0. The FSL layer translates selected list rows into ordered element runs and +/// delegates the execution strategy to the elements child via `PiecewiseSequenceArray` indices. impl TakeExecute for FixedSizeList { fn take( array: ArrayView<'_, FixedSizeList>, indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let max_element_idx = array.elements().len(); - // Indices are non-negative; dispatch over the 4 unsigned widths (the executed array is - // reinterpreted to unsigned in `take_with_indices`). `E` is already unsigned. - match_each_unsigned_integer_ptype!(indices.dtype().as_ptype().to_unsigned(), |I| { - match_smallest_offset_type!(max_element_idx, |E| { - take_with_indices::(array, indices, ctx) - }) - }) - .map(Some) + if array.is_empty() { + return take_empty_fsl(array, indices, ctx).map(Some); + } + + take_non_empty_fsl(array, indices, ctx).map(Some) } } -/// Dispatches to the appropriate take implementation based on list size and nullability. -fn take_with_indices( +fn take_empty_fsl( array: ArrayView<'_, FixedSizeList>, indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult { - let list_size = array.list_size() as usize; + debug_assert!(array.is_empty()); - let indices_array = indices.clone().execute::(ctx)?; - // Reinterpret to unsigned so `as_slice::` (with unsigned `I`) matches; values are unchanged. - let indices_array = indices_array.reinterpret_cast(indices_array.ptype().to_unsigned()); - - // Make sure to handle degenerate case where lists have size 0 (these can take fast paths). - if list_size == 0 { - debug_assert!( - array.elements().is_empty(), - "degenerate list must have empty elements" + let new_len = indices.len(); + if new_len != 0 { + let indices_validity = indices.validity()?.execute_mask(new_len, ctx)?; + vortex_ensure!( + indices_validity.all_false(), + "cannot take valid indices from an empty FixedSizeList" ); + } - // Since there are no elements to take, we just need to take on the validity map. - let new_validity = array.validity()?.take(indices)?; - let new_len = indices_array.len(); - - Ok( - // SAFETY: list_size is 0, elements array is empty, and validity has the correct length. - unsafe { - FixedSizeListArray::new_unchecked( - array.elements().clone(), // Remember that this is an empty array. - array.list_size(), - new_validity, - new_len, - ) - } - .into_array(), + let list_size = array.list_size() as usize; + 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 new_elements = default_elements(array, elements_len); + let new_validity = if new_len == 0 { + array.validity()?.take(indices)? } else { - // The result's nullability is the union of the input nullabilities. - if array.dtype().is_nullable() || indices_array.dtype().is_nullable() { - let indices_array = indices_array.as_view(); - take_nullable_fsl::(array, indices_array, ctx) - } else { - let indices_array = indices_array.as_view(); - take_non_nullable_fsl::(array, indices_array) - } + Validity::AllInvalid + }; + + // SAFETY: empty output needs no child values; otherwise the index validity mask proves every + // output row is null. Placeholder child elements have the exact length required by FSL. + Ok(unsafe { + FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) } + .into_array()) } -/// Takes from an array when both the array and indices are non-nullable. -fn take_non_nullable_fsl( +fn take_non_empty_fsl( array: ArrayView<'_, FixedSizeList>, - indices_array: ArrayView<'_, Primitive>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, ) -> VortexResult { - let list_size = array.list_size() as usize; - let indices: &[I] = indices_array.as_slice::(); - let new_len = indices.len(); - - // Build the element indices directly without validity tracking. - let mut elements_indices = BufferMut::::with_capacity(new_len * list_size); - - // Build the element indices for each list. - for data_idx in indices { - let data_idx = data_idx - .to_usize() - .unwrap_or_else(|| vortex_panic!("Failed to convert index to usize: {}", data_idx)); + debug_assert!(!array.is_empty()); - let list_start = data_idx * list_size; - let list_end = (data_idx + 1) * list_size; + let DType::Primitive(ptype, nullability) = indices.dtype() else { + vortex_bail!("Invalid indices dtype: {}", indices.dtype()) + }; + if !ptype.is_int() { + vortex_bail!("Invalid indices dtype: {}", indices.dtype()); + } - // Expand the list into individual element indices. - for i in list_start..list_end { - // SAFETY: We've allocated enough space for enough indices for all `new_len` lists (that each consist of `list_size = list_end - list_start` elements), so we know we have enough capacity. - unsafe { - elements_indices.push_unchecked(E::from_usize(i).vortex_expect("i < list_end")) - }; - } + if array.list_size() == 0 { + return take_non_empty_degenerate_fsl(array, indices, ctx); } - let elements_indices = elements_indices.freeze(); - debug_assert_eq!(elements_indices.len(), new_len * list_size); + 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(), + *nullability, + ctx, + ) + }) +} + +fn take_non_empty_degenerate_fsl( + array: ArrayView<'_, FixedSizeList>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + debug_assert!(!array.is_empty()); + debug_assert_eq!(array.list_size(), 0); + vortex_ensure!( + array.elements().is_empty(), + "degenerate list must have empty elements" + ); - let elements_indices_array = PrimitiveArray::new(elements_indices, Validity::NonNullable); - let new_elements = array.elements().take(elements_indices_array.into_array())?; - debug_assert_eq!(new_elements.len(), new_len * list_size); + let indices_array = indices.clone().execute::(ctx)?; + match_each_integer_ptype!(indices_array.ptype(), |I| { + bounds_check_valid_indices::(&indices_array.as_view(), array.as_ref().len(), ctx) + })?; + let new_validity = array.validity()?.take(indices)?; + let new_len = indices_array.len(); - // Both inputs are non-nullable, so the result is non-nullable. + // SAFETY: degenerate FSL inputs have no elements, valid index payloads were checked against + // the source length, and `Validity::take` produces validity for `new_len`. Ok(unsafe { FixedSizeListArray::new_unchecked( - new_elements, + array.elements().clone(), array.list_size(), - Validity::NonNullable, + new_validity, new_len, ) } .into_array()) } -/// Takes from an array when either the array or indices are nullable. -fn take_nullable_fsl( +fn take_non_empty_non_degenerate_fsl( array: ArrayView<'_, FixedSizeList>, + indices: &ArrayRef, indices_array: ArrayView<'_, Primitive>, + indices_nullability: Nullability, ctx: &mut ExecutionCtx, ) -> VortexResult { + debug_assert!(!array.is_empty()); + debug_assert_ne!(array.list_size(), 0); + + let (new_elements, new_len) = + take_non_empty_non_degenerate_elements::(array, indices_array, ctx)?; + let new_validity = if array.dtype().is_nullable() || indices_nullability.is_nullable() { + array.validity()?.take(indices)? + } else { + Validity::NonNullable + }; + + // SAFETY: `new_elements` has `new_len * list_size` elements. `new_validity` is either + // non-nullable or was produced by `Validity::take` for `new_len`. + Ok(unsafe { + FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) + } + .into_array()) +} + +fn take_non_empty_non_degenerate_elements( + array: ArrayView<'_, FixedSizeList>, + indices_array: ArrayView<'_, Primitive>, + ctx: &mut ExecutionCtx, +) -> VortexResult<(ArrayRef, usize)> { + debug_assert!(!array.is_empty()); + debug_assert_ne!(array.list_size(), 0); + let list_size = array.list_size() as usize; + let array_len = array.as_ref().len(); let indices: &[I] = indices_array.as_slice::(); let new_len = indices.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 array_validity = array - .fixed_size_list_validity() - .execute_mask(array.as_ref().len(), ctx) - .vortex_expect("Failed to compute validity mask"); - let indices_len = indices_array.as_ref().len(); - let indices_validity = match indices_array - .validity() - .vortex_expect("Failed to compute validity mask") - { - Validity::NonNullable | Validity::AllValid => Mask::new_true(indices_len), - Validity::AllInvalid => Mask::new_false(indices_len), - Validity::Array(a) => a.execute::(ctx)?.execute_mask(ctx), - }; + let indices_validity = indices_array.validity()?.execute_mask(new_len, ctx)?; + let starts = indices + .iter() + .zip_eq(indices_validity.iter()) + .map(|(&data_idx, is_index_valid)| { + if !is_index_valid { + return Ok(0); + } - // We must use placeholder zeros for null lists to maintain the array length without - // propagating nullability to the element array's take operation. - let mut elements_indices = BufferMut::::with_capacity(new_len * list_size); - let mut new_validity_builder = BitBufferMut::with_capacity(new_len); - - // Build the element indices while tracking which lists are null. - for (data_idx, is_index_valid) in indices.iter().zip(indices_validity.iter()) { - let data_idx = data_idx - .to_usize() - .unwrap_or_else(|| vortex_panic!("Failed to convert index to usize: {}", data_idx)); - - // The list is null if the index is null or the indexed element is null. - if !is_index_valid || !array_validity.value(data_idx) { - // Append placeholder zeros for null lists. These will be masked by the validity array. - // We cannot use append_nulls here as explained above. - unsafe { elements_indices.push_n_unchecked(E::zero(), list_size) }; - new_validity_builder.append(false); - } else { - // Append the actual element indices for this list. - let list_start = data_idx * list_size; - let list_end = (data_idx + 1) * list_size; - - // Expand the list into individual element indices. - for i in list_start..list_end { - // SAFETY: We've allocated enough space for enough indices for all `new_len` lists (that each consist of `list_size = list_end - list_start` elements), so we know we have enough capacity. - unsafe { - elements_indices.push_unchecked(E::from_usize(i).vortex_expect("i < list_end")) - }; + let data_idx: usize = data_idx.as_(); + if data_idx >= array_len { + vortex_bail!(OutOfBounds: data_idx, 0, array_len); } + Ok((data_idx * list_size) as u64) + }) + .process_results(|iter| iter.collect::>())?; + + let new_elements = + take_element_runs(array.elements(), starts.freeze(), list_size, elements_len)?; + + Ok((new_elements, new_len)) +} - new_validity_builder.append(true); +fn bounds_check_valid_indices( + indices_array: &ArrayView<'_, Primitive>, + array_len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let indices: &[I] = indices_array.as_slice::(); + let indices_validity = indices_array.validity()?.execute_mask(indices.len(), ctx)?; + + for (&data_idx, is_index_valid) in indices.iter().zip_eq(indices_validity.iter()) { + if is_index_valid { + let data_idx = data_idx.as_(); + if data_idx >= array_len { + vortex_bail!(OutOfBounds: data_idx, 0, array_len); + } } } + Ok(()) +} - let elements_indices = elements_indices.freeze(); - debug_assert_eq!(elements_indices.len(), new_len * list_size); - - let elements_indices_array = PrimitiveArray::new(elements_indices, Validity::NonNullable); - let new_elements = array.elements().take(elements_indices_array.into_array())?; - debug_assert_eq!(new_elements.len(), new_len * list_size); +fn default_elements(array: ArrayView<'_, FixedSizeList>, len: usize) -> ArrayRef { + let mut builder = builder_with_capacity(array.elements().dtype(), len); + builder.append_defaults(len); + builder.finish() +} - // At least one input was nullable, so the result is nullable. - let new_validity = Validity::from(new_validity_builder.freeze()); - debug_assert!(new_validity.maybe_len().is_none_or(|vl| vl == new_len)); +fn take_element_runs( + elements: &ArrayRef, + starts: Buffer, + length: usize, + output_len: usize, +) -> VortexResult { + let run_count = starts.len(); + let starts = PrimitiveArray::new(starts, Validity::NonNullable).into_array(); + let lengths = ConstantArray::new(length as u64, run_count).into_array(); + let multipliers = ConstantArray::new(1u64, run_count).into_array(); - Ok(unsafe { - FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) - } - .into_array()) + // SAFETY: callers produced one start per output row after validating list indices against the + // source FSL length. `length` and multiplier 1 are represented as non-nullable unsigned + // constant arrays, and `output_len` was computed as `run_count * length`. + let indices = + unsafe { PiecewiseSequenceArray::new_unchecked(starts, lengths, multipliers, output_len) }; + elements.take(indices.into_array()) } diff --git a/vortex-array/src/arrays/piecewise_sequence/mod.rs b/vortex-array/src/arrays/piecewise_sequence/mod.rs index b914c5e9651..a5361567c7d 100644 --- a/vortex-array/src/arrays/piecewise_sequence/mod.rs +++ b/vortex-array/src/arrays/piecewise_sequence/mod.rs @@ -21,6 +21,7 @@ use crate::arrays::PrimitiveArray; use crate::arrays::piecewise_sequence::array::PiecewiseSequenceArraySlotsExt; use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; +use crate::scalar::PValue; pub mod array; mod vtable; @@ -65,6 +66,30 @@ pub(crate) fn execute_index_arrays( Ok((starts, lengths, multipliers)) } +pub(crate) fn execute_unit_multiplier_index_arrays( + array: ArrayView<'_, PiecewiseSequence>, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + if !is_constant_multiplier_one(array.multipliers()) { + return Ok(None); + } + + let starts = array.starts().clone().execute::(ctx)?; + let lengths = array.lengths().clone().execute::(ctx)?; + check_index_arrays(starts.as_ref(), lengths.as_ref(), array.multipliers())?; + Ok(Some((starts, lengths))) +} + +pub(crate) fn is_constant_multiplier_one(multipliers: &ArrayRef) -> bool { + let Some(scalar) = multipliers.as_constant() else { + return false; + }; + matches!( + scalar.as_primitive_opt().and_then(|scalar| scalar.pvalue()), + Some(PValue::U8(1) | PValue::U16(1) | PValue::U32(1) | PValue::U64(1)) + ) +} + fn check_index_array(name: &str, array: &ArrayRef) -> VortexResult<()> { vortex_ensure!( array.dtype().is_unsigned_int(), @@ -123,3 +148,36 @@ where } Ok(values) } + +pub(crate) fn validate_index_ranges( + source_len: usize, + starts: &[S], + lengths: &[L], + output_len: usize, +) -> VortexResult<()> +where + S: UnsignedPType, + L: UnsignedPType, +{ + let mut computed_len = 0usize; + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start: usize = start.as_(); + let length: usize = length.as_(); + let end = start.checked_add(length).ok_or_else(|| { + vortex_err!("PiecewiseSequenceArray range overflows usize") + })?; + vortex_ensure!( + end <= source_len, + "PiecewiseSequenceArray range {start}..{end} exceeds source length {source_len}" + ); + computed_len = computed_len.checked_add(length).ok_or_else(|| { + vortex_err!("PiecewiseSequenceArray output length overflows usize") + })?; + } + + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); + Ok(()) +} diff --git a/vortex-array/src/arrays/piecewise_sequence/tests.rs b/vortex-array/src/arrays/piecewise_sequence/tests.rs index ab201943e3d..1abbada7a92 100644 --- a/vortex-array/src/arrays/piecewise_sequence/tests.rs +++ b/vortex-array/src/arrays/piecewise_sequence/tests.rs @@ -7,16 +7,37 @@ use vortex_error::VortexResult; use vortex_session::registry::ReadContext; use crate::ArrayContext; +use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; +use crate::arrays::BoolArray; use crate::arrays::ConstantArray; +use crate::arrays::DecimalArray; +use crate::arrays::FixedSizeListArray; use crate::arrays::PiecewiseSequence; use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; +use crate::arrays::VarBinArray; +use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::Nullability; use crate::serde::SerializeOptions; use crate::serde::SerializedArray; +use crate::validity::Validity; + +fn piecewise_indices( + starts: impl IntoIterator, + lengths: &[u64], +) -> VortexResult { + let len = lengths.iter().map(|&length| length as usize).sum(); + let starts = PrimitiveArray::from_iter(starts).into_array(); + let lengths = PrimitiveArray::from_iter(lengths.iter().copied()).into_array(); + let multipliers = ConstantArray::new(1u64, lengths.len()).into_array(); + Ok(PiecewiseSequenceArray::try_new(starts, lengths, multipliers, len)?.into_array()) +} #[test] fn materializes_piecewise_indices() -> VortexResult<()> { @@ -152,3 +173,98 @@ fn serde_roundtrip_preserves_piecewise_indices() -> VortexResult<()> { ); Ok(()) } + +#[test] +fn primitive_take_consumes_piecewise_indices() -> VortexResult<()> { + let values = PrimitiveArray::from_iter(0i32..20).into_array(); + let taken = values.take(piecewise_indices([3, 10], &[2, 3])?)?; + + assert_arrays_eq!( + taken, + PrimitiveArray::from_iter([3i32, 4, 10, 11, 12]).into_array(), + &mut array_session().create_execution_ctx() + ); + Ok(()) +} + +#[test] +fn bool_take_consumes_piecewise_indices() -> VortexResult<()> { + let values = BoolArray::from_iter([true, false, true, true, false, false]).into_array(); + let taken = values.take(piecewise_indices([1, 4], &[2, 2])?)?; + + assert_arrays_eq!( + taken, + BoolArray::from_iter([false, true, false, false]).into_array(), + &mut array_session().create_execution_ctx() + ); + Ok(()) +} + +#[test] +fn decimal_take_consumes_piecewise_indices() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(19, 2); + let values = DecimalArray::from_iter([10i128, 20, 30, 40, 50, 60], decimal_dtype).into_array(); + let taken = values.take(piecewise_indices([1, 4], &[2, 1])?)?; + + assert_arrays_eq!( + taken, + DecimalArray::from_iter([20i128, 30, 50], decimal_dtype).into_array(), + &mut array_session().create_execution_ctx() + ); + Ok(()) +} + +#[test] +fn varbinview_take_consumes_piecewise_indices() -> VortexResult<()> { + let values = VarBinViewArray::from_iter( + ["a", "bb", "ccc", "dddd", "eeeee"].map(Some), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let taken = values.take(piecewise_indices([1, 3], &[2, 1])?)?; + + assert_arrays_eq!( + taken, + VarBinViewArray::from_iter( + ["bb", "ccc", "dddd"].map(Some), + DType::Utf8(Nullability::NonNullable) + ) + .into_array(), + &mut array_session().create_execution_ctx() + ); + Ok(()) +} + +#[test] +fn varbin_take_consumes_piecewise_indices() -> VortexResult<()> { + let values = VarBinArray::from_iter( + ["a", "bb", "ccc", "dddd", "eeeee"].map(Some), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let taken = values.take(piecewise_indices([1, 3], &[2, 1])?)?; + + assert_arrays_eq!( + taken, + VarBinArray::from_iter( + ["bb", "ccc", "dddd"].map(Some), + DType::Utf8(Nullability::NonNullable) + ) + .into_array(), + &mut array_session().create_execution_ctx() + ); + Ok(()) +} + +#[test] +fn fixed_size_list_take_builds_piecewise_element_indices() -> VortexResult<()> { + let elements = PrimitiveArray::from_iter(0i32..12).into_array(); + let array = FixedSizeListArray::new(elements, 3, Validity::NonNullable, 4).into_array(); + let taken = array.take(buffer![1u64, 3].into_array())?; + + let expected_elements = PrimitiveArray::from_iter([3i32, 4, 5, 9, 10, 11]).into_array(); + let expected = + FixedSizeListArray::new(expected_elements, 3, Validity::NonNullable, 2).into_array(); + assert_arrays_eq!(taken, expected, &mut array_session().create_execution_ctx()); + Ok(()) +} diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index dd562281608..a7b316a4c9b 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -6,6 +6,7 @@ mod avx2; use std::sync::LazyLock; +use itertools::Itertools as _; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; @@ -16,15 +17,19 @@ use crate::ArrayRef; use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::ConstantArray; +use crate::arrays::PiecewiseSequence; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; +use crate::arrays::piecewise_sequence::validate_index_ranges; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::executor::ExecutionCtx; use crate::match_each_integer_ptype; use crate::match_each_native_ptype; +use crate::match_each_unsigned_integer_ptype; use crate::scalar::Scalar; use crate::validity::Validity; @@ -79,6 +84,12 @@ impl TakeExecute for Primitive { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_piecewise_sequence(array, piecewise_indices, indices, ctx)? + { + return Ok(Some(taken)); + } + let DType::Primitive(ptype, null) = indices.dtype() else { vortex_bail!("Invalid indices dtype: {}", indices.dtype()) }; @@ -123,6 +134,31 @@ impl TakeExecute for Primitive { } } +fn take_piecewise_sequence( + array: ArrayView<'_, Primitive>, + indices: ArrayView<'_, PiecewiseSequence>, + indices_ref: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some((starts, lengths)) = execute_unit_multiplier_index_arrays(indices, ctx)? else { + return Ok(None); + }; + match_each_native_ptype!(array.ptype(), |T| { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + let values = primitive_piecewise_values::( + array.as_slice::(), + starts.as_slice::(), + lengths.as_slice::(), + indices_ref.len(), + )?; + let validity = array.validity()?.take(indices_ref)?; + Ok(Some(PrimitiveArray::new(values, validity).into_array())) + }) + }) + }) +} + // Compiler may see this as unused based on enabled features #[inline(always)] fn take_primitive_scalar(buffer: &[T], indices: &[I]) -> Buffer { @@ -144,6 +180,29 @@ fn take_primitive_scalar(buffer: &[T], indices: &[I]) result.freeze() } +fn primitive_piecewise_values( + source: &[T], + starts: &[S], + lengths: &[L], + output_len: usize, +) -> VortexResult> +where + T: Copy, + S: crate::dtype::UnsignedPType, + L: crate::dtype::UnsignedPType, +{ + validate_index_ranges(source.len(), starts, lengths, output_len)?; + + let mut values = BufferMut::::with_capacity(output_len); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = start.as_(); + let length = length.as_(); + values.extend_from_slice(&source[start..start + length]); + } + + Ok(values.freeze()) +} + #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] #[cfg(test)] mod test { diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 6a6454d0603..f3c31bc9f1c 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -1,21 +1,27 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use itertools::Itertools as _; use vortex_buffer::BitBufferMut; use vortex_buffer::BufferMut; use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; use crate::array::ArrayView; +use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::VarBin; use crate::arrays::VarBinArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; +use crate::arrays::piecewise_sequence::validate_index_ranges; use crate::arrays::primitive::PrimitiveArrayExt; use crate::arrays::varbin::VarBinArrayExt; use crate::dtype::DType; @@ -43,6 +49,12 @@ impl TakeExecute for VarBin { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_piecewise_sequence(array, piecewise_indices, indices, ctx)? + { + return Ok(Some(taken)); + } + // TODO(joe): Be lazy with execute let offsets = array.offsets().clone().execute::(ctx)?; let data = array.bytes(); @@ -113,6 +125,80 @@ impl TakeExecute for VarBin { } } +fn take_piecewise_sequence( + array: ArrayView<'_, VarBin>, + indices: ArrayView<'_, PiecewiseSequence>, + indices_ref: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some((starts, lengths)) = execute_unit_multiplier_index_arrays(indices, ctx)? else { + return Ok(None); + }; + let offsets = array.offsets().clone().execute::(ctx)?; + let out_offset_ptype = taken_offset_ptype(offsets.ptype()); + let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); + + let result = match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + match offsets.ptype() { + PType::U8 => gather_piecewise_varbin::( + array.dtype().clone(), + offsets.as_slice::(), + array.bytes().as_slice(), + starts.as_slice::(), + lengths.as_slice::(), + indices_ref.len(), + out_offset_ptype, + ), + PType::U16 => gather_piecewise_varbin::( + array.dtype().clone(), + offsets.as_slice::(), + array.bytes().as_slice(), + starts.as_slice::(), + lengths.as_slice::(), + indices_ref.len(), + out_offset_ptype, + ), + PType::U32 => gather_piecewise_varbin::( + array.dtype().clone(), + offsets.as_slice::(), + array.bytes().as_slice(), + starts.as_slice::(), + lengths.as_slice::(), + indices_ref.len(), + out_offset_ptype, + ), + PType::U64 => gather_piecewise_varbin::( + array.dtype().clone(), + offsets.as_slice::(), + array.bytes().as_slice(), + starts.as_slice::(), + lengths.as_slice::(), + indices_ref.len(), + out_offset_ptype, + ), + _ => unreachable!("offsets were reinterpreted to an unsigned integer ptype"), + } + }) + })?; + + let validity = array.validity()?.take(indices_ref)?; + + // SAFETY: output offsets are built from valid input offsets, start at zero, are monotonically + // non-decreasing, and the copied data buffer has exactly the referenced byte length. + unsafe { + Ok(Some( + VarBinArray::new_unchecked( + result.offsets, + result.data.freeze(), + result.dtype, + validity, + ) + .into_array(), + )) + } +} + fn take( dtype: DType, offsets: &[Offset], @@ -183,6 +269,98 @@ fn take( } } +struct GatheredPiecewiseVarBin { + dtype: DType, + offsets: ArrayRef, + data: ByteBufferMut, +} + +fn gather_piecewise_varbin( + dtype: DType, + offsets: &[Offset], + data: &[u8], + starts: &[S], + lengths: &[L], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult +where + S: crate::dtype::UnsignedPType, + L: crate::dtype::UnsignedPType, + Offset: IntegerPType, + NewOffset: IntegerPType, +{ + validate_index_ranges(offsets.len() - 1, starts, lengths, output_len)?; + + let mut new_offsets = BufferMut::::with_capacity(output_len + 1); + new_offsets.push(NewOffset::zero()); + let mut output_bytes = 0usize; + + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = start.as_(); + let length = length.as_(); + let end = start + length; + if length == 0 { + continue; + } + + let byte_start = offsets[start].as_(); + let byte_end = offsets[end].as_(); + vortex_ensure!( + byte_start <= byte_end && byte_end <= data.len(), + "VarBin offsets range {byte_start}..{byte_end} exceeds data length {}", + data.len() + ); + + for &offset in &offsets[start + 1..=end] { + let offset = offset.as_(); + let relative = offset.checked_sub(byte_start).ok_or_else(|| { + vortex_err!("VarBin offsets are not monotonic at offset {offset}") + })?; + let output_offset = output_bytes.checked_add(relative).ok_or_else(|| { + vortex_err!("PiecewiseSequence VarBin output byte length overflow") + })?; + new_offsets.push(new_offset_value::(output_offset)?); + } + + output_bytes = output_bytes + .checked_add(byte_end - byte_start) + .ok_or_else(|| vortex_err!("PiecewiseSequence VarBin output byte length overflow"))?; + } + + let mut new_data = ByteBufferMut::with_capacity(output_bytes); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = start.as_(); + let length = length.as_(); + let end = start + length; + if length == 0 { + continue; + } + + let byte_start = offsets[start].as_(); + let byte_end = offsets[end].as_(); + new_data.extend_from_slice(&data[byte_start..byte_end]); + } + + let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) + .reinterpret_cast(out_offset_ptype) + .into_array(); + Ok(GatheredPiecewiseVarBin { + dtype, + offsets, + data: new_data, + }) +} + +fn new_offset_value(value: usize) -> VortexResult { + T::from(value).ok_or_else(|| { + vortex_err!( + "PiecewiseSequence VarBin offset value {value} does not fit in {}", + T::PTYPE + ) + }) +} + fn take_nullable( dtype: DType, offsets: &[Offset], diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index ce10abda3d0..89206a17c9f 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -4,8 +4,10 @@ use std::iter; use std::sync::Arc; +use itertools::Itertools as _; use num_traits::AsPrimitive; use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_mask::AllOr; use vortex_mask::Mask; @@ -13,14 +15,18 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; use crate::array::ArrayView; +use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; +use crate::arrays::piecewise_sequence::validate_index_ranges; use crate::arrays::varbinview::BinaryView; use crate::buffer::BufferHandle; use crate::executor::ExecutionCtx; use crate::match_each_integer_ptype; +use crate::match_each_unsigned_integer_ptype; impl TakeExecute for VarBinView { /// Take involves creating a new array that references the old array, just with the given set of views. @@ -29,6 +35,12 @@ impl TakeExecute for VarBinView { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_piecewise_sequence(array, piecewise_indices, indices, ctx)? + { + return Ok(Some(taken)); + } + let validity = array.validity()?.take(indices)?; let indices = indices.clone().execute::(ctx)?; @@ -57,6 +69,40 @@ impl TakeExecute for VarBinView { } } +fn take_piecewise_sequence( + array: ArrayView<'_, VarBinView>, + indices: ArrayView<'_, PiecewiseSequence>, + indices_ref: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some((starts, lengths)) = execute_unit_multiplier_index_arrays(indices, ctx)? else { + return Ok(None); + }; + let views = match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + gather_piecewise_views( + array.views(), + starts.as_slice::(), + lengths.as_slice::(), + indices_ref.len(), + )? + }) + }); + let validity = array.validity()?.take(indices_ref)?; + + // SAFETY: ranges were validated against the source views, and copied views still reference the + // same backing data buffers. + unsafe { + Ok(Some(VarBinViewArray::new_handle_unchecked( + BufferHandle::new_host(views.into_byte_buffer()), + Arc::clone(array.data_buffers()), + array.dtype().clone(), + validity, + ) + .into_array())) + } +} + fn take_views>( views_ref: &[BinaryView], indices: &[I], @@ -85,6 +131,28 @@ fn take_views>( } } +fn gather_piecewise_views( + source: &[BinaryView], + starts: &[S], + lengths: &[L], + output_len: usize, +) -> VortexResult> +where + S: crate::dtype::UnsignedPType, + L: crate::dtype::UnsignedPType, +{ + validate_index_ranges(source.len(), starts, lengths, output_len)?; + + let mut views = BufferMut::::with_capacity(output_len); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = start.as_(); + let length = length.as_(); + views.extend_from_slice(&source[start..start + length]); + } + + Ok(views.freeze()) +} + #[cfg(test)] mod tests { use rstest::rstest; From da98e61c9105eee1527b7f25de47c1ddd290ac73 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 16 Jul 2026 14:38:00 -0400 Subject: [PATCH 08/37] Port FSL take benchmarks to PiecewiseSequential Signed-off-by: Daniel King --- Cargo.lock | 1 + vortex-array/benches/take_fsl.rs | 216 +++++++++++++++++++++++++++++-- vortex-bench/Cargo.toml | 1 + 3 files changed, 209 insertions(+), 9 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e23ac0d40f8..384403520e1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9783,6 +9783,7 @@ dependencies = [ "reqwest 0.13.4", "serde", "serde_json", + "sha2 0.11.0", "spatialbench", "spatialbench-arrow", "sysinfo", diff --git a/vortex-array/benches/take_fsl.rs b/vortex-array/benches/take_fsl.rs index 5dc491c28c5..a47896a7497 100644 --- a/vortex-array/benches/take_fsl.rs +++ b/vortex-array/benches/take_fsl.rs @@ -6,6 +6,7 @@ //! Parameterized over: //! - Number of indices to take //! - Fixed size list length (elements per list) +//! - Element byte width #![expect(clippy::cast_possible_truncation)] #![expect(clippy::unwrap_used)] @@ -13,16 +14,28 @@ use std::sync::LazyLock; use divan::Bencher; +use divan::counter::BytesCount; +use num_traits::FromPrimitive; use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; +use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::array_session; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::PiecewiseSequentialArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; +use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::half::f16; +use vortex_array::match_smallest_offset_type; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; use vortex_session::VortexSession; fn main() { @@ -35,16 +48,25 @@ static SESSION: LazyLock = LazyLock::new(array_session); /// Number of lists in the source array. const NUM_LISTS: usize = 500; -/// Number of indices to take. -const NUM_INDICES: &[usize] = &[100, 1_000]; +/// Number of indices to take. This keeps even the widest, longest cases below one millisecond in +/// CodSpeed's instruction-count simulation. +const NUM_INDICES: &[usize] = &[10]; /// Fixed size list lengths (elements per list). -const LIST_SIZES: &[usize] = &[16, 64, 256, 1024, 4096]; +const LIST_SIZES: &[usize] = &[16, 64, 128, 256, 512, 1024, 2048, 4096]; + +/// F16 list lengths for isolating the per-index, piecewise, and manual range-copy strategies. +const F16_STRATEGY_LIST_SIZES: &[usize] = &[1, 2, 4, 8, 16, 64, 128, 256, 512, 1024, 2048, 4096]; /// Creates a FixedSizeListArray with the given list size and number of lists. -fn create_fsl(list_size: usize, num_lists: usize) -> FixedSizeListArray { +fn create_fsl(list_size: usize, num_lists: usize) -> FixedSizeListArray +where + T: NativePType + FromPrimitive, +{ let total_elements = list_size * num_lists; - let elements: Buffer = (0..total_elements as i64).collect(); + let elements: Buffer = (0..total_elements) + .map(|idx| T::from_u16((idx % 251) as u16).unwrap()) + .collect(); FixedSizeListArray::new( elements.into_array(), list_size as u32, @@ -62,12 +84,35 @@ fn create_random_indices(num_indices: usize, max_index: usize) -> Buffer { } #[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] -fn take_fsl_random(bencher: Bencher, num_indices: usize) { - let fsl = create_fsl(LIST_SIZE, NUM_LISTS); +fn take_fsl_f16_random(bencher: Bencher, num_indices: usize) { + take_fsl_random::(bencher, num_indices); +} + +#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] +fn take_fsl_u8_random(bencher: Bencher, num_indices: usize) { + take_fsl_random::(bencher, num_indices); +} + +#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] +fn take_fsl_u32_random(bencher: Bencher, num_indices: usize) { + take_fsl_random::(bencher, num_indices); +} + +#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] +fn take_fsl_u64_random(bencher: Bencher, num_indices: usize) { + take_fsl_random::(bencher, num_indices); +} + +fn take_fsl_random(bencher: Bencher, num_indices: usize) +where + T: NativePType + FromPrimitive, +{ + let fsl = create_fsl::(LIST_SIZE, NUM_LISTS); let indices = create_random_indices(num_indices, NUM_LISTS); let indices_array = indices.into_array(); bencher + .counter(BytesCount::of_many::(num_indices * LIST_SIZE)) .with_inputs(|| (&fsl, &indices_array, SESSION.create_execution_ctx())) .bench_refs(|(array, indices, execution_ctx)| { array @@ -79,10 +124,162 @@ fn take_fsl_random(bencher: Bencher, num_indices: usize) }); } +#[divan::bench(args = NUM_INDICES, consts = F16_STRATEGY_LIST_SIZES)] +fn take_fsl_f16_force_per_index(bencher: Bencher, num_indices: usize) { + let fsl = create_fsl::(LIST_SIZE, NUM_LISTS); + let indices = create_random_indices(num_indices, NUM_LISTS); + + bencher + .counter(BytesCount::of_many::(num_indices * LIST_SIZE)) + .with_inputs(|| (&fsl, &indices, SESSION.create_execution_ctx())) + .bench_refs(|(array, indices, execution_ctx)| { + match_smallest_offset_type!(array.elements().len(), |E| { + take_fsl_f16_per_index_strategy::(array, indices) + }) + .into_array() + .execute::(execution_ctx) + .unwrap() + }); +} + +#[divan::bench(args = NUM_INDICES, consts = F16_STRATEGY_LIST_SIZES)] +fn take_fsl_f16_force_piecewise_sequential( + bencher: Bencher, + num_indices: usize, +) { + let fsl = create_fsl::(LIST_SIZE, NUM_LISTS); + let indices = create_random_indices(num_indices, NUM_LISTS); + + bencher + .counter(BytesCount::of_many::(num_indices * LIST_SIZE)) + .with_inputs(|| (&fsl, &indices, SESSION.create_execution_ctx())) + .bench_refs(|(array, indices, execution_ctx)| { + take_fsl_f16_piecewise_sequential_strategy::(array, indices) + .into_array() + .execute::(execution_ctx) + .unwrap() + }); +} + +#[divan::bench(args = NUM_INDICES, consts = F16_STRATEGY_LIST_SIZES)] +fn take_fsl_f16_force_manual_range_copy( + bencher: Bencher, + num_indices: usize, +) { + let fsl = create_fsl::(LIST_SIZE, NUM_LISTS); + let indices = create_random_indices(num_indices, NUM_LISTS); + + bencher + .counter(BytesCount::of_many::(num_indices * LIST_SIZE)) + .with_inputs(|| (&fsl, &indices, SESSION.create_execution_ctx())) + .bench_refs(|(array, indices, execution_ctx)| { + take_fsl_f16_manual_range_copy_strategy::(array, indices, execution_ctx) + .into_array() + .execute::(execution_ctx) + .unwrap() + }); +} + +fn take_fsl_f16_per_index_strategy( + array: &FixedSizeListArray, + indices: &Buffer, +) -> FixedSizeListArray { + let mut element_indices = BufferMut::::with_capacity(indices.len() * LIST_SIZE); + for &idx in indices.as_ref() { + let start = idx as usize * LIST_SIZE; + let end = start + LIST_SIZE; + for element_idx in start..end { + // SAFETY: capacity is exactly `indices.len() * LIST_SIZE`, and this loop appends + // exactly `LIST_SIZE` element indices for each input index. + unsafe { element_indices.push_unchecked(E::from_usize(element_idx).unwrap()) }; + } + } + + let element_indices = + PrimitiveArray::new(element_indices.freeze(), Validity::NonNullable).into_array(); + let elements = array.elements().take(element_indices).unwrap(); + + // SAFETY: `elements` was built by taking exactly `LIST_SIZE` elements per input index, so its + // length is `indices.len() * LIST_SIZE`; the output is non-nullable by construction. + unsafe { + FixedSizeListArray::new_unchecked( + elements, + LIST_SIZE as u32, + Validity::NonNullable, + indices.len(), + ) + } +} + +fn take_fsl_f16_piecewise_sequential_strategy( + array: &FixedSizeListArray, + indices: &Buffer, +) -> FixedSizeListArray { + let starts = indices + .as_ref() + .iter() + .map(|&idx| idx * LIST_SIZE as u64) + .collect::>(); + let run_count = starts.len(); + let starts = PrimitiveArray::from_iter(starts).into_array(); + let lengths = ConstantArray::new(LIST_SIZE as u64, run_count).into_array(); + + // SAFETY: benchmark indices are generated in-bounds, lengths is a non-nullable unsigned + // constant, and output length is exactly `indices.len() * LIST_SIZE`. + let element_indices = unsafe { + PiecewiseSequentialArray::new_unchecked(starts, lengths, indices.len() * LIST_SIZE) + } + .into_array(); + let elements = array.elements().take(element_indices).unwrap(); + + // SAFETY: each generated run has width `LIST_SIZE`, and there is one run per input index, + // so `elements.len() == indices.len() * LIST_SIZE`. + unsafe { + FixedSizeListArray::new_unchecked( + elements, + LIST_SIZE as u32, + Validity::NonNullable, + indices.len(), + ) + } +} + +fn take_fsl_f16_manual_range_copy_strategy( + array: &FixedSizeListArray, + indices: &Buffer, + execution_ctx: &mut ExecutionCtx, +) -> FixedSizeListArray { + let elements = array + .elements() + .clone() + .execute::(execution_ctx) + .unwrap(); + let source = elements.as_slice::(); + let mut values = BufferMut::::with_capacity(indices.len() * LIST_SIZE); + + for &idx in indices.as_ref() { + let start = idx as usize * LIST_SIZE; + values.extend_from_slice(&source[start..start + LIST_SIZE]); + } + + // SAFETY: the buffer was filled with exactly `LIST_SIZE` copied values per input index, so it + // has the element length required by the constructed FSL. + unsafe { + FixedSizeListArray::new_unchecked( + PrimitiveArray::new(values.freeze(), Validity::NonNullable).into_array(), + LIST_SIZE as u32, + Validity::NonNullable, + indices.len(), + ) + } +} + #[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] -fn take_fsl_nullable_random(bencher: Bencher, num_indices: usize) { +fn take_fsl_f16_nullable_random(bencher: Bencher, num_indices: usize) { let total_elements = LIST_SIZE * NUM_LISTS; - let elements: Buffer = (0..total_elements as i64).collect(); + let elements: Buffer = (0..total_elements) + .map(|idx| f16::from_u16((idx % 251) as u16).unwrap()) + .collect(); // Create validity with ~10% nulls let mut rng = StdRng::seed_from_u64(123); @@ -94,6 +291,7 @@ fn take_fsl_nullable_random(bencher: Bencher, num_indice let indices_array = indices.into_array(); bencher + .counter(BytesCount::of_many::(num_indices * LIST_SIZE)) .with_inputs(|| (&fsl, &indices_array, SESSION.create_execution_ctx())) .bench_refs(|(array, indices, execution_ctx)| { array diff --git a/vortex-bench/Cargo.toml b/vortex-bench/Cargo.toml index f228e1571d3..dcf445eb99e 100644 --- a/vortex-bench/Cargo.toml +++ b/vortex-bench/Cargo.toml @@ -54,6 +54,7 @@ regex = { workspace = true } reqwest = { workspace = true, features = ["stream"] } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } +sha2 = { workspace = true } spatialbench = { workspace = true } spatialbench-arrow = { workspace = true } spatialbench-parquet = { workspace = true } From ca7d42382541ec702b8d7acb2d504ead57719dce Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 16 Jul 2026 16:24:01 -0400 Subject: [PATCH 09/37] Gate PiecewiseSequence take paths on unit multipliers Signed-off-by: Daniel King --- vortex-array/benches/take_fsl.rs | 20 ++++++++++++------- .../src/arrays/decimal/compute/take.rs | 10 ++++------ .../src/arrays/piecewise_sequence/mod.rs | 12 +++++------ .../src/arrays/piecewise_sequence/tests.rs | 20 +++++++++++++++++++ .../src/arrays/varbinview/compute/take.rs | 16 ++++++++------- 5 files changed, 52 insertions(+), 26 deletions(-) diff --git a/vortex-array/benches/take_fsl.rs b/vortex-array/benches/take_fsl.rs index a47896a7497..e02d9695b82 100644 --- a/vortex-array/benches/take_fsl.rs +++ b/vortex-array/benches/take_fsl.rs @@ -26,7 +26,7 @@ use vortex_array::VortexSessionExecute; use vortex_array::array_session; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeListArray; -use vortex_array::arrays::PiecewiseSequentialArray; +use vortex_array::arrays::PiecewiseSequenceArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; use vortex_array::dtype::IntegerPType; @@ -143,7 +143,7 @@ fn take_fsl_f16_force_per_index(bencher: Bencher, num_in } #[divan::bench(args = NUM_INDICES, consts = F16_STRATEGY_LIST_SIZES)] -fn take_fsl_f16_force_piecewise_sequential( +fn take_fsl_f16_force_piecewise_sequence( bencher: Bencher, num_indices: usize, ) { @@ -154,7 +154,7 @@ fn take_fsl_f16_force_piecewise_sequential( .counter(BytesCount::of_many::(num_indices * LIST_SIZE)) .with_inputs(|| (&fsl, &indices, SESSION.create_execution_ctx())) .bench_refs(|(array, indices, execution_ctx)| { - take_fsl_f16_piecewise_sequential_strategy::(array, indices) + take_fsl_f16_piecewise_sequence_strategy::(array, indices) .into_array() .execute::(execution_ctx) .unwrap() @@ -211,7 +211,7 @@ fn take_fsl_f16_per_index_strategy( } } -fn take_fsl_f16_piecewise_sequential_strategy( +fn take_fsl_f16_piecewise_sequence_strategy( array: &FixedSizeListArray, indices: &Buffer, ) -> FixedSizeListArray { @@ -223,11 +223,17 @@ fn take_fsl_f16_piecewise_sequential_strategy( let run_count = starts.len(); let starts = PrimitiveArray::from_iter(starts).into_array(); let lengths = ConstantArray::new(LIST_SIZE as u64, run_count).into_array(); + let multipliers = ConstantArray::new(1u64, run_count).into_array(); - // SAFETY: benchmark indices are generated in-bounds, lengths is a non-nullable unsigned - // constant, and output length is exactly `indices.len() * LIST_SIZE`. + // SAFETY: benchmark indices are generated in-bounds; lengths and multiplier 1 are + // non-nullable unsigned constants; output length is exactly `indices.len() * LIST_SIZE`. let element_indices = unsafe { - PiecewiseSequentialArray::new_unchecked(starts, lengths, indices.len() * LIST_SIZE) + PiecewiseSequenceArray::new_unchecked( + starts, + lengths, + multipliers, + indices.len() * LIST_SIZE, + ) } .into_array(); let elements = array.elements().take(element_indices).unwrap(); diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index 213d6db39bf..0bd79299d5c 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -75,12 +75,10 @@ fn take_piecewise_sequence( let validity = array.validity()?.take(indices_ref)?; // SAFETY: contiguous gather preserves the decimal dtype and value representation. - Ok( - Some(unsafe { - DecimalArray::new_unchecked(values, array.decimal_dtype(), validity) - } - .into_array()), - ) + Ok(Some( + unsafe { DecimalArray::new_unchecked(values, array.decimal_dtype(), validity) } + .into_array(), + )) }) }) }) diff --git a/vortex-array/src/arrays/piecewise_sequence/mod.rs b/vortex-array/src/arrays/piecewise_sequence/mod.rs index a5361567c7d..37f009ca19f 100644 --- a/vortex-array/src/arrays/piecewise_sequence/mod.rs +++ b/vortex-array/src/arrays/piecewise_sequence/mod.rs @@ -163,16 +163,16 @@ where for (&start, &length) in starts.iter().zip_eq(lengths) { let start: usize = start.as_(); let length: usize = length.as_(); - let end = start.checked_add(length).ok_or_else(|| { - vortex_err!("PiecewiseSequenceArray range overflows usize") - })?; + let end = start + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; vortex_ensure!( end <= source_len, "PiecewiseSequenceArray range {start}..{end} exceeds source length {source_len}" ); - computed_len = computed_len.checked_add(length).ok_or_else(|| { - vortex_err!("PiecewiseSequenceArray output length overflows usize") - })?; + computed_len = computed_len + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; } vortex_ensure!( diff --git a/vortex-array/src/arrays/piecewise_sequence/tests.rs b/vortex-array/src/arrays/piecewise_sequence/tests.rs index 1abbada7a92..bc11bb9dc69 100644 --- a/vortex-array/src/arrays/piecewise_sequence/tests.rs +++ b/vortex-array/src/arrays/piecewise_sequence/tests.rs @@ -187,6 +187,26 @@ fn primitive_take_consumes_piecewise_indices() -> VortexResult<()> { Ok(()) } +#[test] +fn primitive_take_handles_non_unit_multiplier() -> VortexResult<()> { + let values = PrimitiveArray::from_iter(0i32..20).into_array(); + let indices = PiecewiseSequenceArray::try_new( + buffer![3u64].into_array(), + buffer![3u64].into_array(), + buffer![2u64].into_array(), + 3, + )? + .into_array(); + let taken = values.take(indices)?; + + assert_arrays_eq!( + taken, + PrimitiveArray::from_iter([3i32, 5, 7]).into_array(), + &mut array_session().create_execution_ctx() + ); + Ok(()) +} + #[test] fn bool_take_consumes_piecewise_indices() -> VortexResult<()> { let values = BoolArray::from_iter([true, false, true, true, false, false]).into_array(); diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index 89206a17c9f..dadc3041a3c 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -93,13 +93,15 @@ fn take_piecewise_sequence( // SAFETY: ranges were validated against the source views, and copied views still reference the // same backing data buffers. unsafe { - Ok(Some(VarBinViewArray::new_handle_unchecked( - BufferHandle::new_host(views.into_byte_buffer()), - Arc::clone(array.data_buffers()), - array.dtype().clone(), - validity, - ) - .into_array())) + Ok(Some( + VarBinViewArray::new_handle_unchecked( + BufferHandle::new_host(views.into_byte_buffer()), + Arc::clone(array.data_buffers()), + array.dtype().clone(), + validity, + ) + .into_array(), + )) } } From 6ba9fe5e5d4e94a94f01479fc7ff066bf9bc4c7d Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 16 Jul 2026 16:36:18 -0400 Subject: [PATCH 10/37] Preserve existing FSL take benchmarks Signed-off-by: Daniel King --- vortex-array/benches/take_fsl.rs | 106 ++++++++++++++++++++----------- 1 file changed, 69 insertions(+), 37 deletions(-) diff --git a/vortex-array/benches/take_fsl.rs b/vortex-array/benches/take_fsl.rs index e02d9695b82..cfd434b16c4 100644 --- a/vortex-array/benches/take_fsl.rs +++ b/vortex-array/benches/take_fsl.rs @@ -48,9 +48,8 @@ static SESSION: LazyLock = LazyLock::new(array_session); /// Number of lists in the source array. const NUM_LISTS: usize = 500; -/// Number of indices to take. This keeps even the widest, longest cases below one millisecond in -/// CodSpeed's instruction-count simulation. -const NUM_INDICES: &[usize] = &[10]; +/// Number of indices to take. +const NUM_INDICES: &[usize] = &[100, 1_000]; /// Fixed size list lengths (elements per list). const LIST_SIZES: &[usize] = &[16, 64, 128, 256, 512, 1024, 2048, 4096]; @@ -58,8 +57,22 @@ const LIST_SIZES: &[usize] = &[16, 64, 128, 256, 512, 1024, 2048, 4096]; /// F16 list lengths for isolating the per-index, piecewise, and manual range-copy strategies. const F16_STRATEGY_LIST_SIZES: &[usize] = &[1, 2, 4, 8, 16, 64, 128, 256, 512, 1024, 2048, 4096]; +/// F16 strategy benchmarks keep a smaller take width so the forced slow strategies stay cheap. +const F16_STRATEGY_NUM_INDICES: &[usize] = &[10]; + /// Creates a FixedSizeListArray with the given list size and number of lists. fn create_fsl(list_size: usize, num_lists: usize) -> FixedSizeListArray +where + T: NativePType + FromPrimitive, +{ + create_fsl_with_validity::(list_size, num_lists, Validity::NonNullable) +} + +fn create_fsl_with_validity( + list_size: usize, + num_lists: usize, + validity: Validity, +) -> FixedSizeListArray where T: NativePType + FromPrimitive, { @@ -67,12 +80,17 @@ where let elements: Buffer = (0..total_elements) .map(|idx| T::from_u16((idx % 251) as u16).unwrap()) .collect(); - FixedSizeListArray::new( - elements.into_array(), - list_size as u32, - Validity::NonNullable, - num_lists, - ) + FixedSizeListArray::new(elements.into_array(), list_size as u32, validity, num_lists) +} + +fn create_i64_fsl_with_validity( + list_size: usize, + num_lists: usize, + validity: Validity, +) -> FixedSizeListArray { + let total_elements = list_size * num_lists; + let elements: Buffer = (0..total_elements as i64).collect(); + FixedSizeListArray::new(elements.into_array(), list_size as u32, validity, num_lists) } /// Creates random indices for taking from the array. @@ -83,31 +101,47 @@ fn create_random_indices(num_indices: usize, max_index: usize) -> Buffer { .collect() } +#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] +fn take_fsl_random(bencher: Bencher, num_indices: usize) { + let fsl = create_i64_fsl_with_validity(LIST_SIZE, NUM_LISTS, Validity::NonNullable); + bench_take_fsl_random::(bencher, num_indices, fsl); +} + #[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] fn take_fsl_f16_random(bencher: Bencher, num_indices: usize) { - take_fsl_random::(bencher, num_indices); + take_fsl_random_typed::(bencher, num_indices); } #[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] fn take_fsl_u8_random(bencher: Bencher, num_indices: usize) { - take_fsl_random::(bencher, num_indices); + take_fsl_random_typed::(bencher, num_indices); } #[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] fn take_fsl_u32_random(bencher: Bencher, num_indices: usize) { - take_fsl_random::(bencher, num_indices); + take_fsl_random_typed::(bencher, num_indices); } #[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] fn take_fsl_u64_random(bencher: Bencher, num_indices: usize) { - take_fsl_random::(bencher, num_indices); + take_fsl_random_typed::(bencher, num_indices); } -fn take_fsl_random(bencher: Bencher, num_indices: usize) +fn take_fsl_random_typed(bencher: Bencher, num_indices: usize) where T: NativePType + FromPrimitive, { let fsl = create_fsl::(LIST_SIZE, NUM_LISTS); + bench_take_fsl_random::(bencher, num_indices, fsl); +} + +fn bench_take_fsl_random( + bencher: Bencher, + num_indices: usize, + fsl: FixedSizeListArray, +) where + T: NativePType, +{ let indices = create_random_indices(num_indices, NUM_LISTS); let indices_array = indices.into_array(); @@ -124,7 +158,7 @@ where }); } -#[divan::bench(args = NUM_INDICES, consts = F16_STRATEGY_LIST_SIZES)] +#[divan::bench(args = F16_STRATEGY_NUM_INDICES, consts = F16_STRATEGY_LIST_SIZES)] fn take_fsl_f16_force_per_index(bencher: Bencher, num_indices: usize) { let fsl = create_fsl::(LIST_SIZE, NUM_LISTS); let indices = create_random_indices(num_indices, NUM_LISTS); @@ -142,7 +176,7 @@ fn take_fsl_f16_force_per_index(bencher: Bencher, num_in }); } -#[divan::bench(args = NUM_INDICES, consts = F16_STRATEGY_LIST_SIZES)] +#[divan::bench(args = F16_STRATEGY_NUM_INDICES, consts = F16_STRATEGY_LIST_SIZES)] fn take_fsl_f16_force_piecewise_sequence( bencher: Bencher, num_indices: usize, @@ -161,7 +195,7 @@ fn take_fsl_f16_force_piecewise_sequence( }); } -#[divan::bench(args = NUM_INDICES, consts = F16_STRATEGY_LIST_SIZES)] +#[divan::bench(args = F16_STRATEGY_NUM_INDICES, consts = F16_STRATEGY_LIST_SIZES)] fn take_fsl_f16_force_manual_range_copy( bencher: Bencher, num_indices: usize, @@ -281,30 +315,28 @@ fn take_fsl_f16_manual_range_copy_strategy( } #[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] -fn take_fsl_f16_nullable_random(bencher: Bencher, num_indices: usize) { - let total_elements = LIST_SIZE * NUM_LISTS; - let elements: Buffer = (0..total_elements) - .map(|idx| f16::from_u16((idx % 251) as u16).unwrap()) - .collect(); - +fn take_fsl_nullable_random(bencher: Bencher, num_indices: usize) { // Create validity with ~10% nulls let mut rng = StdRng::seed_from_u64(123); let validity = Validity::from_iter((0..NUM_LISTS).map(|_| rng.random_ratio(9, 10))); - let fsl = FixedSizeListArray::new(elements.into_array(), LIST_SIZE as u32, validity, NUM_LISTS); + let fsl = create_i64_fsl_with_validity(LIST_SIZE, NUM_LISTS, validity); + bench_take_fsl_random::(bencher, num_indices, fsl); +} - let indices = create_random_indices(num_indices, NUM_LISTS); - let indices_array = indices.into_array(); +#[divan::bench(args = NUM_INDICES, consts = LIST_SIZES)] +fn take_fsl_f16_nullable_random(bencher: Bencher, num_indices: usize) { + take_fsl_nullable_random_typed::(bencher, num_indices); +} - bencher - .counter(BytesCount::of_many::(num_indices * LIST_SIZE)) - .with_inputs(|| (&fsl, &indices_array, SESSION.create_execution_ctx())) - .bench_refs(|(array, indices, execution_ctx)| { - array - .clone() - .take(indices.clone()) - .unwrap() - .execute::(execution_ctx) - .unwrap() - }); +fn take_fsl_nullable_random_typed(bencher: Bencher, num_indices: usize) +where + T: NativePType + FromPrimitive, +{ + // Create validity with ~10% nulls + let mut rng = StdRng::seed_from_u64(123); + let validity = Validity::from_iter((0..NUM_LISTS).map(|_| rng.random_ratio(9, 10))); + + let fsl = create_fsl_with_validity::(LIST_SIZE, NUM_LISTS, validity); + bench_take_fsl_random::(bencher, num_indices, fsl); } From 9ceff2794354b91dc3b25149ef5c38a74b73e7e1 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 16 Jul 2026 19:40:37 -0400 Subject: [PATCH 11/37] Split PiecewiseSequence take dispatch Signed-off-by: Daniel King --- .../src/arrays/decimal/compute/take.rs | 81 ++++++++++++++----- .../src/arrays/primitive/compute/take/mod.rs | 75 +++++++++++++---- 2 files changed, 121 insertions(+), 35 deletions(-) diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index 0bd79299d5c..ed479b42bbb 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -18,10 +18,12 @@ use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; use crate::arrays::piecewise_sequence::validate_index_ranges; use crate::dtype::IntegerPType; use crate::dtype::NativeDecimalType; +use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_decimal_value_type; use crate::match_each_integer_ptype; use crate::match_each_unsigned_integer_ptype; +use crate::validity::Validity; impl TakeExecute for Decimal { fn take( @@ -63,24 +65,65 @@ fn take_piecewise_sequence( let Some((starts, lengths)) = execute_unit_multiplier_index_arrays(indices, ctx)? else { return Ok(None); }; + let validity = array.validity()?.take(indices_ref)?; + let output_len = indices_ref.len(); + let taken = take_piecewise_sequence_lengths(array, &starts, &lengths, validity, output_len)?; + Ok(Some(taken)) +} + +fn take_piecewise_sequence_lengths( + array: ArrayView<'_, Decimal>, + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + validity: Validity, + output_len: usize, +) -> VortexResult { match_each_decimal_value_type!(array.values_type(), |D| { - match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { - let values = take_piecewise_to_buffer::( - starts.as_slice::(), - lengths.as_slice::(), - array.buffer::().as_slice(), - indices_ref.len(), - )?; - let validity = array.validity()?.take(indices_ref)?; - - // SAFETY: contiguous gather preserves the decimal dtype and value representation. - Ok(Some( - unsafe { DecimalArray::new_unchecked(values, array.decimal_dtype(), validity) } - .into_array(), - )) - }) - }) + take_piecewise_sequence_lengths_typed::(array, starts, lengths, validity, output_len) + }) +} + +fn take_piecewise_sequence_lengths_typed( + array: ArrayView<'_, Decimal>, + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + validity: Validity, + output_len: usize, +) -> VortexResult +where + D: NativeDecimalType, +{ + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + take_piecewise_sequence_lengths_start_typed::( + array, starts, lengths, validity, output_len, + ) + }) +} + +fn take_piecewise_sequence_lengths_start_typed( + array: ArrayView<'_, Decimal>, + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + validity: Validity, + output_len: usize, +) -> VortexResult +where + D: NativeDecimalType, + S: UnsignedPType, +{ + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + let values = take_piecewise_to_buffer::( + starts.as_slice::(), + lengths.as_slice::(), + array.buffer::().as_slice(), + output_len, + )?; + + // SAFETY: contiguous gather preserves the decimal dtype and value representation. + Ok( + unsafe { DecimalArray::new_unchecked(values, array.decimal_dtype(), validity) } + .into_array(), + ) }) } @@ -95,8 +138,8 @@ fn take_piecewise_to_buffer( output_len: usize, ) -> VortexResult> where - S: crate::dtype::UnsignedPType, - L: crate::dtype::UnsignedPType, + S: UnsignedPType, + L: UnsignedPType, T: NativeDecimalType, { validate_index_ranges(values.len(), starts, lengths, output_len)?; diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index a7b316a4c9b..c67a2c5d725 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -26,6 +26,8 @@ use crate::arrays::piecewise_sequence::validate_index_ranges; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::IntegerPType; +use crate::dtype::NativePType; +use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_integer_ptype; use crate::match_each_native_ptype; @@ -143,20 +145,10 @@ fn take_piecewise_sequence( let Some((starts, lengths)) = execute_unit_multiplier_index_arrays(indices, ctx)? else { return Ok(None); }; - match_each_native_ptype!(array.ptype(), |T| { - match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { - let values = primitive_piecewise_values::( - array.as_slice::(), - starts.as_slice::(), - lengths.as_slice::(), - indices_ref.len(), - )?; - let validity = array.validity()?.take(indices_ref)?; - Ok(Some(PrimitiveArray::new(values, validity).into_array())) - }) - }) - }) + let validity = array.validity()?.take(indices_ref)?; + let output_len = indices_ref.len(); + let taken = take_piecewise_sequence_lengths(array, &starts, &lengths, validity, output_len)?; + Ok(Some(taken)) } // Compiler may see this as unused based on enabled features @@ -180,6 +172,57 @@ fn take_primitive_scalar(buffer: &[T], indices: &[I]) result.freeze() } +fn take_piecewise_sequence_lengths( + array: ArrayView<'_, Primitive>, + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + validity: Validity, + output_len: usize, +) -> VortexResult { + match_each_native_ptype!(array.ptype(), |T| { + take_piecewise_sequence_lengths_typed::(array, starts, lengths, validity, output_len) + }) +} + +fn take_piecewise_sequence_lengths_typed( + array: ArrayView<'_, Primitive>, + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + validity: Validity, + output_len: usize, +) -> VortexResult +where + T: NativePType, +{ + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + take_piecewise_sequence_lengths_start_typed::( + array, starts, lengths, validity, output_len, + ) + }) +} + +fn take_piecewise_sequence_lengths_start_typed( + array: ArrayView<'_, Primitive>, + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + validity: Validity, + output_len: usize, +) -> VortexResult +where + T: NativePType, + S: UnsignedPType, +{ + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + let values = primitive_piecewise_values::( + array.as_slice::(), + starts.as_slice::(), + lengths.as_slice::(), + output_len, + )?; + Ok(PrimitiveArray::new(values, validity).into_array()) + }) +} + fn primitive_piecewise_values( source: &[T], starts: &[S], @@ -188,8 +231,8 @@ fn primitive_piecewise_values( ) -> VortexResult> where T: Copy, - S: crate::dtype::UnsignedPType, - L: crate::dtype::UnsignedPType, + S: UnsignedPType, + L: UnsignedPType, { validate_index_ranges(source.len(), starts, lengths, output_len)?; From 7cd3098ac7730f1beed2fbca57bb2c4016752cdf Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 16 Jul 2026 19:41:48 -0400 Subject: [PATCH 12/37] Fix PiecewiseSequence test length cast Signed-off-by: Daniel King --- vortex-array/src/arrays/piecewise_sequence/tests.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/vortex-array/src/arrays/piecewise_sequence/tests.rs b/vortex-array/src/arrays/piecewise_sequence/tests.rs index bc11bb9dc69..1e4bb4329b2 100644 --- a/vortex-array/src/arrays/piecewise_sequence/tests.rs +++ b/vortex-array/src/arrays/piecewise_sequence/tests.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use num_traits::AsPrimitive; use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -32,7 +33,10 @@ fn piecewise_indices( starts: impl IntoIterator, lengths: &[u64], ) -> VortexResult { - let len = lengths.iter().map(|&length| length as usize).sum(); + let len = lengths + .iter() + .map(|&length| -> usize { length.as_() }) + .sum(); let starts = PrimitiveArray::from_iter(starts).into_array(); let lengths = PrimitiveArray::from_iter(lengths.iter().copied()).into_array(); let multipliers = ConstantArray::new(1u64, lengths.len()).into_array(); From d694338e9fb72b653c3d37aab2349266cc837602 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 16 Jul 2026 20:15:07 -0400 Subject: [PATCH 13/37] Use imported UnsignedPType in take consumers Signed-off-by: Daniel King --- vortex-array/src/arrays/bool/compute/take.rs | 5 +++-- vortex-array/src/arrays/varbin/compute/take.rs | 5 +++-- vortex-array/src/arrays/varbinview/compute/take.rs | 5 +++-- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/vortex-array/src/arrays/bool/compute/take.rs b/vortex-array/src/arrays/bool/compute/take.rs index fe7955fbbae..2caca36e5df 100644 --- a/vortex-array/src/arrays/bool/compute/take.rs +++ b/vortex-array/src/arrays/bool/compute/take.rs @@ -23,6 +23,7 @@ use crate::arrays::dict::TakeExecute; use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; use crate::arrays::piecewise_sequence::validate_index_ranges; use crate::builtins::ArrayBuiltins; +use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_integer_ptype; use crate::match_each_unsigned_integer_ptype; @@ -125,8 +126,8 @@ fn take_piecewise_bits( output_len: usize, ) -> VortexResult where - S: crate::dtype::UnsignedPType, - L: crate::dtype::UnsignedPType, + S: UnsignedPType, + L: UnsignedPType, { validate_index_ranges(source.len(), starts, lengths, output_len)?; diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index f3c31bc9f1c..34d19863e20 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -27,6 +27,7 @@ use crate::arrays::varbin::VarBinArrayExt; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::dtype::PType; +use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_unsigned_integer_ptype; use crate::validity::Validity; @@ -285,8 +286,8 @@ fn gather_piecewise_varbin( out_offset_ptype: PType, ) -> VortexResult where - S: crate::dtype::UnsignedPType, - L: crate::dtype::UnsignedPType, + S: UnsignedPType, + L: UnsignedPType, Offset: IntegerPType, NewOffset: IntegerPType, { diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index dadc3041a3c..dd38b45929d 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -24,6 +24,7 @@ use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; use crate::arrays::piecewise_sequence::validate_index_ranges; use crate::arrays::varbinview::BinaryView; use crate::buffer::BufferHandle; +use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_integer_ptype; use crate::match_each_unsigned_integer_ptype; @@ -140,8 +141,8 @@ fn gather_piecewise_views( output_len: usize, ) -> VortexResult> where - S: crate::dtype::UnsignedPType, - L: crate::dtype::UnsignedPType, + S: UnsignedPType, + L: UnsignedPType, { validate_index_ranges(source.len(), starts, lengths, output_len)?; From 7c4811bb06f6af841d0f7f9a261c12d147c87326 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 11:28:32 -0400 Subject: [PATCH 14/37] Specialize constant PiecewiseSequence lengths Signed-off-by: Daniel King --- vortex-array/src/arrays/bool/compute/take.rs | 58 +++- .../src/arrays/decimal/compute/take.rs | 72 +++- .../src/arrays/piecewise_sequence/mod.rs | 66 +++- .../src/arrays/piecewise_sequence/tests.rs | 74 +++++ .../src/arrays/primitive/compute/take/mod.rs | 67 +++- .../src/arrays/varbin/compute/take.rs | 313 +++++++++++++++--- .../src/arrays/varbinview/compute/take.rs | 58 +++- 7 files changed, 630 insertions(+), 78 deletions(-) diff --git a/vortex-array/src/arrays/bool/compute/take.rs b/vortex-array/src/arrays/bool/compute/take.rs index 2caca36e5df..a467e19175f 100644 --- a/vortex-array/src/arrays/bool/compute/take.rs +++ b/vortex-array/src/arrays/bool/compute/take.rs @@ -20,8 +20,10 @@ use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::UnitMultiplierLengths; use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; use crate::arrays::piecewise_sequence::validate_index_ranges; +use crate::arrays::piecewise_sequence::validate_index_ranges_constant; use crate::builtins::ArrayBuiltins; use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; @@ -76,16 +78,32 @@ fn take_piecewise_sequence( let Some((starts, lengths)) = execute_unit_multiplier_index_arrays(indices, ctx)? else { return Ok(None); }; - let buffer = match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { - take_piecewise_bits( - &array.to_bit_buffer(), - starts.as_slice::(), - lengths.as_slice::(), - indices_ref.len(), - )? - }) - }); + let source = array.to_bit_buffer(); + let output_len = indices_ref.len(); + let buffer = match &lengths { + UnitMultiplierLengths::Constant(length) => { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + take_piecewise_bits_constant_length( + &source, + starts.as_slice::(), + *length, + output_len, + )? + }) + } + UnitMultiplierLengths::Array(lengths) => { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + take_piecewise_bits( + &source, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + )? + }) + }) + } + }; Ok(Some( BoolArray::new(buffer, array.validity()?.take(indices_ref)?).into_array(), @@ -119,6 +137,26 @@ fn take_bool_impl>(bools: BitBufferView<'_>, indices: &[I] }) } +fn take_piecewise_bits_constant_length( + source: &BitBuffer, + starts: &[S], + length: usize, + output_len: usize, +) -> VortexResult +where + S: UnsignedPType, +{ + validate_index_ranges_constant(source.len(), starts, length, output_len)?; + + let mut values = BitBufferMut::with_capacity(output_len); + for &start in starts { + let start = start.as_(); + values.append_buffer(&source.slice(start..start + length)); + } + + Ok(values.freeze()) +} + fn take_piecewise_bits( source: &BitBuffer, starts: &[S], diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index ed479b42bbb..a2d508fd31e 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -14,8 +14,10 @@ use crate::arrays::DecimalArray; use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::UnitMultiplierLengths; use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; use crate::arrays::piecewise_sequence::validate_index_ranges; +use crate::arrays::piecewise_sequence::validate_index_ranges_constant; use crate::dtype::IntegerPType; use crate::dtype::NativeDecimalType; use crate::dtype::UnsignedPType; @@ -67,10 +69,57 @@ fn take_piecewise_sequence( }; let validity = array.validity()?.take(indices_ref)?; let output_len = indices_ref.len(); - let taken = take_piecewise_sequence_lengths(array, &starts, &lengths, validity, output_len)?; + let taken = match lengths { + UnitMultiplierLengths::Constant(length) => { + take_piecewise_sequence_constant_length(array, &starts, length, validity, output_len)? + } + UnitMultiplierLengths::Array(lengths) => { + take_piecewise_sequence_lengths(array, &starts, &lengths, validity, output_len)? + } + }; Ok(Some(taken)) } +fn take_piecewise_sequence_constant_length( + array: ArrayView<'_, Decimal>, + starts: &PrimitiveArray, + length: usize, + validity: Validity, + output_len: usize, +) -> VortexResult { + match_each_decimal_value_type!(array.values_type(), |D| { + take_piecewise_sequence_constant_length_typed::( + array, starts, length, validity, output_len, + ) + }) +} + +fn take_piecewise_sequence_constant_length_typed( + array: ArrayView<'_, Decimal>, + starts: &PrimitiveArray, + length: usize, + validity: Validity, + output_len: usize, +) -> VortexResult +where + D: NativeDecimalType, +{ + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + let values = take_piecewise_constant_length_to_buffer::( + starts.as_slice::(), + length, + array.buffer::().as_slice(), + output_len, + )?; + + // SAFETY: contiguous gather preserves the decimal dtype and value representation. + Ok( + unsafe { DecimalArray::new_unchecked(values, array.decimal_dtype(), validity) } + .into_array(), + ) + }) +} + fn take_piecewise_sequence_lengths( array: ArrayView<'_, Decimal>, starts: &PrimitiveArray, @@ -131,6 +180,27 @@ fn take_to_buffer(indices: &[I], values: indices.iter().map(|idx| values[idx.as_()]).collect() } +fn take_piecewise_constant_length_to_buffer( + starts: &[S], + length: usize, + values: &[T], + output_len: usize, +) -> VortexResult> +where + S: UnsignedPType, + T: NativeDecimalType, +{ + validate_index_ranges_constant(values.len(), starts, length, output_len)?; + + let mut result = BufferMut::::with_capacity(output_len); + for &start in starts { + let start = start.as_(); + result.extend_from_slice(&values[start..start + length]); + } + + Ok(result.freeze()) +} + fn take_piecewise_to_buffer( starts: &[S], lengths: &[L], diff --git a/vortex-array/src/arrays/piecewise_sequence/mod.rs b/vortex-array/src/arrays/piecewise_sequence/mod.rs index 37f009ca19f..e5043e42e16 100644 --- a/vortex-array/src/arrays/piecewise_sequence/mod.rs +++ b/vortex-array/src/arrays/piecewise_sequence/mod.rs @@ -9,6 +9,7 @@ //! element. use itertools::Itertools; +use num_traits::AsPrimitive; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -66,18 +67,29 @@ pub(crate) fn execute_index_arrays( Ok((starts, lengths, multipliers)) } +pub(crate) enum UnitMultiplierLengths { + Constant(usize), + Array(PrimitiveArray), +} + pub(crate) fn execute_unit_multiplier_index_arrays( array: ArrayView<'_, PiecewiseSequence>, ctx: &mut ExecutionCtx, -) -> VortexResult> { +) -> VortexResult> { if !is_constant_multiplier_one(array.multipliers()) { return Ok(None); } + check_index_arrays(array.starts(), array.lengths(), array.multipliers())?; let starts = array.starts().clone().execute::(ctx)?; + if let Some(length) = constant_unsigned_usize(array.lengths())? { + check_index_arrays(starts.as_ref(), array.lengths(), array.multipliers())?; + return Ok(Some((starts, UnitMultiplierLengths::Constant(length)))); + } + let lengths = array.lengths().clone().execute::(ctx)?; check_index_arrays(starts.as_ref(), lengths.as_ref(), array.multipliers())?; - Ok(Some((starts, lengths))) + Ok(Some((starts, UnitMultiplierLengths::Array(lengths)))) } pub(crate) fn is_constant_multiplier_one(multipliers: &ArrayRef) -> bool { @@ -90,6 +102,24 @@ pub(crate) fn is_constant_multiplier_one(multipliers: &ArrayRef) -> bool { ) } +fn constant_unsigned_usize(array: &ArrayRef) -> VortexResult> { + let Some(scalar) = array.as_constant() else { + return Ok(None); + }; + + let Some(pvalue) = scalar.as_primitive_opt().and_then(|scalar| scalar.pvalue()) else { + vortex_bail!("PiecewiseSequenceArray constant length must be an unsigned integer"); + }; + + Ok(Some(match pvalue { + PValue::U8(value) => value as usize, + PValue::U16(value) => value as usize, + PValue::U32(value) => value as usize, + PValue::U64(value) => value.as_(), + _ => vortex_bail!("PiecewiseSequenceArray constant length must be an unsigned integer"), + })) +} + fn check_index_array(name: &str, array: &ArrayRef) -> VortexResult<()> { vortex_ensure!( array.dtype().is_unsigned_int(), @@ -104,6 +134,38 @@ fn check_index_array(name: &str, array: &ArrayRef) -> VortexResult<()> { Ok(()) } +pub(crate) fn validate_index_ranges_constant( + source_len: usize, + starts: &[S], + length: usize, + output_len: usize, +) -> VortexResult<()> +where + S: UnsignedPType, +{ + let computed_len = starts + .len() + .checked_mul(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); + + for &start in starts { + let start: usize = start.as_(); + let end = start + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + vortex_ensure!( + end <= source_len, + "PiecewiseSequenceArray range {start}..{end} exceeds source length {source_len}" + ); + } + + Ok(()) +} + pub(crate) fn materialize_ranges( starts: &PrimitiveArray, lengths: &PrimitiveArray, diff --git a/vortex-array/src/arrays/piecewise_sequence/tests.rs b/vortex-array/src/arrays/piecewise_sequence/tests.rs index 1e4bb4329b2..82ead343f96 100644 --- a/vortex-array/src/arrays/piecewise_sequence/tests.rs +++ b/vortex-array/src/arrays/piecewise_sequence/tests.rs @@ -43,6 +43,20 @@ fn piecewise_indices( Ok(PiecewiseSequenceArray::try_new(starts, lengths, multipliers, len)?.into_array()) } +fn piecewise_indices_constant_length( + starts: impl IntoIterator, + length: u64, +) -> VortexResult { + let starts = starts.into_iter().collect::>(); + let length_usize: usize = length.as_(); + let len = starts.len() * length_usize; + let piece_count = starts.len(); + let starts = PrimitiveArray::from_iter(starts).into_array(); + let lengths = ConstantArray::new(length, piece_count).into_array(); + let multipliers = ConstantArray::new(1u64, piece_count).into_array(); + Ok(PiecewiseSequenceArray::try_new(starts, lengths, multipliers, len)?.into_array()) +} + #[test] fn materializes_piecewise_indices() -> VortexResult<()> { let starts = buffer![3u64, 15, 21].into_array(); @@ -280,6 +294,66 @@ fn varbin_take_consumes_piecewise_indices() -> VortexResult<()> { Ok(()) } +#[test] +fn contiguous_take_consumers_support_constant_piecewise_lengths() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let indices = piecewise_indices_constant_length([1, 3], 2)?; + + let primitive = PrimitiveArray::from_iter(0i32..10).into_array(); + assert_arrays_eq!( + primitive.take(indices.clone())?, + PrimitiveArray::from_iter([1i32, 2, 3, 4]).into_array(), + &mut ctx + ); + + let bools = BoolArray::from_iter([false, true, false, true, false]).into_array(); + assert_arrays_eq!( + bools.take(indices.clone())?, + BoolArray::from_iter([true, false, true, false]).into_array(), + &mut ctx + ); + + let decimal_dtype = DecimalDType::new(19, 2); + let decimals = DecimalArray::from_iter([10i128, 20, 30, 40, 50], decimal_dtype).into_array(); + assert_arrays_eq!( + decimals.take(indices.clone())?, + DecimalArray::from_iter([20i128, 30, 40, 50], decimal_dtype).into_array(), + &mut ctx + ); + + let varbinview = VarBinViewArray::from_iter( + ["a", "bb", "ccc", "dddd", "eeeee"].map(Some), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + assert_arrays_eq!( + varbinview.take(indices.clone())?, + VarBinViewArray::from_iter( + ["bb", "ccc", "dddd", "eeeee"].map(Some), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(), + &mut ctx + ); + + let varbin = VarBinArray::from_iter( + ["a", "bb", "ccc", "dddd", "eeeee"].map(Some), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + assert_arrays_eq!( + varbin.take(indices)?, + VarBinArray::from_iter( + ["bb", "ccc", "dddd", "eeeee"].map(Some), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(), + &mut ctx + ); + + Ok(()) +} + #[test] fn fixed_size_list_take_builds_piecewise_element_indices() -> VortexResult<()> { let elements = PrimitiveArray::from_iter(0i32..12).into_array(); diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index c67a2c5d725..d0a49d971e5 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -21,8 +21,10 @@ use crate::arrays::PiecewiseSequence; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::UnitMultiplierLengths; use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; use crate::arrays::piecewise_sequence::validate_index_ranges; +use crate::arrays::piecewise_sequence::validate_index_ranges_constant; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::IntegerPType; @@ -147,7 +149,14 @@ fn take_piecewise_sequence( }; let validity = array.validity()?.take(indices_ref)?; let output_len = indices_ref.len(); - let taken = take_piecewise_sequence_lengths(array, &starts, &lengths, validity, output_len)?; + let taken = match lengths { + UnitMultiplierLengths::Constant(length) => { + take_piecewise_sequence_constant_length(array, &starts, length, validity, output_len)? + } + UnitMultiplierLengths::Array(lengths) => { + take_piecewise_sequence_lengths(array, &starts, &lengths, validity, output_len)? + } + }; Ok(Some(taken)) } @@ -246,6 +255,62 @@ where Ok(values.freeze()) } +fn take_piecewise_sequence_constant_length( + array: ArrayView<'_, Primitive>, + starts: &PrimitiveArray, + length: usize, + validity: Validity, + output_len: usize, +) -> VortexResult { + match_each_native_ptype!(array.ptype(), |T| { + take_piecewise_sequence_constant_length_typed::( + array, starts, length, validity, output_len, + ) + }) +} + +fn take_piecewise_sequence_constant_length_typed( + array: ArrayView<'_, Primitive>, + starts: &PrimitiveArray, + length: usize, + validity: Validity, + output_len: usize, +) -> VortexResult +where + T: NativePType, +{ + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + let values = primitive_piecewise_constant_length_values::( + array.as_slice::(), + starts.as_slice::(), + length, + output_len, + )?; + Ok(PrimitiveArray::new(values, validity).into_array()) + }) +} + +fn primitive_piecewise_constant_length_values( + source: &[T], + starts: &[S], + length: usize, + output_len: usize, +) -> VortexResult> +where + T: Copy, + S: UnsignedPType, +{ + validate_index_ranges_constant(source.len(), starts, length, output_len)?; + + let mut values = BufferMut::::with_capacity(output_len); + for &start in starts { + let start = start.as_(); + values.extend_from_slice(&source[start..start + length]); + } + + Ok(values.freeze()) +} + #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] #[cfg(test)] mod test { diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 34d19863e20..a6326b6f61e 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -20,8 +20,10 @@ use crate::arrays::PrimitiveArray; use crate::arrays::VarBin; use crate::arrays::VarBinArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::UnitMultiplierLengths; use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; use crate::arrays::piecewise_sequence::validate_index_ranges; +use crate::arrays::piecewise_sequence::validate_index_ranges_constant; use crate::arrays::primitive::PrimitiveArrayExt; use crate::arrays::varbin::VarBinArrayExt; use crate::dtype::DType; @@ -138,50 +140,29 @@ fn take_piecewise_sequence( let offsets = array.offsets().clone().execute::(ctx)?; let out_offset_ptype = taken_offset_ptype(offsets.ptype()); let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); - - let result = match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { - match offsets.ptype() { - PType::U8 => gather_piecewise_varbin::( - array.dtype().clone(), - offsets.as_slice::(), - array.bytes().as_slice(), - starts.as_slice::(), - lengths.as_slice::(), - indices_ref.len(), - out_offset_ptype, - ), - PType::U16 => gather_piecewise_varbin::( - array.dtype().clone(), - offsets.as_slice::(), - array.bytes().as_slice(), - starts.as_slice::(), - lengths.as_slice::(), - indices_ref.len(), - out_offset_ptype, - ), - PType::U32 => gather_piecewise_varbin::( - array.dtype().clone(), - offsets.as_slice::(), - array.bytes().as_slice(), - starts.as_slice::(), - lengths.as_slice::(), - indices_ref.len(), - out_offset_ptype, - ), - PType::U64 => gather_piecewise_varbin::( - array.dtype().clone(), - offsets.as_slice::(), - array.bytes().as_slice(), - starts.as_slice::(), - lengths.as_slice::(), - indices_ref.len(), - out_offset_ptype, - ), - _ => unreachable!("offsets were reinterpreted to an unsigned integer ptype"), - } - }) - })?; + let bytes = array.bytes(); + let data = bytes.as_slice(); + let dtype = array.dtype().clone(); + let output_len = indices_ref.len(); + + let result = match &lengths { + UnitMultiplierLengths::Constant(length) => gather_piecewise_varbin_constant_dispatch( + &starts, + *length, + &offsets, + data, + output_len, + out_offset_ptype, + )?, + UnitMultiplierLengths::Array(lengths) => gather_piecewise_varbin_dispatch( + &starts, + lengths, + &offsets, + data, + output_len, + out_offset_ptype, + )?, + }; let validity = array.validity()?.take(indices_ref)?; @@ -189,13 +170,8 @@ fn take_piecewise_sequence( // non-decreasing, and the copied data buffer has exactly the referenced byte length. unsafe { Ok(Some( - VarBinArray::new_unchecked( - result.offsets, - result.data.freeze(), - result.dtype, - validity, - ) - .into_array(), + VarBinArray::new_unchecked(result.offsets, result.data.freeze(), dtype, validity) + .into_array(), )) } } @@ -271,13 +247,243 @@ fn take( } struct GatheredPiecewiseVarBin { - dtype: DType, offsets: ArrayRef, data: ByteBufferMut, } +fn gather_piecewise_varbin_constant_dispatch( + starts: &PrimitiveArray, + length: usize, + offsets: &PrimitiveArray, + data: &[u8], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + gather_piecewise_varbin_constant_start_dispatch::( + starts, + length, + offsets, + data, + output_len, + out_offset_ptype, + ) + }) +} + +fn gather_piecewise_varbin_constant_start_dispatch( + starts: &PrimitiveArray, + length: usize, + offsets: &PrimitiveArray, + data: &[u8], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult +where + S: UnsignedPType, +{ + match offsets.ptype() { + PType::U8 => gather_piecewise_varbin_constant_length::( + offsets.as_slice::(), + data, + starts.as_slice::(), + length, + output_len, + out_offset_ptype, + ), + PType::U16 => gather_piecewise_varbin_constant_length::( + offsets.as_slice::(), + data, + starts.as_slice::(), + length, + output_len, + out_offset_ptype, + ), + PType::U32 => gather_piecewise_varbin_constant_length::( + offsets.as_slice::(), + data, + starts.as_slice::(), + length, + output_len, + out_offset_ptype, + ), + PType::U64 => gather_piecewise_varbin_constant_length::( + offsets.as_slice::(), + data, + starts.as_slice::(), + length, + output_len, + out_offset_ptype, + ), + _ => unreachable!("offsets were reinterpreted to an unsigned integer ptype"), + } +} + +fn gather_piecewise_varbin_dispatch( + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + offsets: &PrimitiveArray, + data: &[u8], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + gather_piecewise_varbin_start_dispatch::( + starts, + lengths, + offsets, + data, + output_len, + out_offset_ptype, + ) + }) +} + +fn gather_piecewise_varbin_start_dispatch( + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + offsets: &PrimitiveArray, + data: &[u8], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult +where + S: UnsignedPType, +{ + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + gather_piecewise_varbin_start_length_dispatch::( + starts, + lengths, + offsets, + data, + output_len, + out_offset_ptype, + ) + }) +} + +fn gather_piecewise_varbin_start_length_dispatch( + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + offsets: &PrimitiveArray, + data: &[u8], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, +{ + match offsets.ptype() { + PType::U8 => gather_piecewise_varbin::( + offsets.as_slice::(), + data, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + out_offset_ptype, + ), + PType::U16 => gather_piecewise_varbin::( + offsets.as_slice::(), + data, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + out_offset_ptype, + ), + PType::U32 => gather_piecewise_varbin::( + offsets.as_slice::(), + data, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + out_offset_ptype, + ), + PType::U64 => gather_piecewise_varbin::( + offsets.as_slice::(), + data, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + out_offset_ptype, + ), + _ => unreachable!("offsets were reinterpreted to an unsigned integer ptype"), + } +} + +fn gather_piecewise_varbin_constant_length( + offsets: &[Offset], + data: &[u8], + starts: &[S], + length: usize, + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult +where + S: UnsignedPType, + Offset: IntegerPType, + NewOffset: IntegerPType, +{ + validate_index_ranges_constant(offsets.len() - 1, starts, length, output_len)?; + + let mut new_offsets = BufferMut::::with_capacity(output_len + 1); + new_offsets.push(NewOffset::zero()); + let mut output_bytes = 0usize; + + for &start in starts { + let start = start.as_(); + let end = start + length; + if length == 0 { + continue; + } + + let byte_start = offsets[start].as_(); + let byte_end = offsets[end].as_(); + vortex_ensure!( + byte_start <= byte_end && byte_end <= data.len(), + "VarBin offsets range {byte_start}..{byte_end} exceeds data length {}", + data.len() + ); + + for &offset in &offsets[start + 1..=end] { + let offset = offset.as_(); + let relative = offset.checked_sub(byte_start).ok_or_else(|| { + vortex_err!("VarBin offsets are not monotonic at offset {offset}") + })?; + let output_offset = output_bytes.checked_add(relative).ok_or_else(|| { + vortex_err!("PiecewiseSequence VarBin output byte length overflow") + })?; + new_offsets.push(new_offset_value::(output_offset)?); + } + + output_bytes = output_bytes + .checked_add(byte_end - byte_start) + .ok_or_else(|| vortex_err!("PiecewiseSequence VarBin output byte length overflow"))?; + } + + let mut new_data = ByteBufferMut::with_capacity(output_bytes); + for &start in starts { + let start = start.as_(); + let end = start + length; + if length == 0 { + continue; + } + + let byte_start = offsets[start].as_(); + let byte_end = offsets[end].as_(); + new_data.extend_from_slice(&data[byte_start..byte_end]); + } + + let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) + .reinterpret_cast(out_offset_ptype) + .into_array(); + Ok(GatheredPiecewiseVarBin { + offsets, + data: new_data, + }) +} + fn gather_piecewise_varbin( - dtype: DType, offsets: &[Offset], data: &[u8], starts: &[S], @@ -347,7 +553,6 @@ where .reinterpret_cast(out_offset_ptype) .into_array(); Ok(GatheredPiecewiseVarBin { - dtype, offsets, data: new_data, }) diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index dd38b45929d..0197a47722a 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -20,8 +20,10 @@ use crate::arrays::PrimitiveArray; use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::UnitMultiplierLengths; use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; use crate::arrays::piecewise_sequence::validate_index_ranges; +use crate::arrays::piecewise_sequence::validate_index_ranges_constant; use crate::arrays::varbinview::BinaryView; use crate::buffer::BufferHandle; use crate::dtype::UnsignedPType; @@ -79,16 +81,32 @@ fn take_piecewise_sequence( let Some((starts, lengths)) = execute_unit_multiplier_index_arrays(indices, ctx)? else { return Ok(None); }; - let views = match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { - gather_piecewise_views( - array.views(), - starts.as_slice::(), - lengths.as_slice::(), - indices_ref.len(), - )? - }) - }); + let source = array.views(); + let output_len = indices_ref.len(); + let views = match &lengths { + UnitMultiplierLengths::Constant(length) => { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + gather_piecewise_views_constant_length( + source, + starts.as_slice::(), + *length, + output_len, + )? + }) + } + UnitMultiplierLengths::Array(lengths) => { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + gather_piecewise_views( + source, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + )? + }) + }) + } + }; let validity = array.validity()?.take(indices_ref)?; // SAFETY: ranges were validated against the source views, and copied views still reference the @@ -134,6 +152,26 @@ fn take_views>( } } +fn gather_piecewise_views_constant_length( + source: &[BinaryView], + starts: &[S], + length: usize, + output_len: usize, +) -> VortexResult> +where + S: UnsignedPType, +{ + validate_index_ranges_constant(source.len(), starts, length, output_len)?; + + let mut views = BufferMut::::with_capacity(output_len); + for &start in starts { + let start = start.as_(); + views.extend_from_slice(&source[start..start + length]); + } + + Ok(views.freeze()) +} + fn gather_piecewise_views( source: &[BinaryView], starts: &[S], From ad5d6ffee3f23bcc2e1e586a8b883d9c222db80a Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 11:35:29 -0400 Subject: [PATCH 15/37] Rename contiguous range take helpers Signed-off-by: Daniel King --- vortex-array/src/arrays/bool/compute/take.rs | 12 ++--- .../src/arrays/decimal/compute/take.rs | 36 +++++++-------- .../src/arrays/primitive/compute/take/mod.rs | 36 +++++++-------- .../src/arrays/varbin/compute/take.rs | 44 +++++++++---------- .../src/arrays/varbinview/compute/take.rs | 12 ++--- 5 files changed, 66 insertions(+), 74 deletions(-) diff --git a/vortex-array/src/arrays/bool/compute/take.rs b/vortex-array/src/arrays/bool/compute/take.rs index a467e19175f..2accc52be42 100644 --- a/vortex-array/src/arrays/bool/compute/take.rs +++ b/vortex-array/src/arrays/bool/compute/take.rs @@ -38,7 +38,7 @@ impl TakeExecute for Bool { ctx: &mut ExecutionCtx, ) -> VortexResult> { if let Some(piecewise_indices) = indices.as_opt::() - && let Some(taken) = take_piecewise_sequence(array, piecewise_indices, indices, ctx)? + && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)? { return Ok(Some(taken)); } @@ -69,7 +69,7 @@ impl TakeExecute for Bool { } } -fn take_piecewise_sequence( +fn take_contiguous_ranges( array: ArrayView<'_, Bool>, indices: ArrayView<'_, PiecewiseSequence>, indices_ref: &ArrayRef, @@ -83,7 +83,7 @@ fn take_piecewise_sequence( let buffer = match &lengths { UnitMultiplierLengths::Constant(length) => { match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - take_piecewise_bits_constant_length( + take_bit_slices_constant_length( &source, starts.as_slice::(), *length, @@ -94,7 +94,7 @@ fn take_piecewise_sequence( UnitMultiplierLengths::Array(lengths) => { match_each_unsigned_integer_ptype!(starts.ptype(), |S| { match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { - take_piecewise_bits( + take_bit_slices( &source, starts.as_slice::(), lengths.as_slice::(), @@ -137,7 +137,7 @@ fn take_bool_impl>(bools: BitBufferView<'_>, indices: &[I] }) } -fn take_piecewise_bits_constant_length( +fn take_bit_slices_constant_length( source: &BitBuffer, starts: &[S], length: usize, @@ -157,7 +157,7 @@ where Ok(values.freeze()) } -fn take_piecewise_bits( +fn take_bit_slices( source: &BitBuffer, starts: &[S], lengths: &[L], diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index a2d508fd31e..fb4518fdc85 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -34,7 +34,7 @@ impl TakeExecute for Decimal { ctx: &mut ExecutionCtx, ) -> VortexResult> { if let Some(piecewise_indices) = indices.as_opt::() - && let Some(taken) = take_piecewise_sequence(array, piecewise_indices, indices, ctx)? + && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)? { return Ok(Some(taken)); } @@ -58,7 +58,7 @@ impl TakeExecute for Decimal { } } -fn take_piecewise_sequence( +fn take_contiguous_ranges( array: ArrayView<'_, Decimal>, indices: ArrayView<'_, PiecewiseSequence>, indices_ref: &ArrayRef, @@ -71,16 +71,16 @@ fn take_piecewise_sequence( let output_len = indices_ref.len(); let taken = match lengths { UnitMultiplierLengths::Constant(length) => { - take_piecewise_sequence_constant_length(array, &starts, length, validity, output_len)? + take_slices_constant_length(array, &starts, length, validity, output_len)? } UnitMultiplierLengths::Array(lengths) => { - take_piecewise_sequence_lengths(array, &starts, &lengths, validity, output_len)? + take_slices(array, &starts, &lengths, validity, output_len)? } }; Ok(Some(taken)) } -fn take_piecewise_sequence_constant_length( +fn take_slices_constant_length( array: ArrayView<'_, Decimal>, starts: &PrimitiveArray, length: usize, @@ -88,13 +88,11 @@ fn take_piecewise_sequence_constant_length( output_len: usize, ) -> VortexResult { match_each_decimal_value_type!(array.values_type(), |D| { - take_piecewise_sequence_constant_length_typed::( - array, starts, length, validity, output_len, - ) + take_slices_constant_length_typed::(array, starts, length, validity, output_len) }) } -fn take_piecewise_sequence_constant_length_typed( +fn take_slices_constant_length_typed( array: ArrayView<'_, Decimal>, starts: &PrimitiveArray, length: usize, @@ -105,7 +103,7 @@ where D: NativeDecimalType, { match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - let values = take_piecewise_constant_length_to_buffer::( + let values = take_slices_constant_length_to_buffer::( starts.as_slice::(), length, array.buffer::().as_slice(), @@ -120,7 +118,7 @@ where }) } -fn take_piecewise_sequence_lengths( +fn take_slices( array: ArrayView<'_, Decimal>, starts: &PrimitiveArray, lengths: &PrimitiveArray, @@ -128,11 +126,11 @@ fn take_piecewise_sequence_lengths( output_len: usize, ) -> VortexResult { match_each_decimal_value_type!(array.values_type(), |D| { - take_piecewise_sequence_lengths_typed::(array, starts, lengths, validity, output_len) + take_slices_typed::(array, starts, lengths, validity, output_len) }) } -fn take_piecewise_sequence_lengths_typed( +fn take_slices_typed( array: ArrayView<'_, Decimal>, starts: &PrimitiveArray, lengths: &PrimitiveArray, @@ -143,13 +141,11 @@ where D: NativeDecimalType, { match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - take_piecewise_sequence_lengths_start_typed::( - array, starts, lengths, validity, output_len, - ) + take_slices_start_typed::(array, starts, lengths, validity, output_len) }) } -fn take_piecewise_sequence_lengths_start_typed( +fn take_slices_start_typed( array: ArrayView<'_, Decimal>, starts: &PrimitiveArray, lengths: &PrimitiveArray, @@ -161,7 +157,7 @@ where S: UnsignedPType, { match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { - let values = take_piecewise_to_buffer::( + let values = take_slices_to_buffer::( starts.as_slice::(), lengths.as_slice::(), array.buffer::().as_slice(), @@ -180,7 +176,7 @@ fn take_to_buffer(indices: &[I], values: indices.iter().map(|idx| values[idx.as_()]).collect() } -fn take_piecewise_constant_length_to_buffer( +fn take_slices_constant_length_to_buffer( starts: &[S], length: usize, values: &[T], @@ -201,7 +197,7 @@ where Ok(result.freeze()) } -fn take_piecewise_to_buffer( +fn take_slices_to_buffer( starts: &[S], lengths: &[L], values: &[T], diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index d0a49d971e5..2e2aa0e2397 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -89,7 +89,7 @@ impl TakeExecute for Primitive { ctx: &mut ExecutionCtx, ) -> VortexResult> { if let Some(piecewise_indices) = indices.as_opt::() - && let Some(taken) = take_piecewise_sequence(array, piecewise_indices, indices, ctx)? + && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)? { return Ok(Some(taken)); } @@ -138,7 +138,7 @@ impl TakeExecute for Primitive { } } -fn take_piecewise_sequence( +fn take_contiguous_ranges( array: ArrayView<'_, Primitive>, indices: ArrayView<'_, PiecewiseSequence>, indices_ref: &ArrayRef, @@ -151,10 +151,10 @@ fn take_piecewise_sequence( let output_len = indices_ref.len(); let taken = match lengths { UnitMultiplierLengths::Constant(length) => { - take_piecewise_sequence_constant_length(array, &starts, length, validity, output_len)? + take_slices_constant_length(array, &starts, length, validity, output_len)? } UnitMultiplierLengths::Array(lengths) => { - take_piecewise_sequence_lengths(array, &starts, &lengths, validity, output_len)? + take_slices(array, &starts, &lengths, validity, output_len)? } }; Ok(Some(taken)) @@ -181,7 +181,7 @@ fn take_primitive_scalar(buffer: &[T], indices: &[I]) result.freeze() } -fn take_piecewise_sequence_lengths( +fn take_slices( array: ArrayView<'_, Primitive>, starts: &PrimitiveArray, lengths: &PrimitiveArray, @@ -189,11 +189,11 @@ fn take_piecewise_sequence_lengths( output_len: usize, ) -> VortexResult { match_each_native_ptype!(array.ptype(), |T| { - take_piecewise_sequence_lengths_typed::(array, starts, lengths, validity, output_len) + take_slices_typed::(array, starts, lengths, validity, output_len) }) } -fn take_piecewise_sequence_lengths_typed( +fn take_slices_typed( array: ArrayView<'_, Primitive>, starts: &PrimitiveArray, lengths: &PrimitiveArray, @@ -204,13 +204,11 @@ where T: NativePType, { match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - take_piecewise_sequence_lengths_start_typed::( - array, starts, lengths, validity, output_len, - ) + take_slices_start_typed::(array, starts, lengths, validity, output_len) }) } -fn take_piecewise_sequence_lengths_start_typed( +fn take_slices_start_typed( array: ArrayView<'_, Primitive>, starts: &PrimitiveArray, lengths: &PrimitiveArray, @@ -222,7 +220,7 @@ where S: UnsignedPType, { match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { - let values = primitive_piecewise_values::( + let values = take_slices_to_buffer::( array.as_slice::(), starts.as_slice::(), lengths.as_slice::(), @@ -232,7 +230,7 @@ where }) } -fn primitive_piecewise_values( +fn take_slices_to_buffer( source: &[T], starts: &[S], lengths: &[L], @@ -255,7 +253,7 @@ where Ok(values.freeze()) } -fn take_piecewise_sequence_constant_length( +fn take_slices_constant_length( array: ArrayView<'_, Primitive>, starts: &PrimitiveArray, length: usize, @@ -263,13 +261,11 @@ fn take_piecewise_sequence_constant_length( output_len: usize, ) -> VortexResult { match_each_native_ptype!(array.ptype(), |T| { - take_piecewise_sequence_constant_length_typed::( - array, starts, length, validity, output_len, - ) + take_slices_constant_length_typed::(array, starts, length, validity, output_len) }) } -fn take_piecewise_sequence_constant_length_typed( +fn take_slices_constant_length_typed( array: ArrayView<'_, Primitive>, starts: &PrimitiveArray, length: usize, @@ -280,7 +276,7 @@ where T: NativePType, { match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - let values = primitive_piecewise_constant_length_values::( + let values = take_slices_constant_length_to_buffer::( array.as_slice::(), starts.as_slice::(), length, @@ -290,7 +286,7 @@ where }) } -fn primitive_piecewise_constant_length_values( +fn take_slices_constant_length_to_buffer( source: &[T], starts: &[S], length: usize, diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index a6326b6f61e..b24ab319598 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -53,7 +53,7 @@ impl TakeExecute for VarBin { ctx: &mut ExecutionCtx, ) -> VortexResult> { if let Some(piecewise_indices) = indices.as_opt::() - && let Some(taken) = take_piecewise_sequence(array, piecewise_indices, indices, ctx)? + && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)? { return Ok(Some(taken)); } @@ -128,7 +128,7 @@ impl TakeExecute for VarBin { } } -fn take_piecewise_sequence( +fn take_contiguous_ranges( array: ArrayView<'_, VarBin>, indices: ArrayView<'_, PiecewiseSequence>, indices_ref: &ArrayRef, @@ -146,7 +146,7 @@ fn take_piecewise_sequence( let output_len = indices_ref.len(); let result = match &lengths { - UnitMultiplierLengths::Constant(length) => gather_piecewise_varbin_constant_dispatch( + UnitMultiplierLengths::Constant(length) => gather_slices_constant_dispatch( &starts, *length, &offsets, @@ -154,7 +154,7 @@ fn take_piecewise_sequence( output_len, out_offset_ptype, )?, - UnitMultiplierLengths::Array(lengths) => gather_piecewise_varbin_dispatch( + UnitMultiplierLengths::Array(lengths) => gather_slices_dispatch( &starts, lengths, &offsets, @@ -251,7 +251,7 @@ struct GatheredPiecewiseVarBin { data: ByteBufferMut, } -fn gather_piecewise_varbin_constant_dispatch( +fn gather_slices_constant_dispatch( starts: &PrimitiveArray, length: usize, offsets: &PrimitiveArray, @@ -260,7 +260,7 @@ fn gather_piecewise_varbin_constant_dispatch( out_offset_ptype: PType, ) -> VortexResult { match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - gather_piecewise_varbin_constant_start_dispatch::( + gather_slices_constant_start_dispatch::( starts, length, offsets, @@ -271,7 +271,7 @@ fn gather_piecewise_varbin_constant_dispatch( }) } -fn gather_piecewise_varbin_constant_start_dispatch( +fn gather_slices_constant_start_dispatch( starts: &PrimitiveArray, length: usize, offsets: &PrimitiveArray, @@ -283,7 +283,7 @@ where S: UnsignedPType, { match offsets.ptype() { - PType::U8 => gather_piecewise_varbin_constant_length::( + PType::U8 => gather_slices_constant_length::( offsets.as_slice::(), data, starts.as_slice::(), @@ -291,7 +291,7 @@ where output_len, out_offset_ptype, ), - PType::U16 => gather_piecewise_varbin_constant_length::( + PType::U16 => gather_slices_constant_length::( offsets.as_slice::(), data, starts.as_slice::(), @@ -299,7 +299,7 @@ where output_len, out_offset_ptype, ), - PType::U32 => gather_piecewise_varbin_constant_length::( + PType::U32 => gather_slices_constant_length::( offsets.as_slice::(), data, starts.as_slice::(), @@ -307,7 +307,7 @@ where output_len, out_offset_ptype, ), - PType::U64 => gather_piecewise_varbin_constant_length::( + PType::U64 => gather_slices_constant_length::( offsets.as_slice::(), data, starts.as_slice::(), @@ -319,7 +319,7 @@ where } } -fn gather_piecewise_varbin_dispatch( +fn gather_slices_dispatch( starts: &PrimitiveArray, lengths: &PrimitiveArray, offsets: &PrimitiveArray, @@ -328,7 +328,7 @@ fn gather_piecewise_varbin_dispatch( out_offset_ptype: PType, ) -> VortexResult { match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - gather_piecewise_varbin_start_dispatch::( + gather_slices_start_dispatch::( starts, lengths, offsets, @@ -339,7 +339,7 @@ fn gather_piecewise_varbin_dispatch( }) } -fn gather_piecewise_varbin_start_dispatch( +fn gather_slices_start_dispatch( starts: &PrimitiveArray, lengths: &PrimitiveArray, offsets: &PrimitiveArray, @@ -351,7 +351,7 @@ where S: UnsignedPType, { match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { - gather_piecewise_varbin_start_length_dispatch::( + gather_slices_start_length_dispatch::( starts, lengths, offsets, @@ -362,7 +362,7 @@ where }) } -fn gather_piecewise_varbin_start_length_dispatch( +fn gather_slices_start_length_dispatch( starts: &PrimitiveArray, lengths: &PrimitiveArray, offsets: &PrimitiveArray, @@ -375,7 +375,7 @@ where L: UnsignedPType, { match offsets.ptype() { - PType::U8 => gather_piecewise_varbin::( + PType::U8 => gather_slices::( offsets.as_slice::(), data, starts.as_slice::(), @@ -383,7 +383,7 @@ where output_len, out_offset_ptype, ), - PType::U16 => gather_piecewise_varbin::( + PType::U16 => gather_slices::( offsets.as_slice::(), data, starts.as_slice::(), @@ -391,7 +391,7 @@ where output_len, out_offset_ptype, ), - PType::U32 => gather_piecewise_varbin::( + PType::U32 => gather_slices::( offsets.as_slice::(), data, starts.as_slice::(), @@ -399,7 +399,7 @@ where output_len, out_offset_ptype, ), - PType::U64 => gather_piecewise_varbin::( + PType::U64 => gather_slices::( offsets.as_slice::(), data, starts.as_slice::(), @@ -411,7 +411,7 @@ where } } -fn gather_piecewise_varbin_constant_length( +fn gather_slices_constant_length( offsets: &[Offset], data: &[u8], starts: &[S], @@ -483,7 +483,7 @@ where }) } -fn gather_piecewise_varbin( +fn gather_slices( offsets: &[Offset], data: &[u8], starts: &[S], diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index 0197a47722a..d347a72ee4b 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -39,7 +39,7 @@ impl TakeExecute for VarBinView { ctx: &mut ExecutionCtx, ) -> VortexResult> { if let Some(piecewise_indices) = indices.as_opt::() - && let Some(taken) = take_piecewise_sequence(array, piecewise_indices, indices, ctx)? + && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)? { return Ok(Some(taken)); } @@ -72,7 +72,7 @@ impl TakeExecute for VarBinView { } } -fn take_piecewise_sequence( +fn take_contiguous_ranges( array: ArrayView<'_, VarBinView>, indices: ArrayView<'_, PiecewiseSequence>, indices_ref: &ArrayRef, @@ -86,7 +86,7 @@ fn take_piecewise_sequence( let views = match &lengths { UnitMultiplierLengths::Constant(length) => { match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - gather_piecewise_views_constant_length( + gather_view_slices_constant_length( source, starts.as_slice::(), *length, @@ -97,7 +97,7 @@ fn take_piecewise_sequence( UnitMultiplierLengths::Array(lengths) => { match_each_unsigned_integer_ptype!(starts.ptype(), |S| { match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { - gather_piecewise_views( + gather_view_slices( source, starts.as_slice::(), lengths.as_slice::(), @@ -152,7 +152,7 @@ fn take_views>( } } -fn gather_piecewise_views_constant_length( +fn gather_view_slices_constant_length( source: &[BinaryView], starts: &[S], length: usize, @@ -172,7 +172,7 @@ where Ok(views.freeze()) } -fn gather_piecewise_views( +fn gather_view_slices( source: &[BinaryView], starts: &[S], lengths: &[L], From 2a33748296e8a8305785993d787044a3a6cf7d5d Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 16 Jul 2026 16:48:53 -0400 Subject: [PATCH 16/37] Optimize FixedSizeList take results Signed-off-by: Daniel King --- .../src/arrays/fixed_size_list/compute/take.rs | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) 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 b2533d132b9..9a8cc63ec2d 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -27,6 +27,7 @@ use crate::dtype::IntegerPType; use crate::dtype::Nullability; use crate::executor::ExecutionCtx; use crate::match_each_integer_ptype; +use crate::optimizer::ArrayOptimizer; use crate::validity::Validity; /// Take implementation for [`FixedSizeListArray`]. @@ -79,10 +80,11 @@ fn take_empty_fsl( // SAFETY: empty output needs no child values; otherwise the index validity mask proves every // output row is null. Placeholder child elements have the exact length required by FSL. - Ok(unsafe { + unsafe { FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) } - .into_array()) + .into_array() + .optimize_ctx(ctx.session()) } fn take_non_empty_fsl( @@ -136,7 +138,7 @@ fn take_non_empty_degenerate_fsl( // SAFETY: degenerate FSL inputs have no elements, valid index payloads were checked against // the source length, and `Validity::take` produces validity for `new_len`. - Ok(unsafe { + unsafe { FixedSizeListArray::new_unchecked( array.elements().clone(), array.list_size(), @@ -144,7 +146,8 @@ fn take_non_empty_degenerate_fsl( new_len, ) } - .into_array()) + .into_array() + .optimize_ctx(ctx.session()) } fn take_non_empty_non_degenerate_fsl( @@ -167,10 +170,11 @@ fn take_non_empty_non_degenerate_fsl( // SAFETY: `new_elements` has `new_len * list_size` elements. `new_validity` is either // non-nullable or was produced by `Validity::take` for `new_len`. - Ok(unsafe { + unsafe { FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len) } - .into_array()) + .into_array() + .optimize_ctx(ctx.session()) } fn take_non_empty_non_degenerate_elements( From 815e475a05bf30027d9820aa5065264263be7868 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 14:39:07 -0400 Subject: [PATCH 17/37] Inline PiecewiseSequence range length checks Signed-off-by: Daniel King --- vortex-array/src/arrays/bool/compute/take.rs | 33 ++++++++-- .../src/arrays/decimal/compute/take.rs | 33 ++++++++-- .../src/arrays/piecewise_sequence/mod.rs | 65 ------------------- .../src/arrays/primitive/compute/take/mod.rs | 33 ++++++++-- .../src/arrays/varbin/compute/take.rs | 37 ++++++++--- .../src/arrays/varbinview/compute/take.rs | 33 ++++++++-- 6 files changed, 132 insertions(+), 102 deletions(-) diff --git a/vortex-array/src/arrays/bool/compute/take.rs b/vortex-array/src/arrays/bool/compute/take.rs index 2accc52be42..8e43a1a1e0b 100644 --- a/vortex-array/src/arrays/bool/compute/take.rs +++ b/vortex-array/src/arrays/bool/compute/take.rs @@ -8,6 +8,8 @@ use vortex_buffer::BitBufferMut; use vortex_buffer::BitBufferView; use vortex_buffer::get_bit; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_mask::Mask; use crate::ArrayRef; @@ -22,8 +24,6 @@ use crate::arrays::bool::BoolArrayExt; use crate::arrays::dict::TakeExecute; use crate::arrays::piecewise_sequence::UnitMultiplierLengths; use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; -use crate::arrays::piecewise_sequence::validate_index_ranges; -use crate::arrays::piecewise_sequence::validate_index_ranges_constant; use crate::builtins::ArrayBuiltins; use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; @@ -146,12 +146,22 @@ fn take_bit_slices_constant_length( where S: UnsignedPType, { - validate_index_ranges_constant(source.len(), starts, length, output_len)?; + let computed_len = starts + .len() + .checked_mul(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); let mut values = BitBufferMut::with_capacity(output_len); for &start in starts { let start = start.as_(); - values.append_buffer(&source.slice(start..start + length)); + let end = start + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + values.append_buffer(&source.slice(start..end)); } Ok(values.freeze()) @@ -167,15 +177,24 @@ where S: UnsignedPType, L: UnsignedPType, { - validate_index_ranges(source.len(), starts, lengths, output_len)?; - let mut values = BitBufferMut::with_capacity(output_len); + let mut computed_len = 0usize; for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - values.append_buffer(&source.slice(start..start + length)); + let end = start + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + computed_len = computed_len + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + values.append_buffer(&source.slice(start..end)); } + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); Ok(values.freeze()) } diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index fb4518fdc85..0bb354b4a6a 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -5,6 +5,8 @@ use itertools::Itertools as _; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use crate::ArrayRef; use crate::IntoArray; @@ -16,8 +18,6 @@ use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; use crate::arrays::piecewise_sequence::UnitMultiplierLengths; use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; -use crate::arrays::piecewise_sequence::validate_index_ranges; -use crate::arrays::piecewise_sequence::validate_index_ranges_constant; use crate::dtype::IntegerPType; use crate::dtype::NativeDecimalType; use crate::dtype::UnsignedPType; @@ -186,12 +186,22 @@ where S: UnsignedPType, T: NativeDecimalType, { - validate_index_ranges_constant(values.len(), starts, length, output_len)?; + let computed_len = starts + .len() + .checked_mul(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); let mut result = BufferMut::::with_capacity(output_len); for &start in starts { let start = start.as_(); - result.extend_from_slice(&values[start..start + length]); + let end = start + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + result.extend_from_slice(&values[start..end]); } Ok(result.freeze()) @@ -208,15 +218,24 @@ where L: UnsignedPType, T: NativeDecimalType, { - validate_index_ranges(values.len(), starts, lengths, output_len)?; - let mut result = BufferMut::::with_capacity(output_len); + let mut computed_len = 0usize; for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - result.extend_from_slice(&values[start..start + length]); + let end = start + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + computed_len = computed_len + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + result.extend_from_slice(&values[start..end]); } + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); Ok(result.freeze()) } diff --git a/vortex-array/src/arrays/piecewise_sequence/mod.rs b/vortex-array/src/arrays/piecewise_sequence/mod.rs index e5043e42e16..cdc5649c9f2 100644 --- a/vortex-array/src/arrays/piecewise_sequence/mod.rs +++ b/vortex-array/src/arrays/piecewise_sequence/mod.rs @@ -134,38 +134,6 @@ fn check_index_array(name: &str, array: &ArrayRef) -> VortexResult<()> { Ok(()) } -pub(crate) fn validate_index_ranges_constant( - source_len: usize, - starts: &[S], - length: usize, - output_len: usize, -) -> VortexResult<()> -where - S: UnsignedPType, -{ - let computed_len = starts - .len() - .checked_mul(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; - vortex_ensure!( - computed_len == output_len, - "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" - ); - - for &start in starts { - let start: usize = start.as_(); - let end = start - .checked_add(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; - vortex_ensure!( - end <= source_len, - "PiecewiseSequenceArray range {start}..{end} exceeds source length {source_len}" - ); - } - - Ok(()) -} - pub(crate) fn materialize_ranges( starts: &PrimitiveArray, lengths: &PrimitiveArray, @@ -210,36 +178,3 @@ where } Ok(values) } - -pub(crate) fn validate_index_ranges( - source_len: usize, - starts: &[S], - lengths: &[L], - output_len: usize, -) -> VortexResult<()> -where - S: UnsignedPType, - L: UnsignedPType, -{ - let mut computed_len = 0usize; - for (&start, &length) in starts.iter().zip_eq(lengths) { - let start: usize = start.as_(); - let length: usize = length.as_(); - let end = start - .checked_add(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; - vortex_ensure!( - end <= source_len, - "PiecewiseSequenceArray range {start}..{end} exceeds source length {source_len}" - ); - computed_len = computed_len - .checked_add(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; - } - - vortex_ensure!( - computed_len == output_len, - "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" - ); - Ok(()) -} diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index 2e2aa0e2397..58c9d271cff 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -11,6 +11,8 @@ 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; @@ -23,8 +25,6 @@ use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; use crate::arrays::piecewise_sequence::UnitMultiplierLengths; use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; -use crate::arrays::piecewise_sequence::validate_index_ranges; -use crate::arrays::piecewise_sequence::validate_index_ranges_constant; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::IntegerPType; @@ -241,15 +241,24 @@ where S: UnsignedPType, L: UnsignedPType, { - validate_index_ranges(source.len(), starts, lengths, output_len)?; - let mut values = BufferMut::::with_capacity(output_len); + let mut computed_len = 0usize; for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - values.extend_from_slice(&source[start..start + length]); + let end = start + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + computed_len = computed_len + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + values.extend_from_slice(&source[start..end]); } + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); Ok(values.freeze()) } @@ -296,12 +305,22 @@ where T: Copy, S: UnsignedPType, { - validate_index_ranges_constant(source.len(), starts, length, output_len)?; + let computed_len = starts + .len() + .checked_mul(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); let mut values = BufferMut::::with_capacity(output_len); for &start in starts { let start = start.as_(); - values.extend_from_slice(&source[start..start + length]); + let end = start + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + values.extend_from_slice(&source[start..end]); } Ok(values.freeze()) diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index b24ab319598..416a194397a 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -22,8 +22,6 @@ use crate::arrays::VarBinArray; use crate::arrays::dict::TakeExecute; use crate::arrays::piecewise_sequence::UnitMultiplierLengths; use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; -use crate::arrays::piecewise_sequence::validate_index_ranges; -use crate::arrays::piecewise_sequence::validate_index_ranges_constant; use crate::arrays::primitive::PrimitiveArrayExt; use crate::arrays::varbin::VarBinArrayExt; use crate::dtype::DType; @@ -424,7 +422,14 @@ where Offset: IntegerPType, NewOffset: IntegerPType, { - validate_index_ranges_constant(offsets.len() - 1, starts, length, output_len)?; + let computed_len = starts + .len() + .checked_mul(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); let mut new_offsets = BufferMut::::with_capacity(output_len + 1); new_offsets.push(NewOffset::zero()); @@ -432,7 +437,9 @@ where for &start in starts { let start = start.as_(); - let end = start + length; + let end = start + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; if length == 0 { continue; } @@ -464,7 +471,9 @@ where let mut new_data = ByteBufferMut::with_capacity(output_bytes); for &start in starts { let start = start.as_(); - let end = start + length; + let end = start + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; if length == 0 { continue; } @@ -497,16 +506,20 @@ where Offset: IntegerPType, NewOffset: IntegerPType, { - validate_index_ranges(offsets.len() - 1, starts, lengths, output_len)?; - let mut new_offsets = BufferMut::::with_capacity(output_len + 1); new_offsets.push(NewOffset::zero()); let mut output_bytes = 0usize; + let mut computed_len = 0usize; for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - let end = start + length; + let end = start + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + computed_len = computed_len + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; if length == 0 { continue; } @@ -534,12 +547,18 @@ where .checked_add(byte_end - byte_start) .ok_or_else(|| vortex_err!("PiecewiseSequence VarBin output byte length overflow"))?; } + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); let mut new_data = ByteBufferMut::with_capacity(output_bytes); for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - let end = start + length; + let end = start + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; if length == 0 { continue; } diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index d347a72ee4b..a7e565df02f 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -9,6 +9,8 @@ use num_traits::AsPrimitive; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_mask::AllOr; use vortex_mask::Mask; @@ -22,8 +24,6 @@ use crate::arrays::VarBinViewArray; use crate::arrays::dict::TakeExecute; use crate::arrays::piecewise_sequence::UnitMultiplierLengths; use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; -use crate::arrays::piecewise_sequence::validate_index_ranges; -use crate::arrays::piecewise_sequence::validate_index_ranges_constant; use crate::arrays::varbinview::BinaryView; use crate::buffer::BufferHandle; use crate::dtype::UnsignedPType; @@ -161,12 +161,22 @@ fn gather_view_slices_constant_length( where S: UnsignedPType, { - validate_index_ranges_constant(source.len(), starts, length, output_len)?; + let computed_len = starts + .len() + .checked_mul(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); let mut views = BufferMut::::with_capacity(output_len); for &start in starts { let start = start.as_(); - views.extend_from_slice(&source[start..start + length]); + let end = start + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + views.extend_from_slice(&source[start..end]); } Ok(views.freeze()) @@ -182,15 +192,24 @@ where S: UnsignedPType, L: UnsignedPType, { - validate_index_ranges(source.len(), starts, lengths, output_len)?; - let mut views = BufferMut::::with_capacity(output_len); + let mut computed_len = 0usize; for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - views.extend_from_slice(&source[start..start + length]); + let end = start + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + computed_len = computed_len + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + views.extend_from_slice(&source[start..end]); } + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); Ok(views.freeze()) } From 66a5896dbe9b7d5c5c0ddad9129133a3dadf85ff Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 14:48:33 -0400 Subject: [PATCH 18/37] Simplify PiecewiseSequence take range handling Signed-off-by: Daniel King --- vortex-array/src/arrays/bool/compute/take.rs | 10 +---- .../src/arrays/decimal/compute/take.rs | 10 +---- .../arrays/fixed_size_list/compute/take.rs | 22 +++------- .../src/arrays/primitive/compute/take/mod.rs | 10 +---- .../src/arrays/varbin/compute/take.rs | 40 ++++++++----------- .../src/arrays/varbinview/compute/take.rs | 10 +---- 6 files changed, 29 insertions(+), 73 deletions(-) diff --git a/vortex-array/src/arrays/bool/compute/take.rs b/vortex-array/src/arrays/bool/compute/take.rs index 8e43a1a1e0b..4e2596249e3 100644 --- a/vortex-array/src/arrays/bool/compute/take.rs +++ b/vortex-array/src/arrays/bool/compute/take.rs @@ -158,10 +158,7 @@ where let mut values = BitBufferMut::with_capacity(output_len); for &start in starts { let start = start.as_(); - let end = start - .checked_add(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; - values.append_buffer(&source.slice(start..end)); + values.append_buffer(&source.slice(start..).slice(..length)); } Ok(values.freeze()) @@ -182,13 +179,10 @@ where for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - let end = start - .checked_add(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; computed_len = computed_len .checked_add(length) .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; - values.append_buffer(&source.slice(start..end)); + values.append_buffer(&source.slice(start..).slice(..length)); } vortex_ensure!( diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index 0bb354b4a6a..6542c77ccd2 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -198,10 +198,7 @@ where let mut result = BufferMut::::with_capacity(output_len); for &start in starts { let start = start.as_(); - let end = start - .checked_add(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; - result.extend_from_slice(&values[start..end]); + result.extend_from_slice(&values[start..][..length]); } Ok(result.freeze()) @@ -223,13 +220,10 @@ where for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - let end = start - .checked_add(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; computed_len = computed_len .checked_add(length) .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; - result.extend_from_slice(&values[start..end]); + result.extend_from_slice(&values[start..][..length]); } vortex_ensure!( 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 9a8cc63ec2d..5f2e0a50835 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -24,7 +24,6 @@ use crate::arrays::primitive::PrimitiveArrayExt; use crate::builders::builder_with_capacity; use crate::dtype::DType; use crate::dtype::IntegerPType; -use crate::dtype::Nullability; use crate::executor::ExecutionCtx; use crate::match_each_integer_ptype; use crate::optimizer::ArrayOptimizer; @@ -94,7 +93,7 @@ fn take_non_empty_fsl( ) -> VortexResult { debug_assert!(!array.is_empty()); - let DType::Primitive(ptype, nullability) = indices.dtype() else { + let DType::Primitive(ptype, _) = indices.dtype() else { vortex_bail!("Invalid indices dtype: {}", indices.dtype()) }; if !ptype.is_int() { @@ -107,13 +106,7 @@ fn take_non_empty_fsl( 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(), - *nullability, - ctx, - ) + take_non_empty_non_degenerate_fsl::(array, indices, indices_array.as_view(), ctx) }) } @@ -154,7 +147,6 @@ fn take_non_empty_non_degenerate_fsl( array: ArrayView<'_, FixedSizeList>, indices: &ArrayRef, indices_array: ArrayView<'_, Primitive>, - indices_nullability: Nullability, ctx: &mut ExecutionCtx, ) -> VortexResult { debug_assert!(!array.is_empty()); @@ -162,14 +154,10 @@ fn take_non_empty_non_degenerate_fsl( let (new_elements, new_len) = take_non_empty_non_degenerate_elements::(array, indices_array, ctx)?; - let new_validity = if array.dtype().is_nullable() || indices_nullability.is_nullable() { - array.validity()?.take(indices)? - } else { - Validity::NonNullable - }; + let new_validity = array.validity()?.take(indices)?; - // SAFETY: `new_elements` has `new_len * list_size` elements. `new_validity` is either - // non-nullable or was produced by `Validity::take` for `new_len`. + // 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) } diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index 58c9d271cff..25dba1c3b53 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -246,13 +246,10 @@ where for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - let end = start - .checked_add(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; computed_len = computed_len .checked_add(length) .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; - values.extend_from_slice(&source[start..end]); + values.extend_from_slice(&source[start..][..length]); } vortex_ensure!( @@ -317,10 +314,7 @@ where let mut values = BufferMut::::with_capacity(output_len); for &start in starts { let start = start.as_(); - let end = start - .checked_add(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; - values.extend_from_slice(&source[start..end]); + values.extend_from_slice(&source[start..][..length]); } Ok(values.freeze()) diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 416a194397a..cc0aeb1d896 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -437,22 +437,20 @@ where for &start in starts { let start = start.as_(); - let end = start - .checked_add(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; if length == 0 { continue; } - let byte_start = offsets[start].as_(); - let byte_end = offsets[end].as_(); + let offset_range = &offsets[start..][..=length]; + let byte_start = offset_range[0].as_(); + let byte_end = offset_range[length].as_(); vortex_ensure!( byte_start <= byte_end && byte_end <= data.len(), "VarBin offsets range {byte_start}..{byte_end} exceeds data length {}", data.len() ); - for &offset in &offsets[start + 1..=end] { + for &offset in &offset_range[1..] { let offset = offset.as_(); let relative = offset.checked_sub(byte_start).ok_or_else(|| { vortex_err!("VarBin offsets are not monotonic at offset {offset}") @@ -471,16 +469,14 @@ where let mut new_data = ByteBufferMut::with_capacity(output_bytes); for &start in starts { let start = start.as_(); - let end = start - .checked_add(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; if length == 0 { continue; } - let byte_start = offsets[start].as_(); - let byte_end = offsets[end].as_(); - new_data.extend_from_slice(&data[byte_start..byte_end]); + let offset_range = &offsets[start..][..=length]; + let byte_start = offset_range[0].as_(); + let byte_end = offset_range[length].as_(); + new_data.extend_from_slice(&data[byte_start..][..byte_end - byte_start]); } let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) @@ -514,9 +510,6 @@ where for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - let end = start - .checked_add(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; computed_len = computed_len .checked_add(length) .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; @@ -524,15 +517,16 @@ where continue; } - let byte_start = offsets[start].as_(); - let byte_end = offsets[end].as_(); + let offset_range = &offsets[start..][..=length]; + let byte_start = offset_range[0].as_(); + let byte_end = offset_range[length].as_(); vortex_ensure!( byte_start <= byte_end && byte_end <= data.len(), "VarBin offsets range {byte_start}..{byte_end} exceeds data length {}", data.len() ); - for &offset in &offsets[start + 1..=end] { + for &offset in &offset_range[1..] { let offset = offset.as_(); let relative = offset.checked_sub(byte_start).ok_or_else(|| { vortex_err!("VarBin offsets are not monotonic at offset {offset}") @@ -556,16 +550,14 @@ where for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - let end = start - .checked_add(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; if length == 0 { continue; } - let byte_start = offsets[start].as_(); - let byte_end = offsets[end].as_(); - new_data.extend_from_slice(&data[byte_start..byte_end]); + let offset_range = &offsets[start..][..=length]; + let byte_start = offset_range[0].as_(); + let byte_end = offset_range[length].as_(); + new_data.extend_from_slice(&data[byte_start..][..byte_end - byte_start]); } let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index a7e565df02f..f91354d8be6 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -173,10 +173,7 @@ where let mut views = BufferMut::::with_capacity(output_len); for &start in starts { let start = start.as_(); - let end = start - .checked_add(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; - views.extend_from_slice(&source[start..end]); + views.extend_from_slice(&source[start..][..length]); } Ok(views.freeze()) @@ -197,13 +194,10 @@ where for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - let end = start - .checked_add(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; computed_len = computed_len .checked_add(length) .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; - views.extend_from_slice(&source[start..end]); + views.extend_from_slice(&source[start..][..length]); } vortex_ensure!( From 18ec1b6963fbb36d4bd82f6569ae57bdf0b0d0a6 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 15:10:45 -0400 Subject: [PATCH 19/37] Validate PiecewiseSequence buffer lengths after copy Signed-off-by: Daniel King --- vortex-array/src/arrays/bool/compute/take.rs | 9 +++------ vortex-array/src/arrays/decimal/compute/take.rs | 9 +++------ vortex-array/src/arrays/primitive/compute/take/mod.rs | 9 +++------ vortex-array/src/arrays/varbin/compute/take.rs | 9 +++------ vortex-array/src/arrays/varbinview/compute/take.rs | 9 +++------ 5 files changed, 15 insertions(+), 30 deletions(-) diff --git a/vortex-array/src/arrays/bool/compute/take.rs b/vortex-array/src/arrays/bool/compute/take.rs index 4e2596249e3..410382f57a4 100644 --- a/vortex-array/src/arrays/bool/compute/take.rs +++ b/vortex-array/src/arrays/bool/compute/take.rs @@ -175,19 +175,16 @@ where L: UnsignedPType, { let mut values = BitBufferMut::with_capacity(output_len); - let mut computed_len = 0usize; for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - computed_len = computed_len - .checked_add(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; values.append_buffer(&source.slice(start..).slice(..length)); } vortex_ensure!( - computed_len == output_len, - "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + values.len() == output_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + values.len() ); Ok(values.freeze()) } diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index 6542c77ccd2..a4ddd287ee3 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -216,19 +216,16 @@ where T: NativeDecimalType, { let mut result = BufferMut::::with_capacity(output_len); - let mut computed_len = 0usize; for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - computed_len = computed_len - .checked_add(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; result.extend_from_slice(&values[start..][..length]); } vortex_ensure!( - computed_len == output_len, - "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + result.len() == output_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + result.len() ); Ok(result.freeze()) } diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index 25dba1c3b53..207f5cbdefc 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -242,19 +242,16 @@ where L: UnsignedPType, { let mut values = BufferMut::::with_capacity(output_len); - let mut computed_len = 0usize; for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - computed_len = computed_len - .checked_add(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; values.extend_from_slice(&source[start..][..length]); } vortex_ensure!( - computed_len == output_len, - "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + values.len() == output_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + values.len() ); Ok(values.freeze()) } diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index cc0aeb1d896..e3bbb7d98ca 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -505,14 +505,10 @@ where let mut new_offsets = BufferMut::::with_capacity(output_len + 1); new_offsets.push(NewOffset::zero()); let mut output_bytes = 0usize; - let mut computed_len = 0usize; for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - computed_len = computed_len - .checked_add(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; if length == 0 { continue; } @@ -542,8 +538,9 @@ where .ok_or_else(|| vortex_err!("PiecewiseSequence VarBin output byte length overflow"))?; } vortex_ensure!( - computed_len == output_len, - "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + new_offsets.len() == output_len + 1, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + new_offsets.len() - 1 ); let mut new_data = ByteBufferMut::with_capacity(output_bytes); diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index f91354d8be6..41636e87a8f 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -190,19 +190,16 @@ where L: UnsignedPType, { let mut views = BufferMut::::with_capacity(output_len); - let mut computed_len = 0usize; for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - computed_len = computed_len - .checked_add(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; views.extend_from_slice(&source[start..][..length]); } vortex_ensure!( - computed_len == output_len, - "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + views.len() == output_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + views.len() ); Ok(views.freeze()) } From 699f48a2f9a64018af9e767db91e05f122a7adac Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 16:02:53 -0400 Subject: [PATCH 20/37] Cover null FSL take placeholder indices Signed-off-by: Daniel King --- .../src/arrays/fixed_size_list/tests/take.rs | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) 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 af8f94f2d5f..93262a74b37 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -147,6 +147,28 @@ fn test_take_fsl_with_null_indices_preserves_elements() { assert_arrays_eq!(expected, result, &mut ctx); } +#[test] +fn test_take_fsl_null_index_ignores_out_of_bounds_payload() { + let mut ctx = array_session().create_execution_ctx(); + let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); + let fsl = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 3); + + let indices = PrimitiveArray::new( + buffer![1u32, 99, 2], + Validity::from_iter([true, false, true]), + ) + .into_array(); + let result = fsl.take(indices).unwrap(); + + let expected = FixedSizeListArray::new( + buffer![3i32, 4, 0, 0, 5, 6].into_array(), + 2, + Validity::from_iter([true, false, true]), + 3, + ); + assert_arrays_eq!(expected, result, &mut ctx); +} + // 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] From a2cf57156a607ed0e5210231e58a0b3f6f9d11db Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 16:17:34 -0400 Subject: [PATCH 21/37] Rename contiguous PiecewiseSequence helpers Signed-off-by: Daniel King --- vortex-array/src/arrays/bool/compute/take.rs | 10 +++++----- vortex-array/src/arrays/decimal/compute/take.rs | 10 +++++----- vortex-array/src/arrays/piecewise_sequence/mod.rs | 10 +++++----- vortex-array/src/arrays/primitive/compute/take/mod.rs | 10 +++++----- vortex-array/src/arrays/varbin/compute/take.rs | 10 +++++----- vortex-array/src/arrays/varbinview/compute/take.rs | 10 +++++----- 6 files changed, 30 insertions(+), 30 deletions(-) diff --git a/vortex-array/src/arrays/bool/compute/take.rs b/vortex-array/src/arrays/bool/compute/take.rs index 410382f57a4..1297a0a18e7 100644 --- a/vortex-array/src/arrays/bool/compute/take.rs +++ b/vortex-array/src/arrays/bool/compute/take.rs @@ -22,8 +22,8 @@ use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; use crate::arrays::dict::TakeExecute; -use crate::arrays::piecewise_sequence::UnitMultiplierLengths; -use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; +use crate::arrays::piecewise_sequence::ConstantOrArray; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::builtins::ArrayBuiltins; use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; @@ -75,13 +75,13 @@ fn take_contiguous_ranges( indices_ref: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let Some((starts, lengths)) = execute_unit_multiplier_index_arrays(indices, ctx)? else { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { return Ok(None); }; let source = array.to_bit_buffer(); let output_len = indices_ref.len(); let buffer = match &lengths { - UnitMultiplierLengths::Constant(length) => { + ConstantOrArray::Constant(length) => { match_each_unsigned_integer_ptype!(starts.ptype(), |S| { take_bit_slices_constant_length( &source, @@ -91,7 +91,7 @@ fn take_contiguous_ranges( )? }) } - UnitMultiplierLengths::Array(lengths) => { + ConstantOrArray::Array(lengths) => { match_each_unsigned_integer_ptype!(starts.ptype(), |S| { match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { take_bit_slices( diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index a4ddd287ee3..ef34f47ffc7 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -16,8 +16,8 @@ use crate::arrays::DecimalArray; use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; -use crate::arrays::piecewise_sequence::UnitMultiplierLengths; -use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; +use crate::arrays::piecewise_sequence::ConstantOrArray; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::dtype::IntegerPType; use crate::dtype::NativeDecimalType; use crate::dtype::UnsignedPType; @@ -64,16 +64,16 @@ fn take_contiguous_ranges( indices_ref: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let Some((starts, lengths)) = execute_unit_multiplier_index_arrays(indices, ctx)? else { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { return Ok(None); }; let validity = array.validity()?.take(indices_ref)?; let output_len = indices_ref.len(); let taken = match lengths { - UnitMultiplierLengths::Constant(length) => { + ConstantOrArray::Constant(length) => { take_slices_constant_length(array, &starts, length, validity, output_len)? } - UnitMultiplierLengths::Array(lengths) => { + ConstantOrArray::Array(lengths) => { take_slices(array, &starts, &lengths, validity, output_len)? } }; diff --git a/vortex-array/src/arrays/piecewise_sequence/mod.rs b/vortex-array/src/arrays/piecewise_sequence/mod.rs index cdc5649c9f2..18ce7debec1 100644 --- a/vortex-array/src/arrays/piecewise_sequence/mod.rs +++ b/vortex-array/src/arrays/piecewise_sequence/mod.rs @@ -67,15 +67,15 @@ pub(crate) fn execute_index_arrays( Ok((starts, lengths, multipliers)) } -pub(crate) enum UnitMultiplierLengths { +pub(crate) enum ConstantOrArray { Constant(usize), Array(PrimitiveArray), } -pub(crate) fn execute_unit_multiplier_index_arrays( +pub(crate) fn maybe_contiguous_slices( array: ArrayView<'_, PiecewiseSequence>, ctx: &mut ExecutionCtx, -) -> VortexResult> { +) -> VortexResult> { if !is_constant_multiplier_one(array.multipliers()) { return Ok(None); } @@ -84,12 +84,12 @@ pub(crate) fn execute_unit_multiplier_index_arrays( let starts = array.starts().clone().execute::(ctx)?; if let Some(length) = constant_unsigned_usize(array.lengths())? { check_index_arrays(starts.as_ref(), array.lengths(), array.multipliers())?; - return Ok(Some((starts, UnitMultiplierLengths::Constant(length)))); + return Ok(Some((starts, ConstantOrArray::Constant(length)))); } let lengths = array.lengths().clone().execute::(ctx)?; check_index_arrays(starts.as_ref(), lengths.as_ref(), array.multipliers())?; - Ok(Some((starts, UnitMultiplierLengths::Array(lengths)))) + Ok(Some((starts, ConstantOrArray::Array(lengths)))) } pub(crate) fn is_constant_multiplier_one(multipliers: &ArrayRef) -> bool { diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index 207f5cbdefc..a301cf99b29 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -23,8 +23,8 @@ use crate::arrays::PiecewiseSequence; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; -use crate::arrays::piecewise_sequence::UnitMultiplierLengths; -use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; +use crate::arrays::piecewise_sequence::ConstantOrArray; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::IntegerPType; @@ -144,16 +144,16 @@ fn take_contiguous_ranges( indices_ref: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let Some((starts, lengths)) = execute_unit_multiplier_index_arrays(indices, ctx)? else { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { return Ok(None); }; let validity = array.validity()?.take(indices_ref)?; let output_len = indices_ref.len(); let taken = match lengths { - UnitMultiplierLengths::Constant(length) => { + ConstantOrArray::Constant(length) => { take_slices_constant_length(array, &starts, length, validity, output_len)? } - UnitMultiplierLengths::Array(lengths) => { + ConstantOrArray::Array(lengths) => { take_slices(array, &starts, &lengths, validity, output_len)? } }; diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index e3bbb7d98ca..b956359ef0b 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -20,8 +20,8 @@ use crate::arrays::PrimitiveArray; use crate::arrays::VarBin; use crate::arrays::VarBinArray; use crate::arrays::dict::TakeExecute; -use crate::arrays::piecewise_sequence::UnitMultiplierLengths; -use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; +use crate::arrays::piecewise_sequence::ConstantOrArray; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; use crate::arrays::varbin::VarBinArrayExt; use crate::dtype::DType; @@ -132,7 +132,7 @@ fn take_contiguous_ranges( indices_ref: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let Some((starts, lengths)) = execute_unit_multiplier_index_arrays(indices, ctx)? else { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { return Ok(None); }; let offsets = array.offsets().clone().execute::(ctx)?; @@ -144,7 +144,7 @@ fn take_contiguous_ranges( let output_len = indices_ref.len(); let result = match &lengths { - UnitMultiplierLengths::Constant(length) => gather_slices_constant_dispatch( + ConstantOrArray::Constant(length) => gather_slices_constant_dispatch( &starts, *length, &offsets, @@ -152,7 +152,7 @@ fn take_contiguous_ranges( output_len, out_offset_ptype, )?, - UnitMultiplierLengths::Array(lengths) => gather_slices_dispatch( + ConstantOrArray::Array(lengths) => gather_slices_dispatch( &starts, lengths, &offsets, diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index 41636e87a8f..40e05e08737 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -22,8 +22,8 @@ use crate::arrays::PrimitiveArray; use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::arrays::dict::TakeExecute; -use crate::arrays::piecewise_sequence::UnitMultiplierLengths; -use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; +use crate::arrays::piecewise_sequence::ConstantOrArray; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::varbinview::BinaryView; use crate::buffer::BufferHandle; use crate::dtype::UnsignedPType; @@ -78,13 +78,13 @@ fn take_contiguous_ranges( indices_ref: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let Some((starts, lengths)) = execute_unit_multiplier_index_arrays(indices, ctx)? else { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { return Ok(None); }; let source = array.views(); let output_len = indices_ref.len(); let views = match &lengths { - UnitMultiplierLengths::Constant(length) => { + ConstantOrArray::Constant(length) => { match_each_unsigned_integer_ptype!(starts.ptype(), |S| { gather_view_slices_constant_length( source, @@ -94,7 +94,7 @@ fn take_contiguous_ranges( )? }) } - UnitMultiplierLengths::Array(lengths) => { + ConstantOrArray::Array(lengths) => { match_each_unsigned_integer_ptype!(starts.ptype(), |S| { match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { gather_view_slices( From 57388044f3b780109f7e5bda904b0f29a1a96175 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 16:23:07 -0400 Subject: [PATCH 22/37] Move VarBin PiecewiseSequence takes out of base PR Signed-off-by: Daniel King --- .../src/arrays/varbin/compute/take.rs | 397 +----------------- .../src/arrays/varbinview/compute/take.rs | 124 +----- 2 files changed, 2 insertions(+), 519 deletions(-) diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index b956359ef0b..74e0bd6af9d 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -1,33 +1,26 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use itertools::Itertools as _; use vortex_buffer::BitBufferMut; use vortex_buffer::BufferMut; use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; -use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; use crate::array::ArrayView; -use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::VarBin; use crate::arrays::VarBinArray; use crate::arrays::dict::TakeExecute; -use crate::arrays::piecewise_sequence::ConstantOrArray; -use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; use crate::arrays::varbin::VarBinArrayExt; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::dtype::PType; -use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_unsigned_integer_ptype; use crate::validity::Validity; @@ -50,12 +43,6 @@ impl TakeExecute for VarBin { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - if let Some(piecewise_indices) = indices.as_opt::() - && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)? - { - return Ok(Some(taken)); - } - // TODO(joe): Be lazy with execute let offsets = array.offsets().clone().execute::(ctx)?; let data = array.bytes(); @@ -126,54 +113,6 @@ impl TakeExecute for VarBin { } } -fn take_contiguous_ranges( - array: ArrayView<'_, VarBin>, - indices: ArrayView<'_, PiecewiseSequence>, - indices_ref: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult> { - let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { - return Ok(None); - }; - let offsets = array.offsets().clone().execute::(ctx)?; - let out_offset_ptype = taken_offset_ptype(offsets.ptype()); - let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); - let bytes = array.bytes(); - let data = bytes.as_slice(); - let dtype = array.dtype().clone(); - let output_len = indices_ref.len(); - - let result = match &lengths { - ConstantOrArray::Constant(length) => gather_slices_constant_dispatch( - &starts, - *length, - &offsets, - data, - output_len, - out_offset_ptype, - )?, - ConstantOrArray::Array(lengths) => gather_slices_dispatch( - &starts, - lengths, - &offsets, - data, - output_len, - out_offset_ptype, - )?, - }; - - let validity = array.validity()?.take(indices_ref)?; - - // SAFETY: output offsets are built from valid input offsets, start at zero, are monotonically - // non-decreasing, and the copied data buffer has exactly the referenced byte length. - unsafe { - Ok(Some( - VarBinArray::new_unchecked(result.offsets, result.data.freeze(), dtype, validity) - .into_array(), - )) - } -} - fn take( dtype: DType, offsets: &[Offset], @@ -244,337 +183,6 @@ fn take( } } -struct GatheredPiecewiseVarBin { - offsets: ArrayRef, - data: ByteBufferMut, -} - -fn gather_slices_constant_dispatch( - starts: &PrimitiveArray, - length: usize, - offsets: &PrimitiveArray, - data: &[u8], - output_len: usize, - out_offset_ptype: PType, -) -> VortexResult { - match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - gather_slices_constant_start_dispatch::( - starts, - length, - offsets, - data, - output_len, - out_offset_ptype, - ) - }) -} - -fn gather_slices_constant_start_dispatch( - starts: &PrimitiveArray, - length: usize, - offsets: &PrimitiveArray, - data: &[u8], - output_len: usize, - out_offset_ptype: PType, -) -> VortexResult -where - S: UnsignedPType, -{ - match offsets.ptype() { - PType::U8 => gather_slices_constant_length::( - offsets.as_slice::(), - data, - starts.as_slice::(), - length, - output_len, - out_offset_ptype, - ), - PType::U16 => gather_slices_constant_length::( - offsets.as_slice::(), - data, - starts.as_slice::(), - length, - output_len, - out_offset_ptype, - ), - PType::U32 => gather_slices_constant_length::( - offsets.as_slice::(), - data, - starts.as_slice::(), - length, - output_len, - out_offset_ptype, - ), - PType::U64 => gather_slices_constant_length::( - offsets.as_slice::(), - data, - starts.as_slice::(), - length, - output_len, - out_offset_ptype, - ), - _ => unreachable!("offsets were reinterpreted to an unsigned integer ptype"), - } -} - -fn gather_slices_dispatch( - starts: &PrimitiveArray, - lengths: &PrimitiveArray, - offsets: &PrimitiveArray, - data: &[u8], - output_len: usize, - out_offset_ptype: PType, -) -> VortexResult { - match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - gather_slices_start_dispatch::( - starts, - lengths, - offsets, - data, - output_len, - out_offset_ptype, - ) - }) -} - -fn gather_slices_start_dispatch( - starts: &PrimitiveArray, - lengths: &PrimitiveArray, - offsets: &PrimitiveArray, - data: &[u8], - output_len: usize, - out_offset_ptype: PType, -) -> VortexResult -where - S: UnsignedPType, -{ - match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { - gather_slices_start_length_dispatch::( - starts, - lengths, - offsets, - data, - output_len, - out_offset_ptype, - ) - }) -} - -fn gather_slices_start_length_dispatch( - starts: &PrimitiveArray, - lengths: &PrimitiveArray, - offsets: &PrimitiveArray, - data: &[u8], - output_len: usize, - out_offset_ptype: PType, -) -> VortexResult -where - S: UnsignedPType, - L: UnsignedPType, -{ - match offsets.ptype() { - PType::U8 => gather_slices::( - offsets.as_slice::(), - data, - starts.as_slice::(), - lengths.as_slice::(), - output_len, - out_offset_ptype, - ), - PType::U16 => gather_slices::( - offsets.as_slice::(), - data, - starts.as_slice::(), - lengths.as_slice::(), - output_len, - out_offset_ptype, - ), - PType::U32 => gather_slices::( - offsets.as_slice::(), - data, - starts.as_slice::(), - lengths.as_slice::(), - output_len, - out_offset_ptype, - ), - PType::U64 => gather_slices::( - offsets.as_slice::(), - data, - starts.as_slice::(), - lengths.as_slice::(), - output_len, - out_offset_ptype, - ), - _ => unreachable!("offsets were reinterpreted to an unsigned integer ptype"), - } -} - -fn gather_slices_constant_length( - offsets: &[Offset], - data: &[u8], - starts: &[S], - length: usize, - output_len: usize, - out_offset_ptype: PType, -) -> VortexResult -where - S: UnsignedPType, - Offset: IntegerPType, - NewOffset: IntegerPType, -{ - let computed_len = starts - .len() - .checked_mul(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; - vortex_ensure!( - computed_len == output_len, - "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" - ); - - let mut new_offsets = BufferMut::::with_capacity(output_len + 1); - new_offsets.push(NewOffset::zero()); - let mut output_bytes = 0usize; - - for &start in starts { - let start = start.as_(); - if length == 0 { - continue; - } - - let offset_range = &offsets[start..][..=length]; - let byte_start = offset_range[0].as_(); - let byte_end = offset_range[length].as_(); - vortex_ensure!( - byte_start <= byte_end && byte_end <= data.len(), - "VarBin offsets range {byte_start}..{byte_end} exceeds data length {}", - data.len() - ); - - for &offset in &offset_range[1..] { - let offset = offset.as_(); - let relative = offset.checked_sub(byte_start).ok_or_else(|| { - vortex_err!("VarBin offsets are not monotonic at offset {offset}") - })?; - let output_offset = output_bytes.checked_add(relative).ok_or_else(|| { - vortex_err!("PiecewiseSequence VarBin output byte length overflow") - })?; - new_offsets.push(new_offset_value::(output_offset)?); - } - - output_bytes = output_bytes - .checked_add(byte_end - byte_start) - .ok_or_else(|| vortex_err!("PiecewiseSequence VarBin output byte length overflow"))?; - } - - let mut new_data = ByteBufferMut::with_capacity(output_bytes); - for &start in starts { - let start = start.as_(); - if length == 0 { - continue; - } - - let offset_range = &offsets[start..][..=length]; - let byte_start = offset_range[0].as_(); - let byte_end = offset_range[length].as_(); - new_data.extend_from_slice(&data[byte_start..][..byte_end - byte_start]); - } - - let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) - .reinterpret_cast(out_offset_ptype) - .into_array(); - Ok(GatheredPiecewiseVarBin { - offsets, - data: new_data, - }) -} - -fn gather_slices( - offsets: &[Offset], - data: &[u8], - starts: &[S], - lengths: &[L], - output_len: usize, - out_offset_ptype: PType, -) -> VortexResult -where - S: UnsignedPType, - L: UnsignedPType, - Offset: IntegerPType, - NewOffset: IntegerPType, -{ - let mut new_offsets = BufferMut::::with_capacity(output_len + 1); - new_offsets.push(NewOffset::zero()); - let mut output_bytes = 0usize; - - for (&start, &length) in starts.iter().zip_eq(lengths) { - let start = start.as_(); - let length = length.as_(); - if length == 0 { - continue; - } - - let offset_range = &offsets[start..][..=length]; - let byte_start = offset_range[0].as_(); - let byte_end = offset_range[length].as_(); - vortex_ensure!( - byte_start <= byte_end && byte_end <= data.len(), - "VarBin offsets range {byte_start}..{byte_end} exceeds data length {}", - data.len() - ); - - for &offset in &offset_range[1..] { - let offset = offset.as_(); - let relative = offset.checked_sub(byte_start).ok_or_else(|| { - vortex_err!("VarBin offsets are not monotonic at offset {offset}") - })?; - let output_offset = output_bytes.checked_add(relative).ok_or_else(|| { - vortex_err!("PiecewiseSequence VarBin output byte length overflow") - })?; - new_offsets.push(new_offset_value::(output_offset)?); - } - - output_bytes = output_bytes - .checked_add(byte_end - byte_start) - .ok_or_else(|| vortex_err!("PiecewiseSequence VarBin output byte length overflow"))?; - } - vortex_ensure!( - new_offsets.len() == output_len + 1, - "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", - new_offsets.len() - 1 - ); - - let mut new_data = ByteBufferMut::with_capacity(output_bytes); - for (&start, &length) in starts.iter().zip_eq(lengths) { - let start = start.as_(); - let length = length.as_(); - if length == 0 { - continue; - } - - let offset_range = &offsets[start..][..=length]; - let byte_start = offset_range[0].as_(); - let byte_end = offset_range[length].as_(); - new_data.extend_from_slice(&data[byte_start..][..byte_end - byte_start]); - } - - let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) - .reinterpret_cast(out_offset_ptype) - .into_array(); - Ok(GatheredPiecewiseVarBin { - offsets, - data: new_data, - }) -} - -fn new_offset_value(value: usize) -> VortexResult { - T::from(value).ok_or_else(|| { - vortex_err!( - "PiecewiseSequence VarBin offset value {value} does not fit in {}", - T::PTYPE - ) - }) -} - fn take_nullable( dtype: DType, offsets: &[Offset], @@ -695,10 +303,7 @@ mod tests { ))] #[case(VarBinArray::from_iter(["single"].map(Some), DType::Utf8(Nullability::NonNullable)))] fn test_take_varbin_conformance(#[case] array: VarBinArray) { - test_take_conformance( - &array.into_array(), - &mut array_session().create_execution_ctx(), - ); + test_take_conformance(&array.into_array()); } #[test] diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index 40e05e08737..1b7435729f1 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -4,32 +4,23 @@ use std::iter; use std::sync::Arc; -use itertools::Itertools as _; use num_traits::AsPrimitive; use vortex_buffer::Buffer; -use vortex_buffer::BufferMut; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; -use vortex_error::vortex_err; use vortex_mask::AllOr; use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; use crate::array::ArrayView; -use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::arrays::dict::TakeExecute; -use crate::arrays::piecewise_sequence::ConstantOrArray; -use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::varbinview::BinaryView; use crate::buffer::BufferHandle; -use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_integer_ptype; -use crate::match_each_unsigned_integer_ptype; impl TakeExecute for VarBinView { /// Take involves creating a new array that references the old array, just with the given set of views. @@ -38,12 +29,6 @@ impl TakeExecute for VarBinView { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - if let Some(piecewise_indices) = indices.as_opt::() - && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)? - { - return Ok(Some(taken)); - } - let validity = array.validity()?.take(indices)?; let indices = indices.clone().execute::(ctx)?; @@ -72,58 +57,6 @@ impl TakeExecute for VarBinView { } } -fn take_contiguous_ranges( - array: ArrayView<'_, VarBinView>, - indices: ArrayView<'_, PiecewiseSequence>, - indices_ref: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult> { - let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { - return Ok(None); - }; - let source = array.views(); - let output_len = indices_ref.len(); - let views = match &lengths { - ConstantOrArray::Constant(length) => { - match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - gather_view_slices_constant_length( - source, - starts.as_slice::(), - *length, - output_len, - )? - }) - } - ConstantOrArray::Array(lengths) => { - match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { - gather_view_slices( - source, - starts.as_slice::(), - lengths.as_slice::(), - output_len, - )? - }) - }) - } - }; - let validity = array.validity()?.take(indices_ref)?; - - // SAFETY: ranges were validated against the source views, and copied views still reference the - // same backing data buffers. - unsafe { - Ok(Some( - VarBinViewArray::new_handle_unchecked( - BufferHandle::new_host(views.into_byte_buffer()), - Arc::clone(array.data_buffers()), - array.dtype().clone(), - validity, - ) - .into_array(), - )) - } -} - fn take_views>( views_ref: &[BinaryView], indices: &[I], @@ -152,58 +85,6 @@ fn take_views>( } } -fn gather_view_slices_constant_length( - source: &[BinaryView], - starts: &[S], - length: usize, - output_len: usize, -) -> VortexResult> -where - S: UnsignedPType, -{ - let computed_len = starts - .len() - .checked_mul(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; - vortex_ensure!( - computed_len == output_len, - "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" - ); - - let mut views = BufferMut::::with_capacity(output_len); - for &start in starts { - let start = start.as_(); - views.extend_from_slice(&source[start..][..length]); - } - - Ok(views.freeze()) -} - -fn gather_view_slices( - source: &[BinaryView], - starts: &[S], - lengths: &[L], - output_len: usize, -) -> VortexResult> -where - S: UnsignedPType, - L: UnsignedPType, -{ - let mut views = BufferMut::::with_capacity(output_len); - for (&start, &length) in starts.iter().zip_eq(lengths) { - let start = start.as_(); - let length = length.as_(); - views.extend_from_slice(&source[start..][..length]); - } - - vortex_ensure!( - views.len() == output_len, - "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", - views.len() - ); - Ok(views.freeze()) -} - #[cfg(test)] mod tests { use rstest::rstest; @@ -292,9 +173,6 @@ mod tests { ))] #[case(VarBinViewArray::from_iter(["single"].map(Some), DType::Utf8(NonNullable)))] fn test_take_varbinview_conformance(#[case] array: VarBinViewArray) { - test_take_conformance( - &array.into_array(), - &mut array_session().create_execution_ctx(), - ); + test_take_conformance(&array.into_array()); } } From 6ee68b3ce77c9c66daa01ee68d966735f3900431 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 16:36:23 -0400 Subject: [PATCH 23/37] Fix VarBin take conformance tests Signed-off-by: Daniel King --- vortex-array/src/arrays/varbin/compute/take.rs | 5 ++++- vortex-array/src/arrays/varbinview/compute/take.rs | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 74e0bd6af9d..6a6454d0603 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -303,7 +303,10 @@ mod tests { ))] #[case(VarBinArray::from_iter(["single"].map(Some), DType::Utf8(Nullability::NonNullable)))] fn test_take_varbin_conformance(#[case] array: VarBinArray) { - test_take_conformance(&array.into_array()); + test_take_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } #[test] diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index 1b7435729f1..ce10abda3d0 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -173,6 +173,9 @@ mod tests { ))] #[case(VarBinViewArray::from_iter(["single"].map(Some), DType::Utf8(NonNullable)))] fn test_take_varbinview_conformance(#[case] array: VarBinViewArray) { - test_take_conformance(&array.into_array()); + test_take_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); } } From 4b8d2423c98f4941db4a16230463d92f9f09f5d2 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 16 Jul 2026 15:03:02 -0400 Subject: [PATCH 24/37] Port PiecewiseSequence run take consumers Signed-off-by: Daniel King --- .../src/arrays/extension/compute/rules.rs | 2 + .../src/arrays/extension/compute/take.rs | 16 ++ vortex-array/src/arrays/list/compute/take.rs | 249 ++++++++++++++++++ vortex-array/src/arrays/listview/rebuild.rs | 228 ++++++++++------ 4 files changed, 410 insertions(+), 85 deletions(-) diff --git a/vortex-array/src/arrays/extension/compute/rules.rs b/vortex-array/src/arrays/extension/compute/rules.rs index 2d02d1ae7a7..0699f2f68f4 100644 --- a/vortex-array/src/arrays/extension/compute/rules.rs +++ b/vortex-array/src/arrays/extension/compute/rules.rs @@ -11,6 +11,7 @@ use crate::arrays::ConstantArray; use crate::arrays::Extension; use crate::arrays::ExtensionArray; use crate::arrays::Filter; +use crate::arrays::dict::TakeReduceAdaptor; use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::filter::FilterReduceAdaptor; use crate::arrays::slice::SliceReduceAdaptor; @@ -50,6 +51,7 @@ pub(crate) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ ParentRuleSet::lift(&FilterReduceAdaptor(Extension)), ParentRuleSet::lift(&MaskReduceAdaptor(Extension)), ParentRuleSet::lift(&SliceReduceAdaptor(Extension)), + ParentRuleSet::lift(&TakeReduceAdaptor(Extension)), ]); /// Push filter operations into the storage array of an extension array. diff --git a/vortex-array/src/arrays/extension/compute/take.rs b/vortex-array/src/arrays/extension/compute/take.rs index 8fb0da11222..d9ac1b4f600 100644 --- a/vortex-array/src/arrays/extension/compute/take.rs +++ b/vortex-array/src/arrays/extension/compute/take.rs @@ -10,8 +10,24 @@ use crate::array::ArrayView; use crate::arrays::Extension; use crate::arrays::ExtensionArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::dict::TakeReduce; use crate::arrays::extension::ExtensionArrayExt; +impl TakeReduce for Extension { + fn take(array: ArrayView<'_, Extension>, indices: &ArrayRef) -> VortexResult> { + let taken_storage = array.storage_array().take(indices.clone())?; + Ok(Some( + ExtensionArray::new( + array + .ext_dtype() + .with_nullability(taken_storage.dtype().nullability()), + taken_storage, + ) + .into_array(), + )) + } +} + impl TakeExecute for Extension { fn take( array: ArrayView<'_, Extension>, diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index 14e89d4c238..35788796d78 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -1,26 +1,36 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use itertools::Itertools as _; +use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use crate::ArrayRef; use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::List; use crate::arrays::ListArray; +use crate::arrays::PiecewiseSequence; +use crate::arrays::PiecewiseSequenceArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; use crate::arrays::list::ListArrayExt; +use crate::arrays::piecewise_sequence::execute_index_arrays; +use crate::arrays::piecewise_sequence::validate_index_ranges; use crate::arrays::primitive::PrimitiveArrayExt; use crate::builders::ArrayBuilder; use crate::builders::PrimitiveBuilder; use crate::dtype::IntegerPType; use crate::dtype::Nullability; +use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_unsigned_integer_ptype; use crate::match_smallest_offset_type; +use crate::validity::Validity; // TODO(connor)[ListView]: Re-revert to the version where we simply convert to a `ListView` and call // the `ListView::take` compute function once `ListView` is more stable. @@ -37,6 +47,12 @@ impl TakeExecute for List { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_piecewise_sequence(array, piecewise_indices, indices, ctx)? + { + return Ok(Some(taken)); + } + let indices = indices.clone().execute::(ctx)?; let indices = indices.reinterpret_cast(indices.ptype().to_unsigned()); let offsets = array.offsets().clone().execute::(ctx)?; @@ -127,6 +143,191 @@ fn _take( .into_array()) } +fn take_piecewise_sequence( + array: ArrayView<'_, List>, + indices: ArrayView<'_, PiecewiseSequence>, + indices_ref: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let data_validity = array + .list_validity() + .execute_mask(array.as_ref().len(), ctx)?; + if !data_validity.all_true() { + return Ok(None); + } + + let (starts, lengths) = execute_index_arrays(indices, ctx)?; + let offsets = array.offsets().clone().execute::(ctx)?; + let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); + + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { + take_piecewise_sequence_typed::( + array, + starts.as_slice::(), + lengths.as_slice::(), + offsets.as_slice::(), + indices_ref, + ) + }) + }) + }) + .map(Some) +} + +fn take_piecewise_sequence_typed( + array: ArrayView<'_, List>, + starts: &[S], + lengths: &[L], + offsets: &[Offset], + indices_ref: &ArrayRef, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, + Offset: UnsignedPType, +{ + validate_index_ranges(array.len(), starts, lengths, indices_ref.len())?; + let total_elements = + piecewise_list_elements_len(array.elements().len(), offsets, starts, lengths)?; + + match_smallest_offset_type!(total_elements, |OutputOffset| { + let gathered = gather_piecewise_list::( + array.elements(), + offsets, + starts, + lengths, + indices_ref.len(), + total_elements, + )?; + let validity = array.validity()?.take(indices_ref)?; + + // SAFETY: output offsets are rebuilt from valid monotonic source offsets; output elements + // are exactly the gathered child ranges referenced by those offsets; validity has one bit + // per output row. + Ok( + unsafe { ListArray::new_unchecked(gathered.elements, gathered.offsets, validity) } + .into_array(), + ) + }) +} + +struct GatheredList { + elements: ArrayRef, + offsets: ArrayRef, +} + +fn piecewise_list_elements_len( + elements_len: usize, + offsets: &[Offset], + starts: &[S], + lengths: &[L], +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, + Offset: UnsignedPType, +{ + let mut total = 0usize; + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start: usize = start.as_(); + let length: usize = length.as_(); + let end = start + length; + if length == 0 { + continue; + } + + let element_start: usize = offsets[start].as_(); + let element_end: usize = offsets[end].as_(); + vortex_ensure!( + element_start <= element_end && element_end <= elements_len, + "List offsets range {element_start}..{element_end} exceeds elements length {elements_len}", + ); + total = total + .checked_add(element_end - element_start) + .ok_or_else(|| vortex_err!("List take output elements length overflow"))?; + } + Ok(total) +} + +fn gather_piecewise_list( + elements: &ArrayRef, + offsets: &[Offset], + starts: &[S], + lengths: &[L], + output_len: usize, + total_elements: usize, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, + Offset: UnsignedPType, + OutputOffset: IntegerPType, +{ + let offsets_capacity = output_len + .checked_add(1) + .ok_or_else(|| vortex_err!("List take offsets length overflow"))?; + let mut new_offsets = BufferMut::::with_capacity(offsets_capacity); + let mut element_starts = BufferMut::::with_capacity(starts.len()); + let mut element_lengths = BufferMut::::with_capacity(lengths.len()); + let mut output_elements = 0usize; + + new_offsets.push(OutputOffset::zero()); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start: usize = start.as_(); + let length: usize = length.as_(); + let end = start + length; + if length == 0 { + continue; + } + + let element_start: usize = offsets[start].as_(); + let element_end: usize = offsets[end].as_(); + for &offset in &offsets[start + 1..=end] { + let offset: usize = offset.as_(); + let relative = offset + .checked_sub(element_start) + .ok_or_else(|| vortex_err!("List offsets are not monotonic at offset {offset}"))?; + let output_offset = output_elements + .checked_add(relative) + .ok_or_else(|| vortex_err!("List take output elements length overflow"))?; + new_offsets.push(new_offset_value::(output_offset)?); + } + + let element_length = element_end - element_start; + element_starts.push(element_start as u64); + element_lengths.push(element_length as u64); + output_elements = output_elements + .checked_add(element_length) + .ok_or_else(|| vortex_err!("List take output elements length overflow"))?; + } + debug_assert_eq!(output_elements, total_elements); + + let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable).into_array(); + // SAFETY: element ranges are derived from validated source list offsets, and total_elements is + // the sum of the gathered element range lengths. + let element_indices = unsafe { + PiecewiseSequenceArray::new_unchecked( + element_starts.into_array(), + element_lengths.into_array(), + total_elements, + ) + }; + let elements = elements.take(element_indices.into_array())?; + + Ok(GatheredList { elements, offsets }) +} + +fn new_offset_value(value: usize) -> VortexResult { + T::from_usize(value).ok_or_else(|| { + vortex_err!( + "List take offset value {value} does not fit in {}", + T::PTYPE + ) + }) +} + // Kept out-of-line: as a single-callsite generic helper it would otherwise be inlined into every // monomorphization of `_take`, duplicating the entire nullable path across all specializations. #[inline(never)] @@ -217,6 +418,7 @@ mod test { use crate::arrays::BoolArray; use crate::arrays::ListArray; use crate::arrays::ListViewArray; + use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; use crate::compute::conformance::take::test_take_conformance; use crate::dtype::DType; @@ -403,6 +605,53 @@ mod test { ); } + #[test] + fn piecewise_sequence_take() { + let mut ctx = array_session().create_execution_ctx(); + let list = ListArray::try_new( + buffer![0i32, 1, 2, 3, 4, 5, 6].into_array(), + buffer![0u32, 2, 5, 5, 7].into_array(), + Validity::NonNullable, + ) + .unwrap() + .into_array(); + let idx = PiecewiseSequenceArray::try_new( + buffer![1u64, 0].into_array(), + buffer![2u64, 1].into_array(), + 3, + ) + .unwrap() + .into_array(); + + let result = list + .take(idx) + .unwrap() + .execute::(&mut ctx) + .unwrap(); + + let element_dtype: Arc = Arc::new(I32.into()); + assert_eq!( + result.execute_scalar(0, &mut ctx).unwrap(), + Scalar::list( + Arc::clone(&element_dtype), + vec![2i32.into(), 3.into(), 4.into()], + Nullability::NonNullable + ) + ); + assert_eq!( + result.execute_scalar(1, &mut ctx).unwrap(), + Scalar::list(Arc::clone(&element_dtype), vec![], Nullability::NonNullable) + ); + assert_eq!( + result.execute_scalar(2, &mut ctx).unwrap(), + Scalar::list( + element_dtype, + vec![0i32.into(), 1.into()], + Nullability::NonNullable + ) + ); + } + #[test] fn test_take_empty_array() { let list = ListArray::try_new( diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 342c9eccf19..1ce8f69d12e 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -5,16 +5,18 @@ use num_traits::FromPrimitive; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::Mask; -use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; +use crate::RecursiveCanonical; use crate::arrays::ConstantArray; use crate::arrays::ListViewArray; +use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; use crate::arrays::listview::ListViewArrayExt; use crate::arrays::primitive::PrimitiveArrayExt; -use crate::builders::builder_with_capacity; use crate::builtins::ArrayBuiltins; use crate::dtype::IntegerPType; use crate::dtype::Nullability; @@ -140,19 +142,18 @@ impl ListViewArray { // for sizes as well. match_each_unsigned_integer_ptype!(sizes_ptype.to_unsigned(), |S| { match offsets_ptype.to_unsigned() { - PType::U8 => self.naive_rebuild::(ctx), - PType::U16 => self.naive_rebuild::(ctx), - PType::U32 => self.naive_rebuild::(ctx), - PType::U64 => self.naive_rebuild::(ctx), + PType::U8 => self.rebuild_with_take_or_piecewise::(ctx), + PType::U16 => self.rebuild_with_take_or_piecewise::(ctx), + PType::U32 => self.rebuild_with_take_or_piecewise::(ctx), + PType::U64 => self.rebuild_with_take_or_piecewise::(ctx), _ => unreachable!("invalid offsets PType"), } }) } /// Picks between [`rebuild_with_take`](Self::rebuild_with_take) and - /// [`rebuild_list_by_list`](Self::rebuild_list_by_list) based on element dtype and average - /// list size. - fn naive_rebuild( + /// [`rebuild_with_piecewise`](Self::rebuild_with_piecewise) based on average list size. + fn rebuild_with_take_or_piecewise( &self, ctx: &mut ExecutionCtx, ) -> VortexResult { @@ -167,12 +168,12 @@ impl ListViewArray { if Self::should_use_take(total, self.len()) { self.rebuild_with_take::(ctx) } else { - self.rebuild_list_by_list::(ctx) + self.rebuild_with_piecewise::(ctx) } } /// Returns `true` when we are confident that `rebuild_with_take` will - /// outperform `rebuild_list_by_list`. + /// outperform `rebuild_with_piecewise`. /// /// Take is dramatically faster for small lists (often 10-100×) because it /// avoids per-list builder overhead. LBL is the safer default for larger @@ -243,87 +244,65 @@ impl ListViewArray { .reinterpret_cast(size_ptype) .into_array(); - // SAFETY: same invariants as `rebuild_list_by_list` — offsets are sequential and - // non-overlapping, all (offset, size) pairs reference valid elements, and the validity - // array is preserved from the original. + // SAFETY: offsets are sequential and non-overlapping, all (offset, size) pairs reference + // valid elements, and the validity array is preserved from the original. Ok(unsafe { ListViewArray::new_unchecked(elements, offsets, sizes, self.validity()?) .with_zero_copy_to_list(true) }) } - /// Rebuilds elements list-by-list: canonicalize elements upfront, then for each list `slice` - /// the relevant range and `append_to_builder` into a typed builder. - fn rebuild_list_by_list( + /// Rebuilds elements using one contiguous-run take over the element child. + fn rebuild_with_piecewise( &self, ctx: &mut ExecutionCtx, ) -> VortexResult { - let element_dtype = self - .dtype() - .as_list_element_opt() - .vortex_expect("somehow had a canonical list that was not a list"); - let new_offset_ptype = rebuilt_offset_ptype(self.offsets().dtype().as_ptype()); let size_ptype = self.sizes().dtype().as_ptype(); let offsets_canonical = self.offsets().clone().execute::(ctx)?; let offsets_canonical = offsets_canonical.reinterpret_cast(offsets_canonical.ptype().to_unsigned()); - let offsets_slice = offsets_canonical.as_slice::(); let sizes_canonical = self.sizes().clone().execute::(ctx)?; let sizes_canonical = sizes_canonical.reinterpret_cast(sizes_canonical.ptype().to_unsigned()); - let sizes_slice = sizes_canonical.as_slice::(); - let len = offsets_slice.len(); + let len = offsets_canonical.len(); + let validity = self.validity()?; + let validity_mask = validity.execute_mask(len, ctx)?; - let mut new_offsets = BufferMut::::with_capacity(len); - // TODO(connor)[ListView]: Do we really need to do this? - // The only reason we need to rebuild the sizes here is that the validity may indicate that - // a list is null even though it has a non-zero size. This rebuild will set the size of all - // null lists to 0. - let mut new_sizes = BufferMut::::with_capacity(len); - - // Canonicalize the elements up front as we will be slicing the elements quite a lot. - let elements_canonical = self + let ranges = match_each_unsigned_integer_ptype!(offsets_canonical.ptype(), |O| { + rebuild_ranges::( + offsets_canonical.as_slice::(), + sizes_canonical.as_slice::(), + &validity_mask, + ) + })?; + + let RebuildRanges { + new_offsets, + new_sizes, + starts, + lengths, + elements_len, + } = ranges; + + // SAFETY: range starts and lengths are derived from valid ListView metadata; elements_len + // is the sum of all generated range lengths. + let element_indices = unsafe { + PiecewiseSequenceArray::new_unchecked( + starts.into_array(), + lengths.into_array(), + elements_len, + ) + }; + let elements = self .elements() - .clone() - .execute::(ctx)? + .take(element_indices.into_array())? + .execute::(ctx)? + .0 .into_array(); - // Note that we do not know what the exact capacity should be of the new elements since - // there could be overlaps in the existing `ListViewArray`. - let mut new_elements_builder = - builder_with_capacity(element_dtype.as_ref(), self.elements().len()); - - // Resolve validity to a mask once instead of probing it per row (see `rebuild_with_take`). - let validity = self.validity()?.execute_mask(len, ctx)?; - - let mut n_elements = NewOffset::zero(); - for index in 0..len { - if !validity.value(index) { - // For NULL lists, place them after the previous item's data to maintain the - // no-overlap invariant for zero-copy to `ListArray` arrays. - new_offsets.push(n_elements); - new_sizes.push(S::zero()); - continue; - } - - let offset = offsets_slice[index]; - let size = sizes_slice[index]; - - let start = offset.as_(); - let stop = start + size.as_(); - - new_offsets.push(n_elements); - new_sizes.push(size); - elements_canonical - .slice(start..stop)? - .append_to_builder(new_elements_builder.as_mut(), ctx)?; - - n_elements += num_traits::cast(size).vortex_expect("Cast failed"); - } - // Built unsigned; reinterpret back to the signed-preserving result types. let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) .reinterpret_cast(new_offset_ptype) @@ -331,24 +310,11 @@ impl ListViewArray { let sizes = PrimitiveArray::new(new_sizes.freeze(), Validity::NonNullable) .reinterpret_cast(size_ptype) .into_array(); - let elements = new_elements_builder.finish(); - debug_assert_eq!( - n_elements.as_(), - elements.len(), - "The accumulated elements somehow had the wrong length" - ); - - // SAFETY: - // - All offsets are sequential and non-overlapping (`n_elements` tracks running total). - // - Each `offset[i] + size[i]` equals `offset[i+1]` for all valid indices (including null - // lists). - // - All elements referenced by (offset, size) pairs exist within the new `elements` array. - // - The validity array is preserved from the original array unchanged - // - The array satisfies the zero-copy-to-list property by having sorted offsets, no gaps, - // and no overlaps. + // SAFETY: offsets are sequential and non-overlapping, all (offset, size) pairs reference + // valid elements, and validity is preserved from the original array. Ok(unsafe { - ListViewArray::new_unchecked(elements, offsets, sizes, self.validity()?) + ListViewArray::new_unchecked(elements, offsets, sizes, validity) .with_zero_copy_to_list(true) }) } @@ -414,6 +380,69 @@ impl ListViewArray { } } +struct RebuildRanges { + new_offsets: BufferMut, + new_sizes: BufferMut, + starts: BufferMut, + lengths: BufferMut, + elements_len: usize, +} + +fn rebuild_ranges( + offsets: &[O], + sizes: &[S], + validity_mask: &Mask, +) -> VortexResult> +where + NewOffset: IntegerPType, + O: IntegerPType, + S: IntegerPType, +{ + let len = offsets.len(); + let mut new_offsets = BufferMut::::with_capacity(len); + let mut new_sizes = BufferMut::::with_capacity(len); + let mut starts = BufferMut::::with_capacity(len); + let mut lengths = BufferMut::::with_capacity(len); + let mut n_elements = NewOffset::zero(); + let mut elements_len = 0usize; + + for (index, is_valid) in validity_mask.iter().enumerate() { + if !is_valid { + new_offsets.push(n_elements); + new_sizes.push(S::zero()); + starts.push(0); + lengths.push(0); + continue; + } + + let size = sizes[index]; + let start: usize = offsets[index].as_(); + let length: usize = size.as_(); + start.checked_add(length).ok_or_else(|| { + vortex_err!( + "ListView rebuild element range overflow for start {start} and length {length}" + ) + })?; + elements_len = elements_len + .checked_add(length) + .ok_or_else(|| vortex_err!("ListView rebuild elements length overflow"))?; + + new_offsets.push(n_elements); + new_sizes.push(size); + starts.push(start as u64); + lengths.push(length as u64); + n_elements += num_traits::cast(size).vortex_expect("Cast failed"); + } + + Ok(RebuildRanges { + new_offsets, + new_sizes, + starts, + lengths, + elements_len, + }) +} + #[cfg(test)] mod tests { use vortex_buffer::BitBuffer; @@ -561,6 +590,35 @@ mod tests { Ok(()) } + #[test] + fn test_rebuild_flatten_large_lists_with_piecewise_indices() -> VortexResult<()> { + let elements = PrimitiveArray::from_iter(0i32..300).into_array(); + let offsets = PrimitiveArray::from_iter(vec![10u32, 0]).into_array(); + let sizes = PrimitiveArray::from_iter(vec![150u32, 130]).into_array(); + let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable); + + let mut ctx = SESSION.create_execution_ctx(); + let flattened = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx)?; + + assert_eq!(flattened.elements().len(), 280); + assert_eq!(flattened.offset_at(0), 0); + assert_eq!(flattened.size_at(0), 150); + assert_eq!(flattened.offset_at(1), 150); + assert_eq!(flattened.size_at(1), 130); + + assert_arrays_eq!( + flattened.list_elements_at(0)?, + PrimitiveArray::from_iter(10i32..160), + &mut ctx + ); + assert_arrays_eq!( + flattened.list_elements_at(1)?, + PrimitiveArray::from_iter(0i32..130), + &mut ctx + ); + Ok(()) + } + #[test] fn test_rebuild_with_trailing_nulls_regression() -> VortexResult<()> { // Regression test for issue #5412 From 116ee72c5ad2201f94f6f2e017eceb7b3e340d8a Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 16 Jul 2026 16:28:40 -0400 Subject: [PATCH 25/37] Use unit multipliers for PiecewiseSequence consumers Signed-off-by: Daniel King --- vortex-array/src/arrays/list/compute/take.rs | 13 ++++++++++--- vortex-array/src/arrays/listview/rebuild.rs | 4 +++- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index 35788796d78..9a6c447b730 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -11,6 +11,7 @@ use vortex_error::vortex_err; use crate::ArrayRef; use crate::IntoArray; use crate::array::ArrayView; +use crate::arrays::ConstantArray; use crate::arrays::List; use crate::arrays::ListArray; use crate::arrays::PiecewiseSequence; @@ -19,7 +20,7 @@ use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; use crate::arrays::list::ListArrayExt; -use crate::arrays::piecewise_sequence::execute_index_arrays; +use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; use crate::arrays::piecewise_sequence::validate_index_ranges; use crate::arrays::primitive::PrimitiveArrayExt; use crate::builders::ArrayBuilder; @@ -156,7 +157,9 @@ fn take_piecewise_sequence( return Ok(None); } - let (starts, lengths) = execute_index_arrays(indices, ctx)?; + let Some((starts, lengths)) = execute_unit_multiplier_index_arrays(indices, ctx)? else { + return Ok(None); + }; let offsets = array.offsets().clone().execute::(ctx)?; let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); @@ -305,12 +308,14 @@ where debug_assert_eq!(output_elements, total_elements); let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable).into_array(); + let multipliers = ConstantArray::new(1u64, element_starts.len()).into_array(); // SAFETY: element ranges are derived from validated source list offsets, and total_elements is - // the sum of the gathered element range lengths. + // the sum of the gathered element range lengths. Multiplier 1 preserves contiguous ranges. let element_indices = unsafe { PiecewiseSequenceArray::new_unchecked( element_starts.into_array(), element_lengths.into_array(), + multipliers, total_elements, ) }; @@ -416,6 +421,7 @@ mod test { use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::BoolArray; + use crate::arrays::ConstantArray; use crate::arrays::ListArray; use crate::arrays::ListViewArray; use crate::arrays::PiecewiseSequenceArray; @@ -618,6 +624,7 @@ mod test { let idx = PiecewiseSequenceArray::try_new( buffer![1u64, 0].into_array(), buffer![2u64, 1].into_array(), + ConstantArray::new(1u64, 2).into_array(), 3, ) .unwrap() diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 1ce8f69d12e..5cf8bcb95b1 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -288,11 +288,13 @@ impl ListViewArray { } = ranges; // SAFETY: range starts and lengths are derived from valid ListView metadata; elements_len - // is the sum of all generated range lengths. + // is the sum of all generated range lengths. Multiplier 1 preserves contiguous ranges. + let multipliers = ConstantArray::new(1u64, starts.len()).into_array(); let element_indices = unsafe { PiecewiseSequenceArray::new_unchecked( starts.into_array(), lengths.into_array(), + multipliers, elements_len, ) }; From 107d2fbec990bb62c5eb517eeadf26458a880682 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 16 Jul 2026 19:51:55 -0400 Subject: [PATCH 26/37] Specialize constant PiecewiseSequence runs Signed-off-by: Daniel King --- vortex-array/src/arrays/list/compute/take.rs | 283 +++++++++++++++++- vortex-array/src/arrays/listview/rebuild.rs | 18 +- .../src/arrays/piecewise_sequence/tests.rs | 22 ++ 3 files changed, 301 insertions(+), 22 deletions(-) diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index 9a6c447b730..369a743d4dc 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -20,8 +20,10 @@ use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; use crate::arrays::list::ListArrayExt; +use crate::arrays::piecewise_sequence::UnitMultiplierLengths; use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; use crate::arrays::piecewise_sequence::validate_index_ranges; +use crate::arrays::piecewise_sequence::validate_index_ranges_constant; use crate::arrays::primitive::PrimitiveArrayExt; use crate::builders::ArrayBuilder; use crate::builders::PrimitiveBuilder; @@ -162,21 +164,174 @@ fn take_piecewise_sequence( }; let offsets = array.offsets().clone().execute::(ctx)?; let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); + let output_len = indices_ref.len(); + + let taken = match &lengths { + UnitMultiplierLengths::Constant(length) => take_piecewise_sequence_constant_dispatch( + array, + &starts, + *length, + &offsets, + indices_ref, + output_len, + )?, + UnitMultiplierLengths::Array(lengths) => take_piecewise_sequence_lengths_dispatch( + array, + &starts, + lengths, + &offsets, + indices_ref, + output_len, + )?, + }; + Ok(Some(taken)) +} +fn take_piecewise_sequence_constant_dispatch( + array: ArrayView<'_, List>, + starts: &PrimitiveArray, + length: usize, + offsets: &PrimitiveArray, + indices_ref: &ArrayRef, + output_len: usize, +) -> VortexResult { match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { - match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { - take_piecewise_sequence_typed::( - array, - starts.as_slice::(), - lengths.as_slice::(), - offsets.as_slice::(), - indices_ref, - ) - }) - }) + take_piecewise_sequence_constant_start_dispatch::( + array, + starts, + length, + offsets, + indices_ref, + output_len, + ) + }) +} + +fn take_piecewise_sequence_constant_start_dispatch( + array: ArrayView<'_, List>, + starts: &PrimitiveArray, + length: usize, + offsets: &PrimitiveArray, + indices_ref: &ArrayRef, + output_len: usize, +) -> VortexResult +where + S: UnsignedPType, +{ + match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { + take_piecewise_sequence_constant_length::( + array, + starts.as_slice::(), + length, + offsets.as_slice::(), + indices_ref, + output_len, + ) + }) +} + +fn take_piecewise_sequence_lengths_dispatch( + array: ArrayView<'_, List>, + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + offsets: &PrimitiveArray, + indices_ref: &ArrayRef, + output_len: usize, +) -> VortexResult { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + take_piecewise_sequence_lengths_start_dispatch::( + array, + starts, + lengths, + offsets, + indices_ref, + output_len, + ) + }) +} + +fn take_piecewise_sequence_lengths_start_dispatch( + array: ArrayView<'_, List>, + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + offsets: &PrimitiveArray, + indices_ref: &ArrayRef, + output_len: usize, +) -> VortexResult +where + S: UnsignedPType, +{ + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + take_piecewise_sequence_lengths_start_length_dispatch::( + array, + starts, + lengths, + offsets, + indices_ref, + output_len, + ) + }) +} + +fn take_piecewise_sequence_lengths_start_length_dispatch( + array: ArrayView<'_, List>, + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + offsets: &PrimitiveArray, + indices_ref: &ArrayRef, + output_len: usize, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, +{ + match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { + take_piecewise_sequence_typed::( + array, + starts.as_slice::(), + lengths.as_slice::(), + offsets.as_slice::(), + indices_ref, + output_len, + ) + }) +} + +fn take_piecewise_sequence_constant_length( + array: ArrayView<'_, List>, + starts: &[S], + length: usize, + offsets: &[Offset], + indices_ref: &ArrayRef, + output_len: usize, +) -> VortexResult +where + S: UnsignedPType, + Offset: UnsignedPType, +{ + validate_index_ranges_constant(array.len(), starts, length, output_len)?; + let total_elements = + piecewise_list_elements_len_constant(array.elements().len(), offsets, starts, length)?; + let validity = array.validity()?.take(indices_ref)?; + + match_smallest_offset_type!(total_elements, |OutputOffset| { + let gathered = gather_piecewise_list_constant_length::( + array.elements(), + offsets, + starts, + length, + output_len, + total_elements, + )?; + + // SAFETY: output offsets are rebuilt from valid monotonic source offsets; output elements + // are exactly the gathered child ranges referenced by those offsets; validity has one bit + // per output row. + Ok( + unsafe { ListArray::new_unchecked(gathered.elements, gathered.offsets, validity) } + .into_array(), + ) }) - .map(Some) } fn take_piecewise_sequence_typed( @@ -185,13 +340,14 @@ fn take_piecewise_sequence_typed( lengths: &[L], offsets: &[Offset], indices_ref: &ArrayRef, + output_len: usize, ) -> VortexResult where S: UnsignedPType, L: UnsignedPType, Offset: UnsignedPType, { - validate_index_ranges(array.len(), starts, lengths, indices_ref.len())?; + validate_index_ranges(array.len(), starts, lengths, output_len)?; let total_elements = piecewise_list_elements_len(array.elements().len(), offsets, starts, lengths)?; @@ -201,7 +357,7 @@ where offsets, starts, lengths, - indices_ref.len(), + output_len, total_elements, )?; let validity = array.validity()?.take(indices_ref)?; @@ -221,6 +377,37 @@ struct GatheredList { offsets: ArrayRef, } +fn piecewise_list_elements_len_constant( + elements_len: usize, + offsets: &[Offset], + starts: &[S], + length: usize, +) -> VortexResult +where + S: UnsignedPType, + Offset: UnsignedPType, +{ + let mut total = 0usize; + for &start in starts { + let start: usize = start.as_(); + let end = start + length; + if length == 0 { + continue; + } + + let element_start: usize = offsets[start].as_(); + let element_end: usize = offsets[end].as_(); + vortex_ensure!( + element_start <= element_end && element_end <= elements_len, + "List offsets range {element_start}..{element_end} exceeds elements length {elements_len}", + ); + total = total + .checked_add(element_end - element_start) + .ok_or_else(|| vortex_err!("List take output elements length overflow"))?; + } + Ok(total) +} + fn piecewise_list_elements_len( elements_len: usize, offsets: &[Offset], @@ -254,6 +441,74 @@ where Ok(total) } +fn gather_piecewise_list_constant_length( + elements: &ArrayRef, + offsets: &[Offset], + starts: &[S], + length: usize, + output_len: usize, + total_elements: usize, +) -> VortexResult +where + S: UnsignedPType, + Offset: UnsignedPType, + OutputOffset: IntegerPType, +{ + let offsets_capacity = output_len + .checked_add(1) + .ok_or_else(|| vortex_err!("List take offsets length overflow"))?; + let mut new_offsets = BufferMut::::with_capacity(offsets_capacity); + let mut element_starts = BufferMut::::with_capacity(starts.len()); + let mut element_lengths = BufferMut::::with_capacity(starts.len()); + let mut output_elements = 0usize; + + new_offsets.push(OutputOffset::zero()); + for &start in starts { + let start: usize = start.as_(); + let end = start + length; + if length == 0 { + continue; + } + + let element_start: usize = offsets[start].as_(); + let element_end: usize = offsets[end].as_(); + for &offset in &offsets[start + 1..=end] { + let offset: usize = offset.as_(); + let relative = offset + .checked_sub(element_start) + .ok_or_else(|| vortex_err!("List offsets are not monotonic at offset {offset}"))?; + let output_offset = output_elements + .checked_add(relative) + .ok_or_else(|| vortex_err!("List take output elements length overflow"))?; + new_offsets.push(new_offset_value::(output_offset)?); + } + + let element_length = element_end - element_start; + element_starts.push(element_start as u64); + element_lengths.push(element_length as u64); + output_elements = output_elements + .checked_add(element_length) + .ok_or_else(|| vortex_err!("List take output elements length overflow"))?; + } + debug_assert_eq!(output_elements, total_elements); + + let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable).into_array(); + let multipliers = ConstantArray::new(1u64, element_starts.len()).into_array(); + // SAFETY: element ranges are derived from validated source list offsets, and total_elements is + // the sum of the gathered element range lengths. Multiplier 1 preserves contiguous ranges. + let element_indices = unsafe { + PiecewiseSequenceArray::new_unchecked( + element_starts.into_array(), + element_lengths.into_array(), + multipliers, + total_elements, + ) + }; + let elements = elements.take(element_indices.into_array())?; + + Ok(GatheredList { elements, offsets }) +} + fn gather_piecewise_list( elements: &ArrayRef, offsets: &[Offset], diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 5cf8bcb95b1..15eb9feab38 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -10,7 +10,6 @@ use vortex_mask::Mask; use crate::ExecutionCtx; use crate::IntoArray; -use crate::RecursiveCanonical; use crate::arrays::ConstantArray; use crate::arrays::ListViewArray; use crate::arrays::PiecewiseSequenceArray; @@ -286,6 +285,14 @@ impl ListViewArray { lengths, elements_len, } = ranges; + let constant_length = lengths + .first() + .copied() + .filter(|first| lengths.iter().all(|length| *length == *first)); + let lengths = match constant_length { + Some(length) => ConstantArray::new(length, starts.len()).into_array(), + None => lengths.into_array(), + }; // SAFETY: range starts and lengths are derived from valid ListView metadata; elements_len // is the sum of all generated range lengths. Multiplier 1 preserves contiguous ranges. @@ -293,17 +300,12 @@ impl ListViewArray { let element_indices = unsafe { PiecewiseSequenceArray::new_unchecked( starts.into_array(), - lengths.into_array(), + lengths, multipliers, elements_len, ) }; - let elements = self - .elements() - .take(element_indices.into_array())? - .execute::(ctx)? - .0 - .into_array(); + let elements = self.elements().take(element_indices.into_array())?; // Built unsigned; reinterpret back to the signed-preserving result types. let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) diff --git a/vortex-array/src/arrays/piecewise_sequence/tests.rs b/vortex-array/src/arrays/piecewise_sequence/tests.rs index 82ead343f96..7a2a9401ec5 100644 --- a/vortex-array/src/arrays/piecewise_sequence/tests.rs +++ b/vortex-array/src/arrays/piecewise_sequence/tests.rs @@ -16,6 +16,7 @@ use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::DecimalArray; use crate::arrays::FixedSizeListArray; +use crate::arrays::ListArray; use crate::arrays::PiecewiseSequence; use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; @@ -354,6 +355,27 @@ fn contiguous_take_consumers_support_constant_piecewise_lengths() -> VortexResul Ok(()) } +#[test] +fn list_take_consumes_constant_piecewise_lengths() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let list = ListArray::try_new( + buffer![0i32, 1, 2, 3, 4, 5, 6].into_array(), + buffer![0u32, 2, 5, 5, 7].into_array(), + Validity::NonNullable, + )? + .into_array(); + let list_indices = piecewise_indices_constant_length([1, 0], 2)?; + let expected = ListArray::try_new( + buffer![2i32, 3, 4, 0, 1, 2, 3, 4].into_array(), + buffer![0u32, 3, 3, 5, 8].into_array(), + Validity::NonNullable, + )? + .into_array(); + assert_arrays_eq!(list.take(list_indices)?, expected, &mut ctx); + + Ok(()) +} + #[test] fn fixed_size_list_take_builds_piecewise_element_indices() -> VortexResult<()> { let elements = PrimitiveArray::from_iter(0i32..12).into_array(); From b0cfdc2d50fbaa589702703929c892f26d868a30 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 14:42:25 -0400 Subject: [PATCH 27/37] Inline List PiecewiseSequence range length checks Signed-off-by: Daniel King --- vortex-array/src/arrays/list/compute/take.rs | 51 +++++++++++++------- 1 file changed, 33 insertions(+), 18 deletions(-) diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index 369a743d4dc..9d83cca3a48 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -22,8 +22,6 @@ use crate::arrays::dict::TakeExecute; use crate::arrays::list::ListArrayExt; use crate::arrays::piecewise_sequence::UnitMultiplierLengths; use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; -use crate::arrays::piecewise_sequence::validate_index_ranges; -use crate::arrays::piecewise_sequence::validate_index_ranges_constant; use crate::arrays::primitive::PrimitiveArrayExt; use crate::builders::ArrayBuilder; use crate::builders::PrimitiveBuilder; @@ -309,7 +307,14 @@ where S: UnsignedPType, Offset: UnsignedPType, { - validate_index_ranges_constant(array.len(), starts, length, output_len)?; + let computed_len = starts + .len() + .checked_mul(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); let total_elements = piecewise_list_elements_len_constant(array.elements().len(), offsets, starts, length)?; let validity = array.validity()?.take(indices_ref)?; @@ -347,7 +352,17 @@ where L: UnsignedPType, Offset: UnsignedPType, { - validate_index_ranges(array.len(), starts, lengths, output_len)?; + let mut computed_len = 0usize; + for &length in lengths { + let length: usize = length.as_(); + computed_len = computed_len + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + } + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); let total_elements = piecewise_list_elements_len(array.elements().len(), offsets, starts, lengths)?; @@ -390,13 +405,13 @@ where let mut total = 0usize; for &start in starts { let start: usize = start.as_(); - let end = start + length; if length == 0 { continue; } - let element_start: usize = offsets[start].as_(); - let element_end: usize = offsets[end].as_(); + let offset_range = &offsets[start..][..=length]; + let element_start: usize = offset_range[0].as_(); + let element_end: usize = offset_range[length].as_(); vortex_ensure!( element_start <= element_end && element_end <= elements_len, "List offsets range {element_start}..{element_end} exceeds elements length {elements_len}", @@ -423,13 +438,13 @@ where for (&start, &length) in starts.iter().zip_eq(lengths) { let start: usize = start.as_(); let length: usize = length.as_(); - let end = start + length; if length == 0 { continue; } - let element_start: usize = offsets[start].as_(); - let element_end: usize = offsets[end].as_(); + let offset_range = &offsets[start..][..=length]; + let element_start: usize = offset_range[0].as_(); + let element_end: usize = offset_range[length].as_(); vortex_ensure!( element_start <= element_end && element_end <= elements_len, "List offsets range {element_start}..{element_end} exceeds elements length {elements_len}", @@ -465,14 +480,14 @@ where new_offsets.push(OutputOffset::zero()); for &start in starts { let start: usize = start.as_(); - let end = start + length; if length == 0 { continue; } - let element_start: usize = offsets[start].as_(); - let element_end: usize = offsets[end].as_(); - for &offset in &offsets[start + 1..=end] { + let offset_range = &offsets[start..][..=length]; + let element_start: usize = offset_range[0].as_(); + let element_end: usize = offset_range[length].as_(); + for &offset in &offset_range[1..] { let offset: usize = offset.as_(); let relative = offset .checked_sub(element_start) @@ -535,14 +550,14 @@ where for (&start, &length) in starts.iter().zip_eq(lengths) { let start: usize = start.as_(); let length: usize = length.as_(); - let end = start + length; if length == 0 { continue; } - let element_start: usize = offsets[start].as_(); - let element_end: usize = offsets[end].as_(); - for &offset in &offsets[start + 1..=end] { + let offset_range = &offsets[start..][..=length]; + let element_start: usize = offset_range[0].as_(); + let element_end: usize = offset_range[length].as_(); + for &offset in &offset_range[1..] { let offset: usize = offset.as_(); let relative = offset .checked_sub(element_start) From 69b929a8b7a729e2b04504b817fe49a3d59d2c8e Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 16:04:53 -0400 Subject: [PATCH 28/37] Cover null ListView rebuild placeholders Signed-off-by: Daniel King --- vortex-array/src/arrays/listview/rebuild.rs | 38 +++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 15eb9feab38..6e36958522e 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -544,6 +544,44 @@ mod tests { Ok(()) } + #[test] + fn test_rebuild_flatten_null_row_ignores_invalid_range_payload() -> VortexResult<()> { + let elements = PrimitiveArray::from_iter(vec![1i32, 2, 3]).into_array(); + let offsets = PrimitiveArray::from_iter(vec![0u32, 999, 2]).into_array(); + let sizes = PrimitiveArray::from_iter(vec![2u32, 999, 1]).into_array(); + let validity = Validity::from_iter([true, false, true]); + + // SAFETY: this intentionally models a null row whose physical offset and size payloads are + // invalid. Rebuild must ignore those payloads and only emit safe placeholder ranges for the + // null row. + let listview = unsafe { ListViewArray::new_unchecked(elements, offsets, sizes, validity) }; + + let mut ctx = SESSION.create_execution_ctx(); + let flattened = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx)?; + + assert_eq!(flattened.offset_at(0), 0); + assert_eq!(flattened.size_at(0), 2); + assert_eq!(flattened.offset_at(1), 2); + assert_eq!(flattened.size_at(1), 0); + assert_eq!(flattened.offset_at(2), 2); + assert_eq!(flattened.size_at(2), 1); + assert!(flattened.validity()?.execute_is_valid(0, &mut ctx)?); + assert!(!flattened.validity()?.execute_is_valid(1, &mut ctx)?); + assert!(flattened.validity()?.execute_is_valid(2, &mut ctx)?); + + assert_arrays_eq!( + flattened.list_elements_at(0)?, + PrimitiveArray::from_iter([1i32, 2]), + &mut ctx + ); + assert_arrays_eq!( + flattened.list_elements_at(2)?, + PrimitiveArray::from_iter([3i32]), + &mut ctx + ); + Ok(()) + } + #[test] fn test_rebuild_trim_elements_basic() -> VortexResult<()> { // Test trimming both leading and trailing unused elements while preserving gaps in the From 7248dd2cf33badd2dbc65fd6229da1e290f2775b Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 16:12:18 -0400 Subject: [PATCH 29/37] Mask null ListView take metadata Signed-off-by: Daniel King --- .../src/arrays/listview/compute/take.rs | 14 +++++++---- .../src/arrays/listview/tests/take.rs | 23 +++++++++++++++++++ 2 files changed, 33 insertions(+), 4 deletions(-) diff --git a/vortex-array/src/arrays/listview/compute/take.rs b/vortex-array/src/arrays/listview/compute/take.rs index 2b6c016d2c3..2389715c2d6 100644 --- a/vortex-array/src/arrays/listview/compute/take.rs +++ b/vortex-array/src/arrays/listview/compute/take.rs @@ -50,14 +50,20 @@ fn apply_take(array: ArrayView<'_, ListView>, indices: &ArrayRef) -> VortexResul // duplicates. let nullable_new_offsets = offsets.take(indices.clone())?; let nullable_new_sizes = sizes.take(indices.clone())?; + let validity_array = new_validity.to_array(indices.len()); - // `take` returns nullable arrays; cast back to non-nullable (filling with zeros to represent - // the null lists — the validity mask tracks nullness separately). + // Null output rows may carry arbitrary physical offset/size payloads from either the indices + // or the source rows. Mask with the final validity before filling so metadata placeholders are + // safe and non-nullable. let new_offsets = match_each_integer_ptype!(nullable_new_offsets.dtype().as_ptype(), |O| { - nullable_new_offsets.fill_null(Scalar::primitive(O::zero(), Nullability::NonNullable))? + nullable_new_offsets + .mask(validity_array.clone())? + .fill_null(Scalar::primitive(O::zero(), Nullability::NonNullable))? }); let new_sizes = match_each_integer_ptype!(nullable_new_sizes.dtype().as_ptype(), |S| { - nullable_new_sizes.fill_null(Scalar::primitive(S::zero(), Nullability::NonNullable))? + nullable_new_sizes + .mask(validity_array.clone())? + .fill_null(Scalar::primitive(S::zero(), Nullability::NonNullable))? }); // SAFETY: Take operation maintains all `ListViewArray` invariants: diff --git a/vortex-array/src/arrays/listview/tests/take.rs b/vortex-array/src/arrays/listview/tests/take.rs index 2af397e262f..8fb7a4af1bb 100644 --- a/vortex-array/src/arrays/listview/tests/take.rs +++ b/vortex-array/src/arrays/listview/tests/take.rs @@ -101,6 +101,29 @@ fn test_take_with_gaps() { ); } +#[test] +fn test_take_null_source_row_zeros_offset_size_payloads() { + let elements = buffer![1i32, 2].into_array(); + let offsets = buffer![0u32, 999].into_array(); + let sizes = buffer![2u32, 999].into_array(); + let validity = Validity::from_iter([true, false]); + let listview = + unsafe { ListViewArray::new_unchecked(elements, offsets, sizes, validity) }.into_array(); + + let result = listview.take(buffer![1u32].into_array()).unwrap(); + let result_list = result + .execute::(&mut SESSION.create_execution_ctx()) + .unwrap(); + + assert_eq!(result_list.offset_at(0), 0); + assert_eq!(result_list.size_at(0), 0); + assert!( + result_list + .is_invalid(0, &mut SESSION.create_execution_ctx()) + .unwrap() + ); +} + #[test] fn test_take_constant_arrays() { // ListView-specific: Test with ConstantArray for offsets/sizes. From e50a07e5acc06d9e9a59e296aa2123d1a51be006 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 16:19:04 -0400 Subject: [PATCH 30/37] Rename List contiguous slice helper uses Signed-off-by: Daniel King --- vortex-array/src/arrays/list/compute/take.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index 9d83cca3a48..ff9549d6a34 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -20,8 +20,8 @@ use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; use crate::arrays::list::ListArrayExt; -use crate::arrays::piecewise_sequence::UnitMultiplierLengths; -use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays; +use crate::arrays::piecewise_sequence::ConstantOrArray; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; use crate::builders::ArrayBuilder; use crate::builders::PrimitiveBuilder; @@ -157,7 +157,7 @@ fn take_piecewise_sequence( return Ok(None); } - let Some((starts, lengths)) = execute_unit_multiplier_index_arrays(indices, ctx)? else { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { return Ok(None); }; let offsets = array.offsets().clone().execute::(ctx)?; @@ -165,7 +165,7 @@ fn take_piecewise_sequence( let output_len = indices_ref.len(); let taken = match &lengths { - UnitMultiplierLengths::Constant(length) => take_piecewise_sequence_constant_dispatch( + ConstantOrArray::Constant(length) => take_piecewise_sequence_constant_dispatch( array, &starts, *length, @@ -173,7 +173,7 @@ fn take_piecewise_sequence( indices_ref, output_len, )?, - UnitMultiplierLengths::Array(lengths) => take_piecewise_sequence_lengths_dispatch( + ConstantOrArray::Array(lengths) => take_piecewise_sequence_lengths_dispatch( array, &starts, lengths, From e4a05437e02640dff06eaf22d4e3ee223ce6eb52 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 16:37:33 -0400 Subject: [PATCH 31/37] Remove redundant ListView validity clone Signed-off-by: Daniel King --- vortex-array/src/arrays/listview/compute/take.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-array/src/arrays/listview/compute/take.rs b/vortex-array/src/arrays/listview/compute/take.rs index 2389715c2d6..01ffb80dab7 100644 --- a/vortex-array/src/arrays/listview/compute/take.rs +++ b/vortex-array/src/arrays/listview/compute/take.rs @@ -62,7 +62,7 @@ fn apply_take(array: ArrayView<'_, ListView>, indices: &ArrayRef) -> VortexResul }); let new_sizes = match_each_integer_ptype!(nullable_new_sizes.dtype().as_ptype(), |S| { nullable_new_sizes - .mask(validity_array.clone())? + .mask(validity_array)? .fill_null(Scalar::primitive(S::zero(), Nullability::NonNullable))? }); From 79d9e10dee7cc10301176a99cd0e4f7f8e7d3368 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 16:24:09 -0400 Subject: [PATCH 32/37] Specialize VarBin PiecewiseSequence takes Signed-off-by: Daniel King --- .../src/arrays/varbin/compute/take.rs | 392 ++++++++++++++++++ .../src/arrays/varbinview/compute/take.rs | 119 ++++++ 2 files changed, 511 insertions(+) diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 6a6454d0603..b956359ef0b 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -1,26 +1,33 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use itertools::Itertools as _; use vortex_buffer::BitBufferMut; use vortex_buffer::BufferMut; use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; use crate::array::ArrayView; +use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::VarBin; use crate::arrays::VarBinArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::ConstantOrArray; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; use crate::arrays::varbin::VarBinArrayExt; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::dtype::PType; +use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_unsigned_integer_ptype; use crate::validity::Validity; @@ -43,6 +50,12 @@ impl TakeExecute for VarBin { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)? + { + return Ok(Some(taken)); + } + // TODO(joe): Be lazy with execute let offsets = array.offsets().clone().execute::(ctx)?; let data = array.bytes(); @@ -113,6 +126,54 @@ impl TakeExecute for VarBin { } } +fn take_contiguous_ranges( + array: ArrayView<'_, VarBin>, + indices: ArrayView<'_, PiecewiseSequence>, + indices_ref: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { + return Ok(None); + }; + let offsets = array.offsets().clone().execute::(ctx)?; + let out_offset_ptype = taken_offset_ptype(offsets.ptype()); + let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); + let bytes = array.bytes(); + let data = bytes.as_slice(); + let dtype = array.dtype().clone(); + let output_len = indices_ref.len(); + + let result = match &lengths { + ConstantOrArray::Constant(length) => gather_slices_constant_dispatch( + &starts, + *length, + &offsets, + data, + output_len, + out_offset_ptype, + )?, + ConstantOrArray::Array(lengths) => gather_slices_dispatch( + &starts, + lengths, + &offsets, + data, + output_len, + out_offset_ptype, + )?, + }; + + let validity = array.validity()?.take(indices_ref)?; + + // SAFETY: output offsets are built from valid input offsets, start at zero, are monotonically + // non-decreasing, and the copied data buffer has exactly the referenced byte length. + unsafe { + Ok(Some( + VarBinArray::new_unchecked(result.offsets, result.data.freeze(), dtype, validity) + .into_array(), + )) + } +} + fn take( dtype: DType, offsets: &[Offset], @@ -183,6 +244,337 @@ fn take( } } +struct GatheredPiecewiseVarBin { + offsets: ArrayRef, + data: ByteBufferMut, +} + +fn gather_slices_constant_dispatch( + starts: &PrimitiveArray, + length: usize, + offsets: &PrimitiveArray, + data: &[u8], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + gather_slices_constant_start_dispatch::( + starts, + length, + offsets, + data, + output_len, + out_offset_ptype, + ) + }) +} + +fn gather_slices_constant_start_dispatch( + starts: &PrimitiveArray, + length: usize, + offsets: &PrimitiveArray, + data: &[u8], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult +where + S: UnsignedPType, +{ + match offsets.ptype() { + PType::U8 => gather_slices_constant_length::( + offsets.as_slice::(), + data, + starts.as_slice::(), + length, + output_len, + out_offset_ptype, + ), + PType::U16 => gather_slices_constant_length::( + offsets.as_slice::(), + data, + starts.as_slice::(), + length, + output_len, + out_offset_ptype, + ), + PType::U32 => gather_slices_constant_length::( + offsets.as_slice::(), + data, + starts.as_slice::(), + length, + output_len, + out_offset_ptype, + ), + PType::U64 => gather_slices_constant_length::( + offsets.as_slice::(), + data, + starts.as_slice::(), + length, + output_len, + out_offset_ptype, + ), + _ => unreachable!("offsets were reinterpreted to an unsigned integer ptype"), + } +} + +fn gather_slices_dispatch( + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + offsets: &PrimitiveArray, + data: &[u8], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + gather_slices_start_dispatch::( + starts, + lengths, + offsets, + data, + output_len, + out_offset_ptype, + ) + }) +} + +fn gather_slices_start_dispatch( + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + offsets: &PrimitiveArray, + data: &[u8], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult +where + S: UnsignedPType, +{ + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + gather_slices_start_length_dispatch::( + starts, + lengths, + offsets, + data, + output_len, + out_offset_ptype, + ) + }) +} + +fn gather_slices_start_length_dispatch( + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + offsets: &PrimitiveArray, + data: &[u8], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, +{ + match offsets.ptype() { + PType::U8 => gather_slices::( + offsets.as_slice::(), + data, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + out_offset_ptype, + ), + PType::U16 => gather_slices::( + offsets.as_slice::(), + data, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + out_offset_ptype, + ), + PType::U32 => gather_slices::( + offsets.as_slice::(), + data, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + out_offset_ptype, + ), + PType::U64 => gather_slices::( + offsets.as_slice::(), + data, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + out_offset_ptype, + ), + _ => unreachable!("offsets were reinterpreted to an unsigned integer ptype"), + } +} + +fn gather_slices_constant_length( + offsets: &[Offset], + data: &[u8], + starts: &[S], + length: usize, + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult +where + S: UnsignedPType, + Offset: IntegerPType, + NewOffset: IntegerPType, +{ + let computed_len = starts + .len() + .checked_mul(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); + + let mut new_offsets = BufferMut::::with_capacity(output_len + 1); + new_offsets.push(NewOffset::zero()); + let mut output_bytes = 0usize; + + for &start in starts { + let start = start.as_(); + if length == 0 { + continue; + } + + let offset_range = &offsets[start..][..=length]; + let byte_start = offset_range[0].as_(); + let byte_end = offset_range[length].as_(); + vortex_ensure!( + byte_start <= byte_end && byte_end <= data.len(), + "VarBin offsets range {byte_start}..{byte_end} exceeds data length {}", + data.len() + ); + + for &offset in &offset_range[1..] { + let offset = offset.as_(); + let relative = offset.checked_sub(byte_start).ok_or_else(|| { + vortex_err!("VarBin offsets are not monotonic at offset {offset}") + })?; + let output_offset = output_bytes.checked_add(relative).ok_or_else(|| { + vortex_err!("PiecewiseSequence VarBin output byte length overflow") + })?; + new_offsets.push(new_offset_value::(output_offset)?); + } + + output_bytes = output_bytes + .checked_add(byte_end - byte_start) + .ok_or_else(|| vortex_err!("PiecewiseSequence VarBin output byte length overflow"))?; + } + + let mut new_data = ByteBufferMut::with_capacity(output_bytes); + for &start in starts { + let start = start.as_(); + if length == 0 { + continue; + } + + let offset_range = &offsets[start..][..=length]; + let byte_start = offset_range[0].as_(); + let byte_end = offset_range[length].as_(); + new_data.extend_from_slice(&data[byte_start..][..byte_end - byte_start]); + } + + let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) + .reinterpret_cast(out_offset_ptype) + .into_array(); + Ok(GatheredPiecewiseVarBin { + offsets, + data: new_data, + }) +} + +fn gather_slices( + offsets: &[Offset], + data: &[u8], + starts: &[S], + lengths: &[L], + output_len: usize, + out_offset_ptype: PType, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, + Offset: IntegerPType, + NewOffset: IntegerPType, +{ + let mut new_offsets = BufferMut::::with_capacity(output_len + 1); + new_offsets.push(NewOffset::zero()); + let mut output_bytes = 0usize; + + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = start.as_(); + let length = length.as_(); + if length == 0 { + continue; + } + + let offset_range = &offsets[start..][..=length]; + let byte_start = offset_range[0].as_(); + let byte_end = offset_range[length].as_(); + vortex_ensure!( + byte_start <= byte_end && byte_end <= data.len(), + "VarBin offsets range {byte_start}..{byte_end} exceeds data length {}", + data.len() + ); + + for &offset in &offset_range[1..] { + let offset = offset.as_(); + let relative = offset.checked_sub(byte_start).ok_or_else(|| { + vortex_err!("VarBin offsets are not monotonic at offset {offset}") + })?; + let output_offset = output_bytes.checked_add(relative).ok_or_else(|| { + vortex_err!("PiecewiseSequence VarBin output byte length overflow") + })?; + new_offsets.push(new_offset_value::(output_offset)?); + } + + output_bytes = output_bytes + .checked_add(byte_end - byte_start) + .ok_or_else(|| vortex_err!("PiecewiseSequence VarBin output byte length overflow"))?; + } + vortex_ensure!( + new_offsets.len() == output_len + 1, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + new_offsets.len() - 1 + ); + + let mut new_data = ByteBufferMut::with_capacity(output_bytes); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = start.as_(); + let length = length.as_(); + if length == 0 { + continue; + } + + let offset_range = &offsets[start..][..=length]; + let byte_start = offset_range[0].as_(); + let byte_end = offset_range[length].as_(); + new_data.extend_from_slice(&data[byte_start..][..byte_end - byte_start]); + } + + let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) + .reinterpret_cast(out_offset_ptype) + .into_array(); + Ok(GatheredPiecewiseVarBin { + offsets, + data: new_data, + }) +} + +fn new_offset_value(value: usize) -> VortexResult { + T::from(value).ok_or_else(|| { + vortex_err!( + "PiecewiseSequence VarBin offset value {value} does not fit in {}", + T::PTYPE + ) + }) +} + fn take_nullable( dtype: DType, offsets: &[Offset], diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index ce10abda3d0..40e05e08737 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -4,23 +4,32 @@ use std::iter; use std::sync::Arc; +use itertools::Itertools as _; use num_traits::AsPrimitive; use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_mask::AllOr; use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; use crate::array::ArrayView; +use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::piecewise_sequence::ConstantOrArray; +use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::varbinview::BinaryView; use crate::buffer::BufferHandle; +use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_integer_ptype; +use crate::match_each_unsigned_integer_ptype; impl TakeExecute for VarBinView { /// Take involves creating a new array that references the old array, just with the given set of views. @@ -29,6 +38,12 @@ impl TakeExecute for VarBinView { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)? + { + return Ok(Some(taken)); + } + let validity = array.validity()?.take(indices)?; let indices = indices.clone().execute::(ctx)?; @@ -57,6 +72,58 @@ impl TakeExecute for VarBinView { } } +fn take_contiguous_ranges( + array: ArrayView<'_, VarBinView>, + indices: ArrayView<'_, PiecewiseSequence>, + indices_ref: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { + return Ok(None); + }; + let source = array.views(); + let output_len = indices_ref.len(); + let views = match &lengths { + ConstantOrArray::Constant(length) => { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + gather_view_slices_constant_length( + source, + starts.as_slice::(), + *length, + output_len, + )? + }) + } + ConstantOrArray::Array(lengths) => { + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + gather_view_slices( + source, + starts.as_slice::(), + lengths.as_slice::(), + output_len, + )? + }) + }) + } + }; + let validity = array.validity()?.take(indices_ref)?; + + // SAFETY: ranges were validated against the source views, and copied views still reference the + // same backing data buffers. + unsafe { + Ok(Some( + VarBinViewArray::new_handle_unchecked( + BufferHandle::new_host(views.into_byte_buffer()), + Arc::clone(array.data_buffers()), + array.dtype().clone(), + validity, + ) + .into_array(), + )) + } +} + fn take_views>( views_ref: &[BinaryView], indices: &[I], @@ -85,6 +152,58 @@ fn take_views>( } } +fn gather_view_slices_constant_length( + source: &[BinaryView], + starts: &[S], + length: usize, + output_len: usize, +) -> VortexResult> +where + S: UnsignedPType, +{ + let computed_len = starts + .len() + .checked_mul(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); + + let mut views = BufferMut::::with_capacity(output_len); + for &start in starts { + let start = start.as_(); + views.extend_from_slice(&source[start..][..length]); + } + + Ok(views.freeze()) +} + +fn gather_view_slices( + source: &[BinaryView], + starts: &[S], + lengths: &[L], + output_len: usize, +) -> VortexResult> +where + S: UnsignedPType, + L: UnsignedPType, +{ + let mut views = BufferMut::::with_capacity(output_len); + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = start.as_(); + let length = length.as_(); + views.extend_from_slice(&source[start..][..length]); + } + + vortex_ensure!( + views.len() == output_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + views.len() + ); + Ok(views.freeze()) +} + #[cfg(test)] mod tests { use rstest::rstest; From 534b70f11b01b147107861672dd711ace1bd018c Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 15:33:07 -0400 Subject: [PATCH 33/37] Use cursor copies for PiecewiseSequence slices Signed-off-by: Daniel King --- .../src/arrays/decimal/compute/take.rs | 21 ++++++++--- .../src/arrays/piecewise_sequence/mod.rs | 36 +++++++++++++++++++ .../src/arrays/primitive/compute/take/mod.rs | 21 ++++++++--- .../src/arrays/varbin/compute/take.rs | 31 ++++++++++++++-- .../src/arrays/varbinview/compute/take.rs | 21 ++++++++--- 5 files changed, 116 insertions(+), 14 deletions(-) diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index ef34f47ffc7..1b0a6aff7d6 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -17,6 +17,7 @@ use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; use crate::arrays::piecewise_sequence::ConstantOrArray; +use crate::arrays::piecewise_sequence::copy_slice_to_spare; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::dtype::IntegerPType; use crate::dtype::NativeDecimalType; @@ -196,11 +197,19 @@ where ); let mut result = BufferMut::::with_capacity(output_len); + let mut cursor = 0usize; for &start in starts { let start = start.as_(); - result.extend_from_slice(&values[start..][..length]); + cursor = copy_slice_to_spare(&mut result, cursor, &values[start..][..length], output_len)?; } + vortex_ensure!( + cursor == output_len, + "PiecewiseSequenceArray expanded length {cursor} does not match declared length {output_len}" + ); + // SAFETY: `copy_slice_to_spare` checked every write against `output_len`, and `cursor == + // output_len` proves all slots were initialized. + unsafe { result.set_len(output_len) }; Ok(result.freeze()) } @@ -216,17 +225,21 @@ where T: NativeDecimalType, { let mut result = BufferMut::::with_capacity(output_len); + let mut cursor = 0usize; for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - result.extend_from_slice(&values[start..][..length]); + cursor = copy_slice_to_spare(&mut result, cursor, &values[start..][..length], output_len)?; } vortex_ensure!( - result.len() == output_len, + cursor == output_len, "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", - result.len() + cursor ); + // SAFETY: `copy_slice_to_spare` checked every write against `output_len`, and `cursor == + // output_len` proves all slots were initialized. + unsafe { result.set_len(output_len) }; Ok(result.freeze()) } diff --git a/vortex-array/src/arrays/piecewise_sequence/mod.rs b/vortex-array/src/arrays/piecewise_sequence/mod.rs index 18ce7debec1..eba974604bc 100644 --- a/vortex-array/src/arrays/piecewise_sequence/mod.rs +++ b/vortex-array/src/arrays/piecewise_sequence/mod.rs @@ -8,6 +8,8 @@ //! intended for take operations that can gather regular runs without materializing one index per //! element. +use std::ptr; + use itertools::Itertools; use num_traits::AsPrimitive; use vortex_buffer::BufferMut; @@ -102,6 +104,40 @@ pub(crate) fn is_constant_multiplier_one(multipliers: &ArrayRef) -> bool { ) } +pub(crate) fn copy_slice_to_spare( + buffer: &mut BufferMut, + cursor: usize, + source: &[T], + output_len: usize, +) -> VortexResult { + let end = cursor + .checked_add(source.len()) + .ok_or_else(|| vortex_err!("slice copy output length overflows usize"))?; + vortex_ensure!( + end <= output_len, + "slice copy length {end} exceeds declared output length {output_len}" + ); + vortex_ensure!( + buffer.is_empty(), + "slice copy buffer already has {} initialized values", + buffer.len() + ); + vortex_ensure!( + output_len <= buffer.capacity(), + "slice copy output length {output_len} exceeds buffer capacity {}", + buffer.capacity() + ); + + // SAFETY: the checks above prove `cursor..end` is inside the spare capacity of an + // uninitialized buffer allocated for at least `output_len` values. + let dst = unsafe { buffer.spare_capacity_mut().get_unchecked_mut(cursor..end) }; + // SAFETY: `dst` has exactly `source.len()` writable slots and does not overlap with `source`. + unsafe { + ptr::copy_nonoverlapping(source.as_ptr(), dst.as_mut_ptr().cast(), source.len()); + } + Ok(end) +} + fn constant_unsigned_usize(array: &ArrayRef) -> VortexResult> { let Some(scalar) = array.as_constant() else { return Ok(None); diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index a301cf99b29..43a079f0a3c 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -24,6 +24,7 @@ use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; use crate::arrays::piecewise_sequence::ConstantOrArray; +use crate::arrays::piecewise_sequence::copy_slice_to_spare; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; @@ -242,17 +243,21 @@ where L: UnsignedPType, { let mut values = BufferMut::::with_capacity(output_len); + let mut cursor = 0usize; for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - values.extend_from_slice(&source[start..][..length]); + cursor = copy_slice_to_spare(&mut values, cursor, &source[start..][..length], output_len)?; } vortex_ensure!( - values.len() == output_len, + cursor == output_len, "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", - values.len() + cursor ); + // SAFETY: `copy_slice_to_spare` checked every write against `output_len`, and `cursor == + // output_len` proves all slots were initialized. + unsafe { values.set_len(output_len) }; Ok(values.freeze()) } @@ -309,11 +314,19 @@ where ); let mut values = BufferMut::::with_capacity(output_len); + let mut cursor = 0usize; for &start in starts { let start = start.as_(); - values.extend_from_slice(&source[start..][..length]); + cursor = copy_slice_to_spare(&mut values, cursor, &source[start..][..length], output_len)?; } + vortex_ensure!( + cursor == output_len, + "PiecewiseSequenceArray expanded length {cursor} does not match declared length {output_len}" + ); + // SAFETY: `copy_slice_to_spare` checked every write against `output_len`, and `cursor == + // output_len` proves all slots were initialized. + unsafe { values.set_len(output_len) }; Ok(values.freeze()) } diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index b956359ef0b..7a81f05a4d3 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -21,6 +21,7 @@ use crate::arrays::VarBin; use crate::arrays::VarBinArray; use crate::arrays::dict::TakeExecute; use crate::arrays::piecewise_sequence::ConstantOrArray; +use crate::arrays::piecewise_sequence::copy_slice_to_spare; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; use crate::arrays::varbin::VarBinArrayExt; @@ -467,6 +468,7 @@ where } let mut new_data = ByteBufferMut::with_capacity(output_bytes); + let mut cursor = 0usize; for &start in starts { let start = start.as_(); if length == 0 { @@ -476,8 +478,20 @@ where let offset_range = &offsets[start..][..=length]; let byte_start = offset_range[0].as_(); let byte_end = offset_range[length].as_(); - new_data.extend_from_slice(&data[byte_start..][..byte_end - byte_start]); + cursor = copy_slice_to_spare( + &mut new_data, + cursor, + &data[byte_start..][..byte_end - byte_start], + output_bytes, + )?; } + vortex_ensure!( + cursor == output_bytes, + "VarBin byte copy length {cursor} does not match computed byte length {output_bytes}" + ); + // SAFETY: `copy_slice_to_spare` checked every write against `output_bytes`, and `cursor == + // output_bytes` proves all bytes were initialized. + unsafe { new_data.set_len(output_bytes) }; let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) .reinterpret_cast(out_offset_ptype) @@ -544,6 +558,7 @@ where ); let mut new_data = ByteBufferMut::with_capacity(output_bytes); + let mut cursor = 0usize; for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); @@ -554,8 +569,20 @@ where let offset_range = &offsets[start..][..=length]; let byte_start = offset_range[0].as_(); let byte_end = offset_range[length].as_(); - new_data.extend_from_slice(&data[byte_start..][..byte_end - byte_start]); + cursor = copy_slice_to_spare( + &mut new_data, + cursor, + &data[byte_start..][..byte_end - byte_start], + output_bytes, + )?; } + vortex_ensure!( + cursor == output_bytes, + "VarBin byte copy length {cursor} does not match computed byte length {output_bytes}" + ); + // SAFETY: `copy_slice_to_spare` checked every write against `output_bytes`, and `cursor == + // output_bytes` proves all bytes were initialized. + unsafe { new_data.set_len(output_bytes) }; let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable) .reinterpret_cast(out_offset_ptype) diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index 40e05e08737..d6256b33078 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -23,6 +23,7 @@ use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; use crate::arrays::dict::TakeExecute; use crate::arrays::piecewise_sequence::ConstantOrArray; +use crate::arrays::piecewise_sequence::copy_slice_to_spare; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::varbinview::BinaryView; use crate::buffer::BufferHandle; @@ -171,11 +172,19 @@ where ); let mut views = BufferMut::::with_capacity(output_len); + let mut cursor = 0usize; for &start in starts { let start = start.as_(); - views.extend_from_slice(&source[start..][..length]); + cursor = copy_slice_to_spare(&mut views, cursor, &source[start..][..length], output_len)?; } + vortex_ensure!( + cursor == output_len, + "PiecewiseSequenceArray expanded length {cursor} does not match declared length {output_len}" + ); + // SAFETY: `copy_slice_to_spare` checked every write against `output_len`, and `cursor == + // output_len` proves all slots were initialized. + unsafe { views.set_len(output_len) }; Ok(views.freeze()) } @@ -190,17 +199,21 @@ where L: UnsignedPType, { let mut views = BufferMut::::with_capacity(output_len); + let mut cursor = 0usize; for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); let length = length.as_(); - views.extend_from_slice(&source[start..][..length]); + cursor = copy_slice_to_spare(&mut views, cursor, &source[start..][..length], output_len)?; } vortex_ensure!( - views.len() == output_len, + cursor == output_len, "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", - views.len() + cursor ); + // SAFETY: `copy_slice_to_spare` checked every write against `output_len`, and `cursor == + // output_len` proves all slots were initialized. + unsafe { views.set_len(output_len) }; Ok(views.freeze()) } From 9c1de8d1844bf918f180fa1d03ef79c218c5d5f3 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 15:59:02 -0400 Subject: [PATCH 34/37] Use PiecewiseSequence for List take elements Signed-off-by: Daniel King --- vortex-array/src/arrays/list/compute/take.rs | 184 +++++++------------ 1 file changed, 65 insertions(+), 119 deletions(-) diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index ff9549d6a34..b78b5847630 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -23,10 +23,7 @@ use crate::arrays::list::ListArrayExt; use crate::arrays::piecewise_sequence::ConstantOrArray; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; -use crate::builders::ArrayBuilder; -use crate::builders::PrimitiveBuilder; use crate::dtype::IntegerPType; -use crate::dtype::Nullability; use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_unsigned_integer_ptype; @@ -64,7 +61,7 @@ impl TakeExecute for List { match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { match_each_unsigned_integer_ptype!(indices.ptype(), |I| { match_smallest_offset_type!(total_approx, |OutputOffsetType| { - _take::( + take_with_piecewise_elements::( array, offsets.as_view(), indices.as_view(), @@ -77,7 +74,11 @@ impl TakeExecute for List { } } -fn _take( +fn take_with_piecewise_elements< + I: IntegerPType, + O: IntegerPType, + OutputOffsetType: IntegerPType, +>( array: ArrayView<'_, List>, offsets_array: ArrayView<'_, Primitive>, indices_array: ArrayView<'_, Primitive>, @@ -91,56 +92,78 @@ fn _take( .vortex_expect("Failed to compute validity mask") .execute_mask(indices_array.as_ref().len(), ctx)?; - if !indices_validity.all_true() || !data_validity.all_true() { - return _take_nullable::(array, offsets_array, indices_array, ctx); - } - let offsets: &[O] = offsets_array.as_slice(); let indices: &[I] = indices_array.as_slice(); - let mut new_offsets = PrimitiveBuilder::::with_capacity( - Nullability::NonNullable, - indices.len(), - ); - let mut elements_to_take = - PrimitiveBuilder::with_capacity(Nullability::NonNullable, 2 * indices.len()); + let offsets_capacity = indices + .len() + .checked_add(1) + .ok_or_else(|| vortex_err!("List take offsets length overflow"))?; + let mut new_offsets = BufferMut::::with_capacity(offsets_capacity); + let mut element_starts = BufferMut::::with_capacity(indices.len()); + let mut element_lengths = BufferMut::::with_capacity(indices.len()); + + let mut current_offset = 0usize; + new_offsets.push(OutputOffsetType::zero()); - let mut current_offset = OutputOffsetType::zero(); - new_offsets.append_zero(); + for (&data_idx, index_valid) in indices.iter().zip_eq(indices_validity.iter()) { + if !index_valid { + new_offsets.push(new_offset_value::(current_offset)?); + element_starts.push(0); + element_lengths.push(0); + continue; + } - for &data_idx in indices { let data_idx: usize = data_idx.as_(); + if !data_validity.value(data_idx) { + new_offsets.push(new_offset_value::(current_offset)?); + element_starts.push(0); + element_lengths.push(0); + continue; + } + let start = offsets[data_idx]; let stop = offsets[data_idx + 1]; + let start: usize = start.as_(); + let stop: usize = stop.as_(); + let length = stop + .checked_sub(start) + .ok_or_else(|| vortex_err!("List offsets are not monotonic at offset {stop}"))?; - // Annoyingly, we can't turn (start..end) into a range, so we're doing that manually. - // - // We could convert start and end to usize, but that would impose a potentially - // harder constraint - now we don't care if they fit into usize as long as their - // difference does. - let additional: usize = (stop - start).as_(); - - // TODO(0ax1): optimize this - elements_to_take.reserve_exact(additional); - for i in 0..additional { - elements_to_take.append_value(start + O::from_usize(i).vortex_expect("i < additional")); - } - current_offset += - OutputOffsetType::from_usize((stop - start).as_()).vortex_expect("offset conversion"); - new_offsets.append_value(current_offset); + current_offset = current_offset + .checked_add(length) + .ok_or_else(|| vortex_err!("List take output elements length overflow"))?; + new_offsets.push(new_offset_value::(current_offset)?); + element_starts.push(start as u64); + element_lengths.push(length as u64); } - let elements_to_take = elements_to_take.finish(); - let new_offsets = new_offsets.finish(); - - let new_elements = array.elements().take(elements_to_take)?; + let new_offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable).into_array(); + let multipliers = ConstantArray::new(1u64, element_starts.len()).into_array(); - Ok(ListArray::try_new( - new_elements, - new_offsets, - array.validity()?.take(indices_array.array())?, - )? + // SAFETY: valid source rows contribute ranges derived from list offsets; null index/source + // rows contribute zero-length placeholder ranges. `current_offset` is the sum of all generated + // element lengths, and multiplier 1 preserves contiguous ranges. + let element_indices = unsafe { + PiecewiseSequenceArray::new_unchecked( + element_starts.into_array(), + element_lengths.into_array(), + multipliers, + current_offset, + ) + }; + let new_elements = array.elements().take(element_indices.into_array())?; + + // SAFETY: offsets are rebuilt from the gathered element ranges and have one entry per output + // row plus the leading zero; validity is produced by the usual take-validity path. + Ok(unsafe { + ListArray::new_unchecked( + new_elements, + new_offsets, + array.validity()?.take(indices_array.array())?, + ) + } .into_array()) } @@ -603,83 +626,6 @@ fn new_offset_value(value: usize) -> VortexResult { }) } -// Kept out-of-line: as a single-callsite generic helper it would otherwise be inlined into every -// monomorphization of `_take`, duplicating the entire nullable path across all specializations. -#[inline(never)] -fn _take_nullable( - array: ArrayView<'_, List>, - offsets_array: ArrayView<'_, Primitive>, - indices_array: ArrayView<'_, Primitive>, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let offsets: &[O] = offsets_array.as_slice(); - let indices: &[I] = indices_array.as_slice(); - let data_validity = array - .list_validity() - .execute_mask(array.as_ref().len(), ctx)?; - let indices_validity = indices_array - .validity() - .vortex_expect("Failed to compute validity mask") - .execute_mask(indices_array.as_ref().len(), ctx)?; - - let mut new_offsets = PrimitiveBuilder::::with_capacity( - Nullability::NonNullable, - indices.len(), - ); - - // This will be the indices we push down to the child array to call `take` with. - // - // There are 2 things to note here: - // - We do not know how many elements we need to take from our child since lists are variable - // size: thus we arbitrarily choose a capacity of `2 * # of indices`. - // - The type of the primitive builder needs to fit the largest offset of the (parent) - // `ListArray`, so we make this `PrimitiveBuilder` generic over `O` (instead of `I`). - let mut elements_to_take = - PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 2 * indices.len()); - - let mut current_offset = OutputOffsetType::zero(); - new_offsets.append_zero(); - - for (data_idx, index_valid) in indices.iter().zip(indices_validity.iter()) { - if !index_valid { - new_offsets.append_value(current_offset); - continue; - } - - let data_idx: usize = data_idx.as_(); - - if !data_validity.value(data_idx) { - new_offsets.append_value(current_offset); - continue; - } - - let start = offsets[data_idx]; - let stop = offsets[data_idx + 1]; - - // See the note in `_take` on the reasoning. - let additional: usize = (stop - start).as_(); - - elements_to_take.reserve_exact(additional); - for i in 0..additional { - elements_to_take.append_value(start + O::from_usize(i).vortex_expect("i < additional")); - } - current_offset += - OutputOffsetType::from_usize((stop - start).as_()).vortex_expect("offset conversion"); - new_offsets.append_value(current_offset); - } - - let elements_to_take = elements_to_take.finish(); - let new_offsets = new_offsets.finish(); - let new_elements = array.elements().take(elements_to_take)?; - - Ok(ListArray::try_new( - new_elements, - new_offsets, - array.validity()?.take(indices_array.array())?, - )? - .into_array()) -} - #[cfg(test)] mod test { use std::sync::Arc; From ff27447e611e1ffc82d116868d60a56262f3d778 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 16:07:41 -0400 Subject: [PATCH 35/37] Cover null List take placeholders Signed-off-by: Daniel King --- vortex-array/src/arrays/list/compute/take.rs | 50 ++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index b78b5847630..6260f86a681 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -642,6 +642,7 @@ mod test { use crate::arrays::ListViewArray; use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; + use crate::assert_arrays_eq; use crate::compute::conformance::take::test_take_conformance; use crate::dtype::DType; use crate::dtype::Nullability; @@ -730,6 +731,55 @@ mod test { ); } + #[test] + fn null_index_ignores_out_of_bounds_payload() { + let mut ctx = array_session().create_execution_ctx(); + let list = ListArray::try_new( + buffer![1i32, 2, 3, 4].into_array(), + buffer![0u32, 2, 4].into_array(), + Validity::NonNullable, + ) + .unwrap() + .into_array(); + + let idx = PrimitiveArray::new( + buffer![1u32, 99, 0], + Validity::from_iter([true, false, true]), + ) + .into_array(); + let result = list.take(idx).unwrap(); + + let expected = ListArray::new( + buffer![3i32, 4, 1, 2].into_array(), + buffer![0u32, 2, 2, 4].into_array(), + Validity::from_iter([true, false, true]), + ); + assert_arrays_eq!(expected, result, &mut ctx); + } + + #[test] + fn null_source_row_ignores_invalid_offset_payload() { + let mut ctx = array_session().create_execution_ctx(); + let list = unsafe { + ListArray::new_unchecked( + buffer![1i32, 2].into_array(), + buffer![0u32, 2, 999].into_array(), + Validity::from_iter([true, false]), + ) + } + .into_array(); + + let idx = buffer![0u32, 1].into_array(); + let result = list.take(idx).unwrap(); + + let expected = ListArray::new( + buffer![1i32, 2].into_array(), + buffer![0u32, 2, 2].into_array(), + Validity::from_iter([true, false]), + ); + assert_arrays_eq!(expected, result, &mut ctx); + } + #[test] fn change_validity() { let list = ListArray::try_new( From e781b619830a21ba73d698fe043024024bbef91e Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 16:15:24 -0400 Subject: [PATCH 36/37] Normalize List take indices before range rebuild Signed-off-by: Daniel King --- vortex-array/src/arrays/list/compute/take.rs | 46 ++++++++------------ 1 file changed, 17 insertions(+), 29 deletions(-) diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index 6260f86a681..aa583fe886e 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -3,10 +3,10 @@ use itertools::Itertools as _; use vortex_buffer::BufferMut; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; +use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; @@ -23,11 +23,13 @@ use crate::arrays::list::ListArrayExt; use crate::arrays::piecewise_sequence::ConstantOrArray; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; +use crate::builtins::ArrayBuiltins; use crate::dtype::IntegerPType; use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_unsigned_integer_ptype; use crate::match_smallest_offset_type; +use crate::scalar::Scalar; use crate::validity::Validity; // TODO(connor)[ListView]: Re-revert to the version where we simply convert to a `ListView` and call @@ -51,10 +53,16 @@ impl TakeExecute for List { return Ok(Some(taken)); } - let indices = indices.clone().execute::(ctx)?; + let new_validity = array.validity()?.take(indices)?; + let normalized_indices = indices + .clone() + .mask(new_validity.to_array(indices.len()))? + .fill_null(Scalar::from(0).cast(indices.dtype())?)?; + let indices = normalized_indices.execute::(ctx)?; let indices = indices.reinterpret_cast(indices.ptype().to_unsigned()); let offsets = array.offsets().clone().execute::(ctx)?; let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); + let validity_mask = new_validity.execute_mask(indices.len(), ctx)?; // This is an over-approximation of the total number of elements in the resulting array. let total_approx = array.elements().len().saturating_mul(indices.len()); @@ -65,7 +73,8 @@ impl TakeExecute for List { array, offsets.as_view(), indices.as_view(), - ctx, + new_validity, + &validity_mask, ) .map(Some) }) @@ -82,16 +91,9 @@ fn take_with_piecewise_elements< array: ArrayView<'_, List>, offsets_array: ArrayView<'_, Primitive>, indices_array: ArrayView<'_, Primitive>, - ctx: &mut ExecutionCtx, + new_validity: Validity, + validity_mask: &Mask, ) -> VortexResult { - let data_validity = array - .list_validity() - .execute_mask(array.as_ref().len(), ctx)?; - let indices_validity = indices_array - .validity() - .vortex_expect("Failed to compute validity mask") - .execute_mask(indices_array.as_ref().len(), ctx)?; - let offsets: &[O] = offsets_array.as_slice(); let indices: &[I] = indices_array.as_slice(); @@ -106,8 +108,8 @@ fn take_with_piecewise_elements< let mut current_offset = 0usize; new_offsets.push(OutputOffsetType::zero()); - for (&data_idx, index_valid) in indices.iter().zip_eq(indices_validity.iter()) { - if !index_valid { + for (&data_idx, is_valid) in indices.iter().zip_eq(validity_mask.iter()) { + if !is_valid { new_offsets.push(new_offset_value::(current_offset)?); element_starts.push(0); element_lengths.push(0); @@ -116,13 +118,6 @@ fn take_with_piecewise_elements< let data_idx: usize = data_idx.as_(); - if !data_validity.value(data_idx) { - new_offsets.push(new_offset_value::(current_offset)?); - element_starts.push(0); - element_lengths.push(0); - continue; - } - let start = offsets[data_idx]; let stop = offsets[data_idx + 1]; let start: usize = start.as_(); @@ -157,14 +152,7 @@ fn take_with_piecewise_elements< // SAFETY: offsets are rebuilt from the gathered element ranges and have one entry per output // row plus the leading zero; validity is produced by the usual take-validity path. - Ok(unsafe { - ListArray::new_unchecked( - new_elements, - new_offsets, - array.validity()?.take(indices_array.array())?, - ) - } - .into_array()) + Ok(unsafe { ListArray::new_unchecked(new_elements, new_offsets, new_validity) }.into_array()) } fn take_piecewise_sequence( From 243ed198b6bfa4e54eae5ce0ab209ee89bb2d976 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 17 Jul 2026 16:48:36 -0400 Subject: [PATCH 37/37] Avoid normalizing unused null list indices Signed-off-by: Daniel King --- vortex-array/src/arrays/list/compute/take.rs | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index aa583fe886e..56b6ce22511 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -23,13 +23,11 @@ use crate::arrays::list::ListArrayExt; use crate::arrays::piecewise_sequence::ConstantOrArray; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::arrays::primitive::PrimitiveArrayExt; -use crate::builtins::ArrayBuiltins; use crate::dtype::IntegerPType; use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_unsigned_integer_ptype; use crate::match_smallest_offset_type; -use crate::scalar::Scalar; use crate::validity::Validity; // TODO(connor)[ListView]: Re-revert to the version where we simply convert to a `ListView` and call @@ -54,11 +52,7 @@ impl TakeExecute for List { } let new_validity = array.validity()?.take(indices)?; - let normalized_indices = indices - .clone() - .mask(new_validity.to_array(indices.len()))? - .fill_null(Scalar::from(0).cast(indices.dtype())?)?; - let indices = normalized_indices.execute::(ctx)?; + let indices = indices.clone().execute::(ctx)?; let indices = indices.reinterpret_cast(indices.ptype().to_unsigned()); let offsets = array.offsets().clone().execute::(ctx)?; let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned());