vortex_array/arrays/varbin/compute/
cast.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5
6use crate::ArrayRef;
7use crate::ExecutionCtx;
8use crate::IntoArray;
9use crate::array::ArrayView;
10use crate::arrays::VarBin;
11use crate::arrays::VarBinArray;
12use crate::arrays::varbin::VarBinArraySlotsExt;
13use crate::dtype::DType;
14use crate::scalar_fn::fns::cast::CastKernel;
15use crate::scalar_fn::fns::cast::CastReduce;
16use crate::validity::Validity;
17
18fn build_with_validity(
19    array: ArrayView<'_, VarBin>,
20    new_dtype: DType,
21    new_validity: Validity,
22) -> VortexResult<ArrayRef> {
23    // SAFETY: the offsets and bytes are unchanged from a valid array, and the validity was cast
24    // from the array's own validity, so it has the same length.
25    Ok(unsafe {
26        VarBinArray::new_unchecked_from_handle(
27            array.offsets().clone(),
28            array.bytes_handle().clone(),
29            new_dtype,
30            new_validity,
31        )
32    }
33    .into_array())
34}
35
36impl CastReduce for VarBin {
37    fn cast(array: ArrayView<'_, VarBin>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
38        if !array.dtype().eq_ignore_nullability(dtype) {
39            return Ok(None);
40        }
41
42        let new_nullability = dtype.nullability();
43        let Some(new_validity) = array
44            .validity()?
45            .trivially_cast_nullability(new_nullability, array.len())?
46        else {
47            return Ok(None);
48        };
49        let new_dtype = array.dtype().with_nullability(new_nullability);
50        Ok(Some(build_with_validity(array, new_dtype, new_validity)?))
51    }
52}
53
54impl CastKernel for VarBin {
55    fn cast(
56        array: ArrayView<'_, VarBin>,
57        dtype: &DType,
58        ctx: &mut ExecutionCtx,
59    ) -> VortexResult<Option<ArrayRef>> {
60        if !array.dtype().eq_ignore_nullability(dtype) {
61            return Ok(None);
62        }
63
64        let new_nullability = dtype.nullability();
65        let new_validity = array
66            .validity()?
67            .cast_nullability(new_nullability, array.len(), ctx)?;
68        let new_dtype = array.dtype().with_nullability(new_nullability);
69        Ok(Some(build_with_validity(array, new_dtype, new_validity)?))
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use std::sync::LazyLock;
76
77    use rstest::rstest;
78    use vortex_session::VortexSession;
79
80    use crate::Canonical;
81    use crate::IntoArray;
82    use crate::VortexSessionExecute;
83    use crate::arrays::VarBinArray;
84    use crate::arrays::VarBinViewArray;
85    use crate::assert_arrays_eq;
86    use crate::builtins::ArrayBuiltins;
87    use crate::compute::conformance::cast::test_cast_conformance;
88    use crate::dtype::DType;
89    use crate::dtype::Nullability;
90
91    static SESSION: LazyLock<VortexSession> = LazyLock::new(crate::array_session);
92
93    #[rstest]
94    #[case(
95        DType::Utf8(Nullability::Nullable),
96        DType::Utf8(Nullability::NonNullable)
97    )]
98    #[case(
99        DType::Binary(Nullability::Nullable),
100        DType::Binary(Nullability::NonNullable)
101    )]
102    #[case(
103        DType::Utf8(Nullability::NonNullable),
104        DType::Utf8(Nullability::Nullable)
105    )]
106    #[case(
107        DType::Binary(Nullability::NonNullable),
108        DType::Binary(Nullability::Nullable)
109    )]
110    fn try_cast_varbin_nullable(#[case] source: DType, #[case] target: DType) {
111        let varbin = VarBinArray::from_iter(vec![Some("a"), Some("b"), Some("c")], source);
112
113        let res = varbin.into_array().cast(target.clone());
114        assert_arrays_eq!(
115            res.unwrap(),
116            VarBinViewArray::from_iter([Some("a"), Some("b"), Some("c")], target),
117            &mut SESSION.create_execution_ctx()
118        );
119    }
120
121    #[rstest]
122    #[case(DType::Utf8(Nullability::Nullable))]
123    #[case(DType::Binary(Nullability::Nullable))]
124    fn try_cast_varbin_fail(#[case] source: DType) {
125        // Failure surfaces during execution via the kernel.
126        let non_nullable_source = source.as_nonnullable();
127        let varbin = VarBinArray::from_iter(vec![Some("a"), Some("b"), None], source);
128        let mut ctx = SESSION.create_execution_ctx();
129        let result = varbin
130            .into_array()
131            .cast(non_nullable_source)
132            .and_then(|a| a.execute::<Canonical>(&mut ctx).map(|c| c.into_array()));
133        assert!(result.is_err(), "Expected error, got: {result:?}");
134    }
135
136    #[rstest]
137    #[case(VarBinArray::from_iter(vec![Some("hello"), Some("world"), Some("test")], DType::Utf8(Nullability::NonNullable)))]
138    #[case(VarBinArray::from_iter(vec![Some("hello"), None, Some("world")], DType::Utf8(Nullability::Nullable)))]
139    #[case(VarBinArray::from_iter(vec![Some(b"binary".as_slice()), Some(b"data".as_slice())], DType::Binary(Nullability::NonNullable)))]
140    #[case(VarBinArray::from_iter(vec![Some(b"test".as_slice()), None], DType::Binary(Nullability::Nullable)))]
141    #[case(VarBinArray::from_iter(vec![Some("single")], DType::Utf8(Nullability::NonNullable)))]
142    fn test_cast_varbin_conformance(#[case] array: VarBinArray) {
143        test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
144    }
145}
vortex_array/arrays/varbin/compute/cast.rs

vortex_array/arrays/varbin/compute/
cast.rs
