Skip to main content

vortex_array/scalar/convert/
primitive.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Conversions for [`PrimitiveScalar`]s.
5
6use vortex_error::VortexError;
7use vortex_error::VortexResult;
8use vortex_error::vortex_err;
9
10use crate::dtype::DType;
11use crate::dtype::NativePType;
12use crate::dtype::Nullability;
13use crate::dtype::half::f16;
14use crate::scalar::PValue;
15use crate::scalar::PrimitiveScalar;
16use crate::scalar::Scalar;
17use crate::scalar::ScalarValue;
18
19// TODO(connor): Ideally we remove this.
20impl From<PrimitiveScalar<'_>> for Scalar {
21    fn from(ps: PrimitiveScalar<'_>) -> Self {
22        // SAFETY: `PrimitiveScalar` is already a valid `Scalar`.
23        unsafe {
24            Scalar::new_unchecked(ps.dtype().clone(), ps.pvalue().map(ScalarValue::Primitive))
25        }
26    }
27}
28
29impl From<PValue> for ScalarValue {
30    fn from(value: PValue) -> Self {
31        ScalarValue::Primitive(value)
32    }
33}
34
35/// Implements scalar conversions for a primitive type.
36macro_rules! primitive_scalar {
37    ($T:ty) => {
38        ////////////////////////////////////////////////////////////////////////////////////////////
39        // `Scalar` INTO $T conversions.
40        ////////////////////////////////////////////////////////////////////////////////////////////
41
42        /// Fallible conversion from a [`ScalarValue`] into an `T`.
43        ///
44        /// # Errors
45        ///
46        /// Returns an error if unable to convert the scalar value into the target type,
47        impl TryFrom<&ScalarValue> for $T {
48            type Error = VortexError;
49
50            fn try_from(value: &ScalarValue) -> VortexResult<Self> {
51                value.as_primitive().cast::<$T>()
52            }
53        }
54
55        /// Fallible conversion from a [`Scalar`] into an `T`.
56        ///
57        /// # Errors
58        ///
59        /// Returns an error if unable to convert the scalar into the target type, or if the
60        /// `Scalar` itself is null.
61        impl TryFrom<&Scalar> for $T {
62            type Error = VortexError;
63
64            fn try_from(value: &Scalar) -> VortexResult<Self> {
65                match value.value() {
66                    Some(ScalarValue::Primitive(pv)) => pv.cast::<$T>(),
67                    Some(_) => Err(vortex_err!(
68                        "Expected primitive scalar, found {}",
69                        value.dtype()
70                    )),
71                    None => Err(vortex_err!("Can't extract present value from null scalar")),
72                }
73            }
74        }
75
76        /// Fallible conversion from a [`Scalar`] into an `Option<T>`.
77        ///
78        /// # Errors
79        ///
80        /// Returns an error if the [`Scalar`] is not primitive, or if it is unable to convert the
81        /// [`Scalar`] into the target type.
82        impl TryFrom<&Scalar> for Option<$T> {
83            type Error = VortexError;
84
85            fn try_from(value: &Scalar) -> VortexResult<Self> {
86                let primitive_scalar = PrimitiveScalar::try_new(value.dtype(), value.value())?;
87                primitive_scalar.try_typed_value::<$T>()
88            }
89        }
90
91        ////////////////////////////////////////////////////////////////////////////////////////////
92        // `Scalar` FROM $T conversions.
93        ////////////////////////////////////////////////////////////////////////////////////////////
94
95        /// `Into<ScalarValue>` for T.
96        impl From<$T> for ScalarValue {
97            fn from(value: $T) -> Self {
98                ScalarValue::Primitive(value.into())
99            }
100        }
101
102        /// Non-nullable `Into<Scalar>` implementation for T.
103        impl From<$T> for Scalar {
104            fn from(value: $T) -> Self {
105                // SAFETY: The `PValue` made from a `$T` has the ptype of `$T`.
106                unsafe {
107                    Scalar::new_unchecked(
108                        DType::Primitive(<$T>::PTYPE, Nullability::NonNullable),
109                        Some(ScalarValue::Primitive(value.into())),
110                    )
111                }
112            }
113        }
114
115        /// Nullable `Into<Scalar>` implementation for T.
116        impl From<Option<$T>> for Scalar {
117            fn from(value: Option<$T>) -> Self {
118                // SAFETY: The `PValue` made from a `$T` has the ptype of `$T`, and the dtype is
119                // nullable.
120                unsafe {
121                    Scalar::new_unchecked(
122                        DType::Primitive(<$T>::PTYPE, Nullability::Nullable),
123                        value.map(|value| ScalarValue::Primitive(value.into())),
124                    )
125                }
126            }
127        }
128    };
129}
130
131primitive_scalar!(u8);
132primitive_scalar!(u16);
133primitive_scalar!(u32);
134primitive_scalar!(u64);
135
136primitive_scalar!(i8);
137primitive_scalar!(i16);
138primitive_scalar!(i32);
139primitive_scalar!(i64);
140
141primitive_scalar!(f16);
142primitive_scalar!(f32);
143primitive_scalar!(f64);
144
145////////////////////////////////////////////////////////////////////////////////////////////
146// `Scalar` <---> `usize` conversions.
147////////////////////////////////////////////////////////////////////////////////////////////
148
149// NB: We cast `usize` to `u64` (which should always succeed) for better ergonomics.
150
151/// Fallible conversion from a [`ScalarValue`] into an `T`.
152///
153/// # Errors
154///
155/// Returns an error if unable to convert the scalar value into the target type,
156impl TryFrom<&ScalarValue> for usize {
157    type Error = VortexError;
158
159    fn try_from(value: &ScalarValue) -> VortexResult<Self> {
160        let val = value.as_primitive().cast::<u64>()?;
161        Ok(usize::try_from(val)?)
162    }
163}
164
165impl TryFrom<&Scalar> for usize {
166    type Error = VortexError;
167
168    fn try_from(value: &Scalar) -> Result<Self, Self::Error> {
169        let prim_scalar = value
170            .as_primitive_opt()
171            .ok_or_else(|| vortex_err!("Expected primitive scalar, found {}", value.dtype()))?;
172
173        let prim = prim_scalar
174            .as_::<u64>()
175            .ok_or_else(|| vortex_err!("cannot convert Null to usize"))?;
176
177        Ok(usize::try_from(prim)?)
178    }
179}
180
181impl TryFrom<&Scalar> for Option<usize> {
182    type Error = VortexError;
183
184    fn try_from(value: &Scalar) -> Result<Self, Self::Error> {
185        let prim_scalar = value
186            .as_primitive_opt()
187            .ok_or_else(|| vortex_err!("Expected primitive scalar, found {}", value.dtype()))?;
188
189        Ok(prim_scalar.as_::<u64>().map(usize::try_from).transpose()?)
190    }
191}
192
193impl From<usize> for ScalarValue {
194    fn from(value: usize) -> Self {
195        ScalarValue::Primitive((value as u64).into())
196    }
197}
198
199impl From<usize> for Scalar {
200    fn from(value: usize) -> Self {
201        Scalar::from(value as u64)
202    }
203}
204
205impl From<Option<usize>> for Scalar {
206    fn from(value: Option<usize>) -> Self {
207        Scalar::from(value.map(|value| value as u64))
208    }
209}