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::IntoArray;
6use vortex_array::builtins::ArrayBuiltins;
7use vortex_array::compute::CastReduce;
8use vortex_dtype::DType;
9use vortex_error::VortexResult;
10
11use crate::RunEndArray;
12use crate::RunEndVTable;
13
14impl CastReduce for RunEndVTable {
15    fn cast(array: &RunEndArray, 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                RunEndArray::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::Array;
38    use vortex_array::IntoArray;
39    use vortex_array::ToCanonical;
40    use vortex_array::arrays::BoolArray;
41    use vortex_array::arrays::PrimitiveArray;
42    use vortex_array::assert_arrays_eq;
43    use vortex_array::builtins::ArrayBuiltins;
44    use vortex_array::compute::conformance::cast::test_cast_conformance;
45    use vortex_buffer::buffer;
46    use vortex_dtype::DType;
47    use vortex_dtype::Nullability;
48    use vortex_dtype::PType;
49
50    use crate::RunEndArray;
51
52    #[test]
53    fn test_cast_runend_i32_to_i64() {
54        let runend = RunEndArray::try_new(
55            buffer![3u64, 5, 8, 10].into_array(),
56            buffer![100i32, 200, 100, 300].into_array(),
57        )
58        .unwrap();
59
60        let casted = runend
61            .into_array()
62            .cast(DType::Primitive(PType::I64, Nullability::NonNullable))
63            .unwrap();
64        assert_eq!(
65            casted.dtype(),
66            &DType::Primitive(PType::I64, Nullability::NonNullable)
67        );
68
69        // Verify by decoding to canonical form
70        let decoded = casted.to_primitive();
71        // RunEnd encoding should expand to [100, 100, 100, 200, 200, 100, 100, 100, 300, 300]
72        assert_eq!(decoded.len(), 10);
73        assert_eq!(
74            TryInto::<i64>::try_into(&decoded.scalar_at(0).unwrap()).unwrap(),
75            100i64
76        );
77        assert_eq!(
78            TryInto::<i64>::try_into(&decoded.scalar_at(3).unwrap()).unwrap(),
79            200i64
80        );
81        assert_eq!(
82            TryInto::<i64>::try_into(&decoded.scalar_at(5).unwrap()).unwrap(),
83            100i64
84        );
85        assert_eq!(
86            TryInto::<i64>::try_into(&decoded.scalar_at(8).unwrap()).unwrap(),
87            300i64
88        );
89    }
90
91    #[test]
92    fn test_cast_runend_nullable() {
93        let runend = RunEndArray::try_new(
94            buffer![2u64, 4, 7].into_array(),
95            PrimitiveArray::from_option_iter([Some(10i32), None, Some(20)]).into_array(),
96        )
97        .unwrap();
98
99        let casted = runend
100            .into_array()
101            .cast(DType::Primitive(PType::I64, Nullability::Nullable))
102            .unwrap();
103        assert_eq!(
104            casted.dtype(),
105            &DType::Primitive(PType::I64, Nullability::Nullable)
106        );
107    }
108
109    #[test]
110    fn test_cast_runend_with_offset() {
111        // Create a RunEndArray: [100, 100, 100, 200, 200, 300, 300, 300, 300, 300]
112        let runend = RunEndArray::try_new(
113            buffer![3u64, 5, 10].into_array(),
114            buffer![100i32, 200, 300].into_array(),
115        )
116        .unwrap();
117
118        // Slice it to get offset 3, length 5: [200, 200, 300, 300, 300]
119        let sliced = runend.slice(3..8).unwrap();
120
121        // Verify the slice is correct before casting
122        assert_arrays_eq!(sliced, PrimitiveArray::from_iter([200, 200, 300, 300, 300]));
123
124        // Cast the sliced array
125        let casted = sliced
126            .cast(DType::Primitive(PType::I64, Nullability::NonNullable))
127            .unwrap();
128
129        // Verify the cast preserved the offset
130        assert_arrays_eq!(
131            casted,
132            PrimitiveArray::from_iter([200i64, 200, 300, 300, 300])
133        );
134    }
135
136    #[rstest]
137    #[case(RunEndArray::try_new(
138        buffer![3u64, 5, 8].into_array(),
139        buffer![100i32, 200, 300].into_array()
140    ).unwrap())]
141    #[case(RunEndArray::try_new(
142        buffer![1u64, 4, 10].into_array(),
143        buffer![1.5f32, 2.5, 3.5].into_array()
144    ).unwrap())]
145    #[case(RunEndArray::try_new(
146        buffer![2u64, 3, 5].into_array(),
147        PrimitiveArray::from_option_iter([Some(42i32), None, Some(84)]).into_array()
148    ).unwrap())]
149    #[case(RunEndArray::try_new(
150        buffer![10u64].into_array(),
151        buffer![255u8].into_array()
152    ).unwrap())]
153    #[case(RunEndArray::try_new(
154        buffer![2u64, 4, 6, 8, 10].into_array(),
155        BoolArray::from_iter(vec![true, false, true, false, true]).into_array()
156    ).unwrap())]
157    fn test_cast_runend_conformance(#[case] array: RunEndArray) {
158        test_cast_conformance(array.as_ref());
159    }
160}