vortex_array/arrays/varbinview/compute/
cast.rs

1use vortex_dtype::DType;
2use vortex_error::{VortexResult, vortex_bail};
3
4use crate::arrays::{VarBinViewArray, VarBinViewVTable};
5use crate::compute::{CastKernel, CastKernelAdapter};
6use crate::vtable::ValidityHelper;
7use crate::{ArrayRef, IntoArray, register_kernel};
8
9impl CastKernel for VarBinViewVTable {
10    fn cast(&self, array: &VarBinViewArray, dtype: &DType) -> VortexResult<ArrayRef> {
11        if !array.dtype().eq_ignore_nullability(dtype) {
12            vortex_bail!("Cannot cast {} to {}", array.dtype(), dtype);
13        }
14
15        let new_nullability = dtype.nullability();
16        let new_validity = array.validity().clone().cast_nullability(new_nullability)?;
17        let new_dtype = array.dtype().with_nullability(new_nullability);
18        Ok(VarBinViewArray::try_new(
19            array.views().clone(),
20            array.buffers().to_vec(),
21            new_dtype,
22            new_validity,
23        )?
24        .into_array())
25    }
26}
27
28register_kernel!(CastKernelAdapter(VarBinViewVTable).lift());
29
30#[cfg(test)]
31mod tests {
32    use rstest::rstest;
33    use vortex_dtype::{DType, Nullability};
34
35    use crate::arrays::VarBinViewArray;
36    use crate::compute::cast;
37
38    #[rstest]
39    #[case(
40        DType::Utf8(Nullability::Nullable),
41        DType::Utf8(Nullability::NonNullable)
42    )]
43    #[case(
44        DType::Binary(Nullability::Nullable),
45        DType::Binary(Nullability::NonNullable)
46    )]
47    #[case(
48        DType::Utf8(Nullability::NonNullable),
49        DType::Utf8(Nullability::Nullable)
50    )]
51    #[case(
52        DType::Binary(Nullability::NonNullable),
53        DType::Binary(Nullability::Nullable)
54    )]
55    fn try_cast_varbin_nullable(#[case] source: DType, #[case] target: DType) {
56        let varbin = VarBinViewArray::from_iter(vec![Some("a"), Some("b"), Some("c")], source);
57
58        let res = cast(varbin.as_ref(), &target);
59        assert_eq!(res.unwrap().dtype(), &target);
60    }
61
62    #[rstest]
63    #[should_panic]
64    #[case(DType::Utf8(Nullability::Nullable))]
65    #[should_panic]
66    #[case(DType::Binary(Nullability::Nullable))]
67    fn try_cast_varbin_fail(#[case] source: DType) {
68        let non_nullable_source = source.as_nonnullable();
69        let varbin = VarBinViewArray::from_iter(vec![Some("a"), Some("b"), None], source);
70        cast(varbin.as_ref(), &non_nullable_source).unwrap();
71    }
72}