vortex_sequence/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use num_traits::cast::FromPrimitive;
5use vortex_array::arrays::PrimitiveArray;
6use vortex_array::stats::{ArrayStats, StatsSetRef};
7use vortex_array::vtable::{
8    ArrayVTable, CanonicalVTable, NotSupported, OperationsVTable, VTable, ValidityVTable,
9    VisitorVTable,
10};
11use vortex_array::{
12    ArrayBufferVisitor, ArrayChildVisitor, ArrayRef, Canonical, EncodingId, EncodingRef, vtable,
13};
14use vortex_buffer::BufferMut;
15use vortex_dtype::{
16    DType, NativePType, Nullability, PType, match_each_integer_ptype, match_each_native_ptype,
17};
18use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
19use vortex_mask::Mask;
20use vortex_scalar::{PValue, Scalar, ScalarValue};
21
22vtable!(Sequence);
23
24#[derive(Clone, Debug)]
25/// An array representing the equation `A[i] = base + i * multiplier`.
26pub struct SequenceArray {
27    base: PValue,
28    multiplier: PValue,
29    dtype: DType,
30    length: usize,
31    stats_set: ArrayStats,
32}
33
34impl SequenceArray {
35    pub fn typed_new<T: NativePType + Into<PValue>>(
36        base: T,
37        multiplier: T,
38        nullability: Nullability,
39        length: usize,
40    ) -> VortexResult<Self> {
41        Self::new(
42            base.into(),
43            multiplier.into(),
44            T::PTYPE,
45            nullability,
46            length,
47        )
48    }
49
50    /// Constructs a sequence array using two integer values (with the same ptype).
51    pub fn new(
52        base: PValue,
53        multiplier: PValue,
54        ptype: PType,
55        nullability: Nullability,
56        length: usize,
57    ) -> VortexResult<Self> {
58        if !ptype.is_int() {
59            vortex_bail!("only integer ptype are supported in SequenceArray currently")
60        }
61
62        Self::try_last(base, multiplier, ptype, length).map_err(|e| {
63            e.with_context(format!(
64                "final value not expressible, base = {base:?}, multiplier = {multiplier:?}, len = {length} ",
65            ))
66        })?;
67
68        Ok(Self::unchecked_new(
69            base,
70            multiplier,
71            ptype,
72            nullability,
73            length,
74        ))
75    }
76
77    pub(crate) fn unchecked_new(
78        base: PValue,
79        multiplier: PValue,
80        ptype: PType,
81        nullability: Nullability,
82        length: usize,
83    ) -> Self {
84        let dtype = DType::Primitive(ptype, nullability);
85        Self {
86            base,
87            multiplier,
88            dtype,
89            length,
90            // TODO(joe): add stats, on construct or on use?
91            stats_set: Default::default(),
92        }
93    }
94
95    pub fn ptype(&self) -> PType {
96        self.dtype.as_ptype()
97    }
98
99    pub fn base(&self) -> PValue {
100        self.base
101    }
102
103    pub fn multiplier(&self) -> PValue {
104        self.multiplier
105    }
106
107    pub(crate) fn try_last(
108        base: PValue,
109        multiplier: PValue,
110        ptype: PType,
111        length: usize,
112    ) -> VortexResult<PValue> {
113        match_each_integer_ptype!(ptype, |P| {
114            let len_t = <P>::from_usize(length - 1)
115                .ok_or_else(|| vortex_err!("cannot convert length {} into {}", length, ptype))?;
116
117            let base = base.as_primitive::<P>()?;
118            let multiplier = multiplier.as_primitive::<P>()?;
119
120            let last = len_t
121                .checked_mul(multiplier)
122                .and_then(|offset| offset.checked_add(base))
123                .ok_or_else(|| vortex_err!("last value computation overflows"))?;
124            Ok(PValue::from(last))
125        })
126    }
127
128    fn index_value(&self, idx: usize) -> VortexResult<PValue> {
129        if idx > self.length {
130            vortex_bail!("out of bounds")
131        }
132        match_each_native_ptype!(self.ptype(), |P| {
133            let base = self.base.as_primitive::<P>()?;
134            let multiplier = self.multiplier.as_primitive::<P>()?;
135            let value = base + (multiplier * <P>::from_usize(idx).vortex_expect("must fit"));
136
137            Ok(PValue::from(value))
138        })
139    }
140
141    /// Returns the validated final value of a sequence array
142    pub fn last(&self) -> PValue {
143        Self::try_last(self.base, self.multiplier, self.ptype(), self.length)
144            .vortex_expect("validated array")
145    }
146
147    pub fn dtype(&self) -> &DType {
148        &self.dtype
149    }
150}
151
152impl VTable for SequenceVTable {
153    type Array = SequenceArray;
154    type Encoding = SequenceEncoding;
155
156    type ArrayVTable = Self;
157    type CanonicalVTable = Self;
158    type OperationsVTable = Self;
159    type ValidityVTable = Self;
160    type VisitorVTable = Self;
161    type ComputeVTable = NotSupported;
162    type EncodeVTable = Self;
163    type SerdeVTable = Self;
164
165    fn id(_encoding: &Self::Encoding) -> EncodingId {
166        EncodingId::new_ref("vortex.sequence")
167    }
168
169    fn encoding(_array: &Self::Array) -> EncodingRef {
170        EncodingRef::new_ref(SequenceEncoding.as_ref())
171    }
172}
173
174impl ArrayVTable<SequenceVTable> for SequenceVTable {
175    fn len(array: &SequenceArray) -> usize {
176        array.length
177    }
178
179    fn dtype(array: &SequenceArray) -> &DType {
180        &array.dtype
181    }
182
183    fn stats(array: &SequenceArray) -> StatsSetRef<'_> {
184        array.stats_set.to_ref(array.as_ref())
185    }
186}
187
188impl CanonicalVTable<SequenceVTable> for SequenceVTable {
189    fn canonicalize(array: &SequenceArray) -> VortexResult<Canonical> {
190        let prim = match_each_native_ptype!(array.ptype(), |P| {
191            let base = array.base().as_primitive::<P>()?;
192            let multiplier = array.multiplier().as_primitive::<P>()?;
193            let values = BufferMut::from_iter(
194                (0..array.len())
195                    .map(|i| base + <P>::from_usize(i).vortex_expect("must fit") * multiplier),
196            );
197            PrimitiveArray::new(values, array.dtype.nullability().into())
198        });
199
200        Ok(Canonical::Primitive(prim))
201    }
202}
203
204impl OperationsVTable<SequenceVTable> for SequenceVTable {
205    fn slice(array: &SequenceArray, start: usize, stop: usize) -> VortexResult<ArrayRef> {
206        Ok(SequenceArray::unchecked_new(
207            array.index_value(start)?,
208            array.multiplier,
209            array.ptype(),
210            array.dtype().nullability(),
211            stop - start,
212        )
213        .to_array())
214    }
215
216    fn scalar_at(array: &SequenceArray, index: usize) -> VortexResult<Scalar> {
217        // Ok(Scalar::from(array.index_value(index)))
218        Ok(Scalar::new(
219            array.dtype().clone(),
220            ScalarValue::from(array.index_value(index)?),
221        ))
222    }
223}
224
225impl ValidityVTable<SequenceVTable> for SequenceVTable {
226    fn is_valid(_array: &SequenceArray, _index: usize) -> VortexResult<bool> {
227        Ok(true)
228    }
229
230    fn all_valid(_array: &SequenceArray) -> VortexResult<bool> {
231        Ok(true)
232    }
233
234    fn all_invalid(_array: &SequenceArray) -> VortexResult<bool> {
235        Ok(false)
236    }
237
238    fn validity_mask(array: &SequenceArray) -> VortexResult<Mask> {
239        Ok(Mask::AllTrue(array.len()))
240    }
241}
242
243impl VisitorVTable<SequenceVTable> for SequenceVTable {
244    fn visit_buffers(_array: &SequenceArray, _visitor: &mut dyn ArrayBufferVisitor) {
245        // TODO(joe): expose scalar values
246    }
247
248    fn visit_children(_array: &SequenceArray, _visitor: &mut dyn ArrayChildVisitor) {}
249}
250
251#[derive(Clone, Debug)]
252pub struct SequenceEncoding;
253
254#[cfg(test)]
255mod tests {
256    use vortex_array::arrays::PrimitiveArray;
257    use vortex_dtype::Nullability;
258    use vortex_scalar::{Scalar, ScalarValue};
259
260    use crate::array::SequenceArray;
261
262    #[test]
263    fn test_sequence_canonical() {
264        let arr = SequenceArray::typed_new(2i64, 3, Nullability::NonNullable, 4).unwrap();
265
266        let canon = PrimitiveArray::from_iter((0..4).map(|i| 2i64 + i * 3));
267
268        assert_eq!(
269            arr.to_canonical()
270                .unwrap()
271                .into_primitive()
272                .unwrap()
273                .as_slice::<i64>(),
274            canon.as_slice::<i64>()
275        )
276    }
277
278    #[test]
279    fn test_sequence_slice_canonical() {
280        let arr = SequenceArray::typed_new(2i64, 3, Nullability::NonNullable, 4)
281            .unwrap()
282            .slice(2, 3)
283            .unwrap();
284
285        let canon = PrimitiveArray::from_iter((2..3).map(|i| 2i64 + i * 3));
286
287        assert_eq!(
288            arr.to_canonical()
289                .unwrap()
290                .into_primitive()
291                .unwrap()
292                .as_slice::<i64>(),
293            canon.as_slice::<i64>()
294        )
295    }
296
297    #[test]
298    fn test_sequence_scalar_at() {
299        let scalar = SequenceArray::typed_new(2i64, 3, Nullability::NonNullable, 4)
300            .unwrap()
301            .scalar_at(2)
302            .unwrap();
303
304        assert_eq!(
305            scalar,
306            Scalar::new(scalar.dtype().clone(), ScalarValue::from(8i64))
307        )
308    }
309
310    #[test]
311    fn test_sequence_min_max() {
312        assert!(SequenceArray::typed_new(-127i8, -1i8, Nullability::NonNullable, 2).is_ok());
313        assert!(SequenceArray::typed_new(126i8, -1i8, Nullability::NonNullable, 2).is_ok());
314    }
315
316    #[test]
317    fn test_sequence_too_big() {
318        assert!(SequenceArray::typed_new(127i8, 1i8, Nullability::NonNullable, 2).is_err());
319        assert!(SequenceArray::typed_new(-128i8, -1i8, Nullability::NonNullable, 2).is_err());
320    }
321}
vortex_sequence/array.rs

vortex_sequence/
array.rs
