vortex_array/arrays/primitive/compute/
cast.rs

1use vortex_buffer::{Buffer, BufferMut};
2use vortex_dtype::{DType, NativePType, Nullability, match_each_native_ptype};
3use vortex_error::{VortexResult, vortex_bail, vortex_err};
4
5use crate::arrays::PrimitiveEncoding;
6use crate::arrays::primitive::PrimitiveArray;
7use crate::compute::CastFn;
8use crate::validity::Validity;
9use crate::variants::PrimitiveArrayTrait;
10use crate::{Array, ArrayRef};
11
12impl CastFn<&PrimitiveArray> for PrimitiveEncoding {
13    fn cast(&self, array: &PrimitiveArray, dtype: &DType) -> VortexResult<ArrayRef> {
14        let DType::Primitive(new_ptype, new_nullability) = dtype else {
15            vortex_bail!(MismatchedTypes: "primitive type", dtype);
16        };
17        let (new_ptype, new_nullability) = (*new_ptype, *new_nullability);
18
19        // First, check that the cast is compatible with the source array's validity
20        let new_validity = if array.dtype().nullability() == new_nullability {
21            array.validity().clone()
22        } else if new_nullability == Nullability::Nullable {
23            // from non-nullable to nullable
24            array.validity().clone().into_nullable()
25        } else if new_nullability == Nullability::NonNullable && array.validity().all_valid()? {
26            // from nullable but all valid, to non-nullable
27            Validity::NonNullable
28        } else {
29            vortex_bail!(
30                "invalid cast from nullable to non-nullable, since source array actually contains nulls"
31            );
32        };
33
34        // If the bit width is the same, we can short-circuit and simply update the validity
35        if array.ptype() == new_ptype {
36            return Ok(PrimitiveArray::from_byte_buffer(
37                array.byte_buffer().clone(),
38                array.ptype(),
39                new_validity,
40            )
41            .into_array());
42        }
43
44        // Otherwise, we need to cast the values one-by-one
45        match_each_native_ptype!(new_ptype, |$T| {
46            Ok(PrimitiveArray::new(
47                cast::<$T>(array)?,
48                new_validity,
49            ).into_array())
50        })
51    }
52}
53
54fn cast<T: NativePType>(array: &PrimitiveArray) -> VortexResult<Buffer<T>> {
55    let mut buffer = BufferMut::with_capacity(array.len());
56    match_each_native_ptype!(array.ptype(), |$P| {
57        for item in array.as_slice::<$P>() {
58            let item = T::from(*item).ok_or_else(
59                || vortex_err!(ComputeError: "Failed to cast {} to {:?}", item, T::PTYPE),
60            )?;
61            // SAFETY: we've pre-allocated the required capacity
62            unsafe { buffer.push_unchecked(item) }
63        }
64    });
65    Ok(buffer.freeze())
66}
67
68#[cfg(test)]
69mod test {
70    use vortex_buffer::buffer;
71    use vortex_dtype::{DType, Nullability, PType};
72    use vortex_error::VortexError;
73
74    use crate::IntoArray;
75    use crate::arrays::PrimitiveArray;
76    use crate::canonical::ToCanonical;
77    use crate::compute::try_cast;
78    use crate::validity::Validity;
79
80    #[test]
81    fn cast_u32_u8() {
82        let arr = buffer![0u32, 10, 200].into_array();
83
84        // cast from u32 to u8
85        let p = try_cast(&arr, PType::U8.into())
86            .unwrap()
87            .to_primitive()
88            .unwrap();
89        assert_eq!(p.as_slice::<u8>(), vec![0u8, 10, 200]);
90        assert_eq!(p.validity(), &Validity::NonNullable);
91
92        // to nullable
93        let p = try_cast(&p, &DType::Primitive(PType::U8, Nullability::Nullable))
94            .unwrap()
95            .to_primitive()
96            .unwrap();
97        assert_eq!(p.as_slice::<u8>(), vec![0u8, 10, 200]);
98        assert_eq!(p.validity(), &Validity::AllValid);
99
100        // back to non-nullable
101        let p = try_cast(&p, &DType::Primitive(PType::U8, Nullability::NonNullable))
102            .unwrap()
103            .to_primitive()
104            .unwrap();
105        assert_eq!(p.as_slice::<u8>(), vec![0u8, 10, 200]);
106        assert_eq!(p.validity(), &Validity::NonNullable);
107
108        // to nullable u32
109        let p = try_cast(&p, &DType::Primitive(PType::U32, Nullability::Nullable))
110            .unwrap()
111            .to_primitive()
112            .unwrap();
113        assert_eq!(p.as_slice::<u32>(), vec![0u32, 10, 200]);
114        assert_eq!(p.validity(), &Validity::AllValid);
115
116        // to non-nullable u8
117        let p = try_cast(&p, &DType::Primitive(PType::U8, Nullability::NonNullable))
118            .unwrap()
119            .to_primitive()
120            .unwrap();
121        assert_eq!(p.as_slice::<u8>(), vec![0u8, 10, 200]);
122        assert_eq!(p.validity(), &Validity::NonNullable);
123    }
124
125    #[test]
126    fn cast_u32_f32() {
127        let arr = buffer![0u32, 10, 200].into_array();
128        let u8arr = try_cast(&arr, PType::F32.into())
129            .unwrap()
130            .to_primitive()
131            .unwrap();
132        assert_eq!(u8arr.as_slice::<f32>(), vec![0.0f32, 10., 200.]);
133    }
134
135    #[test]
136    fn cast_i32_u32() {
137        let arr = buffer![-1i32].into_array();
138        let error = try_cast(&arr, PType::U32.into()).err().unwrap();
139        let VortexError::ComputeError(s, _) = error else {
140            unreachable!()
141        };
142        assert_eq!(s.to_string(), "Failed to cast -1 to U32");
143    }
144
145    #[test]
146    fn cast_array_with_nulls_to_nonnullable() {
147        let arr = PrimitiveArray::from_option_iter([Some(-1i32), None, Some(10)]);
148        let err = try_cast(&arr, PType::I32.into()).unwrap_err();
149        let VortexError::InvalidArgument(s, _) = err else {
150            unreachable!()
151        };
152        assert_eq!(
153            s.to_string(),
154            "invalid cast from nullable to non-nullable, since source array actually contains nulls"
155        );
156    }
157}
vortex_array/arrays/primitive/compute/cast.rs

vortex_array/arrays/primitive/compute/
cast.rs
