vortex_array/arrays/primitive/array/
cast.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexExpect;
5use vortex_error::vortex_panic;
6
7use super::PrimitiveData;
8use crate::dtype::NativePType;
9
10impl PrimitiveData {
11    /// Return a slice of the array's buffer.
12    ///
13    /// NOTE: these values may be nonsense if the validity buffer indicates that the value is null.
14    ///
15    /// # Panic
16    ///
17    /// This operation will panic if the array is not backed by host memory.
18    pub fn as_slice<T: NativePType>(&self) -> &[T] {
19        if T::PTYPE != self.ptype() {
20            vortex_panic!(
21                "Attempted to get slice of type {} from array of type {}",
22                T::PTYPE,
23                self.ptype()
24            )
25        }
26
27        let byte_buffer = self
28            .buffer
29            .as_host_opt()
30            .vortex_expect("as_slice must be called on host buffer");
31        let raw_slice = byte_buffer.as_ptr();
32
33        // SAFETY: alignment of Buffer is checked on construction
34        unsafe { std::slice::from_raw_parts(raw_slice.cast(), byte_buffer.len() / size_of::<T>()) }
35    }
36}
37
38#[cfg(test)]
39mod tests {
40    use rstest::rstest;
41    use vortex_buffer::Buffer;
42    use vortex_buffer::buffer;
43
44    use crate::arrays::PrimitiveArray;
45    use crate::arrays::primitive::PrimitiveArrayExt;
46    use crate::dtype::DType;
47    use crate::dtype::Nullability;
48    use crate::dtype::PType;
49    use crate::validity::Validity;
50
51    #[test]
52    fn test_downcast_all_invalid() {
53        let array = PrimitiveArray::new(
54            buffer![0_u32, 0, 0, 0, 0, 0, 0, 0, 0, 0],
55            Validity::AllInvalid,
56        );
57
58        let result = array.narrow().unwrap();
59        assert_eq!(
60            result.dtype(),
61            &DType::Primitive(PType::U8, Nullability::Nullable)
62        );
63        assert!(matches!(result.validity(), Ok(Validity::AllInvalid)));
64    }
65
66    #[rstest]
67    #[case(vec![0_i64, 127], PType::U8)]
68    #[case(vec![-128_i64, 127], PType::I8)]
69    #[case(vec![-129_i64, 127], PType::I16)]
70    #[case(vec![-128_i64, 128], PType::I16)]
71    #[case(vec![-32768_i64, 32767], PType::I16)]
72    #[case(vec![-32769_i64, 32767], PType::I32)]
73    #[case(vec![-32768_i64, 32768], PType::I32)]
74    #[case(vec![i32::MIN as i64, i32::MAX as i64], PType::I32)]
75    fn test_downcast_signed(#[case] values: Vec<i64>, #[case] expected_ptype: PType) {
76        let array = PrimitiveArray::from_iter(values);
77        let result = array.narrow().unwrap();
78        assert_eq!(result.ptype(), expected_ptype);
79    }
80
81    #[rstest]
82    #[case(vec![0_u64, 255], PType::U8)]
83    #[case(vec![0_u64, 256], PType::U16)]
84    #[case(vec![0_u64, 65535], PType::U16)]
85    #[case(vec![0_u64, 65536], PType::U32)]
86    #[case(vec![0_u64, u32::MAX as u64], PType::U32)]
87    fn test_downcast_unsigned(#[case] values: Vec<u64>, #[case] expected_ptype: PType) {
88        let array = PrimitiveArray::from_iter(values);
89        let result = array.narrow().unwrap();
90        assert_eq!(result.ptype(), expected_ptype);
91    }
92
93    #[test]
94    fn test_downcast_keeps_original_if_too_large() {
95        let array = PrimitiveArray::from_iter(vec![0_u64, u64::MAX]);
96        let result = array.narrow().unwrap();
97        assert_eq!(result.ptype(), PType::U64);
98    }
99
100    #[test]
101    fn test_downcast_preserves_nullability() {
102        let array = PrimitiveArray::from_option_iter([Some(0_i32), None, Some(127)]);
103        let result = array.narrow().unwrap();
104        assert_eq!(
105            result.dtype(),
106            &DType::Primitive(PType::U8, Nullability::Nullable)
107        );
108        // Check that validity is preserved (the array should still have nullable values)
109        assert!(matches!(result.validity(), Ok(Validity::Array(_))));
110    }
111
112    #[test]
113    fn test_downcast_preserves_values() {
114        let values = vec![-100_i16, 0, 100];
115        let array = PrimitiveArray::from_iter(values);
116        let result = array.narrow().unwrap();
117
118        assert_eq!(result.ptype(), PType::I8);
119        // Check that the values were properly downscaled
120        let downscaled_values: Vec<i8> = result.as_slice::<i8>().to_vec();
121        assert_eq!(downscaled_values, vec![-100_i8, 0, 100]);
122    }
123
124    #[test]
125    fn test_downcast_with_mixed_signs_chooses_signed() {
126        let array = PrimitiveArray::from_iter(vec![-1_i32, 200]);
127        let result = array.narrow().unwrap();
128        assert_eq!(result.ptype(), PType::I16);
129    }
130
131    #[test]
132    fn test_downcast_floats() {
133        let array = PrimitiveArray::from_iter(vec![1.0_f32, 2.0, 3.0]);
134        let result = array.narrow().unwrap();
135        // Floats should remain unchanged since they can't be downscaled to integers
136        assert_eq!(result.ptype(), PType::F32);
137    }
138
139    #[test]
140    fn test_downcast_empty_array() {
141        let array = PrimitiveArray::new(Buffer::<i32>::empty(), Validity::AllInvalid);
142        let result = array.narrow().unwrap();
143        let array2 = PrimitiveArray::new(Buffer::<i64>::empty(), Validity::NonNullable);
144        let result2 = array2.narrow().unwrap();
145        // Empty arrays should not have their validity changed
146        assert!(matches!(result.validity(), Ok(Validity::AllInvalid)));
147        assert!(matches!(result2.validity(), Ok(Validity::NonNullable)));
148    }
149}
vortex_array/arrays/primitive/array/cast.rs

vortex_array/arrays/primitive/array/
cast.rs
