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