vortex_array/arrays/primitive/compute/
cast.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_buffer::{Buffer, BufferMut};
5use vortex_dtype::{DType, NativePType, Nullability, match_each_native_ptype};
6use vortex_error::{VortexResult, vortex_bail, vortex_err};
7
8use crate::arrays::PrimitiveVTable;
9use crate::arrays::primitive::PrimitiveArray;
10use crate::compute::{CastKernel, CastKernelAdapter};
11use crate::validity::Validity;
12use crate::vtable::ValidityHelper;
13use crate::{ArrayRef, IntoArray, register_kernel};
14
15impl CastKernel for PrimitiveVTable {
16    fn cast(&self, array: &PrimitiveArray, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
17        let DType::Primitive(new_ptype, new_nullability) = dtype else {
18            return Ok(None);
19        };
20        let (new_ptype, new_nullability) = (*new_ptype, *new_nullability);
21
22        // First, check that the cast is compatible with the source array's validity
23        let new_validity = if array.dtype().nullability() == new_nullability {
24            array.validity().clone()
25        } else if new_nullability == Nullability::Nullable {
26            // from non-nullable to nullable
27            array.validity().clone().into_nullable()
28        } else if new_nullability == Nullability::NonNullable && array.validity().all_valid()? {
29            // from nullable but all valid, to non-nullable
30            Validity::NonNullable
31        } else {
32            vortex_bail!(
33                "invalid cast from nullable to non-nullable, since source array actually contains nulls"
34            );
35        };
36
37        // If the bit width is the same, we can short-circuit and simply update the validity
38        if array.ptype() == new_ptype {
39            return Ok(Some(
40                PrimitiveArray::from_byte_buffer(
41                    array.byte_buffer().clone(),
42                    array.ptype(),
43                    new_validity,
44                )
45                .into_array(),
46            ));
47        }
48
49        // Otherwise, we need to cast the values one-by-one
50        match_each_native_ptype!(new_ptype, |T| {
51            Ok(Some(
52                PrimitiveArray::new(cast::<T>(array)?, new_validity).into_array(),
53            ))
54        })
55    }
56}
57
58register_kernel!(CastKernelAdapter(PrimitiveVTable).lift());
59
60fn cast<T: NativePType>(array: &PrimitiveArray) -> VortexResult<Buffer<T>> {
61    let mut buffer = BufferMut::with_capacity(array.len());
62    match_each_native_ptype!(array.ptype(), |P| {
63        for item in array.as_slice::<P>() {
64            let item = T::from(*item).ok_or_else(
65                || vortex_err!(ComputeError: "Failed to cast {} to {:?}", item, T::PTYPE),
66            )?;
67            // SAFETY: we've pre-allocated the required capacity
68            unsafe { buffer.push_unchecked(item) }
69        }
70    });
71    Ok(buffer.freeze())
72}
73
74#[cfg(test)]
75mod test {
76    use rstest::rstest;
77    use vortex_buffer::buffer;
78    use vortex_dtype::{DType, Nullability, PType};
79    use vortex_error::VortexError;
80
81    use crate::IntoArray;
82    use crate::arrays::PrimitiveArray;
83    use crate::canonical::ToCanonical;
84    use crate::compute::cast;
85    use crate::compute::conformance::cast::test_cast_conformance;
86    use crate::validity::Validity;
87    use crate::vtable::ValidityHelper;
88
89    #[test]
90    fn cast_u32_u8() {
91        let arr = buffer![0u32, 10, 200].into_array();
92
93        // cast from u32 to u8
94        let p = cast(&arr, PType::U8.into())
95            .unwrap()
96            .to_primitive()
97            .unwrap();
98        assert_eq!(p.as_slice::<u8>(), vec![0u8, 10, 200]);
99        assert_eq!(p.validity(), &Validity::NonNullable);
100
101        // to nullable
102        let p = cast(
103            p.as_ref(),
104            &DType::Primitive(PType::U8, Nullability::Nullable),
105        )
106        .unwrap()
107        .to_primitive()
108        .unwrap();
109        assert_eq!(p.as_slice::<u8>(), vec![0u8, 10, 200]);
110        assert_eq!(p.validity(), &Validity::AllValid);
111
112        // back to non-nullable
113        let p = cast(
114            p.as_ref(),
115            &DType::Primitive(PType::U8, Nullability::NonNullable),
116        )
117        .unwrap()
118        .to_primitive()
119        .unwrap();
120        assert_eq!(p.as_slice::<u8>(), vec![0u8, 10, 200]);
121        assert_eq!(p.validity(), &Validity::NonNullable);
122
123        // to nullable u32
124        let p = cast(
125            p.as_ref(),
126            &DType::Primitive(PType::U32, Nullability::Nullable),
127        )
128        .unwrap()
129        .to_primitive()
130        .unwrap();
131        assert_eq!(p.as_slice::<u32>(), vec![0u32, 10, 200]);
132        assert_eq!(p.validity(), &Validity::AllValid);
133
134        // to non-nullable u8
135        let p = cast(
136            p.as_ref(),
137            &DType::Primitive(PType::U8, Nullability::NonNullable),
138        )
139        .unwrap()
140        .to_primitive()
141        .unwrap();
142        assert_eq!(p.as_slice::<u8>(), vec![0u8, 10, 200]);
143        assert_eq!(p.validity(), &Validity::NonNullable);
144    }
145
146    #[test]
147    fn cast_u32_f32() {
148        let arr = buffer![0u32, 10, 200].into_array();
149        let u8arr = cast(&arr, PType::F32.into())
150            .unwrap()
151            .to_primitive()
152            .unwrap();
153        assert_eq!(u8arr.as_slice::<f32>(), vec![0.0f32, 10., 200.]);
154    }
155
156    #[test]
157    fn cast_i32_u32() {
158        let arr = buffer![-1i32].into_array();
159        let error = cast(&arr, PType::U32.into()).err().unwrap();
160        let VortexError::ComputeError(s, _) = error else {
161            unreachable!()
162        };
163        assert_eq!(s.to_string(), "Failed to cast -1 to U32");
164    }
165
166    #[test]
167    fn cast_array_with_nulls_to_nonnullable() {
168        let arr = PrimitiveArray::from_option_iter([Some(-1i32), None, Some(10)]);
169        let err = cast(arr.as_ref(), PType::I32.into()).unwrap_err();
170        let VortexError::InvalidArgument(s, _) = err else {
171            unreachable!()
172        };
173        assert_eq!(
174            s.to_string(),
175            "invalid cast from nullable to non-nullable, since source array actually contains nulls"
176        );
177    }
178
179    #[rstest]
180    #[case(buffer![0u8, 1, 2, 3, 255].into_array())]
181    #[case(buffer![0u16, 100, 1000, 65535].into_array())]
182    #[case(buffer![0u32, 100, 1000, 1000000].into_array())]
183    #[case(buffer![0u64, 100, 1000, 1000000000].into_array())]
184    #[case(buffer![-128i8, -1, 0, 1, 127].into_array())]
185    #[case(buffer![-1000i16, -1, 0, 1, 1000].into_array())]
186    #[case(buffer![-1000000i32, -1, 0, 1, 1000000].into_array())]
187    #[case(buffer![-1000000000i64, -1, 0, 1, 1000000000].into_array())]
188    #[case(buffer![0.0f32, 1.5, -2.5, 100.0, 1e6].into_array())]
189    #[case(buffer![0.0f64, 1.5, -2.5, 100.0, 1e12].into_array())]
190    #[case(PrimitiveArray::from_option_iter([Some(1u8), None, Some(255), Some(0), None]).into_array())]
191    #[case(PrimitiveArray::from_option_iter([Some(1i32), None, Some(-100), Some(0), None]).into_array())]
192    #[case(buffer![42u32].into_array())]
193    fn test_cast_primitive_conformance(#[case] array: crate::ArrayRef) {
194        test_cast_conformance(array.as_ref());
195    }
196}
vortex_array/arrays/primitive/compute/cast.rs

vortex_array/arrays/primitive/compute/
cast.rs
