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