vortex_vector/primitive/
scalar.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::Deref;
5
6use vortex_dtype::NativePType;
7use vortex_dtype::PType;
8use vortex_dtype::PTypeDowncast;
9use vortex_dtype::PTypeUpcast;
10use vortex_dtype::half::f16;
11use vortex_dtype::match_each_native_ptype;
12use vortex_error::VortexExpect;
13use vortex_error::vortex_panic;
14
15use crate::Scalar;
16use crate::ScalarOps;
17use crate::VectorMut;
18use crate::VectorMutOps;
19use crate::match_each_pscalar;
20use crate::primitive::PVectorMut;
21use crate::primitive::PrimitiveVectorMut;
22
23/// Represents a primitive scalar value.
24#[derive(Clone, Debug, PartialEq)]
25pub enum PrimitiveScalar {
26    /// 8-bit signed integer scalar
27    I8(PScalar<i8>),
28    /// 16-bit signed integer scalar
29    I16(PScalar<i16>),
30    /// 32-bit signed integer scalar
31    I32(PScalar<i32>),
32    /// 64-bit signed integer scalar
33    I64(PScalar<i64>),
34    /// 8-bit unsigned integer scalar
35    U8(PScalar<u8>),
36    /// 16-bit unsigned integer scalar
37    U16(PScalar<u16>),
38    /// 32-bit unsigned integer scalar
39    U32(PScalar<u32>),
40    /// 64-bit unsigned integer scalar
41    U64(PScalar<u64>),
42    /// 16-bit floating point scalar
43    F16(PScalar<f16>),
44    /// 32-bit floating point scalar
45    F32(PScalar<f32>),
46    /// 64-bit floating point scalar
47    F64(PScalar<f64>),
48}
49
50impl PrimitiveScalar {
51    /// Returns the primitive type of the scalar.
52    pub fn ptype(&self) -> PType {
53        match self {
54            PrimitiveScalar::I8(_) => PType::I8,
55            PrimitiveScalar::I16(_) => PType::I16,
56            PrimitiveScalar::I32(_) => PType::I32,
57            PrimitiveScalar::I64(_) => PType::I64,
58            PrimitiveScalar::U8(_) => PType::U8,
59            PrimitiveScalar::U16(_) => PType::U16,
60            PrimitiveScalar::U32(_) => PType::U32,
61            PrimitiveScalar::U64(_) => PType::U64,
62            PrimitiveScalar::F16(_) => PType::F16,
63            PrimitiveScalar::F32(_) => PType::F32,
64            PrimitiveScalar::F64(_) => PType::F64,
65        }
66    }
67
68    /// Creates a zero primitive scalar of the given [`PType`].
69    pub fn zero(ptype: PType) -> Self {
70        match_each_native_ptype!(ptype, |T| { PScalar::<T>::zero().into() })
71    }
72
73    /// Creates a null primitive scalar of the given [`PType`].
74    pub fn null(ptype: PType) -> Self {
75        match_each_native_ptype!(ptype, |T| { PScalar::<T>::null().into() })
76    }
77}
78
79impl ScalarOps for PrimitiveScalar {
80    fn is_valid(&self) -> bool {
81        match self {
82            PrimitiveScalar::I8(v) => v.is_some(),
83            PrimitiveScalar::I16(v) => v.is_some(),
84            PrimitiveScalar::I32(v) => v.is_some(),
85            PrimitiveScalar::I64(v) => v.is_some(),
86            PrimitiveScalar::U8(v) => v.is_some(),
87            PrimitiveScalar::U16(v) => v.is_some(),
88            PrimitiveScalar::U32(v) => v.is_some(),
89            PrimitiveScalar::U64(v) => v.is_some(),
90            PrimitiveScalar::F16(v) => v.is_some(),
91            PrimitiveScalar::F32(v) => v.is_some(),
92            PrimitiveScalar::F64(v) => v.is_some(),
93        }
94    }
95
96    fn mask_validity(&mut self, mask: bool) {
97        match_each_pscalar!(self, |s| { s.mask_validity(mask) })
98    }
99
100    fn repeat(&self, n: usize) -> VectorMut {
101        match_each_pscalar!(self, |s| { s.repeat(n) })
102    }
103}
104
105impl From<PrimitiveScalar> for Scalar {
106    fn from(val: PrimitiveScalar) -> Self {
107        Scalar::Primitive(val)
108    }
109}
110
111/// Represents a primitive scalar value with a specific native primitive type.
112#[derive(Clone, Debug, PartialEq, Eq, Hash)]
113pub struct PScalar<T>(Option<T>);
114
115impl<T: NativePType> PScalar<T> {
116    /// Creates a new primitive scalar with the given value.
117    pub fn new(value: Option<T>) -> Self {
118        Self(value)
119    }
120
121    /// Returns the value of the primitive scalar, or `None` if it is null.
122    pub fn value(&self) -> Option<T> {
123        self.0
124    }
125
126    /// Creates a zero primitive scalar.
127    pub fn zero() -> Self {
128        Self::new(Some(T::default()))
129    }
130
131    /// Creates a null primitive scalar.
132    pub fn null() -> Self {
133        Self::new(None)
134    }
135}
136
137impl<T: NativePType> From<PScalar<T>> for PrimitiveScalar {
138    fn from(value: PScalar<T>) -> Self {
139        T::upcast(value)
140    }
141}
142
143impl<T: NativePType> ScalarOps for PScalar<T> {
144    fn is_valid(&self) -> bool {
145        self.0.is_some()
146    }
147
148    fn mask_validity(&mut self, mask: bool) {
149        if !mask {
150            self.0 = None;
151        }
152    }
153
154    fn repeat(&self, n: usize) -> VectorMut {
155        let mut vec = PVectorMut::<T>::with_capacity(n);
156        match self.0 {
157            None => vec.append_nulls(n),
158            Some(v) => vec.append_values(v, n),
159        }
160        PrimitiveVectorMut::from(vec).into()
161    }
162}
163
164impl<T: NativePType> From<PScalar<T>> for Scalar {
165    fn from(value: PScalar<T>) -> Self {
166        Scalar::Primitive(T::upcast(value))
167    }
168}
169
170impl PTypeUpcast for PrimitiveScalar {
171    type Input<T: NativePType> = PScalar<T>;
172
173    fn from_u8(input: Self::Input<u8>) -> Self {
174        PrimitiveScalar::U8(input)
175    }
176
177    fn from_u16(input: Self::Input<u16>) -> Self {
178        PrimitiveScalar::U16(input)
179    }
180
181    fn from_u32(input: Self::Input<u32>) -> Self {
182        PrimitiveScalar::U32(input)
183    }
184
185    fn from_u64(input: Self::Input<u64>) -> Self {
186        PrimitiveScalar::U64(input)
187    }
188
189    fn from_i8(input: Self::Input<i8>) -> Self {
190        PrimitiveScalar::I8(input)
191    }
192
193    fn from_i16(input: Self::Input<i16>) -> Self {
194        PrimitiveScalar::I16(input)
195    }
196
197    fn from_i32(input: Self::Input<i32>) -> Self {
198        PrimitiveScalar::I32(input)
199    }
200
201    fn from_i64(input: Self::Input<i64>) -> Self {
202        PrimitiveScalar::I64(input)
203    }
204
205    fn from_f16(input: Self::Input<f16>) -> Self {
206        PrimitiveScalar::F16(input)
207    }
208
209    fn from_f32(input: Self::Input<f32>) -> Self {
210        PrimitiveScalar::F32(input)
211    }
212
213    fn from_f64(input: Self::Input<f64>) -> Self {
214        PrimitiveScalar::F64(input)
215    }
216}
217
218impl PTypeDowncast for PrimitiveScalar {
219    type Output<T: NativePType> = PScalar<T>;
220
221    fn into_u8(self) -> Self::Output<u8> {
222        if let Self::U8(v) = self {
223            return v;
224        }
225        vortex_panic!("Expected PrimitiveScalar::U8, got {self:?}");
226    }
227
228    fn into_u16(self) -> Self::Output<u16> {
229        if let Self::U16(v) = self {
230            return v;
231        }
232        vortex_panic!("Expected PrimitiveScalar::U16, got {self:?}");
233    }
234
235    fn into_u32(self) -> Self::Output<u32> {
236        if let Self::U32(v) = self {
237            return v;
238        }
239        vortex_panic!("Expected PrimitiveScalar::U32, got {self:?}");
240    }
241
242    fn into_u64(self) -> Self::Output<u64> {
243        if let Self::U64(v) = self {
244            return v;
245        }
246        vortex_panic!("Expected PrimitiveScalar::U64, got {self:?}");
247    }
248
249    fn into_i8(self) -> Self::Output<i8> {
250        if let Self::I8(v) = self {
251            return v;
252        }
253        vortex_panic!("Expected PrimitiveScalar::I8, got {self:?}");
254    }
255
256    fn into_i16(self) -> Self::Output<i16> {
257        if let Self::I16(v) = self {
258            return v;
259        }
260        vortex_panic!("Expected PrimitiveScalar::I16, got {self:?}");
261    }
262
263    fn into_i32(self) -> Self::Output<i32> {
264        if let Self::I32(v) = self {
265            return v;
266        }
267        vortex_panic!("Expected PrimitiveScalar::I32, got {self:?}");
268    }
269
270    fn into_i64(self) -> Self::Output<i64> {
271        if let Self::I64(v) = self {
272            return v;
273        }
274        vortex_panic!("Expected PrimitiveScalar::I64, got {self:?}");
275    }
276
277    fn into_f16(self) -> Self::Output<f16> {
278        if let Self::F16(v) = self {
279            return v;
280        }
281        vortex_panic!("Expected PrimitiveScalar::F16, got {self:?}");
282    }
283
284    fn into_f32(self) -> Self::Output<f32> {
285        if let Self::F32(v) = self {
286            return v;
287        }
288        vortex_panic!("Expected PrimitiveScalar::F32, got {self:?}");
289    }
290
291    fn into_f64(self) -> Self::Output<f64> {
292        if let Self::F64(v) = self {
293            return v;
294        }
295        vortex_panic!("Expected PrimitiveScalar::F64, got {self:?}");
296    }
297}
298
299impl<T: NativePType> Deref for PScalar<T> {
300    type Target = Option<T>;
301
302    fn deref(&self) -> &Self::Target {
303        &self.0
304    }
305}
306
307impl PrimitiveScalar {
308    /// Returns the scalar value as `usize` if possible.
309    ///
310    /// Returns `None` if the scalar cannot be cast to a usize.
311    ///
312    /// # Panics
313    ///
314    /// If the scalar is null.
315    pub fn to_usize(&self) -> Option<usize> {
316        match self {
317            PrimitiveScalar::I8(v) => usize::try_from(v.vortex_expect("null scalar")).ok(),
318            PrimitiveScalar::I16(v) => usize::try_from(v.vortex_expect("null scalar")).ok(),
319            PrimitiveScalar::I32(v) => usize::try_from(v.vortex_expect("null scalar")).ok(),
320            PrimitiveScalar::I64(v) => usize::try_from(v.vortex_expect("null scalar")).ok(),
321            PrimitiveScalar::U8(v) => Some(v.vortex_expect("null scalar") as usize),
322            PrimitiveScalar::U16(v) => Some(v.vortex_expect("null scalar") as usize),
323            PrimitiveScalar::U32(v) => Some(v.vortex_expect("null scalar") as usize),
324            PrimitiveScalar::U64(v) => usize::try_from(v.vortex_expect("null scalar")).ok(),
325            PrimitiveScalar::F16(_) => None,
326            PrimitiveScalar::F32(_) => None,
327            PrimitiveScalar::F64(_) => None,
328        }
329    }
330}
vortex_vector/primitive/scalar.rs

vortex_vector/primitive/
scalar.rs
