Skip to main content

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