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(ArrayParts::new(
50            PiecewiseSequence,
51            PType::U64.into(),
52            len,
53            EmptyArrayData,
54            smallvec![Some(starts), Some(lengths), Some(multipliers)],
55        ))
56    }
57
58    /// Constructs a new `PiecewiseSequenceArray` without validation.
59    ///
60    /// # Safety
61    ///
62    /// The caller must guarantee the same structural invariants as [`Self::try_new`].
63    pub unsafe fn new_unchecked(
64        starts: ArrayRef,
65        lengths: ArrayRef,
66        multipliers: ArrayRef,
67        len: usize,
68    ) -> Self {
69        unsafe {
70            Array::from_parts_unchecked(ArrayParts::new(
71                PiecewiseSequence,
72                PType::U64.into(),
73                len,
74                EmptyArrayData,
75                smallvec![Some(starts), Some(lengths), Some(multipliers)],
76            ))
77        }
78    }
79}
vortex_array/arrays/piecewise_sequence/array.rs

vortex_array/arrays/piecewise_sequence/
array.rs
