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

vortex_pco/compute/
cast.rs
