vortex_alp/alp/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::ALPArrayExt;
13use crate::ALPArraySlotsExt;
14use crate::alp::ALP;
15
16impl CastReduce for ALP {
17    fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
18        // Check if this is just a nullability change
19        if !array.dtype().eq_ignore_nullability(dtype) {
20            return Ok(None);
21        }
22
23        // For nullability-only changes, we can avoid decoding
24        // Cast the encoded array (integers) to handle nullability
25        let new_encoded = array.encoded().cast(
26            array
27                .encoded()
28                .dtype()
29                .with_nullability(dtype.nullability()),
30        )?;
31
32        let new_patches = array
33            .patches()
34            .map(|p| p.map_values(|v| v.cast(dtype.clone())))
35            .transpose()?;
36
37        // SAFETY: casting nullability doesn't alter the invariants
38        unsafe {
39            Ok(Some(
40                ALP::new_unchecked(new_encoded, array.exponents(), new_patches).into_array(),
41            ))
42        }
43    }
44}
45
46#[cfg(test)]
47mod tests {
48    use std::sync::LazyLock;
49
50    use rstest::rstest;
51    use vortex_array::IntoArray;
52    use vortex_array::VortexSessionExecute;
53    use vortex_array::arrays::PrimitiveArray;
54    use vortex_array::assert_arrays_eq;
55    use vortex_array::builtins::ArrayBuiltins;
56    use vortex_array::compute::conformance::cast::test_cast_conformance;
57    use vortex_array::dtype::DType;
58    use vortex_array::dtype::Nullability;
59    use vortex_array::dtype::PType;
60    use vortex_buffer::buffer;
61    use vortex_error::VortexExpect;
62    use vortex_error::VortexResult;
63    use vortex_session::VortexSession;
64
65    use crate::alp::array::ALPArrayExt;
66    use crate::alp_encode;
67
68    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
69        let session = vortex_array::array_session();
70        crate::initialize(&session);
71        session
72    });
73
74    #[test]
75    fn issue_5766_test_cast_alp_with_patches_to_nullable() -> VortexResult<()> {
76        let mut ctx = SESSION.create_execution_ctx();
77        let values = buffer![1.234f32, f32::NAN, 2.345, f32::INFINITY, 3.456].into_array();
78        let values_primitive = values.clone().execute::<PrimitiveArray>(&mut ctx)?;
79        let alp = alp_encode(values_primitive.as_view(), None, &mut ctx)?;
80
81        assert!(
82            alp.patches().is_some(),
83            "Test requires ALP array with patches"
84        );
85
86        let nullable_dtype = DType::Primitive(PType::F32, Nullability::Nullable);
87        let casted = alp.into_array().cast(nullable_dtype.clone())?;
88
89        let expected = values.cast(nullable_dtype)?;
90
91        let casted_prim = casted.execute::<PrimitiveArray>(&mut ctx)?;
92        assert_arrays_eq!(casted_prim, expected, &mut ctx);
93
94        Ok(())
95    }
96
97    #[test]
98    fn test_cast_alp_f32_to_f64() -> VortexResult<()> {
99        let mut ctx = SESSION.create_execution_ctx();
100        let values = buffer![1.5f32, 2.5, 3.5, 4.5].into_array();
101        let values_primitive = values.execute::<PrimitiveArray>(&mut ctx)?;
102        let alp = alp_encode(values_primitive.as_view(), None, &mut ctx)?;
103
104        let casted = alp
105            .into_array()
106            .cast(DType::Primitive(PType::F64, Nullability::NonNullable))?;
107        assert_eq!(
108            casted.dtype(),
109            &DType::Primitive(PType::F64, Nullability::NonNullable)
110        );
111
112        let decoded = casted.execute::<PrimitiveArray>(&mut ctx)?;
113        let values = decoded.as_slice::<f64>();
114        assert_eq!(values.len(), 4);
115        assert!((values[0] - 1.5).abs() < f64::EPSILON);
116        assert!((values[1] - 2.5).abs() < f64::EPSILON);
117
118        Ok(())
119    }
120
121    #[test]
122    fn test_cast_alp_to_int() -> VortexResult<()> {
123        let mut ctx = SESSION.create_execution_ctx();
124        let values = buffer![1.0f32, 2.0, 3.0, 4.0].into_array();
125        let values_primitive = values.execute::<PrimitiveArray>(&mut ctx)?;
126        let alp = alp_encode(values_primitive.as_view(), None, &mut ctx)?;
127
128        let casted = alp
129            .into_array()
130            .cast(DType::Primitive(PType::I32, Nullability::NonNullable))?;
131        assert_eq!(
132            casted.dtype(),
133            &DType::Primitive(PType::I32, Nullability::NonNullable)
134        );
135
136        let decoded = casted.execute::<PrimitiveArray>(&mut ctx)?;
137        assert_arrays_eq!(
138            decoded,
139            PrimitiveArray::from_iter([1i32, 2, 3, 4]),
140            &mut ctx
141        );
142
143        Ok(())
144    }
145
146    #[rstest]
147    #[case(buffer![1.23f32, 4.56, 7.89, 10.11, 12.13].into_array())]
148    #[case(buffer![100.1f64, 200.2, 300.3, 400.4, 500.5].into_array())]
149    #[case(PrimitiveArray::from_option_iter([Some(1.1f32), None, Some(2.2), Some(3.3), None]).into_array())]
150    #[case(buffer![42.42f64].into_array())]
151    #[case(buffer![0.0f32, -1.5, 2.5, -3.5, 4.5].into_array())]
152    fn test_cast_alp_conformance(#[case] array: vortex_array::ArrayRef) -> VortexResult<()> {
153        let mut ctx = SESSION.create_execution_ctx();
154        let array_primitive = array.execute::<PrimitiveArray>(&mut ctx)?;
155        let alp =
156            alp_encode(array_primitive.as_view(), None, &mut ctx).vortex_expect("cannot fail");
157        test_cast_conformance(&alp.into_array(), &mut ctx);
158
159        Ok(())
160    }
161}
vortex_alp/alp/compute/cast.rs

vortex_alp/alp/compute/
cast.rs
