vortex_array/compute/conformance/
cast.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexExpect;
5use vortex_error::VortexResult;
6use vortex_error::vortex_panic;
7
8use crate::ArrayRef;
9use crate::ExecutionCtx;
10use crate::IntoArray;
11use crate::RecursiveCanonical;
12use crate::aggregate_fn::NumericalAggregateOpts;
13use crate::aggregate_fn::fns::min_max::MinMaxResult;
14use crate::aggregate_fn::fns::min_max::min_max;
15use crate::builtins::ArrayBuiltins;
16use crate::dtype::DType;
17use crate::dtype::Nullability;
18use crate::dtype::PType;
19use crate::scalar::Scalar;
20
21/// Cast and force execution via `to_canonical`, returning the canonical array.
22fn cast_and_execute(
23    array: &ArrayRef,
24    dtype: DType,
25    ctx: &mut ExecutionCtx,
26) -> VortexResult<ArrayRef> {
27    Ok(array
28        .cast(dtype)?
29        .execute::<RecursiveCanonical>(ctx)?
30        .0
31        .into_array())
32}
33
34/// Test conformance of the cast compute function for an array.
35///
36/// This function tests various casting scenarios including:
37/// - Casting between numeric types (widening and narrowing)
38/// - Casting between signed and unsigned types
39/// - Casting between integral and floating-point types
40/// - Casting with nullability changes
41/// - Casting between string types (Utf8/Binary)
42/// - Edge cases like overflow behavior
43pub fn test_cast_conformance(array: &ArrayRef, ctx: &mut ExecutionCtx) {
44    let dtype = array.dtype();
45
46    // Always test identity cast and nullability changes
47    test_cast_identity(array, ctx);
48
49    test_cast_to_non_nullable(array, ctx);
50    test_cast_to_nullable(array, ctx);
51
52    // Test based on the specific DType
53    match dtype {
54        DType::Null => test_cast_from_null(array, ctx),
55        DType::Primitive(ptype, ..) => match ptype {
56            PType::U8
57            | PType::U16
58            | PType::U32
59            | PType::U64
60            | PType::I8
61            | PType::I16
62            | PType::I32
63            | PType::I64 => test_cast_to_integral_types(array, ctx),
64            PType::F16 | PType::F32 | PType::F64 => test_cast_from_floating_point_types(array, ctx),
65        },
66        _ => {}
67    }
68}
69
70fn test_cast_identity(array: &ArrayRef, ctx: &mut ExecutionCtx) {
71    // Casting to the same type should be a no-op
72    let result = cast_and_execute(&array.clone(), array.dtype().clone(), ctx)
73        .vortex_expect("cast should succeed in conformance test");
74    assert_eq!(result.len(), array.len());
75    assert_eq!(result.dtype(), array.dtype());
76
77    // Verify values are unchanged
78    for i in 0..array.len().min(10) {
79        assert_eq!(
80            array
81                .execute_scalar(i, ctx)
82                .vortex_expect("scalar_at should succeed in conformance test"),
83            result
84                .execute_scalar(i, ctx)
85                .vortex_expect("scalar_at should succeed in conformance test")
86        );
87    }
88}
89
90fn test_cast_from_null(array: &ArrayRef, ctx: &mut ExecutionCtx) {
91    // Null can also be cast to any nullable type
92    let nullable_types = vec![
93        DType::Bool(Nullability::Nullable),
94        DType::Primitive(PType::I32, Nullability::Nullable),
95        DType::Primitive(PType::F64, Nullability::Nullable),
96        DType::Utf8(Nullability::Nullable),
97        DType::Binary(Nullability::Nullable),
98    ];
99
100    for dtype in nullable_types {
101        let result = cast_and_execute(&array.clone(), dtype.clone(), ctx)
102            .vortex_expect("cast should succeed in conformance test");
103        assert_eq!(result.len(), array.len());
104        assert_eq!(result.dtype(), &dtype);
105
106        // Verify all values are null
107        for i in 0..array.len().min(10) {
108            assert!(
109                result
110                    .execute_scalar(i, ctx)
111                    .vortex_expect("scalar_at should succeed in conformance test")
112                    .is_null()
113            );
114        }
115    }
116
117    // Casting to non-nullable types should fail
118    let non_nullable_types = vec![
119        DType::Bool(Nullability::NonNullable),
120        DType::Primitive(PType::I32, Nullability::NonNullable),
121    ];
122
123    for dtype in non_nullable_types {
124        assert!(cast_and_execute(&array.clone(), dtype.clone(), ctx).is_err());
125    }
126}
127
128fn test_cast_to_non_nullable(array: &ArrayRef, ctx: &mut ExecutionCtx) {
129    // DType::Null has no non-nullable form
130    if &DType::Null == array.dtype() || !array.dtype().is_nullable() {
131        return;
132    }
133
134    if array
135        .invalid_count(ctx)
136        .vortex_expect("invalid_count should succeed in conformance test")
137        == 0
138    {
139        let non_nullable = cast_and_execute(&array.clone(), array.dtype().as_nonnullable(), ctx)
140            .vortex_expect("arrays without nulls can cast to non-nullable");
141        assert_eq!(non_nullable.dtype(), &array.dtype().as_nonnullable());
142        assert_eq!(non_nullable.len(), array.len());
143
144        for i in 0..array.len().min(10) {
145            assert_eq!(
146                array
147                    .execute_scalar(i, ctx)
148                    .vortex_expect("scalar_at should succeed in conformance test"),
149                non_nullable
150                    .execute_scalar(i, ctx)
151                    .vortex_expect("scalar_at should succeed in conformance test")
152            );
153        }
154
155        let back_to_nullable = cast_and_execute(&non_nullable, array.dtype().clone(), ctx)
156            .vortex_expect("non-nullable arrays can cast to nullable");
157        assert_eq!(back_to_nullable.dtype(), array.dtype());
158        assert_eq!(back_to_nullable.len(), array.len());
159
160        for i in 0..array.len().min(10) {
161            assert_eq!(
162                array
163                    .execute_scalar(i, ctx)
164                    .vortex_expect("scalar_at should succeed in conformance test"),
165                back_to_nullable
166                    .execute_scalar(i, ctx)
167                    .vortex_expect("scalar_at should succeed in conformance test")
168            );
169        }
170    } else {
171        cast_and_execute(&array.clone(), array.dtype().as_nonnullable(), ctx)
172            .err()
173            .unwrap_or_else(|| {
174                vortex_panic!(
175                    "arrays with nulls should error when casting to non-nullable {}",
176                    array,
177                )
178            });
179    }
180}
181
182fn test_cast_to_nullable(array: &ArrayRef, ctx: &mut ExecutionCtx) {
183    if array.dtype().is_nullable() {
184        return;
185    }
186
187    let nullable = cast_and_execute(&array.clone(), array.dtype().as_nullable(), ctx)
188        .vortex_expect("arrays without nulls can cast to nullable");
189    assert_eq!(nullable.dtype(), &array.dtype().as_nullable());
190    assert_eq!(nullable.len(), array.len());
191
192    for i in 0..array.len().min(10) {
193        assert_eq!(
194            array
195                .execute_scalar(i, ctx)
196                .vortex_expect("scalar_at should succeed in conformance test"),
197            nullable
198                .execute_scalar(i, ctx)
199                .vortex_expect("scalar_at should succeed in conformance test")
200        );
201    }
202
203    let back = cast_and_execute(&nullable, array.dtype().clone(), ctx)
204        .vortex_expect("casting to nullable and back should be a no-op");
205    assert_eq!(back.dtype(), array.dtype());
206    assert_eq!(back.len(), array.len());
207
208    for i in 0..array.len().min(10) {
209        assert_eq!(
210            array
211                .execute_scalar(i, ctx)
212                .vortex_expect("scalar_at should succeed in conformance test"),
213            back.execute_scalar(i, ctx)
214                .vortex_expect("scalar_at should succeed in conformance test")
215        );
216    }
217}
218
219fn test_cast_from_floating_point_types(array: &ArrayRef, ctx: &mut ExecutionCtx) {
220    let ptype = array.dtype().as_ptype();
221    test_cast_to_primitive(array, PType::I8, false, ctx);
222    test_cast_to_primitive(array, PType::U8, false, ctx);
223    test_cast_to_primitive(array, PType::I16, false, ctx);
224    test_cast_to_primitive(array, PType::U16, false, ctx);
225    test_cast_to_primitive(array, PType::I32, false, ctx);
226    test_cast_to_primitive(array, PType::U32, false, ctx);
227    test_cast_to_primitive(array, PType::I64, false, ctx);
228    test_cast_to_primitive(array, PType::U64, false, ctx);
229    test_cast_to_primitive(array, PType::F16, matches!(ptype, PType::F16), ctx);
230    test_cast_to_primitive(
231        array,
232        PType::F32,
233        matches!(ptype, PType::F16 | PType::F32),
234        ctx,
235    );
236    test_cast_to_primitive(array, PType::F64, true, ctx);
237}
238
239fn test_cast_to_integral_types(array: &ArrayRef, ctx: &mut ExecutionCtx) {
240    test_cast_to_primitive(array, PType::I8, true, ctx);
241    test_cast_to_primitive(array, PType::U8, true, ctx);
242    test_cast_to_primitive(array, PType::I16, true, ctx);
243    test_cast_to_primitive(array, PType::U16, true, ctx);
244    test_cast_to_primitive(array, PType::I32, true, ctx);
245    test_cast_to_primitive(array, PType::U32, true, ctx);
246    test_cast_to_primitive(array, PType::I64, true, ctx);
247    test_cast_to_primitive(array, PType::U64, true, ctx);
248}
249
250/// Does this scalar fit in this type?
251fn fits(value: &Scalar, ptype: PType) -> bool {
252    let dtype = DType::Primitive(ptype, value.dtype().nullability());
253    value.cast(&dtype).is_ok()
254}
255
256fn test_cast_to_primitive(
257    array: &ArrayRef,
258    target_ptype: PType,
259    test_round_trip: bool,
260    ctx: &mut ExecutionCtx,
261) {
262    let maybe_min_max = min_max(array, ctx, NumericalAggregateOpts::default())
263        .vortex_expect("cast should succeed in conformance test");
264
265    if let Some(MinMaxResult { min, max }) = maybe_min_max
266        && (!fits(&min, target_ptype) || !fits(&max, target_ptype))
267    {
268        cast_and_execute(
269            &array.clone(),
270            DType::Primitive(target_ptype, array.dtype().nullability()),
271            ctx,
272        )
273        .err()
274        .unwrap_or_else(|| {
275            vortex_panic!(
276                "Cast must fail because some values are out of bounds. {} {:?} {:?} {} {}",
277                target_ptype,
278                min,
279                max,
280                array,
281                array.display_values(ctx),
282            )
283        });
284        return;
285    }
286
287    // Otherwise, all values must fit.
288    let target_dtype = DType::Primitive(target_ptype, array.dtype().nullability());
289    let casted = cast_and_execute(array, target_dtype.clone(), ctx).unwrap_or_else(|e| {
290        vortex_panic!(
291            "Cast must succeed because all values are within bounds. {} {}: {e}",
292            target_ptype,
293            array.display_values(ctx),
294        )
295    });
296    assert_eq!(casted.dtype(), &target_dtype);
297    assert_eq!(casted.len(), array.len());
298    assert_eq!(
299        array
300            .validity()
301            .vortex_expect("validity_mask should succeed in conformance test")
302            .execute_mask(array.len(), ctx)
303            .vortex_expect("Failed to compute validity mask"),
304        casted
305            .validity()
306            .vortex_expect("validity_mask should succeed in conformance test")
307            .execute_mask(casted.len(), ctx)
308            .vortex_expect("Failed to compute validity mask")
309    );
310    for i in 0..array.len().min(10) {
311        let original = array
312            .execute_scalar(i, ctx)
313            .vortex_expect("scalar_at should succeed in conformance test");
314        let casted = casted
315            .execute_scalar(i, ctx)
316            .vortex_expect("scalar_at should succeed in conformance test");
317        assert_eq!(
318            original
319                .cast(&target_dtype)
320                .vortex_expect("cast should succeed in conformance test"),
321            casted,
322            "{i} {original} {casted}"
323        );
324        if test_round_trip {
325            assert_eq!(
326                original,
327                casted
328                    .cast(original.dtype())
329                    .vortex_expect("cast should succeed in conformance test"),
330                "{i} {original} {casted}"
331            );
332        }
333    }
334}
335
336#[cfg(test)]
337mod tests {
338    use std::sync::LazyLock;
339
340    use vortex_buffer::buffer;
341    use vortex_session::VortexSession;
342
343    use super::*;
344    use crate::IntoArray;
345    use crate::VortexSessionExecute;
346    use crate::array_session;
347    use crate::arrays::BoolArray;
348    use crate::arrays::ListArray;
349    use crate::arrays::NullArray;
350    use crate::arrays::PrimitiveArray;
351    use crate::arrays::StructArray;
352    use crate::arrays::VarBinArray;
353    use crate::dtype::DType;
354    use crate::dtype::FieldNames;
355    use crate::dtype::Nullability;
356
357    static SESSION: LazyLock<VortexSession> = LazyLock::new(array_session);
358
359    #[test]
360    fn test_cast_conformance_u32() {
361        let array = buffer![0u32, 100, 200, 65535, 1000000].into_array();
362        test_cast_conformance(&array, &mut SESSION.create_execution_ctx());
363    }
364
365    #[test]
366    fn test_cast_conformance_i32() {
367        let array = buffer![-100i32, -1, 0, 1, 100].into_array();
368        test_cast_conformance(&array, &mut SESSION.create_execution_ctx());
369    }
370
371    #[test]
372    fn test_cast_conformance_f32() {
373        let array = buffer![0.0f32, 1.5, -2.5, 100.0, 1e6].into_array();
374        test_cast_conformance(&array, &mut SESSION.create_execution_ctx());
375    }
376
377    #[test]
378    fn test_cast_conformance_nullable() {
379        let array = PrimitiveArray::from_option_iter([Some(1u8), None, Some(255), Some(0), None]);
380        test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
381    }
382
383    #[test]
384    fn test_cast_conformance_bool() {
385        let array = BoolArray::from_iter(vec![true, false, true, false]);
386        test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
387    }
388
389    #[test]
390    fn test_cast_conformance_null() {
391        let array = NullArray::new(5);
392        test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
393    }
394
395    #[test]
396    fn test_cast_conformance_utf8() {
397        let array = VarBinArray::from_iter(
398            vec![Some("hello"), None, Some("world")],
399            DType::Utf8(Nullability::Nullable),
400        );
401        test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
402    }
403
404    #[test]
405    fn test_cast_conformance_binary() {
406        let array = VarBinArray::from_iter(
407            vec![Some(b"data".as_slice()), None, Some(b"bytes".as_slice())],
408            DType::Binary(Nullability::Nullable),
409        );
410        test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
411    }
412
413    #[test]
414    fn test_cast_conformance_struct() {
415        let names = FieldNames::from(["a", "b"]);
416
417        let a = buffer![1i32, 2, 3].into_array();
418        let b = VarBinArray::from_iter(
419            vec![Some("x"), None, Some("z")],
420            DType::Utf8(Nullability::Nullable),
421        )
422        .into_array();
423
424        let array =
425            StructArray::try_new(names, vec![a, b], 3, crate::validity::Validity::NonNullable)
426                .unwrap();
427        test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
428    }
429
430    #[test]
431    fn test_cast_conformance_list() {
432        let mut ctx = SESSION.create_execution_ctx();
433        let data = buffer![1i32, 2, 3, 4, 5, 6].into_array();
434        let offsets = buffer![0i64, 2, 2, 5, 6].into_array();
435
436        let array = ListArray::try_new(
437            data,
438            offsets,
439            crate::validity::Validity::NonNullable,
440            &mut ctx,
441        )
442        .unwrap();
443        test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
444    }
445}
vortex_array/compute/conformance/cast.rs

vortex_array/compute/conformance/
cast.rs
