Add PiecewiseSequence index array

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;

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<PiecewiseSequence> + PiecewiseSequenceArraySlotsExt
+{
+}
+impl<T: TypedArrayRef<PiecewiseSequence>> PiecewiseSequenceArrayExt for T {}
+
+impl Array<PiecewiseSequence> {
+    /// 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<Self> {
+        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)]),
+            )
+        }
+    }
+}

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::<PrimitiveArray>(ctx)?;
+    let lengths = array.lengths().clone().execute::<PrimitiveArray>(ctx)?;
+    let multipliers = array.multipliers().clone().execute::<PrimitiveArray>(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<S, L, M>(
+    starts: &PrimitiveArray,
+    lengths: &PrimitiveArray,
+    multipliers: &PrimitiveArray,
+    output_len: usize,
+) -> VortexResult<BufferMut<u64>>
+where
+    S: UnsignedPType,
+    L: UnsignedPType,
+    M: UnsignedPType,
+{
+    let starts = starts.as_slice::<S>();
+    let lengths = lengths.as_slice::<L>();
+    let multipliers = multipliers.as_slice::<M>();
+    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)
+}

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::<PrimitiveArray>(&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::<PrimitiveArray>(&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::<PiecewiseSequence>());
+    assert_arrays_eq!(
+        decoded,
+        PrimitiveArray::from_iter([3u64, 5, 21, 24]).into_array(),
+        &mut array_session().create_execution_ctx()
+    );
+    Ok(())
+}