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::PTypeUpcast;
9use vortex_dtype::half::f16;
10use vortex_error::VortexExpect;
11
12use crate::Scalar;
13use crate::ScalarOps;
14use crate::VectorMut;
15use crate::VectorMutOps;
16use crate::match_each_pscalar;
17use crate::primitive::PVectorMut;
18use crate::primitive::PrimitiveVectorMut;
19
20/// Represents a primitive scalar value.
21#[derive(Clone, Debug)]
22pub enum PrimitiveScalar {
23    /// 8-bit signed integer scalar
24    I8(PScalar<i8>),
25    /// 16-bit signed integer scalar
26    I16(PScalar<i16>),
27    /// 32-bit signed integer scalar
28    I32(PScalar<i32>),
29    /// 64-bit signed integer scalar
30    I64(PScalar<i64>),
31    /// 8-bit unsigned integer scalar
32    U8(PScalar<u8>),
33    /// 16-bit unsigned integer scalar
34    U16(PScalar<u16>),
35    /// 32-bit unsigned integer scalar
36    U32(PScalar<u32>),
37    /// 64-bit unsigned integer scalar
38    U64(PScalar<u64>),
39    /// 16-bit floating point scalar
40    F16(PScalar<f16>),
41    /// 32-bit floating point scalar
42    F32(PScalar<f32>),
43    /// 64-bit floating point scalar
44    F64(PScalar<f64>),
45}
46
47impl PrimitiveScalar {
48    /// Returns the primitive type of the scalar.
49    pub fn ptype(&self) -> PType {
50        match self {
51            PrimitiveScalar::I8(_) => PType::I8,
52            PrimitiveScalar::I16(_) => PType::I16,
53            PrimitiveScalar::I32(_) => PType::I32,
54            PrimitiveScalar::I64(_) => PType::I64,
55            PrimitiveScalar::U8(_) => PType::U8,
56            PrimitiveScalar::U16(_) => PType::U16,
57            PrimitiveScalar::U32(_) => PType::U32,
58            PrimitiveScalar::U64(_) => PType::U64,
59            PrimitiveScalar::F16(_) => PType::F16,
60            PrimitiveScalar::F32(_) => PType::F32,
61            PrimitiveScalar::F64(_) => PType::F64,
62        }
63    }
64}
65
66impl ScalarOps for PrimitiveScalar {
67    fn is_valid(&self) -> bool {
68        match self {
69            PrimitiveScalar::I8(v) => v.is_some(),
70            PrimitiveScalar::I16(v) => v.is_some(),
71            PrimitiveScalar::I32(v) => v.is_some(),
72            PrimitiveScalar::I64(v) => v.is_some(),
73            PrimitiveScalar::U8(v) => v.is_some(),
74            PrimitiveScalar::U16(v) => v.is_some(),
75            PrimitiveScalar::U32(v) => v.is_some(),
76            PrimitiveScalar::U64(v) => v.is_some(),
77            PrimitiveScalar::F16(v) => v.is_some(),
78            PrimitiveScalar::F32(v) => v.is_some(),
79            PrimitiveScalar::F64(v) => v.is_some(),
80        }
81    }
82
83    fn mask_validity(&mut self, mask: bool) {
84        match_each_pscalar!(self, |s| { s.mask_validity(mask) })
85    }
86
87    fn repeat(&self, n: usize) -> VectorMut {
88        match_each_pscalar!(self, |s| { s.repeat(n) })
89    }
90}
91
92impl From<PrimitiveScalar> for Scalar {
93    fn from(val: PrimitiveScalar) -> Self {
94        Scalar::Primitive(val)
95    }
96}
97
98/// Represents a primitive scalar value with a specific native primitive type.
99#[derive(Clone, Debug, PartialEq, Eq, Hash)]
100pub struct PScalar<T>(Option<T>);
101
102impl<T: NativePType> PScalar<T> {
103    /// Creates a new primitive scalar with the given value.
104    pub fn new(value: Option<T>) -> Self {
105        Self(value)
106    }
107
108    /// Returns the value of the primitive scalar, or `None` if it is null.
109    pub fn value(&self) -> Option<T> {
110        self.0
111    }
112}
113
114impl<T: NativePType> From<PScalar<T>> for PrimitiveScalar {
115    fn from(value: PScalar<T>) -> Self {
116        T::upcast(value)
117    }
118}
119
120impl<T: NativePType> ScalarOps for PScalar<T> {
121    fn is_valid(&self) -> bool {
122        self.0.is_some()
123    }
124
125    fn mask_validity(&mut self, mask: bool) {
126        if !mask {
127            self.0 = None;
128        }
129    }
130
131    fn repeat(&self, n: usize) -> VectorMut {
132        let mut vec = PVectorMut::<T>::with_capacity(n);
133        match self.0 {
134            None => vec.append_nulls(n),
135            Some(v) => vec.append_values(v, n),
136        }
137        PrimitiveVectorMut::from(vec).into()
138    }
139}
140
141impl<T: NativePType> From<PScalar<T>> for Scalar {
142    fn from(value: PScalar<T>) -> Self {
143        Scalar::Primitive(T::upcast(value))
144    }
145}
146
147impl PTypeUpcast for PrimitiveScalar {
148    type Input<T: NativePType> = PScalar<T>;
149
150    fn from_u8(input: Self::Input<u8>) -> Self {
151        PrimitiveScalar::U8(input)
152    }
153
154    fn from_u16(input: Self::Input<u16>) -> Self {
155        PrimitiveScalar::U16(input)
156    }
157
158    fn from_u32(input: Self::Input<u32>) -> Self {
159        PrimitiveScalar::U32(input)
160    }
161
162    fn from_u64(input: Self::Input<u64>) -> Self {
163        PrimitiveScalar::U64(input)
164    }
165
166    fn from_i8(input: Self::Input<i8>) -> Self {
167        PrimitiveScalar::I8(input)
168    }
169
170    fn from_i16(input: Self::Input<i16>) -> Self {
171        PrimitiveScalar::I16(input)
172    }
173
174    fn from_i32(input: Self::Input<i32>) -> Self {
175        PrimitiveScalar::I32(input)
176    }
177
178    fn from_i64(input: Self::Input<i64>) -> Self {
179        PrimitiveScalar::I64(input)
180    }
181
182    fn from_f16(input: Self::Input<f16>) -> Self {
183        PrimitiveScalar::F16(input)
184    }
185
186    fn from_f32(input: Self::Input<f32>) -> Self {
187        PrimitiveScalar::F32(input)
188    }
189
190    fn from_f64(input: Self::Input<f64>) -> Self {
191        PrimitiveScalar::F64(input)
192    }
193}
194
195impl<T: NativePType> Deref for PScalar<T> {
196    type Target = Option<T>;
197
198    fn deref(&self) -> &Self::Target {
199        &self.0
200    }
201}
202
203impl PrimitiveScalar {
204    /// Returns the scalar value as `usize` if possible.
205    ///
206    /// Returns `None` if the scalar cannot be cast to a usize.
207    ///
208    /// # Panics
209    ///
210    /// If the scalar is null.
211    pub fn to_usize(&self) -> Option<usize> {
212        match self {
213            PrimitiveScalar::I8(v) => usize::try_from(v.vortex_expect("null scalar")).ok(),
214            PrimitiveScalar::I16(v) => usize::try_from(v.vortex_expect("null scalar")).ok(),
215            PrimitiveScalar::I32(v) => usize::try_from(v.vortex_expect("null scalar")).ok(),
216            PrimitiveScalar::I64(v) => usize::try_from(v.vortex_expect("null scalar")).ok(),
217            PrimitiveScalar::U8(v) => Some(v.vortex_expect("null scalar") as usize),
218            PrimitiveScalar::U16(v) => Some(v.vortex_expect("null scalar") as usize),
219            PrimitiveScalar::U32(v) => Some(v.vortex_expect("null scalar") as usize),
220            PrimitiveScalar::U64(v) => usize::try_from(v.vortex_expect("null scalar")).ok(),
221            PrimitiveScalar::F16(_) => None,
222            PrimitiveScalar::F32(_) => None,
223            PrimitiveScalar::F64(_) => None,
224        }
225    }
226}
vortex_vector/primitive/scalar.rs

vortex_vector/primitive/
scalar.rs
