vortex_runend/compute/
cast.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::compute::{CastKernel, CastKernelAdapter, cast};
5use vortex_array::{ArrayRef, IntoArray, register_kernel};
6use vortex_dtype::DType;
7use vortex_error::VortexResult;
8
9use crate::{RunEndArray, RunEndVTable};
10
11impl CastKernel for RunEndVTable {
12    fn cast(&self, array: &RunEndArray, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
13        // Cast the values array to the target type
14        let casted_values = cast(array.values(), dtype)?;
15
16        // SAFETY: casting does not affect the ends being valid
17        unsafe {
18            Ok(Some(
19                RunEndArray::new_unchecked(
20                    array.ends().clone(),
21                    casted_values,
22                    array.offset(),
23                    array.len(),
24                )
25                .into_array(),
26            ))
27        }
28    }
29}
30
31register_kernel!(CastKernelAdapter(RunEndVTable).lift());
32
33#[cfg(test)]
34mod tests {
35    use rstest::rstest;
36    use vortex_array::arrays::{BoolArray, PrimitiveArray};
37    use vortex_array::compute::cast;
38    use vortex_array::compute::conformance::cast::test_cast_conformance;
39    use vortex_array::{Array, IntoArray};
40    use vortex_buffer::buffer;
41    use vortex_dtype::{DType, Nullability, PType};
42
43    use crate::RunEndArray;
44
45    #[test]
46    fn test_cast_runend_i32_to_i64() {
47        let runend = RunEndArray::try_new(
48            buffer![3u64, 5, 8, 10].into_array(),
49            buffer![100i32, 200, 100, 300].into_array(),
50        )
51        .unwrap();
52
53        let casted = cast(
54            runend.as_ref(),
55            &DType::Primitive(PType::I64, Nullability::NonNullable),
56        )
57        .unwrap();
58        assert_eq!(
59            casted.dtype(),
60            &DType::Primitive(PType::I64, Nullability::NonNullable)
61        );
62
63        // Verify by decoding to canonical form
64        let decoded = casted.to_canonical().unwrap().into_primitive().unwrap();
65        // RunEnd encoding should expand to [100, 100, 100, 200, 200, 100, 100, 100, 300, 300]
66        assert_eq!(decoded.len(), 10);
67        assert_eq!(decoded.as_slice::<i64>()[0], 100);
68        assert_eq!(decoded.as_slice::<i64>()[3], 200);
69        assert_eq!(decoded.as_slice::<i64>()[5], 100);
70        assert_eq!(decoded.as_slice::<i64>()[8], 300);
71    }
72
73    #[test]
74    fn test_cast_runend_nullable() {
75        let runend = RunEndArray::try_new(
76            buffer![2u64, 4, 7].into_array(),
77            PrimitiveArray::from_option_iter([Some(10i32), None, Some(20)]).into_array(),
78        )
79        .unwrap();
80
81        let casted = cast(
82            runend.as_ref(),
83            &DType::Primitive(PType::I64, Nullability::Nullable),
84        )
85        .unwrap();
86        assert_eq!(
87            casted.dtype(),
88            &DType::Primitive(PType::I64, Nullability::Nullable)
89        );
90    }
91
92    #[test]
93    fn test_cast_runend_with_offset() {
94        // Create a RunEndArray: [100, 100, 100, 200, 200, 300, 300, 300, 300, 300]
95        let runend = RunEndArray::try_new(
96            buffer![3u64, 5, 10].into_array(),
97            buffer![100i32, 200, 300].into_array(),
98        )
99        .unwrap();
100
101        // Slice it to get offset 3, length 5: [200, 200, 300, 300, 300]
102        let sliced = runend.slice(3, 8);
103
104        // Verify the slice is correct before casting
105        let sliced_decoded = sliced.to_canonical().unwrap().into_primitive().unwrap();
106        assert_eq!(sliced_decoded.len(), 5);
107        assert_eq!(sliced_decoded.as_slice::<i32>(), &[200, 200, 300, 300, 300]);
108
109        // Cast the sliced array
110        let casted = cast(
111            sliced.as_ref(),
112            &DType::Primitive(PType::I64, Nullability::NonNullable),
113        )
114        .unwrap();
115
116        // Verify the cast preserved the offset
117        let casted_decoded = casted.to_canonical().unwrap().into_primitive().unwrap();
118        assert_eq!(casted_decoded.len(), 5);
119        assert_eq!(
120            casted_decoded.as_slice::<i64>(),
121            &[200i64, 200, 300, 300, 300],
122            "Cast failed to preserve offset - got wrong values after cast"
123        );
124    }
125
126    #[rstest]
127    #[case(RunEndArray::try_new(
128        buffer![3u64, 5, 8].into_array(),
129        buffer![100i32, 200, 300].into_array()
130    ).unwrap())]
131    #[case(RunEndArray::try_new(
132        buffer![1u64, 4, 10].into_array(),
133        buffer![1.5f32, 2.5, 3.5].into_array()
134    ).unwrap())]
135    #[case(RunEndArray::try_new(
136        buffer![2u64, 3, 5].into_array(),
137        PrimitiveArray::from_option_iter([Some(42i32), None, Some(84)]).into_array()
138    ).unwrap())]
139    #[case(RunEndArray::try_new(
140        buffer![10u64].into_array(),
141        buffer![255u8].into_array()
142    ).unwrap())]
143    #[case(RunEndArray::try_new(
144        buffer![2u64, 4, 6, 8, 10].into_array(),
145        BoolArray::from_iter(vec![true, false, true, false, true]).into_array()
146    ).unwrap())]
147    fn test_cast_runend_conformance(#[case] array: RunEndArray) {
148        test_cast_conformance(array.as_ref());
149    }
150}