vortex_array/arrays/varbinview/compute/
mod.rs

1mod cast;
2mod is_constant;
3mod is_sorted;
4mod mask;
5mod min_max;
6mod optimize;
7mod take;
8mod to_arrow;
9
10use vortex_error::VortexResult;
11use vortex_scalar::Scalar;
12
13use super::BinaryView;
14use crate::arrays::VarBinViewEncoding;
15use crate::arrays::varbin::varbin_scalar;
16use crate::arrays::varbinview::VarBinViewArray;
17use crate::compute::{
18    IsConstantFn, IsSortedFn, MinMaxFn, ScalarAtFn, SliceFn, TakeFn, ToArrowFn, UncompressedSizeFn,
19};
20use crate::vtable::ComputeVTable;
21use crate::{Array, ArrayRef};
22
23impl ComputeVTable for VarBinViewEncoding {
24    fn is_constant_fn(&self) -> Option<&dyn IsConstantFn<&dyn Array>> {
25        Some(self)
26    }
27
28    fn is_sorted_fn(&self) -> Option<&dyn IsSortedFn<&dyn Array>> {
29        Some(self)
30    }
31
32    fn min_max_fn(&self) -> Option<&dyn MinMaxFn<&dyn Array>> {
33        Some(self)
34    }
35
36    fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<&dyn Array>> {
37        Some(self)
38    }
39
40    fn slice_fn(&self) -> Option<&dyn SliceFn<&dyn Array>> {
41        Some(self)
42    }
43
44    fn take_fn(&self) -> Option<&dyn TakeFn<&dyn Array>> {
45        Some(self)
46    }
47
48    fn to_arrow_fn(&self) -> Option<&dyn ToArrowFn<&dyn Array>> {
49        Some(self)
50    }
51
52    fn uncompressed_size_fn(&self) -> Option<&dyn UncompressedSizeFn<&dyn Array>> {
53        Some(self)
54    }
55}
56
57impl ScalarAtFn<&VarBinViewArray> for VarBinViewEncoding {
58    fn scalar_at(&self, array: &VarBinViewArray, index: usize) -> VortexResult<Scalar> {
59        Ok(varbin_scalar(array.bytes_at(index), array.dtype()))
60    }
61}
62
63impl SliceFn<&VarBinViewArray> for VarBinViewEncoding {
64    fn slice(&self, array: &VarBinViewArray, start: usize, stop: usize) -> VortexResult<ArrayRef> {
65        let views = array.views().slice(start..stop);
66
67        Ok(VarBinViewArray::try_new(
68            views,
69            array.buffers().to_vec(),
70            array.dtype().clone(),
71            array.validity().slice(start, stop)?,
72        )?
73        .into_array())
74    }
75}
76
77impl UncompressedSizeFn<&VarBinViewArray> for VarBinViewEncoding {
78    fn uncompressed_size(&self, array: &VarBinViewArray) -> VortexResult<usize> {
79        let views = array.views().len() * size_of::<BinaryView>();
80        let mut buffers_size = 0;
81        for buffer in array.buffers() {
82            buffers_size += buffer.len();
83        }
84
85        Ok(views + buffers_size + array.validity().uncompressed_size())
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use vortex_buffer::buffer;
92
93    use crate::IntoArray;
94    use crate::accessor::ArrayAccessor;
95    use crate::array::Array;
96    use crate::arrays::VarBinViewArray;
97    use crate::builders::{ArrayBuilder, VarBinViewBuilder};
98    use crate::canonical::ToCanonical;
99    use crate::compute::conformance::mask::test_mask;
100    use crate::compute::{take, take_into};
101
102    #[test]
103    fn take_nullable() {
104        let arr = VarBinViewArray::from_iter_nullable_str([
105            Some("one"),
106            None,
107            Some("three"),
108            Some("four"),
109            None,
110            Some("six"),
111        ]);
112
113        let taken = take(&arr, &buffer![0, 3].into_array()).unwrap();
114
115        assert!(taken.dtype().is_nullable());
116        assert_eq!(
117            taken
118                .to_varbinview()
119                .unwrap()
120                .with_iterator(|it| it
121                    .map(|v| v.map(|b| unsafe { String::from_utf8_unchecked(b.to_vec()) }))
122                    .collect::<Vec<_>>())
123                .unwrap(),
124            [Some("one".to_string()), Some("four".to_string())]
125        );
126    }
127
128    #[test]
129    fn take_mask_var_bin_view_array() {
130        test_mask(&VarBinViewArray::from_iter_str([
131            "one", "two", "three", "four", "five",
132        ]));
133
134        test_mask(&VarBinViewArray::from_iter_nullable_str([
135            Some("one"),
136            None,
137            Some("three"),
138            Some("four"),
139            Some("five"),
140        ]));
141    }
142
143    #[test]
144    fn take_into_nullable() {
145        let arr = VarBinViewArray::from_iter_nullable_str([
146            Some("one"),
147            None,
148            Some("three"),
149            Some("four"),
150            None,
151            Some("six"),
152        ]);
153
154        let mut builder = VarBinViewBuilder::with_capacity(arr.dtype().clone(), arr.len());
155
156        take_into(&arr, &buffer![0, 3].into_array(), &mut builder).unwrap();
157
158        let taken = builder.finish();
159        assert!(taken.dtype().is_nullable());
160        assert_eq!(
161            taken
162                .to_varbinview()
163                .unwrap()
164                .with_iterator(|it| it
165                    .map(|v| v.map(|b| unsafe { String::from_utf8_unchecked(b.to_vec()) }))
166                    .collect::<Vec<_>>())
167                .unwrap(),
168            [Some("one".to_string()), Some("four".to_string())]
169        );
170    }
171}