From 98e82f4c69c74aba9c39318d788fad71d39cd9de Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 16 Jul 2026 13:29:50 -0400 Subject: [PATCH 1/6] 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 2/6] 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 3/6] 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 4/6] 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 5/6] 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 6/6] 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 {