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 std::sync::LazyLock;
37
38    use rstest::rstest;
39    use vortex_array::IntoArray;
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_array::session::ArraySession;
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> =
57        LazyLock::new(|| VortexSession::empty().with::<ArraySession>());
58
59    #[test]
60    fn test_cast_runend_i32_to_i64() {
61        let mut ctx = SESSION.create_execution_ctx();
62        let runend = RunEnd::try_new(
63            buffer![3u64, 5, 8, 10].into_array(),
64            buffer![100i32, 200, 100, 300].into_array(),
65            &mut ctx,
66        )
67        .unwrap();
68
69        let casted = runend
70            .into_array()
71            .cast(DType::Primitive(PType::I64, Nullability::NonNullable))
72            .unwrap();
73        assert_eq!(
74            casted.dtype(),
75            &DType::Primitive(PType::I64, Nullability::NonNullable)
76        );
77
78        // Verify by decoding to canonical form
79        let decoded = casted.execute::<PrimitiveArray>(&mut ctx).unwrap();
80        // RunEnd encoding should expand to [100, 100, 100, 200, 200, 100, 100, 100, 300, 300]
81        assert_eq!(decoded.len(), 10);
82        assert_eq!(
83            TryInto::<i64>::try_into(&decoded.execute_scalar(0, &mut ctx).unwrap()).unwrap(),
84            100i64
85        );
86        assert_eq!(
87            TryInto::<i64>::try_into(&decoded.execute_scalar(3, &mut ctx).unwrap()).unwrap(),
88            200i64
89        );
90        assert_eq!(
91            TryInto::<i64>::try_into(&decoded.execute_scalar(5, &mut ctx).unwrap()).unwrap(),
92            100i64
93        );
94        assert_eq!(
95            TryInto::<i64>::try_into(&decoded.execute_scalar(8, &mut ctx).unwrap()).unwrap(),
96            300i64
97        );
98    }
99
100    #[test]
101    fn test_cast_runend_nullable() {
102        let mut ctx = SESSION.create_execution_ctx();
103        let runend = RunEnd::try_new(
104            buffer![2u64, 4, 7].into_array(),
105            PrimitiveArray::from_option_iter([Some(10i32), None, Some(20)]).into_array(),
106            &mut ctx,
107        )
108        .unwrap();
109
110        let casted = runend
111            .into_array()
112            .cast(DType::Primitive(PType::I64, Nullability::Nullable))
113            .unwrap();
114        assert_eq!(
115            casted.dtype(),
116            &DType::Primitive(PType::I64, Nullability::Nullable)
117        );
118    }
119
120    #[test]
121    fn test_cast_runend_with_offset() {
122        let mut ctx = SESSION.create_execution_ctx();
123        // Create a RunEndArray: [100, 100, 100, 200, 200, 300, 300, 300, 300, 300]
124        let runend = RunEnd::try_new(
125            buffer![3u64, 5, 10].into_array(),
126            buffer![100i32, 200, 300].into_array(),
127            &mut ctx,
128        )
129        .unwrap();
130
131        // Slice it to get offset 3, length 5: [200, 200, 300, 300, 300]
132        let sliced = runend.slice(3..8).unwrap();
133
134        // Verify the slice is correct before casting
135        assert_arrays_eq!(sliced, PrimitiveArray::from_iter([200, 200, 300, 300, 300]));
136
137        // Cast the sliced array
138        let casted = sliced
139            .cast(DType::Primitive(PType::I64, Nullability::NonNullable))
140            .unwrap();
141
142        // Verify the cast preserved the offset
143        assert_arrays_eq!(
144            casted,
145            PrimitiveArray::from_iter([200i64, 200, 300, 300, 300])
146        );
147    }
148
149    type RunEndBuilder = fn(&mut vortex_array::ExecutionCtx) -> RunEndArray;
150
151    #[rstest]
152    #[case(|ctx: &mut vortex_array::ExecutionCtx| RunEnd::try_new(
153        buffer![3u64, 5, 8].into_array(),
154        buffer![100i32, 200, 300].into_array(),
155        ctx,
156    ).unwrap())]
157    #[case(|ctx: &mut vortex_array::ExecutionCtx| RunEnd::try_new(
158        buffer![1u64, 4, 10].into_array(),
159        buffer![1.5f32, 2.5, 3.5].into_array(),
160        ctx,
161    ).unwrap())]
162    #[case(|ctx: &mut vortex_array::ExecutionCtx| RunEnd::try_new(
163        buffer![2u64, 3, 5].into_array(),
164        PrimitiveArray::from_option_iter([Some(42i32), None, Some(84)]).into_array(),
165        ctx,
166    ).unwrap())]
167    #[case(|ctx: &mut vortex_array::ExecutionCtx| RunEnd::try_new(
168        buffer![10u64].into_array(),
169        buffer![255u8].into_array(),
170        ctx,
171    ).unwrap())]
172    #[case(|ctx: &mut vortex_array::ExecutionCtx| RunEnd::try_new(
173        buffer![2u64, 4, 6, 8, 10].into_array(),
174        BoolArray::from_iter(vec![true, false, true, false, true]).into_array(),
175        ctx,
176    ).unwrap())]
177    fn test_cast_runend_conformance(#[case] build: RunEndBuilder) {
178        let mut ctx = SESSION.create_execution_ctx();
179        let array = build(&mut ctx);
180        test_cast_conformance(&array.into_array());
181    }
182}