vortex_array/arrays/varbinview/compute/
mod.rs1mod 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}