vortex_array/arrays/varbinview/vtable/
operator.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use vortex_buffer::ByteBuffer;
7use vortex_compute::filter::Filter;
8use vortex_dtype::DType;
9use vortex_error::VortexResult;
10use vortex_vector::Vector;
11use vortex_vector::binaryview::{BinaryVector, BinaryViewTypeUpcast, StringVector};
12
13use crate::ArrayRef;
14use crate::arrays::{VarBinViewArray, VarBinViewVTable};
15use crate::execution::{BatchKernelRef, BindCtx, kernel};
16use crate::vtable::{OperatorVTable, ValidityHelper};
17
18impl OperatorVTable<VarBinViewVTable> for VarBinViewVTable {
19    fn bind(
20        array: &VarBinViewArray,
21        selection: Option<&ArrayRef>,
22        ctx: &mut dyn BindCtx,
23    ) -> VortexResult<BatchKernelRef> {
24        let mask = ctx.bind_selection(array.len(), selection)?;
25        let validity = ctx.bind_validity(array.validity(), array.len(), selection)?;
26        let dtype = array.dtype().clone();
27
28        let views = array.views().clone();
29        let buffers: Vec<ByteBuffer> = array.buffers().iter().cloned().collect();
30        let buffers = Arc::new(buffers.into_boxed_slice());
31
32        Ok(kernel(move || {
33            let selection = mask.execute()?;
34            let validity = validity.execute()?;
35
36            // We only filter the views buffer
37            let views = views.filter(&selection);
38
39            match dtype {
40                // SAFETY: the incoming array has the same validation as the vector
41                DType::Utf8(_) => Ok(Vector::from_string(unsafe {
42                    StringVector::new_unchecked(views, buffers, validity)
43                })),
44
45                // SAFETY: the incoming array has the same validation as the vector
46                DType::Binary(_) => Ok(Vector::from_binary(unsafe {
47                    BinaryVector::new_unchecked(views, buffers, validity)
48                })),
49                _ => unreachable!("invalid dtype for VarBinViewArray {dtype}"),
50            }
51        }))
52    }
53}
54
55#[cfg(test)]
56mod tests {
57    use rstest::{fixture, rstest};
58    use vortex_dtype::{DType, Nullability};
59
60    use crate::IntoArray;
61    use crate::arrays::{BoolArray, VarBinViewArray};
62    use crate::builders::{ArrayBuilder, VarBinViewBuilder};
63
64    #[fixture]
65    fn strings() -> VarBinViewArray {
66        let mut strings = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 5);
67        strings.append_value("inlined");
68        strings.append_nulls(1);
69        strings.append_value("large string 1");
70        strings.append_value("large string 2");
71        strings.append_value("large string 3");
72        strings.finish_into_varbinview()
73    }
74
75    #[rstest]
76    fn test_bind(strings: VarBinViewArray) {
77        // Attempt to bind with a full selection.
78        let strings_vec = strings
79            .bind(None, &mut ())
80            .unwrap()
81            .execute()
82            .unwrap()
83            .into_string();
84        assert_eq!(strings_vec.get_ref(0), Some("inlined"));
85        assert_eq!(strings_vec.get_ref(1), None);
86        assert_eq!(strings_vec.get_ref(2), Some("large string 1"));
87        assert_eq!(strings_vec.get_ref(3), Some("large string 2"));
88        assert_eq!(strings_vec.get_ref(4), Some("large string 3"));
89    }
90
91    #[rstest]
92    fn test_bind_with_selection(strings: VarBinViewArray) {
93        let selection = BoolArray::from_iter([false, true, false, true, true]).into_array();
94        let strings_vec = strings
95            .bind(Some(&selection), &mut ())
96            .unwrap()
97            .execute()
98            .unwrap()
99            .into_string();
100
101        assert_eq!(strings_vec.get_ref(0), None);
102        assert_eq!(strings_vec.get_ref(1), Some("large string 2"));
103        assert_eq!(strings_vec.get_ref(2), Some("large string 3"));
104    }
105}