vortex_array/arrays/varbinview/compute/
zip.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5
6use crate::arrays::{VarBinViewArray, VarBinViewVTable};
7use crate::builders::VarBinViewBuilder;
8use crate::compute::{ZipKernel, ZipKernelAdapter, zip_impl_with_builder, zip_return_dtype};
9use crate::{Array, ArrayRef, register_kernel};
10
11impl ZipKernel for VarBinViewVTable {
12    fn zip(
13        &self,
14        if_true: &VarBinViewArray,
15        if_false: &dyn Array,
16        mask: &vortex_mask::Mask,
17    ) -> VortexResult<Option<ArrayRef>> {
18        let Some(if_false) = if_false.as_opt::<VarBinViewVTable>() else {
19            return Ok(None);
20        };
21        Ok(Some(zip_impl_with_builder(
22            if_true.as_ref(),
23            if_false.as_ref(),
24            mask,
25            Box::new(VarBinViewBuilder::with_buffer_deduplication(
26                zip_return_dtype(if_true.as_ref(), if_false.as_ref()),
27                if_true.len(),
28            )),
29        )?))
30    }
31}
32
33register_kernel!(ZipKernelAdapter(VarBinViewVTable).lift());
34
35#[cfg(test)]
36mod tests {
37    use arrow_array::cast::AsArray;
38    use arrow_select::zip::zip as arrow_zip;
39    use vortex_dtype::{DType, Nullability};
40    use vortex_mask::Mask;
41
42    use crate::IntoArray;
43    use crate::arrays::VarBinViewVTable;
44    use crate::arrow::IntoArrowArray;
45    use crate::builders::{ArrayBuilder as _, VarBinViewBuilder};
46    use crate::compute::zip;
47
48    #[test]
49    fn test_varbinview_zip() {
50        let if_true = {
51            let mut builder =
52                VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::NonNullable), 10);
53            for _ in 0..100 {
54                builder.append_value("Hello");
55                builder.append_value("Hello this is a long string that won't be inlined.");
56            }
57            builder.finish()
58        };
59
60        let if_false = {
61            let mut builder =
62                VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::NonNullable), 10);
63            for _ in 0..100 {
64                builder.append_value("Hello2");
65                builder.append_value("Hello2 this is a long string that won't be inlined.");
66            }
67            builder.finish()
68        };
69
70        // [1,2,4,5,7,8,..]
71        let mask = Mask::from_indices(200, (0..100).filter(|i| i % 3 != 0).collect());
72
73        let zipped = zip(&if_true, &if_false, &mask).unwrap();
74        let zipped = zipped.as_opt::<VarBinViewVTable>().unwrap();
75        assert_eq!(zipped.nbuffers(), 2);
76
77        // assert the result is the same as arrow
78        let expected = arrow_zip(
79            mask.into_array()
80                .into_arrow_preferred()
81                .unwrap()
82                .as_boolean(),
83            &if_true.into_arrow_preferred().unwrap(),
84            &if_false.into_arrow_preferred().unwrap(),
85        )
86        .unwrap();
87
88        let actual = zipped.clone().into_array().into_arrow_preferred().unwrap();
89        assert_eq!(actual.as_ref(), expected.as_ref());
90    }
91}
vortex_array/arrays/varbinview/compute/zip.rs

vortex_array/arrays/varbinview/compute/
zip.rs
