vortex_array/arrays/varbinview/compute/
zip.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::Range;
5
6use vortex_buffer::BufferMut;
7use vortex_error::VortexResult;
8use vortex_error::vortex_bail;
9use vortex_mask::AllOr;
10use vortex_mask::Mask;
11
12use crate::ArrayRef;
13use crate::ExecutionCtx;
14use crate::IntoArray;
15use crate::array::ArrayView;
16use crate::arrays::VarBinView;
17use crate::arrays::VarBinViewArray;
18use crate::arrays::varbinview::BinaryView;
19use crate::arrays::varbinview::VarBinViewArrayExt;
20use crate::builders::DeduplicatedBuffers;
21use crate::builders::LazyBitBufferBuilder;
22use crate::builtins::ArrayBuiltins;
23use crate::scalar_fn::fns::zip::ZipKernel;
24
25// A dedicated VarBinView zip kernel that builds the result directly by adjusting views and validity,
26// instead of routing through the generic builder (which would redo buffer lookups per mask slice).
27impl ZipKernel for VarBinView {
28    fn zip(
29        if_true: ArrayView<'_, VarBinView>,
30        if_false: &ArrayRef,
31        mask: &ArrayRef,
32        ctx: &mut ExecutionCtx,
33    ) -> VortexResult<Option<ArrayRef>> {
34        let Some(if_false) = if_false.as_opt::<VarBinView>() else {
35            return Ok(None);
36        };
37
38        if !if_true.dtype().eq_ignore_nullability(if_false.dtype()) {
39            vortex_bail!("input arrays to zip must have the same dtype");
40        }
41
42        let len = if_true.len();
43        let dtype = if_true
44            .dtype()
45            .union_nullability(if_false.dtype().nullability());
46
47        // build buffer lookup tables for both arrays, these map from the original buffer idx
48        // to the new buffer index in the result array
49        let mut buffers = DeduplicatedBuffers::default();
50        let true_lookup =
51            buffers.extend_from_iter(if_true.data_buffers().iter().map(|b| b.as_host().clone()));
52        let false_lookup =
53            buffers.extend_from_iter(if_false.data_buffers().iter().map(|b| b.as_host().clone()));
54
55        let mut views_builder = BufferMut::<BinaryView>::with_capacity(len);
56        let mut validity_builder = LazyBitBufferBuilder::new(len, ctx.allocator().clone());
57
58        let true_validity = if_true.varbinview_validity().execute_mask(len, ctx)?;
59        let false_validity = if_false.varbinview_validity().execute_mask(len, ctx)?;
60
61        let mask = mask.clone().fill_null(false)?.execute::<Mask>(ctx)?;
62        let if_false_view = if_false;
63        match mask.slices() {
64            AllOr::All => push_range(
65                if_true,
66                &true_lookup,
67                &true_validity,
68                0..len,
69                &mut views_builder,
70                &mut validity_builder,
71            ),
72            AllOr::None => push_range(
73                if_false_view,
74                &false_lookup,
75                &false_validity,
76                0..len,
77                &mut views_builder,
78                &mut validity_builder,
79            ),
80            AllOr::Some(slices) => {
81                let mut pos = 0;
82                for (start, end) in slices {
83                    if pos < *start {
84                        push_range(
85                            if_false_view,
86                            &false_lookup,
87                            &false_validity,
88                            pos..*start,
89                            &mut views_builder,
90                            &mut validity_builder,
91                        );
92                    }
93                    push_range(
94                        if_true,
95                        &true_lookup,
96                        &true_validity,
97                        *start..*end,
98                        &mut views_builder,
99                        &mut validity_builder,
100                    );
101                    pos = *end;
102                }
103                if pos < len {
104                    push_range(
105                        if_false_view,
106                        &false_lookup,
107                        &false_validity,
108                        pos..len,
109                        &mut views_builder,
110                        &mut validity_builder,
111                    );
112                }
113            }
114        }
115
116        let validity = validity_builder.finish_with_nullability(dtype.nullability());
117
118        // SAFETY: views are built with adjusted buffer indices, validity tracked alongside;
119        // buffers come from `DeduplicatedBuffers`, dtype/nullability preserved.
120        let array = unsafe {
121            VarBinViewArray::new_unchecked(
122                views_builder.freeze(),
123                buffers.finish(),
124                dtype,
125                validity,
126            )
127        };
128
129        Ok(Some(array.into_array()))
130    }
131}
132
133fn push_range(
134    array: ArrayView<'_, VarBinView>,
135    buffer_lookup: &[u32],
136    validity: &Mask,
137    range: Range<usize>,
138    views_builder: &mut BufferMut<BinaryView>,
139    validity_builder: &mut LazyBitBufferBuilder,
140) {
141    let views = array.views();
142
143    match validity.bit_buffer() {
144        AllOr::All => {
145            for idx in range {
146                push_view(
147                    views[idx],
148                    buffer_lookup,
149                    true,
150                    views_builder,
151                    validity_builder,
152                );
153            }
154        }
155        AllOr::None => {
156            for _ in range {
157                push_view(
158                    BinaryView::empty_view(),
159                    buffer_lookup,
160                    false,
161                    views_builder,
162                    validity_builder,
163                );
164            }
165        }
166        AllOr::Some(bit_buffer) => {
167            for idx in range {
168                let is_valid = bit_buffer.value(idx);
169                push_view(
170                    views[idx],
171                    buffer_lookup,
172                    is_valid,
173                    views_builder,
174                    validity_builder,
175                );
176            }
177        }
178    }
179}
180
181#[inline]
182fn push_view(
183    view: BinaryView,
184    buffer_lookup: &[u32],
185    is_valid: bool,
186    views_builder: &mut BufferMut<BinaryView>,
187    validity_builder: &mut LazyBitBufferBuilder,
188) {
189    if !is_valid {
190        views_builder.push(BinaryView::empty_view());
191        validity_builder.append_null();
192        return;
193    }
194
195    let adjusted = if view.is_inlined() {
196        view
197    } else {
198        let view_ref = view.as_view();
199        view_ref
200            .with_buffer_and_offset(
201                buffer_lookup[view_ref.buffer_index as usize],
202                view_ref.offset,
203            )
204            .into()
205    };
206
207    views_builder.push(adjusted);
208    validity_builder.append_non_null();
209}
210
211#[cfg(test)]
212mod tests {
213    use vortex_error::VortexResult;
214    use vortex_mask::Mask;
215
216    use crate::IntoArray;
217    use crate::VortexSessionExecute;
218    use crate::array_session;
219    use crate::arrays::VarBinViewArray;
220    use crate::builtins::ArrayBuiltins;
221    use crate::dtype::DType;
222    use crate::dtype::Nullability;
223
224    #[test]
225    fn zip_varbinview_kernel_zips() -> VortexResult<()> {
226        let a = VarBinViewArray::from_iter(
227            [
228                Some("aaaaaaaaaaaaa_long"), // outlined
229                Some("short"),
230                None,
231                Some("bbbbbbbbbbbbbbbb_long"),
232                Some("tiny"),
233                Some("cccccccccccccccc_long"),
234            ],
235            DType::Utf8(Nullability::Nullable),
236        );
237
238        let b = VarBinViewArray::from_iter(
239            [
240                Some("dddddddddddddddd_long"),
241                Some("eeeeeeeeeeeeeeee_long"),
242                Some("ffff"),
243                Some("gggggggggggggggg_long"),
244                None,
245                Some("hhhhhhhhhhhhhhhh_long"),
246            ],
247            DType::Utf8(Nullability::Nullable),
248        );
249
250        let mask = Mask::from_iter([true, false, true, false, false, true]);
251
252        let mut ctx = array_session().create_execution_ctx();
253        let zipped = mask
254            .clone()
255            .into_array()
256            .zip(a.into_array(), b.into_array())?
257            .execute::<VarBinViewArray>(&mut ctx)?;
258
259        let mut ctx = array_session().create_execution_ctx();
260        let validity_mask = zipped.validity()?.execute_mask(zipped.len(), &mut ctx)?;
261        let values = (0..zipped.len())
262            .map(|i| {
263                validity_mask
264                    .value(i)
265                    .then(|| String::from_utf8(zipped.bytes_at(i).to_vec()).unwrap())
266            })
267            .collect::<Vec<_>>();
268
269        assert_eq!(
270            values,
271            vec![
272                Some("aaaaaaaaaaaaa_long".to_string()),
273                Some("eeeeeeeeeeeeeeee_long".to_string()),
274                None,
275                Some("gggggggggggggggg_long".to_string()),
276                None,
277                Some("cccccccccccccccc_long".to_string())
278            ]
279        );
280        assert_eq!(zipped.len(), mask.len());
281        assert_eq!(zipped.dtype(), &DType::Utf8(Nullability::Nullable));
282        Ok(())
283    }
284}
vortex_array/arrays/varbinview/compute/zip.rs

vortex_array/arrays/varbinview/compute/
zip.rs
