varbinview zip kernel (#5526)

Signed-off-by: Onur Satici <[email protected]>

Author: onursatici

vortex-array/src/arrays/struct_/compute/zip.rs

@@ -29,15 +29,10 @@ impl ZipKernel for StructVTable {
         let Some(if_false) = if_false.as_opt::<StructVTable>() else {
             return Ok(None);
         };
-        assert_eq!(
-            if_true.len(),
-            if_false.len(),
-            "ComputeFn::invoke checks that arrays have the same size"
-        );
         assert_eq!(
             if_true.names(),
             if_false.names(),
-            "Zip checks that arrays type"
+            "input arrays to zip must have the same field names",
         );
 
         let fields = if_true

vortex-array/src/arrays/varbinview/compute/mod.rs

@@ -8,6 +8,7 @@ mod is_sorted;
 mod mask;
 mod min_max;
 mod take;
+mod zip;
 
 #[cfg(test)]
 mod tests {

vortex-array/src/arrays/varbinview/compute/zip.rs

@@ -0,0 +1,268 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::ops::Range;
+
+use vortex_buffer::BufferMut;
+use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
+use vortex_mask::AllOr;
+use vortex_mask::Mask;
+use vortex_vector::binaryview::BinaryView;
+
+use crate::Array;
+use crate::ArrayRef;
+use crate::arrays::VarBinViewArray;
+use crate::arrays::VarBinViewVTable;
+use crate::builders::DeduplicatedBuffers;
+use crate::builders::LazyBitBufferBuilder;
+use crate::compute::ZipKernel;
+use crate::compute::ZipKernelAdapter;
+use crate::register_kernel;
+
+// A dedicated VarBinView zip kernel that builds the result directly by adjusting views and validity,
+// instead of routing through the generic builder (which would redo buffer lookups per mask slice).
+impl ZipKernel for VarBinViewVTable {
+    fn zip(
+        &self,
+        if_true: &VarBinViewArray,
+        if_false: &dyn Array,
+        mask: &Mask,
+    ) -> VortexResult<Option<ArrayRef>> {
+        let Some(if_false) = if_false.as_opt::<VarBinViewVTable>() else {
+            return Ok(None);
+        };
+
+        if !if_true.dtype().eq_ignore_nullability(if_false.dtype()) {
+            vortex_bail!("input arrays to zip must have the same dtype");
+        }
+
+        // compute fn already asserts if_true.len() == if_false.len()
+        let len = if_true.len();
+        let dtype = if_true
+            .dtype()
+            .union_nullability(if_false.dtype().nullability());
+
+        // build buffer lookup tables for both arrays, these map from the original buffer idx
+        // to the new buffer index in the result array
+        let mut buffers = DeduplicatedBuffers::default();
+        let true_lookup = buffers.extend_from_slice(if_true.buffers());
+        let false_lookup = buffers.extend_from_slice(if_false.buffers());
+
+        let mut views_builder = BufferMut::<BinaryView>::with_capacity(len);
+        let mut validity_builder = LazyBitBufferBuilder::new(len);
+
+        let true_validity = if_true.validity_mask();
+        let false_validity = if_false.validity_mask();
+
+        match mask.slices() {
+            AllOr::All => push_range(
+                if_true,
+                &true_lookup,
+                &true_validity,
+                0..len,
+                &mut views_builder,
+                &mut validity_builder,
+            ),
+            AllOr::None => push_range(
+                if_false,
+                &false_lookup,
+                &false_validity,
+                0..len,
+                &mut views_builder,
+                &mut validity_builder,
+            ),
+            AllOr::Some(slices) => {
+                let mut pos = 0;
+                for (start, end) in slices {
+                    if pos < *start {
+                        push_range(
+                            if_false,
+                            &false_lookup,
+                            &false_validity,
+                            pos..*start,
+                            &mut views_builder,
+                            &mut validity_builder,
+                        );
+                    }
+                    push_range(
+                        if_true,
+                        &true_lookup,
+                        &true_validity,
+                        *start..*end,
+                        &mut views_builder,
+                        &mut validity_builder,
+                    );
+                    pos = *end;
+                }
+                if pos < len {
+                    push_range(
+                        if_false,
+                        &false_lookup,
+                        &false_validity,
+                        pos..len,
+                        &mut views_builder,
+                        &mut validity_builder,
+                    );
+                }
+            }
+        }
+
+        let validity = validity_builder.finish_with_nullability(dtype.nullability());
+
+        // SAFETY: views are built with adjusted buffer indices, validity tracked alongside;
+        // buffers come from `DeduplicatedBuffers`, dtype/nullability preserved.
+        let array = unsafe {
+            VarBinViewArray::new_unchecked(
+                views_builder.freeze(),
+                buffers.finish(),
+                dtype,
+                validity,
+            )
+        };
+
+        Ok(Some(array.to_array()))
+    }
+}
+
+fn push_range(
+    array: &VarBinViewArray,
+    buffer_lookup: &[u32],
+    validity: &Mask,
+    range: Range<usize>,
+    views_builder: &mut BufferMut<BinaryView>,
+    validity_builder: &mut LazyBitBufferBuilder,
+) {
+    let views = array.views();
+
+    match validity.bit_buffer() {
+        AllOr::All => {
+            for idx in range {
+                push_view(
+                    views[idx],
+                    buffer_lookup,
+                    true,
+                    views_builder,
+                    validity_builder,
+                );
+            }
+        }
+        AllOr::None => {
+            for _ in range {
+                push_view(
+                    BinaryView::empty_view(),
+                    buffer_lookup,
+                    false,
+                    views_builder,
+                    validity_builder,
+                );
+            }
+        }
+        AllOr::Some(bit_buffer) => {
+            for idx in range {
+                let is_valid = bit_buffer.value(idx);
+                push_view(
+                    views[idx],
+                    buffer_lookup,
+                    is_valid,
+                    views_builder,
+                    validity_builder,
+                );
+            }
+        }
+    }
+}
+
+#[inline]
+fn push_view(
+    view: BinaryView,
+    buffer_lookup: &[u32],
+    is_valid: bool,
+    views_builder: &mut BufferMut<BinaryView>,
+    validity_builder: &mut LazyBitBufferBuilder,
+) {
+    if !is_valid {
+        views_builder.push(BinaryView::empty_view());
+        validity_builder.append_null();
+        return;
+    }
+
+    let adjusted = if view.is_inlined() {
+        view
+    } else {
+        let view_ref = view.as_view();
+        view_ref
+            .with_buffer_and_offset(
+                buffer_lookup[view_ref.buffer_index as usize],
+                view_ref.offset,
+            )
+            .into()
+    };
+
+    views_builder.push(adjusted);
+    validity_builder.append_non_null();
+}
+
+register_kernel!(ZipKernelAdapter(VarBinViewVTable).lift());
+
+#[cfg(test)]
+mod tests {
+    use vortex_dtype::DType;
+    use vortex_dtype::Nullability;
+    use vortex_mask::Mask;
+
+    use crate::accessor::ArrayAccessor;
+    use crate::arrays::VarBinViewArray;
+    use crate::canonical::ToCanonical;
+    use crate::compute::zip;
+
+    #[test]
+    fn zip_varbinview_kernel_zips() {
+        let a = VarBinViewArray::from_iter(
+            [
+                Some("aaaaaaaaaaaaa_long"), // outlined
+                Some("short"),
+                None,
+                Some("bbbbbbbbbbbbbbbb_long"),
+                Some("tiny"),
+                Some("cccccccccccccccc_long"),
+            ],
+            DType::Utf8(Nullability::Nullable),
+        );
+
+        let b = VarBinViewArray::from_iter(
+            [
+                Some("dddddddddddddddd_long"),
+                Some("eeeeeeeeeeeeeeee_long"),
+                Some("ffff"),
+                Some("gggggggggggggggg_long"),
+                None,
+                Some("hhhhhhhhhhhhhhhh_long"),
+            ],
+            DType::Utf8(Nullability::Nullable),
+        );
+
+        let mask = Mask::from_iter([true, false, true, false, false, true]);
+
+        let zipped = zip(a.as_ref(), b.as_ref(), &mask).unwrap().to_varbinview();
+
+        let values = zipped.with_iterator(|it| {
+            it.map(|v| v.map(|bytes| String::from_utf8(bytes.to_vec()).unwrap()))
+                .collect::<Vec<_>>()
+        });
+
+        assert_eq!(
+            values,
+            vec![
+                Some("aaaaaaaaaaaaa_long".to_string()),
+                Some("eeeeeeeeeeeeeeee_long".to_string()),
+                None,
+                Some("gggggggggggggggg_long".to_string()),
+                None,
+                Some("cccccccccccccccc_long".to_string())
+            ]
+        );
+        assert_eq!(zipped.len(), mask.len());
+        assert_eq!(zipped.dtype(), &DType::Utf8(Nullability::Nullable));
+    }
+}

vortex-array/src/builders/mod.rs

@@ -43,7 +43,7 @@ use crate::ArrayRef;
 use crate::canonical::Canonical;
 
 mod lazy_null_builder;
-use lazy_null_builder::LazyBitBufferBuilder;
+pub(crate) use lazy_null_builder::LazyBitBufferBuilder;
 
 mod bool;
 mod decimal;

vortex-array/src/builders/varbinview.rs

@@ -436,7 +436,7 @@ impl DeduplicatedBuffers {
     }
 
     /// Push a new block if not seen before. Returns the idx of the block.
-    fn push(&mut self, block: ByteBuffer) -> u32 {
+    pub(crate) fn push(&mut self, block: ByteBuffer) -> u32 {
         assert!(self.buffers.len() < u32::MAX as usize, "Too many blocks");
 
         let initial_len = self.len();
@@ -452,21 +452,24 @@ impl DeduplicatedBuffers {
         }
     }
 
-    fn extend_from_option_slice(&mut self, buffers: &[Option<ByteBuffer>]) -> Vec<Option<u32>> {
+    pub(crate) fn extend_from_option_slice(
+        &mut self,
+        buffers: &[Option<ByteBuffer>],
+    ) -> Vec<Option<u32>> {
         buffers
             .iter()
             .map(|buffer| buffer.as_ref().map(|buf| self.push(buf.clone())))
             .collect()
     }
 
-    fn extend_from_slice(&mut self, buffers: &[ByteBuffer]) -> Vec<u32> {
+    pub(crate) fn extend_from_slice(&mut self, buffers: &[ByteBuffer]) -> Vec<u32> {
         buffers
             .iter()
             .map(|buffer| self.push(buffer.clone()))
             .collect()
     }
 
-    fn finish(self) -> Arc<[ByteBuffer]> {
+    pub(crate) fn finish(self) -> Arc<[ByteBuffer]> {
         Arc::from(self.buffers)
     }
 }

vortex-array/src/compute/zip.rs

@@ -19,7 +19,6 @@ use super::cast;
 use crate::Array;
 use crate::ArrayRef;
 use crate::builders::ArrayBuilder;
-use crate::builders::VarBinViewBuilder;
 use crate::builders::builder_with_capacity;
 use crate::compute::ComputeFn;
 use crate::compute::Kernel;
@@ -210,27 +209,15 @@ fn zip_impl(if_true: &dyn Array, if_false: &dyn Array, mask: &Mask) -> VortexRes
     );
 
     let return_type = zip_return_dtype(if_true, if_false);
-    let capacity = if_true.len();
-
-    let builder = match return_type {
-        // TODO(blaginin): once https://github.com/vortex-data/vortex/pull/4695 is merged, we can kill
-        //  these two special cases, but before that we need to manually use deduplicated buffers.
-        //  Otherwise, the same buffer will be appended multiple times causing fragmentation.
-        DType::Utf8(n) => Box::new(VarBinViewBuilder::with_buffer_deduplication(
-            DType::Utf8(n),
-            capacity,
-        )),
-        DType::Binary(n) => Box::new(VarBinViewBuilder::with_buffer_deduplication(
-            DType::Binary(n),
-            capacity,
-        )),
-        _ => builder_with_capacity(&return_type, if_true.len()),
-    };
-
-    zip_impl_with_builder(if_true, if_false, mask, builder)
+    zip_impl_with_builder(
+        if_true,
+        if_false,
+        mask,
+        builder_with_capacity(&return_type, if_true.len()),
+    )
 }
 
-pub(crate) fn zip_impl_with_builder(
+fn zip_impl_with_builder(
     if_true: &dyn Array,
     if_false: &dyn Array,
     mask: &Mask,
@@ -272,8 +259,8 @@ mod tests {
     use crate::arrow::IntoArrowArray;
     use crate::builders::ArrayBuilder;
     use crate::builders::BufferGrowthStrategy;
+    use crate::builders::VarBinViewBuilder;
     use crate::compute::zip;
-    use crate::compute::zip::VarBinViewBuilder;
 
     #[test]
     fn test_zip_basic() {