vortex_pco/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::dtype::DType;
8use vortex_array::scalar_fn::fns::cast::CastReduce;
9use vortex_error::VortexResult;
10
11use crate::Pco;
12use crate::PcoArrayExt;
13use crate::PcoData;
14
15impl CastReduce for Pco {
16    fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
17        // PCO (Pcodec) stores compressed data and uses validity bits to decode (the validity
18        // tells PCO which logical positions correspond to compressed values). Casting away
19        // nullability would change the validity-to-compressed-value mapping, so we cannot
20        // construct a non-nullable Pco without re-encoding — we only handle nullability changes
21        // toward `Nullable`. Non-nullable targets fall through to canonicalization.
22        //
23        // No `CastKernel` is provided for the same reason: even with execution context, we
24        // cannot cast away nullability on a PCO array in place.
25        //
26        // PCO supports: F16, F32, F64, I16, I32, I64, U16, U32, U64.
27        if !array.dtype().eq_ignore_nullability(dtype) {
28            return Ok(None);
29        }
30
31        let unsliced_validity = array.unsliced_validity();
32        let Some(new_validity) =
33            unsliced_validity.trivially_cast_nullability(dtype.nullability(), array.len())?
34        else {
35            return Ok(None);
36        };
37
38        let data = PcoData::new(
39            array.chunk_metas.clone(),
40            array.pages.clone(),
41            dtype.as_ptype(),
42            array.metadata.clone(),
43            array.unsliced_n_rows(),
44        )
45        ._slice(array.slice_start(), array.slice_stop());
46
47        Ok(Some(
48            Pco::try_new(dtype.clone(), data, new_validity)?.into_array(),
49        ))
50    }
51}
52
53#[cfg(test)]
54mod tests {
55    use std::sync::LazyLock;
56
57    use rstest::rstest;
58    use vortex_array::IntoArray;
59    use vortex_array::VortexSessionExecute;
60    use vortex_array::arrays::PrimitiveArray;
61    use vortex_array::assert_arrays_eq;
62    use vortex_array::builtins::ArrayBuiltins;
63    use vortex_array::compute::conformance::cast::test_cast_conformance;
64    use vortex_array::dtype::DType;
65    use vortex_array::dtype::Nullability;
66    use vortex_array::dtype::PType;
67    use vortex_array::validity::Validity;
68    use vortex_buffer::buffer;
69    use vortex_session::VortexSession;
70
71    use crate::Pco;
72
73    static SESSION: LazyLock<VortexSession> = LazyLock::new(vortex_array::array_session);
74
75    #[test]
76    fn test_cast_pco_f32_to_f64() {
77        let mut ctx = SESSION.create_execution_ctx();
78        let values = PrimitiveArray::from_iter([1.0f32, 2.0, 3.0, 4.0, 5.0]);
79        let pco = Pco::from_primitive(values.as_view(), 0, 128, &mut ctx).unwrap();
80
81        let casted = pco
82            .into_array()
83            .cast(DType::Primitive(PType::F64, Nullability::NonNullable))
84            .unwrap();
85        assert_eq!(
86            casted.dtype(),
87            &DType::Primitive(PType::F64, Nullability::NonNullable)
88        );
89
90        assert_arrays_eq!(
91            casted,
92            PrimitiveArray::from_iter([1.0f64, 2.0, 3.0, 4.0, 5.0]),
93            &mut ctx
94        );
95    }
96
97    #[test]
98    fn test_cast_pco_nullability_change() {
99        let mut ctx = SESSION.create_execution_ctx();
100        // Test casting from NonNullable to Nullable
101        let values = PrimitiveArray::from_iter([10u32, 20, 30, 40]);
102        let pco = Pco::from_primitive(values.as_view(), 0, 128, &mut ctx).unwrap();
103
104        let casted = pco
105            .into_array()
106            .cast(DType::Primitive(PType::U32, Nullability::Nullable))
107            .unwrap();
108        assert_arrays_eq!(
109            casted,
110            PrimitiveArray::new(buffer![10u32, 20, 30, 40], Validity::AllValid,),
111            &mut ctx
112        );
113    }
114
115    #[test]
116    fn test_cast_sliced_pco_nullable_to_nonnullable() {
117        let mut ctx = SESSION.create_execution_ctx();
118        let values = PrimitiveArray::new(
119            buffer![10u32, 20, 30, 40, 50, 60],
120            Validity::from_iter([true, true, true, true, true, true]),
121        );
122        let pco = Pco::from_primitive(values.as_view(), 0, 128, &mut ctx).unwrap();
123        let sliced = pco.slice(1..5).unwrap();
124        let casted = sliced
125            .cast(DType::Primitive(PType::U32, Nullability::NonNullable))
126            .unwrap();
127        assert_eq!(
128            casted.dtype(),
129            &DType::Primitive(PType::U32, Nullability::NonNullable)
130        );
131        // Verify the values are correct
132        assert_arrays_eq!(
133            casted,
134            PrimitiveArray::from_iter([20u32, 30, 40, 50]),
135            &mut ctx
136        );
137    }
138
139    #[test]
140    fn test_cast_sliced_pco_part_valid_to_nonnullable() {
141        let mut ctx = SESSION.create_execution_ctx();
142        let values = PrimitiveArray::from_option_iter([
143            None,
144            Some(20u32),
145            Some(30),
146            Some(40),
147            Some(50),
148            Some(60),
149        ]);
150        let pco = Pco::from_primitive(values.as_view(), 0, 128, &mut ctx).unwrap();
151        let sliced = pco.slice(1..5).unwrap();
152        let casted = sliced
153            .cast(DType::Primitive(PType::U32, Nullability::NonNullable))
154            .unwrap();
155        assert_eq!(
156            casted.dtype(),
157            &DType::Primitive(PType::U32, Nullability::NonNullable)
158        );
159        assert_arrays_eq!(
160            casted,
161            PrimitiveArray::from_iter([20u32, 30, 40, 50]),
162            &mut ctx
163        );
164    }
165
166    #[rstest]
167    #[case::f32(PrimitiveArray::new(
168        buffer![1.23f32, 4.56, 7.89, 10.11, 12.13],
169        Validity::NonNullable,
170    ))]
171    #[case::f64(PrimitiveArray::new(
172        buffer![100.1f64, 200.2, 300.3, 400.4, 500.5],
173        Validity::NonNullable,
174    ))]
175    #[case::i32(PrimitiveArray::new(
176        buffer![100i32, 200, 300, 400, 500],
177        Validity::NonNullable,
178    ))]
179    #[case::u64(PrimitiveArray::new(
180        buffer![1000u64, 2000, 3000, 4000],
181        Validity::NonNullable,
182    ))]
183    #[case::single(PrimitiveArray::new(
184        buffer![42.42f64],
185        Validity::NonNullable,
186    ))]
187    fn test_cast_pco_conformance(#[case] values: PrimitiveArray) {
188        let mut ctx = SESSION.create_execution_ctx();
189        let pco = Pco::from_primitive(values.as_view(), 0, 128, &mut ctx).unwrap();
190        test_cast_conformance(&pco.into_array(), &mut ctx);
191    }
192}
vortex_pco/compute/cast.rs

vortex_pco/compute/
cast.rs
