Add PiecewiseSequential index array

Signed-off-by: Daniel King <dan@spiraldb.com>

Author: danking

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;

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<PiecewiseSequential> + PiecewiseSequentialArraySlotsExt
+{
+}
+impl<T: TypedArrayRef<PiecewiseSequential>> PiecewiseSequentialArrayExt for T {}
+
+impl Array<PiecewiseSequential> {
+    /// 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<Self> {
+        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)]),
+            )
+        }
+    }
+}

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<T: UnsignedPType>(value: T) -> usize {
+    value.as_()
+}
+
+#[inline]
+pub(crate) fn index_value_to_u64<T: UnsignedPType + AsPrimitive<u64>>(value: T) -> u64 {
+    value.as_()
+}
+
+#[inline]
+pub(crate) fn checked_range_end(start: u64, length: usize) -> VortexResult<u64> {
+    start
+        .checked_add(length as u64)
+        .ok_or_else(|| vortex_error::vortex_err!("PiecewiseSequentialArray range overflows u64"))
+}
+
+pub(crate) fn materialize_ranges<S, L>(
+    starts: &PrimitiveArray,
+    lengths: &PrimitiveArray,
+    output_len: usize,
+) -> VortexResult<BufferMut<u64>>
+where
+    S: UnsignedPType + AsPrimitive<u64>,
+    L: UnsignedPType,
+{
+    let starts = starts.as_slice::<S>();
+    let lengths = lengths.as_slice::<L>();
+    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)
+}

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