Skip to main content

vortex_alp/alp_rd/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::ALPRDArrayExt;
13use crate::ALPRDArraySlotsExt;
14use crate::alp_rd::ALPRD;
15
16impl CastReduce for ALPRD {
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 need to cast the left_parts array
24        // since it carries the validity information
25        let new_left_parts = array.left_parts().cast(
26            array
27                .left_parts()
28                .dtype()
29                .with_nullability(dtype.nullability()),
30        )?;
31
32        Ok(Some(
33            unsafe {
34                ALPRD::new_unchecked(
35                    dtype.clone(),
36                    new_left_parts,
37                    array.left_parts_dictionary().clone(),
38                    array.right_parts().clone(),
39                    array.right_bit_width(),
40                    array.left_parts_patches(),
41                )
42            }
43            .into_array(),
44        ))
45    }
46}
47
48#[cfg(test)]
49mod tests {
50    use rstest::rstest;
51    use vortex_array::IntoArray;
52    use vortex_array::VortexSessionExecute;
53    use vortex_array::array_session;
54    use vortex_array::arrays::PrimitiveArray;
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
61    use crate::RDEncoder;
62
63    #[test]
64    fn test_cast_alprd_f32_to_f64() {
65        let mut ctx = array_session().create_execution_ctx();
66        let values = vec![1.0f32, 1.1, 1.2, 1.3, 1.4];
67        let arr = PrimitiveArray::from_iter(values.clone());
68        let encoder = RDEncoder::new(&values);
69        let alprd = encoder.encode(arr.as_view());
70
71        let casted = alprd
72            .into_array()
73            .cast(DType::Primitive(PType::F64, Nullability::NonNullable))
74            .unwrap();
75        assert_eq!(
76            casted.dtype(),
77            &DType::Primitive(PType::F64, Nullability::NonNullable)
78        );
79
80        let decoded = casted.execute::<PrimitiveArray>(&mut ctx).unwrap();
81        let f64_values = decoded.as_slice::<f64>();
82        assert_eq!(f64_values.len(), 5);
83        assert!((f64_values[0] - 1.0).abs() < f64::EPSILON);
84        assert!((f64_values[1] - 1.1).abs() < 1e-6); // Use larger epsilon for f32->f64 conversion
85    }
86
87    #[test]
88    fn test_cast_alprd_nullable() {
89        let mut ctx = array_session().create_execution_ctx();
90        let arr =
91            PrimitiveArray::from_option_iter([Some(10.0f64), None, Some(10.1), Some(10.2), None]);
92        let values = vec![10.0f64, 10.1, 10.2];
93        let encoder = RDEncoder::new(&values);
94        let alprd = encoder.encode(arr.as_view());
95
96        // Cast to NonNullable should fail since we have nulls. The failure surfaces during
97        // execution since the reduce path defers when the validity stat is not cached.
98        let result = alprd
99            .clone()
100            .into_array()
101            .cast(DType::Primitive(PType::F64, Nullability::NonNullable))
102            .and_then(|a| {
103                a.execute::<PrimitiveArray>(&mut ctx)
104                    .map(|p| p.into_array())
105            });
106        assert!(result.is_err(), "Expected error, got: {result:?}");
107
108        // Cast to same type with Nullable should succeed
109        let casted = alprd
110            .into_array()
111            .cast(DType::Primitive(PType::F64, Nullability::Nullable))
112            .unwrap();
113        assert_eq!(
114            casted.dtype(),
115            &DType::Primitive(PType::F64, Nullability::Nullable)
116        );
117    }
118
119    #[rstest]
120    #[case::f32({
121        let values = vec![1.23f32, 4.56, 7.89, 10.11, 12.13];
122        let arr = PrimitiveArray::from_iter(values.clone());
123        let encoder = RDEncoder::new(&values);
124        encoder.encode(arr.as_view())
125    })]
126    #[case::f64({
127        let values = vec![100.1f64, 200.2, 300.3, 400.4, 500.5];
128        let arr = PrimitiveArray::from_iter(values.clone());
129        let encoder = RDEncoder::new(&values);
130        encoder.encode(arr.as_view())
131    })]
132    #[case::single({
133        let values = vec![42.42f64];
134        let arr = PrimitiveArray::from_iter(values.clone());
135        let encoder = RDEncoder::new(&values);
136        encoder.encode(arr.as_view())
137    })]
138    #[case::negative({
139        let values = vec![0.0f32, -1.5, 2.5, -3.5, 4.5];
140        let arr = PrimitiveArray::from_iter(values.clone());
141        let encoder = RDEncoder::new(&values);
142        encoder.encode(arr.as_view())
143    })]
144    #[case::nullable({
145        let arr = PrimitiveArray::from_option_iter([Some(1.1f32), None, Some(2.2), Some(3.3), None]);
146        let values = vec![1.1f32, 2.2, 3.3];
147        let encoder = RDEncoder::new(&values);
148        encoder.encode(arr.as_view())
149    })]
150    fn test_cast_alprd_conformance(#[case] alprd: crate::alp_rd::ALPRDArray) {
151        test_cast_conformance(
152            &alprd.into_array(),
153            &mut array_session().create_execution_ctx(),
154        );
155    }
156}