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::half::f16;
10use vortex_dtype::{NativePType, PType, PTypeDowncast, PTypeUpcast};
11use vortex_error::vortex_panic;
12use vortex_mask::Mask;
13
14use crate::primitive::{PVector, PrimitiveVectorMut};
15use crate::{Scalar, VectorOps, match_each_pvector};
16
17/// An immutable vector of primitive values.
18///
19/// The mutable equivalent of this type is [`PrimitiveVectorMut`].
20///
21/// `PrimitiveVector` is represented by an enum over all possible [`PVector`] types (which are
22/// templated by the types that implement [`NativePType`]).
23///
24/// See the documentation for [`PVector`] for more information.
25#[derive(Debug, Clone)]
26pub enum PrimitiveVector {
27    /// U8
28    U8(PVector<u8>),
29    /// U16
30    U16(PVector<u16>),
31    /// U32
32    U32(PVector<u32>),
33    /// U64
34    U64(PVector<u64>),
35    /// I8
36    I8(PVector<i8>),
37    /// I16
38    I16(PVector<i16>),
39    /// I32
40    I32(PVector<i32>),
41    /// I64
42    I64(PVector<i64>),
43    /// F16
44    F16(PVector<f16>),
45    /// F32
46    F32(PVector<f32>),
47    /// F64
48    F64(PVector<f64>),
49}
50
51impl PrimitiveVector {
52    /// Returns the [`PType`] of this [`PrimitiveVector`].
53    pub fn ptype(&self) -> PType {
54        match self {
55            Self::U8(_) => PType::U8,
56            Self::U16(_) => PType::U16,
57            Self::U32(_) => PType::U32,
58            Self::U64(_) => PType::U64,
59            Self::I8(_) => PType::I8,
60            Self::I16(_) => PType::I16,
61            Self::I32(_) => PType::I32,
62            Self::I64(_) => PType::I64,
63            Self::F16(_) => PType::F16,
64            Self::F32(_) => PType::F32,
65            Self::F64(_) => PType::F64,
66        }
67    }
68}
69
70impl VectorOps for PrimitiveVector {
71    type Mutable = PrimitiveVectorMut;
72
73    fn len(&self) -> usize {
74        match_each_pvector!(self, |v| { v.len() })
75    }
76
77    fn validity(&self) -> &Mask {
78        match_each_pvector!(self, |v| { v.validity() })
79    }
80
81    fn scalar_at(&self, index: usize) -> Scalar {
82        match_each_pvector!(self, |v| { v.scalar_at(index) })
83    }
84
85    fn slice(&self, range: impl RangeBounds<usize> + Clone + Debug) -> Self {
86        match_each_pvector!(self, |v| { v.slice(range).into() })
87    }
88
89    fn try_into_mut(self) -> Result<PrimitiveVectorMut, Self> {
90        match_each_pvector!(self, |v| {
91            v.try_into_mut()
92                .map(PrimitiveVectorMut::from)
93                .map_err(Self::from)
94        })
95    }
96
97    fn into_mut(self) -> PrimitiveVectorMut {
98        match_each_pvector!(self, |v| { v.into_mut().into() })
99    }
100}
101
102impl PTypeUpcast for PrimitiveVector {
103    type Input<T: NativePType> = PVector<T>;
104
105    fn from_u8(input: Self::Input<u8>) -> Self {
106        Self::U8(input)
107    }
108
109    fn from_u16(input: Self::Input<u16>) -> Self {
110        Self::U16(input)
111    }
112
113    fn from_u32(input: Self::Input<u32>) -> Self {
114        Self::U32(input)
115    }
116
117    fn from_u64(input: Self::Input<u64>) -> Self {
118        Self::U64(input)
119    }
120
121    fn from_i8(input: Self::Input<i8>) -> Self {
122        Self::I8(input)
123    }
124
125    fn from_i16(input: Self::Input<i16>) -> Self {
126        Self::I16(input)
127    }
128
129    fn from_i32(input: Self::Input<i32>) -> Self {
130        Self::I32(input)
131    }
132
133    fn from_i64(input: Self::Input<i64>) -> Self {
134        Self::I64(input)
135    }
136
137    fn from_f16(input: Self::Input<f16>) -> Self {
138        Self::F16(input)
139    }
140
141    fn from_f32(input: Self::Input<f32>) -> Self {
142        Self::F32(input)
143    }
144
145    fn from_f64(input: Self::Input<f64>) -> Self {
146        Self::F64(input)
147    }
148}
149
150impl PTypeDowncast for PrimitiveVector {
151    type Output<T: NativePType> = PVector<T>;
152
153    fn into_u8(self) -> Self::Output<u8> {
154        if let Self::U8(v) = self {
155            return v;
156        }
157        vortex_panic!("Expected PrimitiveVector::U8, got {self:?}");
158    }
159
160    fn into_u16(self) -> Self::Output<u16> {
161        if let Self::U16(v) = self {
162            return v;
163        }
164        vortex_panic!("Expected PrimitiveVector::U16, got {self:?}");
165    }
166
167    fn into_u32(self) -> Self::Output<u32> {
168        if let Self::U32(v) = self {
169            return v;
170        }
171        vortex_panic!("Expected PrimitiveVector::U32, got {self:?}");
172    }
173
174    fn into_u64(self) -> Self::Output<u64> {
175        if let Self::U64(v) = self {
176            return v;
177        }
178        vortex_panic!("Expected PrimitiveVector::U64, got {self:?}");
179    }
180
181    fn into_i8(self) -> Self::Output<i8> {
182        if let Self::I8(v) = self {
183            return v;
184        }
185        vortex_panic!("Expected PrimitiveVector::I8, got {self:?}");
186    }
187
188    fn into_i16(self) -> Self::Output<i16> {
189        if let Self::I16(v) = self {
190            return v;
191        }
192        vortex_panic!("Expected PrimitiveVector::I16, got {self:?}");
193    }
194
195    fn into_i32(self) -> Self::Output<i32> {
196        if let Self::I32(v) = self {
197            return v;
198        }
199        vortex_panic!("Expected PrimitiveVector::I32, got {self:?}");
200    }
201
202    fn into_i64(self) -> Self::Output<i64> {
203        if let Self::I64(v) = self {
204            return v;
205        }
206        vortex_panic!("Expected PrimitiveVector::I64, got {self:?}");
207    }
208
209    fn into_f16(self) -> Self::Output<f16> {
210        if let Self::F16(v) = self {
211            return v;
212        }
213        vortex_panic!("Expected PrimitiveVector::F16, got {self:?}");
214    }
215
216    fn into_f32(self) -> Self::Output<f32> {
217        if let Self::F32(v) = self {
218            return v;
219        }
220        vortex_panic!("Expected PrimitiveVector::F32, got {self:?}");
221    }
222
223    fn into_f64(self) -> Self::Output<f64> {
224        if let Self::F64(v) = self {
225            return v;
226        }
227        vortex_panic!("Expected PrimitiveVector::F64, got {self:?}");
228    }
229}
230
231impl<'a> PTypeDowncast for &'a PrimitiveVector {
232    type Output<T: NativePType> = &'a PVector<T>;
233
234    fn into_u8(self) -> Self::Output<u8> {
235        if let PrimitiveVector::U8(v) = self {
236            return v;
237        }
238        vortex_panic!("Expected PrimitiveVector::U8, got {self:?}");
239    }
240
241    fn into_u16(self) -> Self::Output<u16> {
242        if let PrimitiveVector::U16(v) = self {
243            return v;
244        }
245        vortex_panic!("Expected PrimitiveVector::U16, got {self:?}");
246    }
247
248    fn into_u32(self) -> Self::Output<u32> {
249        if let PrimitiveVector::U32(v) = self {
250            return v;
251        }
252        vortex_panic!("Expected PrimitiveVector::U32, got {self:?}");
253    }
254
255    fn into_u64(self) -> Self::Output<u64> {
256        if let PrimitiveVector::U64(v) = self {
257            return v;
258        }
259        vortex_panic!("Expected PrimitiveVector::U64, got {self:?}");
260    }
261
262    fn into_i8(self) -> Self::Output<i8> {
263        if let PrimitiveVector::I8(v) = self {
264            return v;
265        }
266        vortex_panic!("Expected PrimitiveVector::I8, got {self:?}");
267    }
268
269    fn into_i16(self) -> Self::Output<i16> {
270        if let PrimitiveVector::I16(v) = self {
271            return v;
272        }
273        vortex_panic!("Expected PrimitiveVector::I16, got {self:?}");
274    }
275
276    fn into_i32(self) -> Self::Output<i32> {
277        if let PrimitiveVector::I32(v) = self {
278            return v;
279        }
280        vortex_panic!("Expected PrimitiveVector::I32, got {self:?}");
281    }
282
283    fn into_i64(self) -> Self::Output<i64> {
284        if let PrimitiveVector::I64(v) = self {
285            return v;
286        }
287        vortex_panic!("Expected PrimitiveVector::I64, got {self:?}");
288    }
289
290    fn into_f16(self) -> Self::Output<f16> {
291        if let PrimitiveVector::F16(v) = self {
292            return v;
293        }
294        vortex_panic!("Expected PrimitiveVector::F16, got {self:?}");
295    }
296
297    fn into_f32(self) -> Self::Output<f32> {
298        if let PrimitiveVector::F32(v) = self {
299            return v;
300        }
301        vortex_panic!("Expected PrimitiveVector::F32, got {self:?}");
302    }
303
304    fn into_f64(self) -> Self::Output<f64> {
305        if let PrimitiveVector::F64(v) = self {
306            return v;
307        }
308        vortex_panic!("Expected PrimitiveVector::F64, got {self:?}");
309    }
310}