vortex_array/arrays/piecewise_sequence/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use smallvec::smallvec;
5use vortex_error::VortexResult;
6
7use crate::ArrayRef;
8use crate::array::Array;
9use crate::array::ArrayParts;
10use crate::array::EmptyArrayData;
11use crate::array::TypedArrayRef;
12use crate::array_slots;
13use crate::arrays::PiecewiseSequence;
14use crate::dtype::PType;
15
16#[array_slots(PiecewiseSequence)]
17pub struct PiecewiseSequenceSlots {
18    /// The inclusive start index of each sequential piece.
19    #[slot(0)]
20    pub starts: ArrayRef,
21    /// The length of each sequential piece.
22    #[slot(1)]
23    pub lengths: ArrayRef,
24    /// The distance between consecutive indices in each piece.
25    #[slot(2)]
26    pub multipliers: ArrayRef,
27}
28
29/// Extension methods for [`Array`] values using the [`PiecewiseSequence`] encoding.
30pub trait PiecewiseSequenceArrayExt:
31    TypedArrayRef<PiecewiseSequence> + PiecewiseSequenceArraySlotsExt
32{
33}
34impl<T: TypedArrayRef<PiecewiseSequence>> PiecewiseSequenceArrayExt for T {}
35
36impl Array<PiecewiseSequence> {
37    /// Constructs a new `PiecewiseSequenceArray` from start, length, and multiplier arrays.
38    ///
39    /// This validates only structural invariants: all children must be non-nullable unsigned
40    /// integer arrays with matching lengths, and the outer array length is the declared expanded
41    /// index length. Individual ranges are checked when the index array is executed or consumed by
42    /// a take implementation.
43    pub fn try_new(
44        starts: ArrayRef,
45        lengths: ArrayRef,
46        multipliers: ArrayRef,
47        len: usize,
48    ) -> VortexResult<Self> {
49        Array::try_from_parts(
50            ArrayParts::new(PiecewiseSequence, PType::U64.into(), len, EmptyArrayData)
51                .with_slots(smallvec![Some(starts), Some(lengths), Some(multipliers)]),
52        )
53    }
54
55    /// Constructs a new `PiecewiseSequenceArray` without validation.
56    ///
57    /// # Safety
58    ///
59    /// The caller must guarantee the same structural invariants as [`Self::try_new`].
60    pub unsafe fn new_unchecked(
61        starts: ArrayRef,
62        lengths: ArrayRef,
63        multipliers: ArrayRef,
64        len: usize,
65    ) -> Self {
66        unsafe {
67            Array::from_parts_unchecked(
68                ArrayParts::new(PiecewiseSequence, PType::U64.into(), len, EmptyArrayData)
69                    .with_slots(smallvec![Some(starts), Some(lengths), Some(multipliers)]),
70            )
71        }
72    }
73}
vortex_array/arrays/piecewise_sequence/array.rs

vortex_array/arrays/piecewise_sequence/
array.rs
