vortex_array/arrays/varbinview/compute/
mod.rs

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