vortex_zstd/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::dtype::Nullability;
9use vortex_array::scalar_fn::fns::cast::CastReduce;
10use vortex_array::vtable::child_to_validity;
11use vortex_error::VortexResult;
12
13use crate::Zstd;
14use crate::ZstdData;
15
16impl CastReduce for Zstd {
17    fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
18        if !dtype.eq_ignore_nullability(array.dtype()) {
19            // Type changes can't be handled in ZSTD, need to decode and tweak.
20            // TODO(aduffy): handle trivial conversions like Binary -> UTF8, integer widening, etc.
21            return Ok(None);
22        }
23
24        let src_nullability = array.dtype().nullability();
25        let target_nullability = dtype.nullability();
26
27        let new_validity = match (src_nullability, target_nullability) {
28            // Same type case. This should be handled in the layer above but for
29            // completeness of the match arms we also handle it here.
30            (Nullability::Nullable, Nullability::Nullable)
31            | (Nullability::NonNullable, Nullability::NonNullable) => {
32                return Ok(Some(array.array().clone()));
33            }
34            (Nullability::NonNullable, Nullability::Nullable) => {
35                // nonnull => null, trivial cast by altering the validity
36                child_to_validity(array.slots()[0].as_ref(), array.dtype().nullability())
37            }
38            (Nullability::Nullable, Nullability::NonNullable) => {
39                // null => non-null works if there are no nulls in the sliced range
40                let unsliced_validity =
41                    child_to_validity(array.slots()[0].as_ref(), array.dtype().nullability());
42                let has_nulls = !unsliced_validity
43                    .slice(array.slice_start()..array.slice_stop())?
44                    .definitely_no_nulls();
45
46                // We don't attempt to handle casting when there are nulls.
47                if has_nulls {
48                    return Ok(None);
49                }
50                unsliced_validity
51            }
52        };
53
54        // If there are no nulls, the cast is trivial
55        Ok(Some(
56            Zstd::try_new(
57                dtype.clone(),
58                ZstdData::new(
59                    array.dictionary.clone(),
60                    array.frames.clone(),
61                    array.metadata.clone(),
62                    array.unsliced_n_rows(),
63                ),
64                new_validity,
65            )?
66            .into_array()
67            .slice(array.slice_start()..array.slice_stop())?,
68        ))
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use std::sync::LazyLock;
75
76    use rstest::rstest;
77    use vortex_array::IntoArray;
78    use vortex_array::VortexSessionExecute;
79    use vortex_array::arrays::PrimitiveArray;
80    use vortex_array::assert_arrays_eq;
81    use vortex_array::builtins::ArrayBuiltins;
82    use vortex_array::compute::conformance::cast::test_cast_conformance;
83    use vortex_array::dtype::DType;
84    use vortex_array::dtype::Nullability;
85    use vortex_array::dtype::PType;
86    use vortex_array::validity::Validity;
87    use vortex_buffer::buffer;
88    use vortex_session::VortexSession;
89
90    use crate::Zstd;
91
92    static SESSION: LazyLock<VortexSession> = LazyLock::new(vortex_array::array_session);
93
94    #[test]
95    fn test_cast_zstd_i32_to_i64() {
96        let mut ctx = SESSION.create_execution_ctx();
97        let values = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]);
98        let zstd = Zstd::from_primitive(&values, 0, 0, &mut ctx).unwrap();
99
100        let casted = zstd
101            .into_array()
102            .cast(DType::Primitive(PType::I64, Nullability::NonNullable))
103            .unwrap();
104        assert_eq!(
105            casted.dtype(),
106            &DType::Primitive(PType::I64, Nullability::NonNullable)
107        );
108
109        let decoded = casted.execute::<PrimitiveArray>(&mut ctx).unwrap();
110        assert_arrays_eq!(
111            decoded,
112            PrimitiveArray::from_iter([1i64, 2, 3, 4, 5]),
113            &mut ctx
114        );
115    }
116
117    #[test]
118    fn test_cast_zstd_nullability_change() {
119        let mut ctx = SESSION.create_execution_ctx();
120        let values = PrimitiveArray::from_iter([10u32, 20, 30, 40]);
121        let zstd = Zstd::from_primitive(&values, 0, 0, &mut ctx).unwrap();
122
123        let casted = zstd
124            .into_array()
125            .cast(DType::Primitive(PType::U32, Nullability::Nullable))
126            .unwrap();
127        assert_eq!(
128            casted.dtype(),
129            &DType::Primitive(PType::U32, Nullability::Nullable)
130        );
131    }
132
133    #[test]
134    fn test_cast_sliced_zstd_nullable_to_nonnullable() {
135        let mut ctx = SESSION.create_execution_ctx();
136        let values = PrimitiveArray::new(
137            buffer![10u32, 20, 30, 40, 50, 60],
138            Validity::from_iter([true, true, true, true, true, true]),
139        );
140        let zstd = Zstd::from_primitive(&values, 0, 128, &mut ctx).unwrap();
141        let sliced = zstd.slice(1..5).unwrap();
142        let casted = sliced
143            .cast(DType::Primitive(PType::U32, Nullability::NonNullable))
144            .unwrap();
145        assert_eq!(
146            casted.dtype(),
147            &DType::Primitive(PType::U32, Nullability::NonNullable)
148        );
149        // Verify the values are correct
150        let decoded = casted.execute::<PrimitiveArray>(&mut ctx).unwrap();
151        assert_arrays_eq!(
152            decoded,
153            PrimitiveArray::from_iter([20u32, 30, 40, 50]),
154            &mut ctx
155        );
156    }
157
158    #[test]
159    fn test_cast_sliced_zstd_part_valid_to_nonnullable() {
160        let mut ctx = SESSION.create_execution_ctx();
161        let values = PrimitiveArray::from_option_iter([
162            None,
163            Some(20u32),
164            Some(30),
165            Some(40),
166            Some(50),
167            Some(60),
168        ]);
169        let zstd = Zstd::from_primitive(&values, 0, 128, &mut ctx).unwrap();
170        let sliced = zstd.slice(1..5).unwrap();
171        let casted = sliced
172            .cast(DType::Primitive(PType::U32, Nullability::NonNullable))
173            .unwrap();
174        assert_eq!(
175            casted.dtype(),
176            &DType::Primitive(PType::U32, Nullability::NonNullable)
177        );
178        let decoded = casted.execute::<PrimitiveArray>(&mut ctx).unwrap();
179        let expected = PrimitiveArray::from_iter([20u32, 30, 40, 50]);
180        assert_arrays_eq!(decoded, expected, &mut ctx);
181    }
182
183    #[rstest]
184    #[case::i32(PrimitiveArray::new(
185        buffer![100i32, 200, 300, 400, 500],
186        Validity::NonNullable,
187    ))]
188    #[case::f64(PrimitiveArray::new(
189        buffer![1.1f64, 2.2, 3.3, 4.4, 5.5],
190        Validity::NonNullable,
191    ))]
192    #[case::single(PrimitiveArray::new(
193        buffer![42i64],
194        Validity::NonNullable,
195    ))]
196    #[case::large(PrimitiveArray::new(
197        buffer![0u32..1000],
198        Validity::NonNullable,
199    ))]
200    fn test_cast_zstd_conformance(#[case] values: PrimitiveArray) {
201        let zstd =
202            Zstd::from_primitive(&values, 0, 0, &mut SESSION.create_execution_ctx()).unwrap();
203        test_cast_conformance(&zstd.into_array(), &mut SESSION.create_execution_ctx());
204    }
205}
vortex_zstd/compute/cast.rs

vortex_zstd/compute/
cast.rs
