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::assert_arrays_eq;
51    use crate::dtype::DType;
52    use crate::dtype::Nullability;
53    use crate::dtype::PType;
54    use crate::validity::Validity;
55
56    static SESSION: LazyLock<VortexSession> = LazyLock::new(crate::array_session);
57
58    #[test]
59    fn test_downcast_all_invalid() {
60        let array = PrimitiveArray::new(
61            buffer![0_u32, 0, 0, 0, 0, 0, 0, 0, 0, 0],
62            Validity::AllInvalid,
63        );
64
65        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
66        assert_eq!(
67            result.dtype(),
68            &DType::Primitive(PType::U8, Nullability::Nullable)
69        );
70        assert!(matches!(result.validity(), Ok(Validity::AllInvalid)));
71        assert_arrays_eq!(
72            result,
73            PrimitiveArray::from_option_iter([None::<u8>; 10]),
74            &mut SESSION.create_execution_ctx()
75        );
76    }
77
78    #[rstest]
79    #[case(vec![0_i64, 127], PType::U8)]
80    #[case(vec![-128_i64, 127], PType::I8)]
81    #[case(vec![-129_i64, 127], PType::I16)]
82    #[case(vec![-128_i64, 128], PType::I16)]
83    #[case(vec![-32768_i64, 32767], PType::I16)]
84    #[case(vec![-32769_i64, 32767], PType::I32)]
85    #[case(vec![-32768_i64, 32768], PType::I32)]
86    #[case(vec![i32::MIN as i64, i32::MAX as i64], PType::I32)]
87    fn test_downcast_signed(#[case] values: Vec<i64>, #[case] expected_ptype: PType) {
88        let array = PrimitiveArray::from_iter(values.iter().copied());
89        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
90        assert_eq!(
91            result.dtype(),
92            &DType::Primitive(expected_ptype, Nullability::NonNullable)
93        );
94        assert_eq!(result.len(), values.len());
95        let mut ctx = SESSION.create_execution_ctx();
96        for (index, expected) in values.into_iter().enumerate() {
97            assert_eq!(
98                result
99                    .execute_scalar(index, &mut ctx)
100                    .unwrap()
101                    .as_primitive()
102                    .as_::<i64>(),
103                Some(expected)
104            );
105        }
106    }
107
108    #[rstest]
109    #[case(vec![0_u64, 255], PType::U8)]
110    #[case(vec![0_u64, 256], PType::U16)]
111    #[case(vec![0_u64, 65535], PType::U16)]
112    #[case(vec![0_u64, 65536], PType::U32)]
113    #[case(vec![0_u64, u32::MAX as u64], PType::U32)]
114    fn test_downcast_unsigned(#[case] values: Vec<u64>, #[case] expected_ptype: PType) {
115        let array = PrimitiveArray::from_iter(values.iter().copied());
116        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
117        assert_eq!(
118            result.dtype(),
119            &DType::Primitive(expected_ptype, Nullability::NonNullable)
120        );
121        assert_eq!(result.len(), values.len());
122        let mut ctx = SESSION.create_execution_ctx();
123        for (index, expected) in values.into_iter().enumerate() {
124            assert_eq!(
125                result
126                    .execute_scalar(index, &mut ctx)
127                    .unwrap()
128                    .as_primitive()
129                    .as_::<u64>(),
130                Some(expected)
131            );
132        }
133    }
134
135    #[test]
136    fn test_downcast_keeps_original_if_too_large() {
137        let array = PrimitiveArray::from_iter(vec![0_u64, u64::MAX]);
138        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
139        assert_eq!(result.ptype(), PType::U64);
140        assert_arrays_eq!(
141            result,
142            PrimitiveArray::from_iter([0u64, u64::MAX]),
143            &mut SESSION.create_execution_ctx()
144        );
145    }
146
147    #[test]
148    fn test_downcast_preserves_nullability() {
149        let array = PrimitiveArray::from_option_iter([Some(0_i32), None, Some(127)]);
150        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
151        assert_eq!(
152            result.dtype(),
153            &DType::Primitive(PType::U8, Nullability::Nullable)
154        );
155        // Check that validity is preserved (the array should still have nullable values)
156        assert!(matches!(result.validity(), Ok(Validity::Array(_))));
157        assert_arrays_eq!(
158            result,
159            PrimitiveArray::from_option_iter([Some(0u8), None, Some(127)]),
160            &mut SESSION.create_execution_ctx()
161        );
162    }
163
164    #[test]
165    fn test_downcast_preserves_values() {
166        let values = vec![-100_i16, 0, 100];
167        let array = PrimitiveArray::from_iter(values);
168        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
169
170        assert_eq!(result.ptype(), PType::I8);
171        // Check that the values were properly downscaled
172        let downscaled_values: Vec<i8> = result.as_slice::<i8>().to_vec();
173        assert_eq!(downscaled_values, vec![-100_i8, 0, 100]);
174    }
175
176    #[test]
177    fn test_downcast_with_mixed_signs_chooses_signed() {
178        let array = PrimitiveArray::from_iter(vec![-1_i32, 200]);
179        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
180        assert_eq!(result.ptype(), PType::I16);
181        assert_arrays_eq!(
182            result,
183            PrimitiveArray::from_iter([-1i16, 200]),
184            &mut SESSION.create_execution_ctx()
185        );
186    }
187
188    #[test]
189    fn test_downcast_floats() {
190        let array = PrimitiveArray::from_iter(vec![1.0_f32, 2.0, 3.0]);
191        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
192        // Floats should remain unchanged since they can't be downscaled to integers
193        assert_eq!(result.ptype(), PType::F32);
194        assert_arrays_eq!(
195            result,
196            PrimitiveArray::from_iter([1f32, 2., 3.]),
197            &mut SESSION.create_execution_ctx()
198        );
199    }
200
201    #[test]
202    fn test_downcast_empty_array() {
203        let array = PrimitiveArray::new(Buffer::<i32>::empty(), Validity::AllInvalid);
204        let result = array.narrow(&mut SESSION.create_execution_ctx()).unwrap();
205        let array2 = PrimitiveArray::new(Buffer::<i64>::empty(), Validity::NonNullable);
206        let result2 = array2.narrow(&mut SESSION.create_execution_ctx()).unwrap();
207        assert!(result.is_empty());
208        assert!(result2.is_empty());
209        // Empty arrays should not have their validity changed
210        assert!(matches!(result.validity(), Ok(Validity::AllInvalid)));
211        assert!(matches!(result2.validity(), Ok(Validity::NonNullable)));
212    }
213}
vortex_array/arrays/primitive/array/cast.rs

vortex_array/arrays/primitive/array/
cast.rs
