Skip to main content

vortex_sparse/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::Scalar;
10use vortex_array::scalar_fn::fns::cast::CastReduce;
11use vortex_error::VortexResult;
12
13use crate::Sparse;
14use crate::SparseExt as _;
15
16impl CastReduce for Sparse {
17    fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
18        let casted_patches = array
19            .patches()
20            .map_values(|values| values.cast(dtype.clone()))?;
21
22        let casted_fill = if array.patches().num_patches() == array.len() {
23            // When every position is patched the fill scalar is unused and can be undefined.
24            // We skip casting it entirely and substitute a default value for the target dtype.
25            Scalar::default_value(dtype)
26        } else {
27            array.fill_scalar().cast(dtype)?
28        };
29
30        Ok(Some(
31            Sparse::try_new_from_patches(casted_patches, casted_fill)?.into_array(),
32        ))
33    }
34}
35
36#[cfg(test)]
37mod tests {
38    use std::sync::LazyLock;
39
40    use rstest::rstest;
41    use vortex_array::IntoArray;
42    use vortex_array::VortexSessionExecute;
43    use vortex_array::arrays::PrimitiveArray;
44    use vortex_array::assert_arrays_eq;
45    use vortex_array::builtins::ArrayBuiltins;
46    use vortex_array::compute::conformance::cast::test_cast_conformance;
47    use vortex_array::dtype::DType;
48    use vortex_array::dtype::Nullability;
49    use vortex_array::dtype::PType;
50    use vortex_array::scalar::Scalar;
51    use vortex_buffer::buffer;
52    use vortex_session::VortexSession;
53
54    use crate::Sparse;
55    use crate::SparseArray;
56
57    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
58        let session = vortex_array::array_session();
59        crate::initialize(&session);
60        session
61    });
62
63    #[test]
64    fn test_cast_sparse_i32_to_i64() {
65        let mut ctx = SESSION.create_execution_ctx();
66        let sparse = Sparse::try_new(
67            buffer![2u64, 5, 8].into_array(),
68            buffer![100i32, 200, 300].into_array(),
69            10,
70            Scalar::from(0i32),
71        )
72        .unwrap();
73
74        let casted = sparse
75            .into_array()
76            .cast(DType::Primitive(PType::I64, Nullability::NonNullable))
77            .unwrap();
78        assert_eq!(
79            casted.dtype(),
80            &DType::Primitive(PType::I64, Nullability::NonNullable)
81        );
82
83        let expected = PrimitiveArray::from_iter([0i64, 0, 100, 0, 0, 200, 0, 0, 300, 0]);
84        let casted_primitive = casted.execute::<PrimitiveArray>(&mut ctx).unwrap();
85        assert_arrays_eq!(casted_primitive, expected, &mut ctx);
86    }
87
88    #[test]
89    fn test_cast_sparse_with_null_fill() {
90        let sparse = Sparse::try_new(
91            buffer![1u64, 3, 5].into_array(),
92            PrimitiveArray::from_option_iter([Some(42i32), Some(84), Some(126)]).into_array(),
93            8,
94            Scalar::null_native::<i32>(),
95        )
96        .unwrap();
97
98        let casted = sparse
99            .into_array()
100            .cast(DType::Primitive(PType::I64, Nullability::Nullable))
101            .unwrap();
102        assert_eq!(
103            casted.dtype(),
104            &DType::Primitive(PType::I64, Nullability::Nullable)
105        );
106    }
107
108    #[rstest]
109    #[case(Sparse::try_new(
110        buffer![2u64, 5, 8].into_array(),
111        buffer![100i32, 200, 300].into_array(),
112        10,
113        Scalar::from(0i32)
114    ).unwrap())]
115    #[case(Sparse::try_new(
116        buffer![0u64, 4, 9].into_array(),
117        buffer![1.5f32, 2.5, 3.5].into_array(),
118        10,
119        Scalar::from(0.0f32)
120    ).unwrap())]
121    #[case(Sparse::try_new(
122        buffer![1u64, 3, 7].into_array(),
123        PrimitiveArray::from_option_iter([Some(100i32), None, Some(300)]).into_array(),
124        10,
125        Scalar::null_native::<i32>()
126    ).unwrap())]
127    #[case(Sparse::try_new(
128        buffer![5u64].into_array(),
129        buffer![42u8].into_array(),
130        10,
131        Scalar::from(0u8)
132    ).unwrap())]
133    fn test_cast_sparse_conformance(#[case] array: SparseArray) {
134        test_cast_conformance(&array.into_array());
135    }
136
137    #[test]
138    fn test_cast_sparse_null_fill_all_patched_to_non_nullable() -> vortex_error::VortexResult<()> {
139        let mut ctx = SESSION.create_execution_ctx();
140        // Regression test for https://github.com/vortex-data/vortex/issues/6932
141        //
142        // When all positions are patched the null fill is unused, so a cast to
143        // non-nullable is valid.  Sparse::cast detects this case, substitutes a
144        // zero fill, and keeps the result in the Sparse encoding.
145        let sparse = Sparse::try_new(
146            buffer![0u64, 1, 2, 3, 4].into_array(),
147            // Nullable values to match the null (nullable) fill.
148            PrimitiveArray::from_option_iter([Some(10u64), Some(20), Some(30), Some(40), Some(50)])
149                .into_array(),
150            5,
151            Scalar::null_native::<u64>(),
152        )?;
153
154        let casted = sparse
155            .into_array()
156            .cast(DType::Primitive(PType::U64, Nullability::NonNullable))?;
157
158        assert_eq!(
159            casted.dtype(),
160            &DType::Primitive(PType::U64, Nullability::NonNullable)
161        );
162
163        let expected = PrimitiveArray::from_iter([10u64, 20, 30, 40, 50]);
164        let casted_primitive = casted.execute::<PrimitiveArray>(&mut ctx)?;
165        assert_arrays_eq!(casted_primitive, expected, &mut ctx);
166        Ok(())
167    }
168
169    #[test]
170    fn test_fill_null_sparse_with_null_fill() -> vortex_error::VortexResult<()> {
171        // Regression test for https://github.com/vortex-data/vortex/issues/6932
172        // fill_null on a sparse array with null fill triggers an internal cast to
173        // non-nullable, which must not panic.
174        let sparse = Sparse::try_new(
175            buffer![1u64, 3].into_array(),
176            // Nullable values to match the null (nullable) fill.
177            PrimitiveArray::from_option_iter([Some(10u64), Some(20)]).into_array(),
178            5,
179            Scalar::null_native::<u64>(),
180        )?;
181
182        let filled = sparse.into_array().fill_null(Scalar::from(0u64))?;
183
184        assert_eq!(
185            filled.dtype(),
186            &DType::Primitive(PType::U64, Nullability::NonNullable)
187        );
188        Ok(())
189    }
190}