diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index 15b6b5aa6be..001e20dee65 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_sequence; +pub use piecewise_sequence::PiecewiseSequence; +pub use piecewise_sequence::PiecewiseSequenceArray; + pub mod primitive; pub use primitive::Primitive; pub use primitive::PrimitiveArray; diff --git a/vortex-array/src/arrays/piecewise_sequence/array.rs b/vortex-array/src/arrays/piecewise_sequence/array.rs new file mode 100644 index 00000000000..1559cf7d3ad --- /dev/null +++ b/vortex-array/src/arrays/piecewise_sequence/array.rs @@ -0,0 +1,70 @@ +// 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::PiecewiseSequence; +use crate::dtype::PType; + +#[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, + /// The distance between consecutive indices in each piece. + pub multipliers: ArrayRef, +} + +/// Extension methods for [`Array`] values using the [`PiecewiseSequence`] encoding. +pub trait PiecewiseSequenceArrayExt: + TypedArrayRef + PiecewiseSequenceArraySlotsExt +{ +} +impl> PiecewiseSequenceArrayExt for T {} + +impl Array { + /// Constructs a new `PiecewiseSequenceArray` from start, length, and multiplier arrays. + /// + /// 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, + multipliers: ArrayRef, + len: usize, + ) -> VortexResult { + Array::try_from_parts( + ArrayParts::new(PiecewiseSequence, PType::U64.into(), len, EmptyArrayData) + .with_slots(smallvec![Some(starts), Some(lengths), Some(multipliers)]), + ) + } + + /// Constructs a new `PiecewiseSequenceArray` without validation. + /// + /// # Safety + /// + /// The caller must guarantee the same structural invariants as [`Self::try_new`]. + 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), Some(multipliers)]), + ) + } + } +} diff --git a/vortex-array/src/arrays/piecewise_sequence/mod.rs b/vortex-array/src/arrays/piecewise_sequence/mod.rs new file mode 100644 index 00000000000..b914c5e9651 --- /dev/null +++ b/vortex-array/src/arrays/piecewise_sequence/mod.rs @@ -0,0 +1,125 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Index encoding for concatenated sequential ranges. +//! +//! A `PiecewiseSequenceArray` represents the expanded index sequence +//! `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 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; + +#[cfg(test)] +mod tests; + +pub use array::PiecewiseSequenceArrayExt; +pub use vtable::*; + +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(), + "PiecewiseSequenceArray {name} must have unsigned integer dtype, got {}", + array.dtype() + ); + vortex_ensure!( + !array.dtype().is_nullable(), + "PiecewiseSequenceArray {name} must be non-nullable, got {}", + array.dtype() + ); + Ok(()) +} + +pub(crate) fn materialize_ranges( + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + multipliers: &PrimitiveArray, + output_len: usize, +) -> VortexResult> +where + 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), &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 * multiplier) as u64)); + } + + if computed_len != output_len { + vortex_bail!( + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); + } + Ok(values) +} diff --git a/vortex-array/src/arrays/piecewise_sequence/tests.rs b/vortex-array/src/arrays/piecewise_sequence/tests.rs new file mode 100644 index 00000000000..ab201943e3d --- /dev/null +++ b/vortex-array/src/arrays/piecewise_sequence/tests.rs @@ -0,0 +1,154 @@ +// 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::PiecewiseSequence; +use crate::arrays::PiecewiseSequenceArray; +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 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()); + 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 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 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()); + 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 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()); + 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 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 usize"), + "{err}" + ); + Ok(()) +} + +#[test] +fn execution_checks_declared_length() -> VortexResult<()> { + let starts = buffer![0u64, 3].into_array(); + let lengths = buffer![2u64, 2].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()) + .unwrap_err(); + assert!( + 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 = PiecewiseSequenceArray::try_new( + buffer![3u32, 15, 21].into_array(), + buffer![2u16, 0, 2].into_array(), + buffer![2u8, 1, 3].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, 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 new file mode 100644 index 00000000000..ac2eba61116 --- /dev/null +++ b/vortex-array/src/arrays/piecewise_sequence/vtable.rs @@ -0,0 +1,275 @@ +// 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_sequence::array::PiecewiseSequenceArraySlotsExt; +use crate::arrays::piecewise_sequence::array::PiecewiseSequenceSlots; +use crate::arrays::piecewise_sequence::check_index_arrays; +use crate::arrays::piecewise_sequence::execute_index_arrays; +use crate::arrays::piecewise_sequence::materialize_ranges; +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; +use crate::validity::Validity; + +/// A [`PiecewiseSequence`]-encoded Vortex index array. +pub type PiecewiseSequenceArray = Array; + +#[derive(Clone, prost::Message)] +struct PiecewiseSequenceMetadata { + #[prost(uint64, tag = "1")] + num_pieces: u64, + #[prost(enumeration = "PType", tag = "2")] + starts_ptype: i32, + #[prost(enumeration = "PType", tag = "3")] + lengths_ptype: i32, + #[prost(enumeration = "PType", tag = "4")] + multipliers_ptype: i32, +} + +#[derive(Clone, Debug)] +pub struct PiecewiseSequence; + +impl VTable for PiecewiseSequence { + type TypedArrayData = EmptyArrayData; + type OperationsVTable = Self; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.piecewise-sequence"); + *ID + } + + fn validate( + &self, + _data: &Self::TypedArrayData, + dtype: &DType, + _len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure!( + dtype == &DType::from(PType::U64), + "PiecewiseSequenceArray dtype must be u64, got {dtype}" + ); + vortex_ensure!( + slots.len() == PiecewiseSequenceSlots::NAMES.len(), + "PiecewiseSequenceArray requires {} slots, got {}", + PiecewiseSequenceSlots::NAMES.len(), + slots.len() + ); + let starts = slots[PiecewiseSequenceSlots::STARTS] + .as_ref() + .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_err!("PiecewiseSequenceArray multipliers slot must be present") + })?; + check_index_arrays(starts, lengths, multipliers) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, _idx: usize) -> BufferHandle { + vortex_panic!("PiecewiseSequenceArray 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 { + PiecewiseSequenceSlots::NAMES[idx].to_string() + } + + fn serialize( + array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(Some( + PiecewiseSequenceMetadata { + num_pieces: u64::try_from(array.starts().len()).map_err(|_| { + vortex_err!( + "PiecewiseSequenceArray 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, + multipliers_ptype: PType::try_from(array.multipliers().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 = PiecewiseSequenceMetadata::decode(metadata)?; + vortex_ensure!( + dtype == &DType::from(PType::U64), + "PiecewiseSequenceArray dtype must be u64, got {dtype}" + ); + vortex_ensure!( + buffers.is_empty(), + "PiecewiseSequenceArray expects no buffers, got {}", + buffers.len() + ); + vortex_ensure!( + 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!( + "PiecewiseSequenceArray piece count {} does not fit in usize", + metadata.num_pieces + ) + })?; + let starts = children.get( + PiecewiseSequenceSlots::STARTS, + &metadata.starts_ptype().into(), + num_pieces, + )?; + let lengths = children.get( + PiecewiseSequenceSlots::LENGTHS, + &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), Some(multipliers)]), + ) + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + 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| { + match_each_unsigned_integer_ptype!(multipliers.ptype(), |M| { + materialize_ranges::(&starts, &lengths, &multipliers, array.len())? + }) + }) + }); + Ok(ExecutionResult::done( + PrimitiveArray::new(values.freeze(), Validity::NonNullable).into_array(), + )) + } +} + +impl OperationsVTable for PiecewiseSequence { + fn scalar_at( + array: ArrayView<'_, PiecewiseSequence>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + 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| { + match_each_unsigned_integer_ptype!(multipliers.ptype(), |M| { + scalar_at::(&starts, &lengths, &multipliers, index)? + }) + }) + }); + Ok(value.into()) + } +} + +impl ValidityVTable for PiecewiseSequence { + fn validity(_array: ArrayView<'_, PiecewiseSequence>) -> VortexResult { + Ok(Validity::NonNullable) + } +} + +fn scalar_at( + starts: &PrimitiveArray, + lengths: &PrimitiveArray, + multipliers: &PrimitiveArray, + index: usize, +) -> VortexResult +where + S: UnsignedPType, + L: UnsignedPType, + M: UnsignedPType, +{ + let mut remaining = index; + for ((&start, &length), &multiplier) in starts + .as_slice::() + .iter() + .zip_eq(lengths.as_slice::()) + .zip_eq(multipliers.as_slice::()) + { + let length: usize = length.as_(); + if remaining < length { + 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; + } + 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 2cb8414cbfd..98a769ceea9 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::PiecewiseSequence; 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(PiecewiseSequence); this.register(VarBin); this