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