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_array::vtable::child_to_validity;
10use vortex_error::VortexResult;
11
12use crate::Pco;
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 =
32            child_to_validity(array.slots()[0].as_ref(), array.dtype().nullability());
33        let Some(new_validity) =
34            unsliced_validity.trivially_cast_nullability(dtype.nullability(), array.len())?
35        else {
36            return Ok(None);
37        };
38
39        let data = PcoData::new(
40            array.chunk_metas.clone(),
41            array.pages.clone(),
42            dtype.as_ptype(),
43            array.metadata.clone(),
44            array.unsliced_n_rows(),
45        )
46        ._slice(array.slice_start(), array.slice_stop());
47
48        Ok(Some(
49            Pco::try_new(dtype.clone(), data, new_validity)?.into_array(),
50        ))
51    }
52}
53
54#[cfg(test)]
55mod tests {
56    use std::sync::LazyLock;
57
58    use rstest::rstest;
59    use vortex_array::IntoArray;
60    use vortex_array::VortexSessionExecute;
61    use vortex_array::arrays::PrimitiveArray;
62    use vortex_array::assert_arrays_eq;
63    use vortex_array::builtins::ArrayBuiltins;
64    use vortex_array::compute::conformance::cast::test_cast_conformance;
65    use vortex_array::dtype::DType;
66    use vortex_array::dtype::Nullability;
67    use vortex_array::dtype::PType;
68    use vortex_array::validity::Validity;
69    use vortex_buffer::buffer;
70    use vortex_session::VortexSession;
71
72    use crate::Pco;
73
74    static SESSION: LazyLock<VortexSession> = LazyLock::new(vortex_array::array_session);
75
76    #[test]
77    fn test_cast_pco_f32_to_f64() {
78        let mut ctx = SESSION.create_execution_ctx();
79        let values = PrimitiveArray::from_iter([1.0f32, 2.0, 3.0, 4.0, 5.0]);
80        let pco = Pco::from_primitive(values.as_view(), 0, 128, &mut ctx).unwrap();
81
82        let casted = pco
83            .into_array()
84            .cast(DType::Primitive(PType::F64, Nullability::NonNullable))
85            .unwrap();
86        assert_eq!(
87            casted.dtype(),
88            &DType::Primitive(PType::F64, Nullability::NonNullable)
89        );
90
91        assert_arrays_eq!(
92            casted,
93            PrimitiveArray::from_iter([1.0f64, 2.0, 3.0, 4.0, 5.0]),
94            &mut ctx
95        );
96    }
97
98    #[test]
99    fn test_cast_pco_nullability_change() {
100        let mut ctx = SESSION.create_execution_ctx();
101        // Test casting from NonNullable to Nullable
102        let values = PrimitiveArray::from_iter([10u32, 20, 30, 40]);
103        let pco = Pco::from_primitive(values.as_view(), 0, 128, &mut ctx).unwrap();
104
105        let casted = pco
106            .into_array()
107            .cast(DType::Primitive(PType::U32, Nullability::Nullable))
108            .unwrap();
109        assert_arrays_eq!(
110            casted,
111            PrimitiveArray::new(buffer![10u32, 20, 30, 40], Validity::AllValid,),
112            &mut ctx
113        );
114    }
115
116    #[test]
117    fn test_cast_sliced_pco_nullable_to_nonnullable() {
118        let mut ctx = SESSION.create_execution_ctx();
119        let values = PrimitiveArray::new(
120            buffer![10u32, 20, 30, 40, 50, 60],
121            Validity::from_iter([true, true, true, true, true, true]),
122        );
123        let pco = Pco::from_primitive(values.as_view(), 0, 128, &mut ctx).unwrap();
124        let sliced = pco.slice(1..5).unwrap();
125        let casted = sliced
126            .cast(DType::Primitive(PType::U32, Nullability::NonNullable))
127            .unwrap();
128        assert_eq!(
129            casted.dtype(),
130            &DType::Primitive(PType::U32, Nullability::NonNullable)
131        );
132        // Verify the values are correct
133        assert_arrays_eq!(
134            casted,
135            PrimitiveArray::from_iter([20u32, 30, 40, 50]),
136            &mut ctx
137        );
138    }
139
140    #[test]
141    fn test_cast_sliced_pco_part_valid_to_nonnullable() {
142        let mut ctx = SESSION.create_execution_ctx();
143        let values = PrimitiveArray::from_option_iter([
144            None,
145            Some(20u32),
146            Some(30),
147            Some(40),
148            Some(50),
149            Some(60),
150        ]);
151        let pco = Pco::from_primitive(values.as_view(), 0, 128, &mut ctx).unwrap();
152        let sliced = pco.slice(1..5).unwrap();
153        let casted = sliced
154            .cast(DType::Primitive(PType::U32, Nullability::NonNullable))
155            .unwrap();
156        assert_eq!(
157            casted.dtype(),
158            &DType::Primitive(PType::U32, Nullability::NonNullable)
159        );
160        assert_arrays_eq!(
161            casted,
162            PrimitiveArray::from_iter([20u32, 30, 40, 50]),
163            &mut ctx
164        );
165    }
166
167    #[rstest]
168    #[case::f32(PrimitiveArray::new(
169        buffer![1.23f32, 4.56, 7.89, 10.11, 12.13],
170        Validity::NonNullable,
171    ))]
172    #[case::f64(PrimitiveArray::new(
173        buffer![100.1f64, 200.2, 300.3, 400.4, 500.5],
174        Validity::NonNullable,
175    ))]
176    #[case::i32(PrimitiveArray::new(
177        buffer![100i32, 200, 300, 400, 500],
178        Validity::NonNullable,
179    ))]
180    #[case::u64(PrimitiveArray::new(
181        buffer![1000u64, 2000, 3000, 4000],
182        Validity::NonNullable,
183    ))]
184    #[case::single(PrimitiveArray::new(
185        buffer![42.42f64],
186        Validity::NonNullable,
187    ))]
188    fn test_cast_pco_conformance(#[case] values: PrimitiveArray) {
189        let mut ctx = SESSION.create_execution_ctx();
190        let pco = Pco::from_primitive(values.as_view(), 0, 128, &mut ctx).unwrap();
191        test_cast_conformance(&pco.into_array());
192    }
193}
vortex_pco/compute/cast.rs

vortex_pco/compute/
cast.rs
