vortex_array/arrays/varbinview/compute/
cast.rs1use std::sync::Arc;
5
6use vortex_error::VortexResult;
7
8use crate::ArrayRef;
9use crate::ExecutionCtx;
10use crate::IntoArray;
11use crate::array::ArrayView;
12use crate::arrays::VarBinView;
13use crate::arrays::VarBinViewArray;
14use crate::dtype::DType;
15use crate::scalar_fn::fns::cast::CastKernel;
16use crate::scalar_fn::fns::cast::CastReduce;
17use crate::validity::Validity;
18
19fn build_with_validity(
20 array: ArrayView<'_, VarBinView>,
21 new_dtype: DType,
22 new_validity: Validity,
23) -> ArrayRef {
24 unsafe {
26 VarBinViewArray::new_handle_unchecked(
27 array.views_handle().clone(),
28 Arc::clone(array.data_buffers()),
29 new_dtype,
30 new_validity,
31 )
32 .into_array()
33 }
34}
35
36impl CastReduce for VarBinView {
37 fn cast(array: ArrayView<'_, VarBinView>, 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 VarBinView {
55 fn cast(
56 array: ArrayView<'_, VarBinView>,
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::VarBinViewArray;
84 use crate::assert_arrays_eq;
85 use crate::builtins::ArrayBuiltins;
86 use crate::compute::conformance::cast::test_cast_conformance;
87 use crate::dtype::DType;
88 use crate::dtype::Nullability;
89
90 static SESSION: LazyLock<VortexSession> = LazyLock::new(crate::array_session);
91
92 #[rstest]
93 #[case(
94 DType::Utf8(Nullability::Nullable),
95 DType::Utf8(Nullability::NonNullable)
96 )]
97 #[case(
98 DType::Binary(Nullability::Nullable),
99 DType::Binary(Nullability::NonNullable)
100 )]
101 #[case(
102 DType::Utf8(Nullability::NonNullable),
103 DType::Utf8(Nullability::Nullable)
104 )]
105 #[case(
106 DType::Binary(Nullability::NonNullable),
107 DType::Binary(Nullability::Nullable)
108 )]
109 fn try_cast_varbin_nullable(#[case] source: DType, #[case] target: DType) {
110 let varbin = VarBinViewArray::from_iter(vec![Some("a"), Some("b"), Some("c")], source);
111
112 let res = varbin.into_array().cast(target.clone());
113 assert_arrays_eq!(
114 res.unwrap(),
115 VarBinViewArray::from_iter([Some("a"), Some("b"), Some("c")], target),
116 &mut SESSION.create_execution_ctx()
117 );
118 }
119
120 #[rstest]
121 #[case(DType::Utf8(Nullability::Nullable))]
122 #[case(DType::Binary(Nullability::Nullable))]
123 fn try_cast_varbin_fail(#[case] source: DType) {
124 let non_nullable_source = source.as_nonnullable();
126 let varbin = VarBinViewArray::from_iter(vec![Some("a"), Some("b"), None], source);
127 let mut ctx = SESSION.create_execution_ctx();
128 let result = varbin
129 .into_array()
130 .cast(non_nullable_source)
131 .and_then(|a| a.execute::<Canonical>(&mut ctx).map(|c| c.into_array()));
132 assert!(result.is_err(), "Expected error, got: {result:?}");
133 }
134
135 #[rstest]
136 #[case(VarBinViewArray::from_iter(vec![Some("hello"), Some("world"), Some("test")], DType::Utf8(Nullability::NonNullable)))]
137 #[case(VarBinViewArray::from_iter(vec![Some("hello"), None, Some("world")], DType::Utf8(Nullability::Nullable)))]
138 #[case(VarBinViewArray::from_iter(vec![Some(b"binary".as_slice()), Some(b"data".as_slice())], DType::Binary(Nullability::NonNullable)))]
139 #[case(VarBinViewArray::from_iter(vec![Some(b"test".as_slice()), None], DType::Binary(Nullability::Nullable)))]
140 #[case(VarBinViewArray::from_iter(vec![Some("single")], DType::Utf8(Nullability::NonNullable)))]
141 #[case(VarBinViewArray::from_iter(vec![Some("very long string that exceeds the inline size to test view functionality with multiple buffers")], DType::Utf8(Nullability::NonNullable)))]
142 fn test_cast_varbinview_conformance(#[case] array: VarBinViewArray) {
143 test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
144 }
145}