vortex_array/arrays/varbinview/compute/
mod.rs

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

vortex_array/arrays/varbinview/compute/
mod.rs
