diff --git a/vortex-array/benches/chunk_array_builder.rs b/vortex-array/benches/chunk_array_builder.rs index c2dedc6ef46..dcc64cf8d7a 100644 --- a/vortex-array/benches/chunk_array_builder.rs +++ b/vortex-array/benches/chunk_array_builder.rs @@ -17,6 +17,7 @@ use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ConstantArray; use vortex_array::builders::ArrayBuilder; +use vortex_array::builders::VarBinBuilder; use vortex_array::builders::VarBinViewBuilder; use vortex_array::builders::builder_with_capacity; use vortex_array::dtype::DType; @@ -133,6 +134,54 @@ fn chunked_varbinview_opt_into_canonical(bencher: Bencher, (len, chunk_count): ( .bench_refs(|(chunk, ctx)| (**chunk).clone().execute::(ctx)) } +// Fewer rows than BENCH_ARGS: decoding VarBin values is the most expensive work in this file. +const VARBIN_BENCH_ARGS: &[(usize, usize)] = &[ + // length, chunk_count + (10, 100), + (500, 2), +]; + +#[divan::bench(args = VARBIN_BENCH_ARGS)] +fn chunked_varbin_to_varbinview_builder(bencher: Bencher, (len, chunk_count): (usize, usize)) { + let chunks = make_varbin_chunks(false, len, chunk_count); + + bencher + .with_inputs(|| (&chunks, SESSION.create_execution_ctx())) + .bench_refs(|(chunk, ctx)| { + let mut builder = + VarBinViewBuilder::with_capacity(chunk.dtype().clone(), len * chunk_count); + chunk + .append_to_builder(&mut builder, ctx) + .vortex_expect("append failed"); + builder.finish() + }) +} + +#[divan::bench(args = VARBIN_BENCH_ARGS)] +fn chunked_varbin_opt_to_varbinview_builder(bencher: Bencher, (len, chunk_count): (usize, usize)) { + let chunks = make_varbin_chunks(true, len, chunk_count); + + bencher + .with_inputs(|| (&chunks, SESSION.create_execution_ctx())) + .bench_refs(|(chunk, ctx)| { + let mut builder = + VarBinViewBuilder::with_capacity(chunk.dtype().clone(), len * chunk_count); + chunk + .append_to_builder(&mut builder, ctx) + .vortex_expect("append failed"); + builder.finish() + }) +} + +#[divan::bench(args = VARBIN_BENCH_ARGS)] +fn chunked_varbin_into_canonical(bencher: Bencher, (len, chunk_count): (usize, usize)) { + let chunks = make_varbin_chunks(false, len, chunk_count); + + bencher + .with_inputs(|| (&chunks, SESSION.create_execution_ctx())) + .bench_refs(|(chunk, ctx)| (**chunk).clone().execute::(ctx)) +} + #[divan::bench(args = BENCH_ARGS)] fn chunked_constant_i32_append_to_builder(bencher: Bencher, (len, chunk_count): (usize, usize)) { let chunk = make_constant_i32_chunks(len, chunk_count); @@ -226,6 +275,30 @@ fn make_bool_chunks(len: usize, chunk_count: usize) -> ArrayRef { .into_array() } +fn make_varbin_chunks(nullable: bool, len: usize, chunk_count: usize) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(123); + let dtype = DType::Utf8(nullable.into()); + + (0..chunk_count) + .map(|_| { + let mut builder = VarBinBuilder::::with_capacity(dtype.clone(), len); + (0..len).for_each(|_| { + if nullable && rng.random_bool(0.2) { + builder.push_null() + } else { + builder.append_value( + (0..rng.random_range(0..=20)) + .map(|_| rng.random_range(b'a'..=b'z')) + .collect::>(), + ) + } + }); + builder.finish() + }) + .collect::() + .into_array() +} + fn make_string_chunks(nullable: bool, len: usize, chunk_count: usize) -> ArrayRef { let mut rng = StdRng::seed_from_u64(123); diff --git a/vortex-array/src/arrays/varbin/vtable/canonical.rs b/vortex-array/src/arrays/varbin/vtable/canonical.rs index acea7bd92f4..6c195fe16ce 100644 --- a/vortex-array/src/arrays/varbin/vtable/canonical.rs +++ b/vortex-array/src/arrays/varbin/vtable/canonical.rs @@ -4,6 +4,8 @@ use std::sync::Arc; use num_traits::AsPrimitive; +use vortex_buffer::Buffer; +use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use crate::ExecutionCtx; @@ -11,9 +13,11 @@ use crate::array::ArrayView; use crate::arrays::PrimitiveArray; use crate::arrays::VarBin; use crate::arrays::VarBinViewArray; +use crate::arrays::varbinview::BinaryView; use crate::arrays::varbinview::build_views::MAX_BUFFER_LEN; use crate::arrays::varbinview::build_views::build_views; use crate::arrays::varbinview::build_views::offsets_to_lengths; +use crate::buffer::BufferHandle; use crate::match_each_integer_ptype; /// Converts a VarBinArray to its canonical form (VarBinViewArray). @@ -24,35 +28,54 @@ pub(crate) fn varbin_to_canonical( ctx: &mut ExecutionCtx, ) -> VortexResult { let parts = array.into_owned().into_data_parts(); - let offsets = parts.offsets.execute::(ctx)?; + let (buffers, views) = varbin_decode_views(&offsets, parts.bytes, 0); + + // SAFETY: views are correctly computed from valid offsets + Ok(unsafe { + VarBinViewArray::new_unchecked(views, Arc::from(buffers), parts.dtype, parts.validity) + }) +} +/// Lays a `VarBin` array's value bytes out as `VarBinView` buffers plus the views over them. +/// +/// `start_buf_index` is the index the first returned buffer will occupy in its destination, so the +/// views come out already referencing the right buffer and never need rebasing. Canonicalization +/// passes `0`; appending into a [`VarBinViewBuilder`](crate::builders::VarBinViewBuilder) passes the +/// index its next buffer will land at. +/// +/// The value bytes are handed over as they are — only the offsets are consumed, to derive the view +/// lengths — so this costs one view per row and no byte copy when the buffer is uniquely held. +pub(crate) fn varbin_decode_views( + offsets: &PrimitiveArray, + bytes: BufferHandle, + start_buf_index: u32, +) -> (Vec, Buffer) { match_each_integer_ptype!(offsets.ptype(), |P| { let offsets_slice = offsets.as_slice::

(); let first: usize = offsets_slice[0].as_(); let last: usize = offsets_slice[offsets_slice.len() - 1].as_(); - let bytes = parts.bytes.unwrap_host().slice(first..last).into_mut(); + let bytes = bytes.unwrap_host().slice(first..last).into_mut(); let lens = offsets_to_lengths(offsets_slice); - let (buffers, views) = build_views(0, MAX_BUFFER_LEN, bytes, lens.as_slice()); - - // SAFETY: views are correctly computed from valid offsets - Ok(unsafe { - VarBinViewArray::new_unchecked(views, Arc::from(buffers), parts.dtype, parts.validity) - }) + build_views(start_buf_index, MAX_BUFFER_LEN, bytes, lens.as_slice()) }) } #[cfg(test)] mod tests { use rstest::rstest; + use vortex_error::VortexResult; + use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; + use crate::arrays::ChunkedArray; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::arrays::varbin::builder::VarBinBuilder; use crate::assert_arrays_eq; + use crate::builders::VarBinViewBuilder; use crate::dtype::DType; use crate::dtype::Nullability; @@ -108,6 +131,80 @@ mod tests { assert_arrays_eq!(canonical, expected, &mut ctx); } + /// Appending a `VarBin` array to a `VarBinViewBuilder` builds views over its bytes directly + /// instead of canonicalizing first, so the views must be numbered against the buffers the + /// builder already holds. Interleaving `VarBin` appends with value appends (which stage an + /// in-progress buffer) and with a `VarBinView` append exercises that numbering. + #[rstest] + #[case(DType::Utf8(Nullability::Nullable))] + #[case(DType::Binary(Nullability::Nullable))] + fn append_varbin_to_varbinview_builder(#[case] dtype: DType) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let long = "a value long enough that its view has to reference a buffer"; + let longer = "another value long enough that its view has to reference a buffer"; + + // Two chunks, each with an inlined value, a buffer-referencing value and a null, so both + // pushed buffers are non-empty and the second must not reuse the first's index. + let first = VarBinArray::from_iter([Some("short"), None, Some(long)], dtype.clone()); + let second = VarBinArray::from_iter([Some(longer), Some("tiny"), None], dtype.clone()); + let view = VarBinViewArray::from_iter([Some(long), None], dtype.clone()); + + let mut builder = VarBinViewBuilder::with_capacity(dtype.clone(), 8); + first + .as_array() + .clone() + .append_to_builder(&mut builder, &mut ctx)?; + // Stages an in-progress buffer, which the next append has to account for. + builder.append_value(longer); + second + .as_array() + .clone() + .append_to_builder(&mut builder, &mut ctx)?; + view.clone() + .into_array() + .append_to_builder(&mut builder, &mut ctx)?; + + let expected = ChunkedArray::try_new( + vec![ + first.as_array().clone(), + VarBinViewArray::from_iter([Some(longer)], dtype.clone()).into_array(), + second.as_array().clone(), + view.into_array(), + ], + dtype, + )?; + assert_arrays_eq!(builder.finish_into_varbinview(), expected, &mut ctx); + Ok(()) + } + + /// A builder configured to compact must not be handed a raw buffer behind its back: the + /// inlined values leave the pushed buffer only partly referenced, and skipping compaction + /// would keep those bytes alive. Appending through the canonical array instead drops them. + #[test] + fn append_varbin_to_a_compacting_builder_still_compacts() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Utf8(Nullability::NonNullable); + // Every value inlines, so nothing references the value bytes at all. + let array = VarBinArray::from_iter_nonnull(["short", "tiny", "small"], dtype.clone()); + + let mut builder = VarBinViewBuilder::with_compaction(dtype, 4, 1.0); + array + .as_array() + .clone() + .append_to_builder(&mut builder, &mut ctx)?; + let compacted = builder.finish_into_varbinview(); + + assert!( + compacted + .data_buffers() + .iter() + .all(|buffer| buffer.is_empty()), + "a fully-inlined append should not retain any value bytes" + ); + assert_arrays_eq!(compacted, array.as_array().clone(), &mut ctx); + Ok(()) + } + // Empty array: offsets has exactly one element; no elements to canonicalize. #[test] fn test_canonical_varbin_empty() { diff --git a/vortex-array/src/arrays/varbin/vtable/mod.rs b/vortex-array/src/arrays/varbin/vtable/mod.rs index 22cace7ac2b..420bbcf9d63 100644 --- a/vortex-array/src/arrays/varbin/vtable/mod.rs +++ b/vortex-array/src/arrays/varbin/vtable/mod.rs @@ -20,11 +20,14 @@ use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::VTable; +use crate::arrays::PrimitiveArray; +use crate::arrays::varbin::VarBinArrayExt; use crate::arrays::varbin::VarBinArraySlotsExt; use crate::arrays::varbin::VarBinData; use crate::arrays::varbin::VarBinSlots; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; +use crate::builders::VarBinViewBuilder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -36,6 +39,7 @@ mod kernel; mod operations; mod validity; +use canonical::varbin_decode_views; use canonical::varbin_to_canonical; use vortex_session::VortexSession; @@ -214,9 +218,27 @@ impl VTable for VarBin { { return result; } - varbin_to_canonical(array, ctx)? - .into_array() - .append_to_builder(builder, ctx) + + // The two arms here are every builder a `Utf8`/`Binary` dtype has: all four + // `VarBinBuilder` widths above, and `VarBinViewBuilder` below. + let Some(view_builder) = builder.as_any().downcast_ref::() else { + vortex_bail!("append_to_builder for VarBin requires a variable-binary builder") + }; + + if view_builder.compacts_buffers() { + // A compacting builder decides per buffer whether to keep, slice or rewrite it, which + // it can only do by measuring the finished views against the buffer. Go through the + // canonical array so that policy still applies. + return varbin_to_canonical(array, ctx)? + .into_array() + .append_to_builder(builder, ctx); + } + + let builder = builder + .as_any_mut() + .downcast_mut::() + .vortex_expect("builder type checked above"); + append_to_varbinview(array, builder, ctx) } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { @@ -226,5 +248,32 @@ impl VTable for VarBin { } } +/// Hands the value bytes to `builder` as a data buffer with views built over them. +/// +/// Canonicalizing first would build the same views, then pay for them twice more: once to wrap +/// them in a `VarBinViewArray` the builder immediately unwraps, and once for +/// `append_varbinview_array` to rewrite every view so its buffer index is rebased onto the +/// builder's. Numbering the buffer up front instead makes the whole append one view per row plus +/// pushing the byte buffer. +fn append_to_varbinview( + array: ArrayView<'_, VarBin>, + builder: &mut VarBinViewBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let len = array.as_ref().len(); + let validity = array.varbin_validity().execute_mask(len, ctx)?; + + // Build the views against the index the pushed buffer will land at, so the builder does not + // have to rebase them afterwards. + let next_buffer_index = builder.completed_block_count() + u32::from(builder.in_progress()); + + let parts = array.into_owned().into_data_parts(); + let offsets = parts.offsets.execute::(ctx)?; + let (buffers, views) = varbin_decode_views(&offsets, parts.bytes, next_buffer_index); + + builder.push_buffer_and_adjusted_views(&buffers, &views, validity); + Ok(()) +} + #[derive(Clone, Debug)] pub struct VarBin; diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index 1d530a384ef..1ec39138ebb 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -194,6 +194,16 @@ impl VarBinViewBuilder { self.in_progress.is_some() } + /// Whether this builder compacts the data buffers it is handed. + /// + /// [`push_buffer_and_adjusted_views`](Self::push_buffer_and_adjusted_views) takes buffers + /// exactly as they are, so an encoding that would push a buffer only partly covered by its + /// views should check this first and fall back to a route that measures utilization — + /// otherwise it silently opts the builder out of the compaction it was configured for. + pub fn compacts_buffers(&self) -> bool { + self.compaction_threshold > 0.0 + } + /// Pushes buffers and pre-adjusted views into the builder. /// /// The provided `buffers` contain sections of data from a `VarBinViewArray`, and the