diff --git a/encodings/fsst/src/compute/mod.rs b/encodings/fsst/src/compute/mod.rs index c8477119007..e23d4272866 100644 --- a/encodings/fsst/src/compute/mod.rs +++ b/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(); - ::take(codes, indices, ctx)? - .vortex_expect("VarBin take kernel always returns Some") - } - .try_downcast::() - .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())? diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 22a8768af5d..d00b811a387 100644 --- a/vortex-array/Cargo.toml +++ b/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 diff --git a/vortex-array/benches/take_varbin.rs b/vortex-array/benches/take_varbin.rs new file mode 100644 index 00000000000..494f09fcbe5 --- /dev/null +++ b/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 = 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 = (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::(ctx) + }); +} diff --git a/vortex-array/src/arrays/varbin/compute/mod.rs b/vortex-array/src/arrays/varbin/compute/mod.rs index 480c07a3031..ea1a387008f 100644 --- a/vortex-array/src/arrays/varbin/compute/mod.rs +++ b/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; diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 93d51466605..59689d9116c 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -1,10 +1,14 @@ // 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; @@ -12,6 +16,7 @@ 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> { - if let Some(piecewise_indices) = indices.as_opt::() - && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)? - { - return Ok(Some(taken)); + let offsets = array.offsets().clone().execute::(ctx)?; + let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); + let last_offset = match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { + offsets.as_slice::().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::(ctx)?; - let data = array.bytes(); - let indices = indices.clone().execute::(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::(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::( - dtype, - offsets.as_slice::(), - data.as_slice(), - indices.as_slice::(), - array_validity, - indices_validity, - out_offset_ptype, - ), - PType::U16 => take::( - dtype, - offsets.as_slice::(), - data.as_slice(), - indices.as_slice::(), - array_validity, - indices_validity, - out_offset_ptype, - ), - PType::U32 => take::( - dtype, - offsets.as_slice::(), + let views = match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { + match_each_integer_ptype!(indices.ptype(), |I| { + take_views( + offsets.as_slice::(), data.as_slice(), indices.as_slice::(), - array_validity, - indices_validity, - out_offset_ptype, - ), - PType::U64 => take::( - dtype, - offsets.as_slice::(), - data.as_slice(), - indices.as_slice::(), - 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>( + offsets: &[O], + data: &[u8], + indices: &[I], + mask: &Mask, +) -> Buffer { + 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 { + if let Some(piecewise_indices) = indices.as_opt::() + && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)? + { + return Ok(taken); + } + + let offsets = array.offsets().clone().execute::(ctx)?; + let data = array.bytes(); + let indices = indices.clone().execute::(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::( + dtype, + offsets.as_slice::(), + data.as_slice(), + indices.as_slice::(), + array_validity, + indices_validity, + out_offset_ptype, + ) + }) + }) +} + fn take_contiguous_ranges( array: ArrayView<'_, VarBin>, indices: ArrayView<'_, PiecewiseSequence>, indices_ref: &ArrayRef, ctx: &mut ExecutionCtx, -) -> VortexResult> { +) -> VortexResult> { 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, + ))) } } diff --git a/vortex-array/src/arrays/varbin/mod.rs b/vortex-array/src/arrays/varbin/mod.rs index f7e159fa37c..e7bdca23e70 100644 --- a/vortex-array/src/arrays/varbin/mod.rs +++ b/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; diff --git a/vortex-array/src/arrays/varbin/vtable/mod.rs b/vortex-array/src/arrays/varbin/vtable/mod.rs index 420bbcf9d63..1e2065e29c2 100644 --- a/vortex-array/src/arrays/varbin/vtable/mod.rs +++ b/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;