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 std::sync::LazyLock;
41
42    use rstest::rstest;
43    use vortex_buffer::Buffer;
44    use vortex_buffer::buffer;
45    use vortex_session::VortexSession;
46
47    use crate::VortexSessionExecute;
48    use crate::arrays::PrimitiveArray;
49    use crate::arrays::primitive::PrimitiveArrayExt;
50    use crate::dtype::DType;
51    use crate::dtype::Nullability;
52    use crate::dtype::PType;
53    use crate::validity::Validity;
54
55    static SESSION: LazyLock<VortexSession> = LazyLock::new(crate::array_session);
56
57    #[test]
58    fn test_downcast_all_invalid() {
59        let array = PrimitiveArray::new(
60            buffer![0_u32, 0, 0, 0, 0, 0, 0, 0, 0, 0],
61            Validity::AllInvalid,
62        );
63
64        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
65        assert_eq!(
66            result.dtype(),
67            &DType::Primitive(PType::U8, Nullability::Nullable)
68        );
69        assert!(matches!(result.validity(), Ok(Validity::AllInvalid)));
70    }
71
72    #[rstest]
73    #[case(vec![0_i64, 127], PType::U8)]
74    #[case(vec![-128_i64, 127], PType::I8)]
75    #[case(vec![-129_i64, 127], PType::I16)]
76    #[case(vec![-128_i64, 128], PType::I16)]
77    #[case(vec![-32768_i64, 32767], PType::I16)]
78    #[case(vec![-32769_i64, 32767], PType::I32)]
79    #[case(vec![-32768_i64, 32768], PType::I32)]
80    #[case(vec![i32::MIN as i64, i32::MAX as i64], PType::I32)]
81    fn test_downcast_signed(#[case] values: Vec<i64>, #[case] expected_ptype: PType) {
82        let array = PrimitiveArray::from_iter(values);
83        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
84        assert_eq!(result.ptype(), expected_ptype);
85    }
86
87    #[rstest]
88    #[case(vec![0_u64, 255], PType::U8)]
89    #[case(vec![0_u64, 256], PType::U16)]
90    #[case(vec![0_u64, 65535], PType::U16)]
91    #[case(vec![0_u64, 65536], PType::U32)]
92    #[case(vec![0_u64, u32::MAX as u64], PType::U32)]
93    fn test_downcast_unsigned(#[case] values: Vec<u64>, #[case] expected_ptype: PType) {
94        let array = PrimitiveArray::from_iter(values);
95        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
96        assert_eq!(result.ptype(), expected_ptype);
97    }
98
99    #[test]
100    fn test_downcast_keeps_original_if_too_large() {
101        let array = PrimitiveArray::from_iter(vec![0_u64, u64::MAX]);
102        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
103        assert_eq!(result.ptype(), PType::U64);
104    }
105
106    #[test]
107    fn test_downcast_preserves_nullability() {
108        let array = PrimitiveArray::from_option_iter([Some(0_i32), None, Some(127)]);
109        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
110        assert_eq!(
111            result.dtype(),
112            &DType::Primitive(PType::U8, Nullability::Nullable)
113        );
114        // Check that validity is preserved (the array should still have nullable values)
115        assert!(matches!(result.validity(), Ok(Validity::Array(_))));
116    }
117
118    #[test]
119    fn test_downcast_preserves_values() {
120        let values = vec![-100_i16, 0, 100];
121        let array = PrimitiveArray::from_iter(values);
122        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
123
124        assert_eq!(result.ptype(), PType::I8);
125        // Check that the values were properly downscaled
126        let downscaled_values: Vec<i8> = result.as_slice::<i8>().to_vec();
127        assert_eq!(downscaled_values, vec![-100_i8, 0, 100]);
128    }
129
130    #[test]
131    fn test_downcast_with_mixed_signs_chooses_signed() {
132        let array = PrimitiveArray::from_iter(vec![-1_i32, 200]);
133        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
134        assert_eq!(result.ptype(), PType::I16);
135    }
136
137    #[test]
138    fn test_downcast_floats() {
139        let array = PrimitiveArray::from_iter(vec![1.0_f32, 2.0, 3.0]);
140        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
141        // Floats should remain unchanged since they can't be downscaled to integers
142        assert_eq!(result.ptype(), PType::F32);
143    }
144
145    #[test]
146    fn test_downcast_empty_array() {
147        let array = PrimitiveArray::new(Buffer::<i32>::empty(), Validity::AllInvalid);
148        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
149        let array2 = PrimitiveArray::new(Buffer::<i64>::empty(), Validity::NonNullable);
150        let result2 = array2.narrow(&mut SESSION.create_execution_ctx()).unwrap();
151        // Empty arrays should not have their validity changed
152        assert!(matches!(result.validity(), Ok(Validity::AllInvalid)));
153        assert!(matches!(result2.validity(), Ok(Validity::NonNullable)));
154    }
155}
vortex_array/arrays/primitive/array/cast.rs

vortex_array/arrays/primitive/array/
cast.rs
