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;
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!(
65            dtype == &DType::from(PType::U64),
66            "PiecewiseSequenceArray dtype must be u64, got {dtype}"
67        );
68        vortex_ensure!(
69            slots.len() == PiecewiseSequenceSlots::NAMES.len(),
70            "PiecewiseSequenceArray requires {} slots, got {}",
71            PiecewiseSequenceSlots::NAMES.len(),
72            slots.len()
73        );
74        let starts = slots[PiecewiseSequenceSlots::STARTS]
75            .as_ref()
76            .ok_or_else(|| vortex_err!("PiecewiseSequenceArray starts slot must be present"))?;
77        let lengths = slots[PiecewiseSequenceSlots::LENGTHS]
78            .as_ref()
79            .ok_or_else(|| vortex_err!("PiecewiseSequenceArray lengths slot must be present"))?;
80        let multipliers = slots[PiecewiseSequenceSlots::MULTIPLIERS]
81            .as_ref()
82            .ok_or_else(|| {
83                vortex_err!("PiecewiseSequenceArray multipliers slot must be present")
84            })?;
85        check_index_arrays(starts, lengths, multipliers)
86    }
87
88    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
89        0
90    }
91
92    fn buffer(_array: ArrayView<'_, Self>, _idx: usize) -> BufferHandle {
93        vortex_panic!("PiecewiseSequenceArray has no buffers")
94    }
95
96    fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option<String> {
97        None
98    }
99
100    fn with_buffers(
101        &self,
102        array: ArrayView<'_, Self>,
103        buffers: &[BufferHandle],
104    ) -> VortexResult<ArrayParts<Self>> {
105        with_empty_buffers(self, array, buffers)
106    }
107
108    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
109        PiecewiseSequenceSlots::NAMES[idx].to_string()
110    }
111
112    fn serialize(
113        _array: ArrayView<'_, Self>,
114        _session: &VortexSession,
115    ) -> VortexResult<Option<Vec<u8>>> {
116        Ok(None)
117    }
118
119    fn deserialize(
120        &self,
121        _dtype: &DType,
122        _len: usize,
123        _metadata: &[u8],
124        _buffers: &[BufferHandle],
125        _children: &dyn ArrayChildren,
126        _session: &VortexSession,
127    ) -> VortexResult<ArrayParts<Self>> {
128        vortex_bail!("PiecewiseSequenceArray is not serializable")
129    }
130
131    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
132        let (starts, lengths, multipliers) = execute_index_arrays(array.as_view(), ctx)?;
133
134        let values = match_each_unsigned_integer_ptype!(starts.ptype(), |S| {
135            match_each_unsigned_integer_ptype!(lengths.ptype(), |L| {
136                match_each_unsigned_integer_ptype!(multipliers.ptype(), |M| {
137                    materialize_ranges::<S, L, M>(&starts, &lengths, &multipliers, array.len())?
138                })
139            })
140        });
141        Ok(ExecutionResult::done(
142            PrimitiveArray::new(values.freeze(), Validity::NonNullable).into_array(),
143        ))
144    }
145}
146
147impl OperationsVTable<PiecewiseSequence> for PiecewiseSequence {
148    type ProbeState = ();
149
150    fn scalar_at(
151        array: ArrayView<'_, PiecewiseSequence>,
152        index: usize,
153        ctx: &mut ExecutionCtx,
154    ) -> VortexResult<Scalar> {
155        let (starts, lengths, multipliers) = execute_index_arrays(array, ctx)?;
156
157        let value = match_each_unsigned_integer_ptype!(starts.ptype(), |S| {
158            match_each_unsigned_integer_ptype!(lengths.ptype(), |L| {
159                match_each_unsigned_integer_ptype!(multipliers.ptype(), |M| {
160                    scalar_at::<S, L, M>(&starts, &lengths, &multipliers, index)?
161                })
162            })
163        });
164        Ok(value.into())
165    }
166}
167
168impl ValidityVTable<PiecewiseSequence> for PiecewiseSequence {
169    fn validity(_array: ArrayView<'_, PiecewiseSequence>) -> VortexResult<Validity> {
170        Ok(Validity::NonNullable)
171    }
172}
173
174fn scalar_at<S, L, M>(
175    starts: &PrimitiveArray,
176    lengths: &PrimitiveArray,
177    multipliers: &PrimitiveArray,
178    index: usize,
179) -> VortexResult<u64>
180where
181    S: UnsignedPType,
182    L: UnsignedPType,
183    M: UnsignedPType,
184{
185    let mut remaining = index;
186    for ((&start, &length), &multiplier) in starts
187        .as_slice::<S>()
188        .iter()
189        .zip_eq(lengths.as_slice::<L>())
190        .zip_eq(multipliers.as_slice::<M>())
191    {
192        let length: usize = length.as_();
193        if remaining < length {
194            let start: usize = start.as_();
195            let multiplier: usize = multiplier.as_();
196            let offset = remaining
197                .checked_mul(multiplier)
198                .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
199            let value = start
200                .checked_add(offset)
201                .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
202            return Ok(value as u64);
203        }
204        remaining -= length;
205    }
206    vortex_bail!("PiecewiseSequenceArray index {index} out of bounds")
207}
vortex_array/arrays/piecewise_sequence/vtable.rs

vortex_array/arrays/piecewise_sequence/
vtable.rs
