vortex_array/arrays/piecewise_sequence/
vtable.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use itertools::Itertools;
5use vortex_error::VortexResult;
6use vortex_error::vortex_bail;
7use vortex_error::vortex_ensure_eq;
8use vortex_error::vortex_err;
9use vortex_error::vortex_panic;
10use vortex_session::VortexSession;
11use vortex_session::registry::CachedId;
12
13use crate::ArrayParts;
14use crate::ArrayRef;
15use crate::ExecutionCtx;
16use crate::ExecutionResult;
17use crate::IntoArray;
18use crate::array::Array;
19use crate::array::ArrayId;
20use crate::array::ArrayView;
21use crate::array::EmptyArrayData;
22use crate::array::OperationsVTable;
23use crate::array::VTable;
24use crate::array::ValidityVTable;
25use crate::array::with_empty_buffers;
26use crate::arrays::PrimitiveArray;
27use crate::arrays::piecewise_sequence::array::PiecewiseSequenceSlots;
28use crate::arrays::piecewise_sequence::check_index_arrays;
29use crate::arrays::piecewise_sequence::execute_index_arrays;
30use crate::arrays::piecewise_sequence::materialize_ranges;
31use crate::arrays::primitive::PrimitiveArrayExt;
32use crate::buffer::BufferHandle;
33use crate::dtype::DType;
34use crate::dtype::PType;
35use crate::dtype::UnsignedPType;
36use crate::match_each_unsigned_integer_ptype;
37use crate::scalar::Scalar;
38use crate::serde::ArrayChildren;
39use crate::validity::Validity;
40
41/// A [`PiecewiseSequence`]-encoded Vortex index array.
42pub type PiecewiseSequenceArray = Array<PiecewiseSequence>;
43
44#[derive(Clone, Debug)]
45pub struct PiecewiseSequence;
46
47impl VTable for PiecewiseSequence {
48    type TypedArrayData = EmptyArrayData;
49    type OperationsVTable = Self;
50    type ValidityVTable = Self;
51
52    fn id(&self) -> ArrayId {
53        static ID: CachedId = CachedId::new("vortex.piecewise-sequence");
54        *ID
55    }
56
57    fn validate(
58        &self,
59        _data: &Self::TypedArrayData,
60        dtype: &DType,
61        _len: usize,
62        slots: &[Option<ArrayRef>],
63    ) -> VortexResult<()> {
64        vortex_ensure_eq!(
65            dtype,
66            &DType::from(PType::U64),
67            "PiecewiseSequenceArray dtype must be u64",
68        );
69        vortex_ensure_eq!(slots.len(), PiecewiseSequenceSlots::NAMES.len());
70        let starts = slots[PiecewiseSequenceSlots::STARTS]
71            .as_ref()
72            .ok_or_else(|| vortex_err!("PiecewiseSequenceArray starts slot must be present"))?;
73        let lengths = slots[PiecewiseSequenceSlots::LENGTHS]
74            .as_ref()
75            .ok_or_else(|| vortex_err!("PiecewiseSequenceArray lengths slot must be present"))?;
76        let multipliers = slots[PiecewiseSequenceSlots::MULTIPLIERS]
77            .as_ref()
78            .ok_or_else(|| {
79                vortex_err!("PiecewiseSequenceArray multipliers slot must be present")
80            })?;
81        check_index_arrays(starts, lengths, multipliers)
82    }
83
84    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
85        0
86    }
87
88    fn buffer(_array: ArrayView<'_, Self>, _idx: usize) -> BufferHandle {
89        vortex_panic!("PiecewiseSequenceArray has no buffers")
90    }
91
92    fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option<String> {
93        None
94    }
95
96    fn with_buffers(
97        &self,
98        array: ArrayView<'_, Self>,
99        buffers: &[BufferHandle],
100    ) -> VortexResult<ArrayParts<Self>> {
101        with_empty_buffers(self, array, buffers)
102    }
103
104    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
105        PiecewiseSequenceSlots::NAMES[idx].to_string()
106    }
107
108    fn serialize(
109        _array: ArrayView<'_, Self>,
110        _session: &VortexSession,
111    ) -> VortexResult<Option<Vec<u8>>> {
112        Ok(None)
113    }
114
115    fn deserialize(
116        &self,
117        _dtype: &DType,
118        _len: usize,
119        _metadata: &[u8],
120        _buffers: &[BufferHandle],
121        _children: &dyn ArrayChildren,
122        _session: &VortexSession,
123    ) -> VortexResult<ArrayParts<Self>> {
124        vortex_bail!("PiecewiseSequenceArray is not serializable")
125    }
126
127    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
128        let (starts, lengths, multipliers) = execute_index_arrays(array.as_view(), ctx)?;
129
130        let values = match_each_unsigned_integer_ptype!(starts.ptype(), |S| {
131            match_each_unsigned_integer_ptype!(lengths.ptype(), |L| {
132                match_each_unsigned_integer_ptype!(multipliers.ptype(), |M| {
133                    materialize_ranges::<S, L, M>(&starts, &lengths, &multipliers, array.len())?
134                })
135            })
136        });
137        Ok(ExecutionResult::done(
138            PrimitiveArray::new(values.freeze(), Validity::NonNullable).into_array(),
139        ))
140    }
141}
142
143impl OperationsVTable<PiecewiseSequence> for PiecewiseSequence {
144    type ProbeState = ();
145
146    fn scalar_at(
147        array: ArrayView<'_, PiecewiseSequence>,
148        index: usize,
149        ctx: &mut ExecutionCtx,
150    ) -> VortexResult<Scalar> {
151        let (starts, lengths, multipliers) = execute_index_arrays(array, ctx)?;
152
153        let value = match_each_unsigned_integer_ptype!(starts.ptype(), |S| {
154            match_each_unsigned_integer_ptype!(lengths.ptype(), |L| {
155                match_each_unsigned_integer_ptype!(multipliers.ptype(), |M| {
156                    scalar_at::<S, L, M>(&starts, &lengths, &multipliers, index)?
157                })
158            })
159        });
160        Ok(value.into())
161    }
162}
163
164impl ValidityVTable<PiecewiseSequence> for PiecewiseSequence {
165    fn validity(_array: ArrayView<'_, PiecewiseSequence>) -> VortexResult<Validity> {
166        Ok(Validity::NonNullable)
167    }
168}
169
170fn scalar_at<S, L, M>(
171    starts: &PrimitiveArray,
172    lengths: &PrimitiveArray,
173    multipliers: &PrimitiveArray,
174    index: usize,
175) -> VortexResult<u64>
176where
177    S: UnsignedPType,
178    L: UnsignedPType,
179    M: UnsignedPType,
180{
181    let mut remaining = index;
182    for ((&start, &length), &multiplier) in starts
183        .as_slice::<S>()
184        .iter()
185        .zip_eq(lengths.as_slice::<L>())
186        .zip_eq(multipliers.as_slice::<M>())
187    {
188        let length: usize = length.as_();
189        if remaining < length {
190            let start: usize = start.as_();
191            let multiplier: usize = multiplier.as_();
192            let offset = remaining
193                .checked_mul(multiplier)
194                .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
195            let value = start
196                .checked_add(offset)
197                .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
198            return Ok(value as u64);
199        }
200        remaining -= length;
201    }
202    vortex_bail!("PiecewiseSequenceArray index {index} out of bounds")
203}
vortex_array/arrays/piecewise_sequence/vtable.rs

vortex_array/arrays/piecewise_sequence/
vtable.rs
