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, ToCanonical, assert_arrays_eq};
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_primitive();
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!(
68            TryInto::<i64>::try_into(decoded.scalar_at(0).as_ref())
69                .ok()
70                .unwrap(),
71            100i64
72        );
73        assert_eq!(
74            TryInto::<i64>::try_into(decoded.scalar_at(3).as_ref())
75                .ok()
76                .unwrap(),
77            200i64
78        );
79        assert_eq!(
80            TryInto::<i64>::try_into(decoded.scalar_at(5).as_ref())
81                .ok()
82                .unwrap(),
83            100i64
84        );
85        assert_eq!(
86            TryInto::<i64>::try_into(decoded.scalar_at(8).as_ref())
87                .ok()
88                .unwrap(),
89            300i64
90        );
91    }
92
93    #[test]
94    fn test_cast_runend_nullable() {
95        let runend = RunEndArray::try_new(
96            buffer![2u64, 4, 7].into_array(),
97            PrimitiveArray::from_option_iter([Some(10i32), None, Some(20)]).into_array(),
98        )
99        .unwrap();
100
101        let casted = cast(
102            runend.as_ref(),
103            &DType::Primitive(PType::I64, Nullability::Nullable),
104        )
105        .unwrap();
106        assert_eq!(
107            casted.dtype(),
108            &DType::Primitive(PType::I64, Nullability::Nullable)
109        );
110    }
111
112    #[test]
113    fn test_cast_runend_with_offset() {
114        // Create a RunEndArray: [100, 100, 100, 200, 200, 300, 300, 300, 300, 300]
115        let runend = RunEndArray::try_new(
116            buffer![3u64, 5, 10].into_array(),
117            buffer![100i32, 200, 300].into_array(),
118        )
119        .unwrap();
120
121        // Slice it to get offset 3, length 5: [200, 200, 300, 300, 300]
122        let sliced = runend.slice(3..8);
123
124        // Verify the slice is correct before casting
125        let sliced_decoded = sliced.to_primitive();
126        assert_eq!(sliced_decoded.len(), 5);
127        assert_arrays_eq!(
128            sliced_decoded,
129            PrimitiveArray::from_iter([200, 200, 300, 300, 300])
130        );
131
132        // Cast the sliced array
133        let casted = cast(
134            sliced.as_ref(),
135            &DType::Primitive(PType::I64, Nullability::NonNullable),
136        )
137        .unwrap();
138
139        // Verify the cast preserved the offset
140        let casted_decoded = casted.to_primitive();
141        assert_eq!(casted_decoded.len(), 5);
142        assert_arrays_eq!(
143            casted_decoded,
144            PrimitiveArray::from_iter([200i64, 200, 300, 300, 300])
145        );
146    }
147
148    #[rstest]
149    #[case(RunEndArray::try_new(
150        buffer![3u64, 5, 8].into_array(),
151        buffer![100i32, 200, 300].into_array()
152    ).unwrap())]
153    #[case(RunEndArray::try_new(
154        buffer![1u64, 4, 10].into_array(),
155        buffer![1.5f32, 2.5, 3.5].into_array()
156    ).unwrap())]
157    #[case(RunEndArray::try_new(
158        buffer![2u64, 3, 5].into_array(),
159        PrimitiveArray::from_option_iter([Some(42i32), None, Some(84)]).into_array()
160    ).unwrap())]
161    #[case(RunEndArray::try_new(
162        buffer![10u64].into_array(),
163        buffer![255u8].into_array()
164    ).unwrap())]
165    #[case(RunEndArray::try_new(
166        buffer![2u64, 4, 6, 8, 10].into_array(),
167        BoolArray::from_iter(vec![true, false, true, false, true]).into_array()
168    ).unwrap())]
169    fn test_cast_runend_conformance(#[case] array: RunEndArray) {
170        test_cast_conformance(array.as_ref());
171    }
172}