1use 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 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 let sparse = Sparse::try_new(
146 buffer![0u64, 1, 2, 3, 4].into_array(),
147 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 let sparse = Sparse::try_new(
175 buffer![1u64, 3].into_array(),
176 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}