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    pub starts: ArrayRef,
20    /// The length of each sequential piece.
21    pub lengths: ArrayRef,
22    /// The distance between consecutive indices in each piece.
23    pub multipliers: ArrayRef,
24}
25
26/// Extension methods for [`Array`] values using the [`PiecewiseSequence`] encoding.
27pub trait PiecewiseSequenceArrayExt:
28    TypedArrayRef<PiecewiseSequence> + PiecewiseSequenceArraySlotsExt
29{
30}
31impl<T: TypedArrayRef<PiecewiseSequence>> PiecewiseSequenceArrayExt for T {}
32
33impl Array<PiecewiseSequence> {
34    /// Constructs a new `PiecewiseSequenceArray` from start, length, and multiplier arrays.
35    ///
36    /// This validates only structural invariants: all children must be non-nullable unsigned
37    /// integer arrays with matching lengths, and the outer array length is the declared expanded
38    /// index length. Individual ranges are checked when the index array is executed or consumed by
39    /// a take implementation.
40    pub fn try_new(
41        starts: ArrayRef,
42        lengths: ArrayRef,
43        multipliers: ArrayRef,
44        len: usize,
45    ) -> VortexResult<Self> {
46        Array::try_from_parts(
47            ArrayParts::new(PiecewiseSequence, PType::U64.into(), len, EmptyArrayData)
48                .with_slots(smallvec![Some(starts), Some(lengths), Some(multipliers)]),
49        )
50    }
51
52    /// Constructs a new `PiecewiseSequenceArray` without validation.
53    ///
54    /// # Safety
55    ///
56    /// The caller must guarantee the same structural invariants as [`Self::try_new`].
57    pub unsafe fn new_unchecked(
58        starts: ArrayRef,
59        lengths: ArrayRef,
60        multipliers: ArrayRef,
61        len: usize,
62    ) -> Self {
63        unsafe {
64            Array::from_parts_unchecked(
65                ArrayParts::new(PiecewiseSequence, PType::U64.into(), len, EmptyArrayData)
66                    .with_slots(smallvec![Some(starts), Some(lengths), Some(multipliers)]),
67            )
68        }
69    }
70}
vortex_array/arrays/piecewise_sequence/array.rs

vortex_array/arrays/piecewise_sequence/
array.rs
