fix[vortex-array]: fix take on varbinviews with NULL indices (#5626)

Author: asubiotto

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

@@ -1,12 +1,15 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::iter;
 use std::ops::Deref;
 
 use num_traits::AsPrimitive;
 use vortex_buffer::Buffer;
 use vortex_dtype::match_each_integer_ptype;
 use vortex_error::VortexResult;
+use vortex_mask::AllOr;
+use vortex_mask::Mask;
 use vortex_vector::binaryview::BinaryView;
 
 use crate::Array;
@@ -23,16 +26,16 @@ use crate::vtable::ValidityHelper;
 /// Take involves creating a new array that references the old array, just with the given set of views.
 impl TakeKernel for VarBinViewVTable {
     fn take(&self, array: &VarBinViewArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
-        // Compute the new validity
-
-        // This is valid since all elements (of all arrays) even null values must be inside
-        // min-max valid range.
+        // Compute the new validity.
         let validity = array.validity().take(indices)?;
         let indices = indices.to_primitive();
 
         let views_buffer = match_each_integer_ptype!(indices.ptype(), |I| {
-            // This is valid since all elements even null values are inside the min-max valid range.
-            take_views(array.views(), indices.as_slice::<I>())
+            take_views(
+                array.views(),
+                indices.as_slice::<I>(),
+                &indices.validity_mask(),
+            )
         });
 
         // SAFETY: taking all components at same indices maintains invariants
@@ -55,15 +58,36 @@ register_kernel!(TakeKernelAdapter(VarBinViewVTable).lift());
 fn take_views<I: AsPrimitive<usize>>(
     views: &Buffer<BinaryView>,
     indices: &[I],
+    mask: &Mask,
 ) -> Buffer<BinaryView> {
     // NOTE(ngates): this deref is not actually trivial, so we run it once.
     let views_ref = views.deref();
-    Buffer::<BinaryView>::from_trusted_len_iter(indices.iter().map(|i| views_ref[i.as_()]))
+    // We do not use iter_bools directly, since the resulting dyn iterator cannot
+    // implement TrustedLen.
+    match mask.bit_buffer() {
+        AllOr::All => {
+            Buffer::<BinaryView>::from_trusted_len_iter(indices.iter().map(|i| views_ref[i.as_()]))
+        }
+        AllOr::None => Buffer::<BinaryView>::from_trusted_len_iter(iter::repeat_n(
+            BinaryView::default(),
+            indices.len(),
+        )),
+        AllOr::Some(buffer) => Buffer::<BinaryView>::from_trusted_len_iter(
+            buffer.iter().zip(indices.iter()).map(|(valid, idx)| {
+                if valid {
+                    views_ref[idx.as_()]
+                } else {
+                    BinaryView::default()
+                }
+            }),
+        ),
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use rstest::rstest;
+    use vortex_buffer::BitBuffer;
     use vortex_buffer::buffer;
     use vortex_dtype::DType;
     use vortex_dtype::Nullability::NonNullable;
@@ -76,6 +100,7 @@ mod tests {
     use crate::canonical::ToCanonical;
     use crate::compute::conformance::take::test_take_conformance;
     use crate::compute::take;
+    use crate::validity::Validity;
 
     #[test]
     fn take_nullable() {
@@ -103,11 +128,13 @@ mod tests {
     fn take_nullable_indices() {
         let arr = VarBinViewArray::from_iter(["one", "two"].map(Some), DType::Utf8(NonNullable));
 
-        let taken = take(
-            arr.as_ref(),
-            PrimitiveArray::from_option_iter(vec![Some(1), None]).as_ref(),
-        )
-        .unwrap();
+        let indices = PrimitiveArray::new(
+            // Verify that garbage values at NULL indices are ignored.
+            buffer![1u64, 999],
+            Validity::from(BitBuffer::from(vec![true, false])),
+        );
+
+        let taken = take(arr.as_ref(), indices.as_ref()).unwrap();
 
         assert!(taken.dtype().is_nullable());
         assert_eq!(

vortex-buffer/src/trusted_len.rs

@@ -161,6 +161,8 @@ where
 {
 }
 
+unsafe impl<T: Clone> TrustedLen for std::iter::RepeatN<T> {}
+
 // Arrow bit iterators
 unsafe impl<'a> TrustedLen for crate::bit::BitIterator<'a> {}
 unsafe impl<'a> TrustedLen for crate::bit::BitChunkIterator<'a> {}