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