vortex_vector/primitive/
vector.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Definition and implementation of [`PrimitiveVector`].
5
6use std::fmt::Debug;
7use std::ops::RangeBounds;
8
9use vortex_dtype::NativePType;
10use vortex_dtype::PType;
11use vortex_dtype::PTypeDowncast;
12use vortex_dtype::PTypeUpcast;
13use vortex_dtype::half::f16;
14use vortex_error::vortex_panic;
15use vortex_mask::Mask;
16
17use crate::VectorOps;
18use crate::match_each_pvector;
19use crate::primitive::PVector;
20use crate::primitive::PrimitiveScalar;
21use crate::primitive::PrimitiveVectorMut;
22
23/// An immutable vector of primitive values.
24///
25/// The mutable equivalent of this type is [`PrimitiveVectorMut`].
26///
27/// `PrimitiveVector` is represented by an enum over all possible [`PVector`] types (which are
28/// templated by the types that implement [`NativePType`]).
29///
30/// See the documentation for [`PVector`] for more information.
31#[derive(Debug, Clone)]
32pub enum PrimitiveVector {
33    /// U8
34    U8(PVector<u8>),
35    /// U16
36    U16(PVector<u16>),
37    /// U32
38    U32(PVector<u32>),
39    /// U64
40    U64(PVector<u64>),
41    /// I8
42    I8(PVector<i8>),
43    /// I16
44    I16(PVector<i16>),
45    /// I32
46    I32(PVector<i32>),
47    /// I64
48    I64(PVector<i64>),
49    /// F16
50    F16(PVector<f16>),
51    /// F32
52    F32(PVector<f32>),
53    /// F64
54    F64(PVector<f64>),
55}
56
57impl PartialEq for PrimitiveVector {
58    fn eq(&self, other: &Self) -> bool {
59        // Validity patterns must match
60        if self.validity() != other.validity() {
61            return false;
62        }
63        // Delegate to the underlying PVector equality
64        match (self, other) {
65            (PrimitiveVector::U8(a), PrimitiveVector::U8(b)) => a == b,
66            (PrimitiveVector::U16(a), PrimitiveVector::U16(b)) => a == b,
67            (PrimitiveVector::U32(a), PrimitiveVector::U32(b)) => a == b,
68            (PrimitiveVector::U64(a), PrimitiveVector::U64(b)) => a == b,
69            (PrimitiveVector::I8(a), PrimitiveVector::I8(b)) => a == b,
70            (PrimitiveVector::I16(a), PrimitiveVector::I16(b)) => a == b,
71            (PrimitiveVector::I32(a), PrimitiveVector::I32(b)) => a == b,
72            (PrimitiveVector::I64(a), PrimitiveVector::I64(b)) => a == b,
73            (PrimitiveVector::F16(a), PrimitiveVector::F16(b)) => a == b,
74            (PrimitiveVector::F32(a), PrimitiveVector::F32(b)) => a == b,
75            (PrimitiveVector::F64(a), PrimitiveVector::F64(b)) => a == b,
76            _ => false, // Different variants are not equal
77        }
78    }
79}
80
81impl PrimitiveVector {
82    /// Returns the [`PType`] of this [`PrimitiveVector`].
83    pub fn ptype(&self) -> PType {
84        match self {
85            Self::U8(_) => PType::U8,
86            Self::U16(_) => PType::U16,
87            Self::U32(_) => PType::U32,
88            Self::U64(_) => PType::U64,
89            Self::I8(_) => PType::I8,
90            Self::I16(_) => PType::I16,
91            Self::I32(_) => PType::I32,
92            Self::I64(_) => PType::I64,
93            Self::F16(_) => PType::F16,
94            Self::F32(_) => PType::F32,
95            Self::F64(_) => PType::F64,
96        }
97    }
98}
99
100impl VectorOps for PrimitiveVector {
101    type Mutable = PrimitiveVectorMut;
102    type Scalar = PrimitiveScalar;
103
104    fn len(&self) -> usize {
105        match_each_pvector!(self, |v| { v.len() })
106    }
107
108    fn validity(&self) -> &Mask {
109        match_each_pvector!(self, |v| { v.validity() })
110    }
111
112    fn mask_validity(&mut self, mask: &Mask) {
113        match_each_pvector!(self, |v| { v.mask_validity(mask) })
114    }
115
116    fn scalar_at(&self, index: usize) -> PrimitiveScalar {
117        match_each_pvector!(self, |v| { v.scalar_at(index).into() })
118    }
119
120    fn slice(&self, range: impl RangeBounds<usize> + Clone + Debug) -> Self {
121        match_each_pvector!(self, |v| { v.slice(range).into() })
122    }
123
124    fn clear(&mut self) {
125        match_each_pvector!(self, |v| { v.clear() })
126    }
127
128    fn try_into_mut(self) -> Result<PrimitiveVectorMut, Self> {
129        match_each_pvector!(self, |v| {
130            v.try_into_mut()
131                .map(PrimitiveVectorMut::from)
132                .map_err(Self::from)
133        })
134    }
135
136    fn into_mut(self) -> PrimitiveVectorMut {
137        match_each_pvector!(self, |v| { v.into_mut().into() })
138    }
139}
140
141impl PTypeUpcast for PrimitiveVector {
142    type Input<T: NativePType> = PVector<T>;
143
144    fn from_u8(input: Self::Input<u8>) -> Self {
145        Self::U8(input)
146    }
147
148    fn from_u16(input: Self::Input<u16>) -> Self {
149        Self::U16(input)
150    }
151
152    fn from_u32(input: Self::Input<u32>) -> Self {
153        Self::U32(input)
154    }
155
156    fn from_u64(input: Self::Input<u64>) -> Self {
157        Self::U64(input)
158    }
159
160    fn from_i8(input: Self::Input<i8>) -> Self {
161        Self::I8(input)
162    }
163
164    fn from_i16(input: Self::Input<i16>) -> Self {
165        Self::I16(input)
166    }
167
168    fn from_i32(input: Self::Input<i32>) -> Self {
169        Self::I32(input)
170    }
171
172    fn from_i64(input: Self::Input<i64>) -> Self {
173        Self::I64(input)
174    }
175
176    fn from_f16(input: Self::Input<f16>) -> Self {
177        Self::F16(input)
178    }
179
180    fn from_f32(input: Self::Input<f32>) -> Self {
181        Self::F32(input)
182    }
183
184    fn from_f64(input: Self::Input<f64>) -> Self {
185        Self::F64(input)
186    }
187}
188
189impl PTypeDowncast for PrimitiveVector {
190    type Output<T: NativePType> = PVector<T>;
191
192    fn into_u8(self) -> Self::Output<u8> {
193        if let Self::U8(v) = self {
194            return v;
195        }
196        vortex_panic!("Expected PrimitiveVector::U8, got {self:?}");
197    }
198
199    fn into_u16(self) -> Self::Output<u16> {
200        if let Self::U16(v) = self {
201            return v;
202        }
203        vortex_panic!("Expected PrimitiveVector::U16, got {self:?}");
204    }
205
206    fn into_u32(self) -> Self::Output<u32> {
207        if let Self::U32(v) = self {
208            return v;
209        }
210        vortex_panic!("Expected PrimitiveVector::U32, got {self:?}");
211    }
212
213    fn into_u64(self) -> Self::Output<u64> {
214        if let Self::U64(v) = self {
215            return v;
216        }
217        vortex_panic!("Expected PrimitiveVector::U64, got {self:?}");
218    }
219
220    fn into_i8(self) -> Self::Output<i8> {
221        if let Self::I8(v) = self {
222            return v;
223        }
224        vortex_panic!("Expected PrimitiveVector::I8, got {self:?}");
225    }
226
227    fn into_i16(self) -> Self::Output<i16> {
228        if let Self::I16(v) = self {
229            return v;
230        }
231        vortex_panic!("Expected PrimitiveVector::I16, got {self:?}");
232    }
233
234    fn into_i32(self) -> Self::Output<i32> {
235        if let Self::I32(v) = self {
236            return v;
237        }
238        vortex_panic!("Expected PrimitiveVector::I32, got {self:?}");
239    }
240
241    fn into_i64(self) -> Self::Output<i64> {
242        if let Self::I64(v) = self {
243            return v;
244        }
245        vortex_panic!("Expected PrimitiveVector::I64, got {self:?}");
246    }
247
248    fn into_f16(self) -> Self::Output<f16> {
249        if let Self::F16(v) = self {
250            return v;
251        }
252        vortex_panic!("Expected PrimitiveVector::F16, got {self:?}");
253    }
254
255    fn into_f32(self) -> Self::Output<f32> {
256        if let Self::F32(v) = self {
257            return v;
258        }
259        vortex_panic!("Expected PrimitiveVector::F32, got {self:?}");
260    }
261
262    fn into_f64(self) -> Self::Output<f64> {
263        if let Self::F64(v) = self {
264            return v;
265        }
266        vortex_panic!("Expected PrimitiveVector::F64, got {self:?}");
267    }
268}
269
270impl<'a> PTypeDowncast for &'a PrimitiveVector {
271    type Output<T: NativePType> = &'a PVector<T>;
272
273    fn into_u8(self) -> Self::Output<u8> {
274        if let PrimitiveVector::U8(v) = self {
275            return v;
276        }
277        vortex_panic!("Expected PrimitiveVector::U8, got {self:?}");
278    }
279
280    fn into_u16(self) -> Self::Output<u16> {
281        if let PrimitiveVector::U16(v) = self {
282            return v;
283        }
284        vortex_panic!("Expected PrimitiveVector::U16, got {self:?}");
285    }
286
287    fn into_u32(self) -> Self::Output<u32> {
288        if let PrimitiveVector::U32(v) = self {
289            return v;
290        }
291        vortex_panic!("Expected PrimitiveVector::U32, got {self:?}");
292    }
293
294    fn into_u64(self) -> Self::Output<u64> {
295        if let PrimitiveVector::U64(v) = self {
296            return v;
297        }
298        vortex_panic!("Expected PrimitiveVector::U64, got {self:?}");
299    }
300
301    fn into_i8(self) -> Self::Output<i8> {
302        if let PrimitiveVector::I8(v) = self {
303            return v;
304        }
305        vortex_panic!("Expected PrimitiveVector::I8, got {self:?}");
306    }
307
308    fn into_i16(self) -> Self::Output<i16> {
309        if let PrimitiveVector::I16(v) = self {
310            return v;
311        }
312        vortex_panic!("Expected PrimitiveVector::I16, got {self:?}");
313    }
314
315    fn into_i32(self) -> Self::Output<i32> {
316        if let PrimitiveVector::I32(v) = self {
317            return v;
318        }
319        vortex_panic!("Expected PrimitiveVector::I32, got {self:?}");
320    }
321
322    fn into_i64(self) -> Self::Output<i64> {
323        if let PrimitiveVector::I64(v) = self {
324            return v;
325        }
326        vortex_panic!("Expected PrimitiveVector::I64, got {self:?}");
327    }
328
329    fn into_f16(self) -> Self::Output<f16> {
330        if let PrimitiveVector::F16(v) = self {
331            return v;
332        }
333        vortex_panic!("Expected PrimitiveVector::F16, got {self:?}");
334    }
335
336    fn into_f32(self) -> Self::Output<f32> {
337        if let PrimitiveVector::F32(v) = self {
338            return v;
339        }
340        vortex_panic!("Expected PrimitiveVector::F32, got {self:?}");
341    }
342
343    fn into_f64(self) -> Self::Output<f64> {
344        if let PrimitiveVector::F64(v) = self {
345            return v;
346        }
347        vortex_panic!("Expected PrimitiveVector::F64, got {self:?}");
348    }
349}
350
351impl<'a> PTypeDowncast for &'a mut PrimitiveVector {
352    type Output<T: NativePType> = &'a mut PVector<T>;
353
354    fn into_u8(self) -> Self::Output<u8> {
355        if let PrimitiveVector::U8(v) = self {
356            return v;
357        }
358        vortex_panic!("Expected PrimitiveVector::U8, got {self:?}");
359    }
360
361    fn into_u16(self) -> Self::Output<u16> {
362        if let PrimitiveVector::U16(v) = self {
363            return v;
364        }
365        vortex_panic!("Expected PrimitiveVector::U16, got {self:?}");
366    }
367
368    fn into_u32(self) -> Self::Output<u32> {
369        if let PrimitiveVector::U32(v) = self {
370            return v;
371        }
372        vortex_panic!("Expected PrimitiveVector::U32, got {self:?}");
373    }
374
375    fn into_u64(self) -> Self::Output<u64> {
376        if let PrimitiveVector::U64(v) = self {
377            return v;
378        }
379        vortex_panic!("Expected PrimitiveVector::U64, got {self:?}");
380    }
381
382    fn into_i8(self) -> Self::Output<i8> {
383        if let PrimitiveVector::I8(v) = self {
384            return v;
385        }
386        vortex_panic!("Expected PrimitiveVector::I8, got {self:?}");
387    }
388
389    fn into_i16(self) -> Self::Output<i16> {
390        if let PrimitiveVector::I16(v) = self {
391            return v;
392        }
393        vortex_panic!("Expected PrimitiveVector::I16, got {self:?}");
394    }
395
396    fn into_i32(self) -> Self::Output<i32> {
397        if let PrimitiveVector::I32(v) = self {
398            return v;
399        }
400        vortex_panic!("Expected PrimitiveVector::I32, got {self:?}");
401    }
402
403    fn into_i64(self) -> Self::Output<i64> {
404        if let PrimitiveVector::I64(v) = self {
405            return v;
406        }
407        vortex_panic!("Expected PrimitiveVector::I64, got {self:?}");
408    }
409
410    fn into_f16(self) -> Self::Output<f16> {
411        if let PrimitiveVector::F16(v) = self {
412            return v;
413        }
414        vortex_panic!("Expected PrimitiveVector::F16, got {self:?}");
415    }
416
417    fn into_f32(self) -> Self::Output<f32> {
418        if let PrimitiveVector::F32(v) = self {
419            return v;
420        }
421        vortex_panic!("Expected PrimitiveVector::F32, got {self:?}");
422    }
423
424    fn into_f64(self) -> Self::Output<f64> {
425        if let PrimitiveVector::F64(v) = self {
426            return v;
427        }
428        vortex_panic!("Expected PrimitiveVector::F64, got {self:?}");
429    }
430}