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;
5use vortex_buffer::BufferMut;
6use vortex_error::VortexResult;
7use vortex_error::vortex_err;
8use vortex_mask::AllOr;
9use vortex_mask::Mask;
10
11use crate::ArrayRef;
12use crate::ExecutionCtx;
13use crate::IntoArray;
14use crate::arrays::Primitive;
15use crate::arrays::PrimitiveArray;
16use crate::dtype::DType;
17use crate::dtype::NativePType;
18use crate::match_each_native_ptype;
19use crate::scalar_fn::fns::cast::CastKernel;
20use crate::vtable::ValidityHelper;
21
22impl CastKernel for Primitive {
23    fn cast(
24        array: &PrimitiveArray,
25        dtype: &DType,
26        _ctx: &mut ExecutionCtx,
27    ) -> VortexResult<Option<ArrayRef>> {
28        let DType::Primitive(new_ptype, new_nullability) = dtype else {
29            return Ok(None);
30        };
31        let (new_ptype, new_nullability) = (*new_ptype, *new_nullability);
32
33        // First, check that the cast is compatible with the source array's validity
34        let new_validity = array
35            .validity()
36            .clone()
37            .cast_nullability(new_nullability, array.len())?;
38
39        // If the bit width is the same, we can short-circuit and simply update the validity
40        if array.ptype() == new_ptype {
41            // SAFETY: validity and data buffer still have same length
42            return Ok(Some(unsafe {
43                PrimitiveArray::new_unchecked_from_handle(
44                    array.buffer_handle().clone(),
45                    array.ptype(),
46                    new_validity,
47                )
48                .into_array()
49            }));
50        }
51
52        let mask = array.validity_mask()?;
53
54        // Otherwise, we need to cast the values one-by-one
55        Ok(Some(match_each_native_ptype!(new_ptype, |T| {
56            match_each_native_ptype!(array.ptype(), |F| {
57                PrimitiveArray::new(cast::<F, T>(array.as_slice(), mask)?, new_validity)
58                    .into_array()
59            })
60        })))
61    }
62}
63
64fn cast<F: NativePType, T: NativePType>(array: &[F], mask: Mask) -> VortexResult<Buffer<T>> {
65    match mask.bit_buffer() {
66        AllOr::All => {
67            let mut buffer = BufferMut::with_capacity(array.len());
68            for item in array {
69                let item = T::from(*item).ok_or_else(
70                    || vortex_err!(Compute: "Failed to cast {} to {:?}", item, T::PTYPE),
71                )?;
72                // SAFETY: we've pre-allocated the required capacity
73                unsafe { buffer.push_unchecked(item) }
74            }
75            Ok(buffer.freeze())
76        }
77        AllOr::None => Ok(Buffer::zeroed(array.len())),
78        AllOr::Some(b) => {
79            // TODO(robert): Depending on density of the buffer might be better to prefill Buffer and only write valid values
80            let mut buffer = BufferMut::with_capacity(array.len());
81            for (item, valid) in array.iter().zip(b.iter()) {
82                if valid {
83                    let item = T::from(*item).ok_or_else(
84                        || vortex_err!(Compute: "Failed to cast {} to {:?}", item, T::PTYPE),
85                    )?;
86                    // SAFETY: we've pre-allocated the required capacity
87                    unsafe { buffer.push_unchecked(item) }
88                } else {
89                    // SAFETY: we've pre-allocated the required capacity
90                    unsafe { buffer.push_unchecked(T::default()) }
91                }
92            }
93            Ok(buffer.freeze())
94        }
95    }
96}
97
98#[cfg(test)]
99mod test {
100    use rstest::rstest;
101    use vortex_buffer::BitBuffer;
102    use vortex_buffer::buffer;
103    use vortex_error::VortexError;
104    use vortex_mask::Mask;
105
106    use crate::IntoArray;
107    use crate::arrays::PrimitiveArray;
108    use crate::assert_arrays_eq;
109    use crate::builtins::ArrayBuiltins;
110    use crate::canonical::ToCanonical;
111    use crate::compute::conformance::cast::test_cast_conformance;
112    use crate::dtype::DType;
113    use crate::dtype::Nullability;
114    use crate::dtype::PType;
115    use crate::validity::Validity;
116    use crate::vtable::ValidityHelper;
117
118    #[allow(clippy::cognitive_complexity)]
119    #[test]
120    fn cast_u32_u8() {
121        let arr = buffer![0u32, 10, 200].into_array();
122
123        // cast from u32 to u8
124        let p = arr.cast(PType::U8.into()).unwrap().to_primitive();
125        assert_arrays_eq!(p, PrimitiveArray::from_iter([0u8, 10, 200]));
126        assert!(matches!(p.validity(), Validity::NonNullable));
127
128        // to nullable
129        let p = p
130            .into_array()
131            .cast(DType::Primitive(PType::U8, Nullability::Nullable))
132            .unwrap()
133            .to_primitive();
134        assert_arrays_eq!(
135            p,
136            PrimitiveArray::new(buffer![0u8, 10, 200], Validity::AllValid)
137        );
138        assert!(matches!(p.validity(), Validity::AllValid));
139
140        // back to non-nullable
141        let p = p
142            .into_array()
143            .cast(DType::Primitive(PType::U8, Nullability::NonNullable))
144            .unwrap()
145            .to_primitive();
146        assert_arrays_eq!(p, PrimitiveArray::from_iter([0u8, 10, 200]));
147        assert!(matches!(p.validity(), Validity::NonNullable));
148
149        // to nullable u32
150        let p = p
151            .into_array()
152            .cast(DType::Primitive(PType::U32, Nullability::Nullable))
153            .unwrap()
154            .to_primitive();
155        assert_arrays_eq!(
156            p,
157            PrimitiveArray::new(buffer![0u32, 10, 200], Validity::AllValid)
158        );
159        assert!(matches!(p.validity(), Validity::AllValid));
160
161        // to non-nullable u8
162        let p = p
163            .into_array()
164            .cast(DType::Primitive(PType::U8, Nullability::NonNullable))
165            .unwrap()
166            .to_primitive();
167        assert_arrays_eq!(p, PrimitiveArray::from_iter([0u8, 10, 200]));
168        assert!(matches!(p.validity(), Validity::NonNullable));
169    }
170
171    #[test]
172    fn cast_u32_f32() {
173        let arr = buffer![0u32, 10, 200].into_array();
174        let u8arr = arr.cast(PType::F32.into()).unwrap().to_primitive();
175        assert_arrays_eq!(u8arr, PrimitiveArray::from_iter([0.0f32, 10., 200.]));
176    }
177
178    #[test]
179    fn cast_i32_u32() {
180        let arr = buffer![-1i32].into_array();
181        let error = arr
182            .cast(PType::U32.into())
183            .and_then(|a| a.to_canonical().map(|c| c.into_array()))
184            .unwrap_err();
185        assert!(matches!(error, VortexError::Compute(..)));
186        assert!(error.to_string().contains("Failed to cast -1 to U32"));
187    }
188
189    #[test]
190    fn cast_array_with_nulls_to_nonnullable() {
191        let arr = PrimitiveArray::from_option_iter([Some(-1i32), None, Some(10)]);
192        let err = arr
193            .into_array()
194            .cast(PType::I32.into())
195            .and_then(|a| a.to_canonical().map(|c| c.into_array()))
196            .unwrap_err();
197
198        assert!(matches!(err, VortexError::InvalidArgument(..)));
199        assert!(
200            err.to_string()
201                .contains("Cannot cast array with invalid values to non-nullable type.")
202        );
203    }
204
205    #[test]
206    fn cast_with_invalid_nulls() {
207        let arr = PrimitiveArray::new(
208            buffer![-1i32, 0, 10],
209            Validity::from_iter([false, true, true]),
210        );
211        let p = arr
212            .into_array()
213            .cast(DType::Primitive(PType::U32, Nullability::Nullable))
214            .unwrap()
215            .to_primitive();
216        assert_arrays_eq!(
217            p,
218            PrimitiveArray::from_option_iter([None, Some(0u32), Some(10)])
219        );
220        assert_eq!(
221            p.validity_mask().unwrap(),
222            Mask::from(BitBuffer::from(vec![false, true, true]))
223        );
224    }
225
226    #[rstest]
227    #[case(buffer![0u8, 1, 2, 3, 255].into_array())]
228    #[case(buffer![0u16, 100, 1000, 65535].into_array())]
229    #[case(buffer![0u32, 100, 1000, 1000000].into_array())]
230    #[case(buffer![0u64, 100, 1000, 1000000000].into_array())]
231    #[case(buffer![-128i8, -1, 0, 1, 127].into_array())]
232    #[case(buffer![-1000i16, -1, 0, 1, 1000].into_array())]
233    #[case(buffer![-1000000i32, -1, 0, 1, 1000000].into_array())]
234    #[case(buffer![-1000000000i64, -1, 0, 1, 1000000000].into_array())]
235    #[case(buffer![0.0f32, 1.5, -2.5, 100.0, 1e6].into_array())]
236    #[case(buffer![0.0f64, 1.5, -2.5, 100.0, 1e12].into_array())]
237    #[case(PrimitiveArray::from_option_iter([Some(1u8), None, Some(255), Some(0), None]).into_array())]
238    #[case(PrimitiveArray::from_option_iter([Some(1i32), None, Some(-100), Some(0), None]).into_array())]
239    #[case(buffer![42u32].into_array())]
240    fn test_cast_primitive_conformance(#[case] array: crate::ArrayRef) {
241        test_cast_conformance(&array);
242    }
243}
vortex_array/arrays/primitive/compute/cast.rs

vortex_array/arrays/primitive/compute/
cast.rs
