vortex_runend/compute/
cast.rs

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