take() on VarBin returns VarBinView (#9229)

take() on VarBin returned a VarBin which was slow because referenced
bytes were
copied. However, FSST version requires this. Make take() return
VarBinView but
keep the old version for FSST.

Add a take benchmark for VarBin

Resolves: #4964

---------

Signed-off-by: Mikhail Kot <mikhail@spiraldb.com>

Author: myrrc

encodings/fsst/src/compute/mod.rs

@@ -11,13 +11,11 @@ use vortex_array::ArrayRef;
 use vortex_array::ArrayView;
 use vortex_array::ExecutionCtx;
 use vortex_array::IntoArray;
-use vortex_array::arrays::VarBin;
 use vortex_array::arrays::dict::TakeExecute;
+use vortex_array::arrays::varbin::take_varbin;
 use vortex_array::builtins::ArrayBuiltins;
 use vortex_array::scalar::Scalar;
-use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
-use vortex_error::vortex_err;
 
 use crate::FSST;
 use crate::FSSTArrayExt;
@@ -36,14 +34,7 @@ impl TakeExecute for FSST {
                     .clone()
                     .union_nullability(indices.dtype().nullability()),
                 array.symbol_table(),
-                {
-                    let codes = array.codes();
-                    let codes = codes.as_view();
-                    <VarBin as TakeExecute>::take(codes, indices, ctx)?
-                        .vortex_expect("VarBin take kernel always returns Some")
-                }
-                .try_downcast::<VarBin>()
-                .map_err(|_| vortex_err!("take for codes must return varbin array"))?,
+                take_varbin(array.codes().as_view(), indices, ctx)?,
                 array
                     .uncompressed_lengths()
                     .take(indices.clone())?

vortex-array/Cargo.toml

@@ -247,6 +247,10 @@ harness = false
 name = "take_fsl"
 harness = false
 
+[[bench]]
+name = "take_varbin"
+harness = false
+
 [[bench]]
 name = "take_filter"
 harness = false

vortex-array/benches/take_varbin.rs

@@ -0,0 +1,54 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#![expect(clippy::unwrap_used)]
+
+use std::sync::LazyLock;
+
+use divan::Bencher;
+use rand::RngExt;
+use rand::SeedableRng;
+use rand::rngs::StdRng;
+use vortex_array::IntoArray;
+use vortex_array::RecursiveCanonical;
+use vortex_array::VortexSessionExecute;
+use vortex_array::array_session;
+use vortex_array::arrays::VarBinArray;
+use vortex_array::dtype::DType;
+use vortex_array::dtype::Nullability;
+use vortex_buffer::Buffer;
+use vortex_session::VortexSession;
+
+fn main() {
+    LazyLock::force(&SESSION);
+    divan::main();
+}
+
+static SESSION: LazyLock<VortexSession> = LazyLock::new(array_session);
+
+const ARRAY_SIZE: usize = 20_000;
+const TAKE_SIZE: usize = 8_000;
+
+#[divan::bench]
+fn take_varbin(bencher: Bencher) {
+    let array = VarBinArray::from_iter(
+        (0..ARRAY_SIZE).map(|i| Some(format!("row-{i:0>40}"))),
+        DType::Utf8(Nullability::NonNullable),
+    )
+    .into_array();
+
+    let mut rng = StdRng::seed_from_u64(0);
+    let indices: Buffer<u64> = (0..TAKE_SIZE)
+        .map(|_| rng.random_range(0..ARRAY_SIZE) as u64)
+        .collect();
+    let indices = indices.into_array();
+
+    bencher
+        .with_inputs(|| (&array, &indices, SESSION.create_execution_ctx()))
+        .bench_refs(|(array, indices, ctx)| {
+            array
+                .take((*indices).clone())
+                .unwrap()
+                .execute::<RecursiveCanonical>(ctx)
+        });
+}

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

@@ -10,6 +10,8 @@ mod filter;
 mod mask;
 mod take;
 
+pub use take::take_varbin;
+
 #[cfg(test)]
 mod tests {
     use rstest::rstest;

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

@@ -1,17 +1,22 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::iter;
 use std::ptr;
+use std::sync::Arc;
 
 use itertools::Itertools as _;
+use num_traits::AsPrimitive;
 use vortex_buffer::BitBufferMut;
+use vortex_buffer::Buffer;
 use vortex_buffer::BufferMut;
 use vortex_buffer::ByteBufferMut;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_ensure;
 use vortex_error::vortex_err;
 use vortex_error::vortex_panic;
+use vortex_mask::AllOr;
 use vortex_mask::Mask;
 
 use crate::ArrayRef;
@@ -22,17 +27,21 @@ use crate::arrays::PiecewiseSequence;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::VarBin;
 use crate::arrays::VarBinArray;
+use crate::arrays::VarBinViewArray;
 use crate::arrays::dict::TakeExecute;
 use crate::arrays::piecewise_sequence::constant_unsigned_usize;
 use crate::arrays::piecewise_sequence::maybe_contiguous_slices;
 use crate::arrays::primitive::PrimitiveArrayExt;
 use crate::arrays::varbin::VarBinArrayExt;
 use crate::arrays::varbin::VarBinArraySlotsExt;
+use crate::arrays::varbinview::BinaryView;
+use crate::arrays::varbinview::build_views::MAX_BUFFER_LEN;
 use crate::dtype::DType;
 use crate::dtype::IntegerPType;
 use crate::dtype::PType;
 use crate::dtype::UnsignedPType;
 use crate::executor::ExecutionCtx;
+use crate::match_each_integer_ptype;
 use crate::match_each_unsigned_integer_ptype;
 use crate::validity::Validity;
 
@@ -138,88 +147,149 @@ impl TakeExecute for VarBin {
         indices: &ArrayRef,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>> {
-        if let Some(piecewise_indices) = indices.as_opt::<PiecewiseSequence>()
-            && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)?
-        {
-            return Ok(Some(taken));
+        let offsets = array.offsets().clone().execute::<PrimitiveArray>(ctx)?;
+        let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned());
+        let last_offset = match_each_unsigned_integer_ptype!(offsets.ptype(), |O| {
+            offsets.as_slice::<O>().last().map_or(0usize, |&o| o.as_())
+        });
+
+        // VarBinView can't hold this buffer, so we can't canonicalize and
+        // take() (take panics). Convert to VarBin
+        if last_offset > MAX_BUFFER_LEN {
+            return Ok(Some(take_varbin(array, indices, ctx)?.into_array()));
         }
 
-        // TODO(joe): Be lazy with execute
-        let offsets = array.offsets().clone().execute::<PrimitiveArray>(ctx)?;
-        let data = array.bytes();
-        let indices = indices.clone().execute::<PrimitiveArray>(ctx)?;
+        let data = array.bytes().clone();
         let dtype = array
             .dtype()
             .clone()
             .union_nullability(indices.dtype().nullability());
-        let array_validity = array
-            .varbin_validity()
-            .execute_mask(array.as_ref().len(), ctx)?;
-        let indices_validity = indices
+        let validity = array.validity()?.take(indices)?;
+
+        let indices = indices.clone().execute::<PrimitiveArray>(ctx)?;
+        let indices_mask = indices
             .as_ref()
             .validity()?
             .execute_mask(indices.as_ref().len(), ctx)?;
 
-        // Offsets and indices are non-negative; read them through their unsigned reinterpretations
-        // so we only monomorphize over the 4 unsigned widths each (4x4 instead of 8x8). On take,
-        // offsets get widened to either 32- or 64-bit (to avoid overflow); the built output offsets
-        // are reinterpreted back to `out_offset_ptype` to preserve the result's offset signedness.
-        let out_offset_ptype = taken_offset_ptype(offsets.ptype());
-        let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned());
-        let indices = indices.reinterpret_cast(indices.ptype().to_unsigned());
-
-        let array = match_each_unsigned_integer_ptype!(indices.ptype(), |I| {
-            match offsets.ptype() {
-                PType::U8 => take::<I, u8>(
-                    dtype,
-                    offsets.as_slice::<u8>(),
-                    data.as_slice(),
-                    indices.as_slice::<I>(),
-                    array_validity,
-                    indices_validity,
-                    out_offset_ptype,
-                ),
-                PType::U16 => take::<I, u16>(
-                    dtype,
-                    offsets.as_slice::<u16>(),
-                    data.as_slice(),
-                    indices.as_slice::<I>(),
-                    array_validity,
-                    indices_validity,
-                    out_offset_ptype,
-                ),
-                PType::U32 => take::<I, u32>(
-                    dtype,
-                    offsets.as_slice::<u32>(),
+        let views = match_each_unsigned_integer_ptype!(offsets.ptype(), |O| {
+            match_each_integer_ptype!(indices.ptype(), |I| {
+                take_views(
+                    offsets.as_slice::<O>(),
                     data.as_slice(),
                     indices.as_slice::<I>(),
-                    array_validity,
-                    indices_validity,
-                    out_offset_ptype,
-                ),
-                PType::U64 => take::<I, u64>(
-                    dtype,
-                    offsets.as_slice::<u64>(),
-                    data.as_slice(),
-                    indices.as_slice::<I>(),
-                    array_validity,
-                    indices_validity,
-                    out_offset_ptype,
-                ),
-                _ => unreachable!("invalid PType for offsets"),
-            }
+                    &indices_mask,
+                )
+            })
         });
 
-        Ok(Some(array?.into_array()))
+        // SAFETY: every view references buffer 0 which is inside shared data buffer
+        unsafe {
+            Ok(Some(
+                VarBinViewArray::new_unchecked(views, Arc::from([data]), dtype, validity)
+                    .into_array(),
+            ))
+        }
     }
 }
 
+fn take_views<O: UnsignedPType, I: IntegerPType + AsPrimitive<usize>>(
+    offsets: &[O],
+    data: &[u8],
+    indices: &[I],
+    mask: &Mask,
+) -> Buffer<BinaryView> {
+    let build = |idx: usize| -> BinaryView {
+        let start: usize = offsets[idx].as_();
+        let stop: usize = offsets[idx + 1].as_();
+        let value = &data[start..stop];
+        let len = stop - start;
+
+        // Caller guarantees every offset is <= MAX_BUFFER_LEN
+        let start: u32 = start.as_();
+        if len > BinaryView::MAX_INLINED_SIZE {
+            let mut prefix = [0u8; 4];
+            prefix.copy_from_slice(&value[..4]);
+            let len: u32 = len.as_();
+            BinaryView::new_ref(len, prefix, 0, start)
+        } else {
+            BinaryView::make_view(value, 0, start)
+        }
+    };
+
+    match mask.bit_buffer() {
+        AllOr::All => Buffer::from_trusted_len_iter(indices.iter().map(|i| build(i.as_()))),
+        AllOr::None => {
+            Buffer::from_trusted_len_iter(iter::repeat_n(BinaryView::default(), indices.len()))
+        }
+        AllOr::Some(buffer) => {
+            Buffer::from_trusted_len_iter(buffer.iter().zip(indices.iter()).map(|(valid, i)| {
+                if valid {
+                    build(i.as_())
+                } else {
+                    BinaryView::default()
+                }
+            }))
+        }
+    }
+}
+
+/// Take from a VarBin. Referenced bytes are copied
+pub fn take_varbin(
+    array: ArrayView<'_, VarBin>,
+    indices: &ArrayRef,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<VarBinArray> {
+    if let Some(piecewise_indices) = indices.as_opt::<PiecewiseSequence>()
+        && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)?
+    {
+        return Ok(taken);
+    }
+
+    let offsets = array.offsets().clone().execute::<PrimitiveArray>(ctx)?;
+    let data = array.bytes();
+    let indices = indices.clone().execute::<PrimitiveArray>(ctx)?;
+    let dtype = array
+        .dtype()
+        .clone()
+        .union_nullability(indices.dtype().nullability());
+    let array_validity = array
+        .varbin_validity()
+        .execute_mask(array.as_ref().len(), ctx)?;
+    let indices_validity = indices
+        .as_ref()
+        .validity()?
+        .execute_mask(indices.as_ref().len(), ctx)?;
+
+    // Offsets and indices are non-negative; read them through their unsigned reinterpretations
+    // so we only monomorphize over the 4 unsigned widths each (4x4 instead of 8x8). On take,
+    // offsets get widened to either 32- or 64-bit (to avoid overflow); the built output offsets
+    // are reinterpreted back to `out_offset_ptype` to preserve the result's offset signedness.
+    let out_offset_ptype = taken_offset_ptype(offsets.ptype());
+    let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned());
+    let indices = indices.reinterpret_cast(indices.ptype().to_unsigned());
+
+    match_each_unsigned_integer_ptype!(indices.ptype(), |I| {
+        match_each_unsigned_integer_ptype!(offsets.ptype(), |O| {
+            take::<I, O>(
+                dtype,
+                offsets.as_slice::<O>(),
+                data.as_slice(),
+                indices.as_slice::<I>(),
+                array_validity,
+                indices_validity,
+                out_offset_ptype,
+            )
+        })
+    })
+}
+
 fn take_contiguous_ranges(
     array: ArrayView<'_, VarBin>,
     indices: ArrayView<'_, PiecewiseSequence>,
     indices_ref: &ArrayRef,
     ctx: &mut ExecutionCtx,
-) -> VortexResult<Option<ArrayRef>> {
+) -> VortexResult<Option<VarBinArray>> {
     let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else {
         return Ok(None);
     };
@@ -261,10 +331,12 @@ fn take_contiguous_ranges(
     // SAFETY: output offsets are built from valid input offsets, start at zero, are monotonically
     // non-decreasing, and the copied data buffer has exactly the referenced byte length.
     unsafe {
-        Ok(Some(
-            VarBinArray::new_unchecked(result.offsets, result.data.freeze(), dtype, validity)
-                .into_array(),
-        ))
+        Ok(Some(VarBinArray::new_unchecked(
+            result.offsets,
+            result.data.freeze(),
+            dtype,
+            validity,
+        )))
     }
 }
 

vortex-array/src/arrays/varbin/mod.rs

@@ -11,6 +11,7 @@ pub use array::VarBinSlotsView;
 pub use vtable::VarBinArray;
 
 pub(crate) mod compute;
+pub use compute::take_varbin;
 
 mod vtable;
 pub use vtable::VarBin;

vortex-array/src/arrays/varbin/vtable/mod.rs

@@ -34,7 +34,7 @@ use crate::dtype::PType;
 use crate::match_each_varbin_builder;
 use crate::serde::ArrayChildren;
 use crate::validity::Validity;
-mod canonical;
+pub(crate) mod canonical;
 mod kernel;
 mod operations;
 mod validity;