diff --git a/Cargo.lock b/Cargo.lock index 51188781c9a..43970638788 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9377,12 +9377,14 @@ dependencies = [ "itertools 0.14.0", "num-traits", "rstest", + "simdutf8", "tracing", "vortex-array", "vortex-buffer", "vortex-error", "vortex-fsst", "vortex-mask", + "vortex-onpair", "vortex-runend", "vortex-session", "vortex-zstd", @@ -10437,6 +10439,7 @@ name = "vortex-zstd" version = "0.1.0" dependencies = [ "itertools 0.14.0", + "num-traits", "prost 0.14.4", "rstest", "vortex-array", diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index 4a740417f7e..d211ac85df9 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -5,6 +5,7 @@ use std::fmt::Debug; use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hasher; +use std::mem::MaybeUninit; use std::sync::Arc; use std::sync::OnceLock; @@ -34,13 +35,15 @@ use vortex_array::arrays::VarBinArray; use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::buffer::BufferHandle; use vortex_array::builders::ArrayBuilder; -use vortex_array::builders::DynVarBinBuilder; +use vortex_array::builders::VarBinBuilder; use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; +use vortex_array::dtype::IntegerPType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::legacy_session; use vortex_array::match_each_integer_ptype; +use vortex_array::match_each_varbin_builder; use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; @@ -58,8 +61,9 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::canonical::FSST_DECODE_SLACK; +use crate::canonical::FsstDecodePlan; use crate::canonical::canonicalize_fsst; -use crate::canonical::fsst_decode_bytes; use crate::canonical::fsst_decode_views; use crate::rules::RULES; @@ -307,23 +311,10 @@ impl VTable for FSST { builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - if let Some(builder) = builder.as_any_mut().downcast_mut::() { - let (bytes, lengths) = fsst_decode_bytes(array, ctx)?; - let validity = array - .array() - .validity()? - .execute_mask(array.array().len(), ctx)?; - match_each_integer_ptype!(lengths.ptype(), |P| { - builder.append_values( - bytes.as_slice(), - lengths.as_slice::

().iter().scan(0usize, |end, length| { - *end += AsPrimitive::::as_(*length); - Some(*end) - }), - &validity, - ) - })?; - return Ok(()); + if let Some(result) = + match_each_varbin_builder!(builder, |builder| append_to_varbin(array, builder, ctx)) + { + return result; } let Some(builder) = builder.as_any_mut().downcast_mut::() else { @@ -361,6 +352,41 @@ impl VTable for FSST { } } +/// Decompresses the code stream straight into `builder`'s byte storage. +/// +/// The offsets are the running sum of the uncompressed lengths the array already stores, so the +/// only work beyond the bulk `decompress_into` is one prefix sum over them. +fn append_to_varbin( + array: ArrayView<'_, FSST>, + builder: &mut VarBinBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> +where + usize: AsPrimitive, +{ + let plan = FsstDecodePlan::new(array, ctx)?; + let validity = array + .array() + .validity()? + .execute_mask(array.array().len(), ctx)?; + let decompressor = array.decompressor(); + // Built once, outside the ptype match: the decoder is `#[inline(always)]`, so creating the + // closure inside each arm would stamp out a copy of the whole decode loop per length type. + let mut decode = |out: &mut [MaybeUninit]| plan.decode_into(&decompressor, out); + match_each_integer_ptype!(plan.lengths.ptype(), |P| { + // SAFETY: `decode_into` initializes exactly the prefix whose length it returns. + unsafe { + builder.append_decoded( + plan.total_size, + FSST_DECODE_SLACK, + plan.lengths.as_slice::

(), + &validity, + &mut decode, + ) + } + }) +} + #[array_slots(FSST)] pub struct FSSTSlots { /// Lengths of the original values before compression, can be compressed. diff --git a/encodings/fsst/src/canonical.rs b/encodings/fsst/src/canonical.rs index 44288f48cf9..8db84b22823 100644 --- a/encodings/fsst/src/canonical.rs +++ b/encodings/fsst/src/canonical.rs @@ -1,8 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::mem::MaybeUninit; use std::sync::Arc; +use fsst::Decompressor; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; @@ -42,36 +44,82 @@ pub(super) fn canonicalize_fsst( }) } +/// Extra headroom that keeps [`Decompressor::decompress_into`] on its fast path. +/// +/// It never writes past the slice it is given — every store is bounded by the end of the output — +/// but it emits whole 8-byte symbols, so it can only use the wide-store loop while at least 8 +/// bytes remain. Handing it 7 spare bytes lets that loop run through the final value instead of +/// finishing byte at a time. This is a performance knob, not a safety requirement. +/// +/// [`Decompressor::decompress_into`]: fsst::Decompressor::decompress_into +pub(crate) const FSST_DECODE_SLACK: usize = 7; + +/// Everything needed to decode an FSST array's values in one bulk `decompress_into` call. +pub(crate) struct FsstDecodePlan { + codes: ByteBuffer, + /// Per-row uncompressed lengths, zero for null rows. + pub(crate) lengths: PrimitiveArray, + /// Total decoded size, i.e. the sum of `lengths`. + pub(crate) total_size: usize, +} + +impl FsstDecodePlan { + pub(crate) fn new( + fsst_array: ArrayView<'_, FSST>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let codes = fsst_array.codes().sliced_bytes(); + let lengths = fsst_array + .uncompressed_lengths() + .clone() + .execute::(ctx)?; + + #[expect(clippy::cast_possible_truncation)] + let total_size: usize = match_each_integer_ptype!(lengths.ptype(), |P| { + lengths.as_slice::

().iter().map(|x| *x as usize).sum() + }); + + Ok(Self { + codes, + lengths, + total_size, + }) + } + + /// Bulk-decompresses the whole code stream into `out`, which must hold at least + /// `total_size + FSST_DECODE_SLACK` bytes. + /// + /// Kept inlinable so the decoder is not called from behind an extra frame in whichever + /// codegen unit the caller lands in; see OnPair's equivalent for why that matters. + #[inline] + pub(crate) fn decode_into( + &self, + decompressor: &Decompressor<'_>, + out: &mut [MaybeUninit], + ) -> VortexResult { + let len = decompressor.decompress_into(self.codes.as_slice(), out); + vortex_ensure!( + len == self.total_size, + "FSST decoded {len} bytes, expected {}", + self.total_size + ); + Ok(len) + } +} + pub(crate) fn fsst_decode_bytes( fsst_array: ArrayView<'_, FSST>, ctx: &mut ExecutionCtx, ) -> VortexResult<(ByteBufferMut, PrimitiveArray)> { - let bytes = fsst_array.codes().sliced_bytes(); - let uncompressed_lens_array = fsst_array - .uncompressed_lengths() - .clone() - .execute::(ctx)?; - - #[expect(clippy::cast_possible_truncation)] - let total_size: usize = match_each_integer_ptype!(uncompressed_lens_array.ptype(), |P| { - uncompressed_lens_array - .as_slice::

() - .iter() - .map(|x| *x as usize) - .sum() - }); - - let decompressor = fsst_array.decompressor(); - let mut uncompressed_bytes = ByteBufferMut::with_capacity(total_size + 7); - let len = - decompressor.decompress_into(bytes.as_slice(), uncompressed_bytes.spare_capacity_mut()); - vortex_ensure!( - len == total_size, - "FSST decoded {len} bytes, expected {total_size}" - ); - // SAFETY: `decompress_into` initialized the first `len` bytes. + let plan = FsstDecodePlan::new(fsst_array, ctx)?; + let mut uncompressed_bytes = ByteBufferMut::with_capacity(plan.total_size + FSST_DECODE_SLACK); + let len = plan.decode_into( + &fsst_array.decompressor(), + uncompressed_bytes.spare_capacity_mut(), + )?; + // SAFETY: `decode_into` initialized the first `len` bytes. unsafe { uncompressed_bytes.set_len(len) }; - Ok((uncompressed_bytes, uncompressed_lens_array)) + Ok((uncompressed_bytes, plan.lengths)) } pub(crate) fn fsst_decode_views( @@ -106,7 +154,7 @@ mod tests { use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::varbin::VarBinArrayExt; use vortex_array::builders::ArrayBuilder; - use vortex_array::builders::DynVarBinBuilder; + use vortex_array::builders::VarBinBuilder; use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -212,7 +260,7 @@ mod tests { { let mut builder = - DynVarBinBuilder::with_capacity(chunked_arr.dtype().clone(), false, data.len()); + VarBinBuilder::::with_capacity(chunked_arr.dtype().clone(), data.len()); chunked_arr .into_array() .append_to_builder(&mut builder, &mut ctx)?; diff --git a/encodings/fsst/src/compress.rs b/encodings/fsst/src/compress.rs index e54b6af26f1..0f10a775b5e 100644 --- a/encodings/fsst/src/compress.rs +++ b/encodings/fsst/src/compress.rs @@ -199,7 +199,11 @@ fn compress_views( where O: IntegerPType + 'static, { - let mut sink = FsstSink::::with_capacity(strings.len(), compressor); + let mut sink = FsstSink::::with_capacity( + DType::Binary(strings.dtype().nullability()), + strings.len(), + compressor, + ); let views = strings.views(); let buffers = strings.data_buffers(); match mask.bit_buffer() { @@ -232,7 +236,11 @@ fn compress_varbin( where O: IntegerPType + 'static, { - let mut sink = FsstSink::::with_capacity(strings.len(), compressor); + let mut sink = FsstSink::::with_capacity( + DType::Binary(strings.dtype().nullability()), + strings.len(), + compressor, + ); let bytes = strings.bytes().as_slice(); match_each_integer_ptype!(offsets.ptype(), |I| { let off = offsets.as_slice::(); @@ -277,10 +285,10 @@ struct FsstSink<'c, O: IntegerPType + 'static> { } impl<'c, O: IntegerPType + 'static> FsstSink<'c, O> { - fn with_capacity(len: usize, compressor: &'c Compressor) -> Self { + fn with_capacity(dtype: DType, len: usize, compressor: &'c Compressor) -> Self { Self { buffer: Vec::with_capacity(DEFAULT_BUFFER_LEN), - builder: VarBinBuilder::::with_capacity(len), + builder: VarBinBuilder::::with_capacity(dtype, len), uncompressed_lengths: BufferMut::with_capacity(len), compressor, } @@ -289,7 +297,7 @@ impl<'c, O: IntegerPType + 'static> FsstSink<'c, O> { #[inline] fn emit(&mut self, row: Option<&[u8]>) { let Some(s) = row else { - self.builder.append_null(); + self.builder.push_null(); self.uncompressed_lengths.push(0); return; }; @@ -310,8 +318,8 @@ impl<'c, O: IntegerPType + 'static> FsstSink<'c, O> { self.builder.append_value(&self.buffer); } - fn finish(self, dtype: DType, ctx: &mut ExecutionCtx) -> VortexResult { - let codes = self.builder.finish(DType::Binary(dtype.nullability())); + fn finish(mut self, dtype: DType, ctx: &mut ExecutionCtx) -> VortexResult { + let codes = self.builder.finish_into_varbin(); FSST::try_new( dtype, Buffer::copy_from(self.compressor.symbol_table()), diff --git a/encodings/fsst/src/kernel.rs b/encodings/fsst/src/kernel.rs index 30f25006195..1ae34dfc800 100644 --- a/encodings/fsst/src/kernel.rs +++ b/encodings/fsst/src/kernel.rs @@ -61,7 +61,8 @@ mod tests { }); fn build_test_fsst_array() -> ArrayRef { - let mut builder = VarBinBuilder::::with_capacity(10); + let mut builder = + VarBinBuilder::::with_capacity(DType::Utf8(Nullability::NonNullable), 10); builder.append_value(b"hello world"); builder.append_value(b"foo bar baz"); builder.append_value(b"testing fsst compression"); @@ -72,7 +73,7 @@ mod tests { builder.append_value(b"qrstuvwxyz"); builder.append_value(b"0123456789"); builder.append_value(b"final string"); - let input = builder.finish(DType::Utf8(Nullability::NonNullable)); + let input = builder.finish_into_varbin(); let mut ctx = SESSION.create_execution_ctx(); let arr = input.into_array(); @@ -131,7 +132,8 @@ mod tests { // Test case with special characters and nulls // Values: ["", "", "", "", "", "", "", "", "", "", "", ",", "A<<<<<<<", "", "", "", "", null, null, null, null, null, null] // Mask: only the last element is selected (true at index 22) - let mut builder = VarBinBuilder::::with_capacity(23); + let mut builder = + VarBinBuilder::::with_capacity(DType::Utf8(Nullability::Nullable), 23); // 11 empty strings for _ in 0..11 { builder.append_value(b""); @@ -146,9 +148,9 @@ mod tests { } // 6 nulls for _ in 0..6 { - builder.append_null(); + builder.push_null(); } - let input = builder.finish(DType::Utf8(Nullability::Nullable)); + let input = builder.finish_into_varbin(); let array = input.clone().into_array(); let mut ctx = SESSION.create_execution_ctx(); @@ -172,12 +174,13 @@ mod tests { #[test] fn filter_only_null() -> VortexResult<()> { - let mut builder = VarBinBuilder::::with_capacity(3); - builder.append_null(); + let mut builder = + VarBinBuilder::::with_capacity(DType::Utf8(Nullability::Nullable), 3); + builder.push_null(); builder.append_value(b"A"); - builder.append_null(); + builder.push_null(); - let input = builder.finish(DType::Utf8(Nullability::Nullable)); + let input = builder.finish_into_varbin(); let array = input.clone().into_array(); let mut ctx = SESSION.create_execution_ctx(); @@ -213,15 +216,14 @@ mod tests { #[test] fn test_fsst_byte_length() -> VortexResult<()> { - let mut builder = VarBinBuilder::::with_capacity(3); + let mut builder = + VarBinBuilder::::with_capacity(DType::Utf8(Nullability::NonNullable), 3); builder.append_value(b"hello"); builder.append_value(b"world!!"); builder.append_value("Пуховички"); // 9 characters, 18 bytes builder.append_value(b""); - let varbin = builder - .finish(DType::Utf8(Nullability::NonNullable)) - .into_array(); + let varbin = builder.finish_into_varbin().into_array(); let mut ctx = SESSION.create_execution_ctx(); let compressor = fsst_train_compressor(&varbin, &mut ctx)?; let fsst = fsst_compress(&varbin, &compressor, &mut ctx)?.into_array(); diff --git a/encodings/fsst/src/tests.rs b/encodings/fsst/src/tests.rs index 43cdeb1a02c..2b908a8997f 100644 --- a/encodings/fsst/src/tests.rs +++ b/encodings/fsst/src/tests.rs @@ -30,15 +30,14 @@ static SESSION: LazyLock = LazyLock::new(|| { /// this function is VERY slow on miri, so we only want to run it once pub(crate) fn build_fsst_array(ctx: &mut ExecutionCtx) -> ArrayRef { - let mut input_array = VarBinBuilder::::with_capacity(3); + let mut input_array = + VarBinBuilder::::with_capacity(DType::Utf8(Nullability::NonNullable), 3); input_array.append_value(b"The Greeks never said that the limit could not be overstepped"); input_array.append_value( b"They said it existed and that whoever dared to exceed it was mercilessly struck down", ); input_array.append_value(b"Nothing in present history can contradict them"); - let input_array = input_array - .finish(DType::Utf8(Nullability::NonNullable)) - .into_array(); + let input_array = input_array.finish_into_varbin().into_array(); let compressor = fsst_train_compressor(&input_array, ctx).unwrap(); fsst_compress(&input_array, &compressor, ctx) @@ -162,13 +161,11 @@ fn fsst_compress_offsets_overflow_i32() { println!("building large VarBinArray"); let string = vec![b'a'; STRING_LEN]; - let mut builder = VarBinBuilder::::with_capacity(N); + let mut builder = VarBinBuilder::::with_capacity(DType::Utf8(Nullability::NonNullable), N); for _ in 0..N { builder.append_value(&string); } - let array = builder - .finish(DType::Utf8(Nullability::NonNullable)) - .into_array(); + let array = builder.finish_into_varbin().into_array(); let compressor = CompressorBuilder::default().build(); let len = array.len(); diff --git a/encodings/onpair/src/array.rs b/encodings/onpair/src/array.rs index b67e09dfbd0..3f39c0114ac 100644 --- a/encodings/onpair/src/array.rs +++ b/encodings/onpair/src/array.rs @@ -5,6 +5,7 @@ use std::fmt::Debug; use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hasher; +use std::mem::MaybeUninit; use std::sync::Arc; use std::sync::OnceLock; @@ -26,12 +27,14 @@ use vortex_array::IntoArray; use vortex_array::array_slots; use vortex_array::buffer::BufferHandle; use vortex_array::builders::ArrayBuilder; -use vortex_array::builders::DynVarBinBuilder; +use vortex_array::builders::VarBinBuilder; use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; +use vortex_array::dtype::IntegerPType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::match_each_integer_ptype; +use vortex_array::match_each_varbin_builder; use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; @@ -48,8 +51,8 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::canonical::OnPairDecodePlan; use crate::canonical::canonicalize_onpair; -use crate::canonical::onpair_decode_bytes; use crate::canonical::onpair_decode_views; use crate::decode::collect_widened; use crate::rules::RULES; @@ -553,23 +556,10 @@ impl VTable for OnPair { builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - if let Some(builder) = builder.as_any_mut().downcast_mut::() { - let (bytes, lengths) = onpair_decode_bytes(array, ctx)?; - let validity = array - .array() - .validity()? - .execute_mask(array.array().len(), ctx)?; - match_each_integer_ptype!(lengths.ptype(), |P| { - builder.append_values( - bytes.as_slice(), - lengths.as_slice::

().iter().scan(0usize, |end, length| { - *end += AsPrimitive::::as_(*length); - Some(*end) - }), - &validity, - ) - })?; - return Ok(()); + if let Some(result) = + match_each_varbin_builder!(builder, |builder| append_to_varbin(array, builder, ctx)) + { + return result; } let Some(builder) = builder.as_any_mut().downcast_mut::() else { @@ -603,6 +593,41 @@ impl VTable for OnPair { } } +/// Decodes the code stream straight into `builder`'s byte storage. +/// +/// The offsets are the running sum of the uncompressed lengths the array already stores, so the +/// only work beyond the bulk `try_decode_into` is one prefix sum over them. +fn append_to_varbin( + array: ArrayView<'_, OnPair>, + builder: &mut VarBinBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> +where + usize: AsPrimitive, +{ + let plan = OnPairDecodePlan::new(array, ctx)?; + let validity = array + .array() + .validity()? + .execute_mask(array.array().len(), ctx)?; + // Built once, outside the ptype match: the decoder is `#[inline(always)]`, so creating the + // closure inside each arm would stamp out a copy of the whole decode loop per length type. + let mut decode = |out: &mut [MaybeUninit]| plan.decode_into(out); + match_each_integer_ptype!(plan.lengths.ptype(), |P| { + // SAFETY: `decode_into` initializes exactly the prefix whose length it returns. It needs + // no slack: it derives its bound from the slice it is handed and writes each value exactly. + unsafe { + builder.append_decoded( + plan.total_size, + 0, + plan.lengths.as_slice::

(), + &validity, + &mut decode, + ) + } + }) +} + impl ValidityVTable for OnPair { fn validity(array: ArrayView<'_, OnPair>) -> VortexResult { Ok(child_to_validity( diff --git a/encodings/onpair/src/canonical.rs b/encodings/onpair/src/canonical.rs index 5179def3bf1..172555e9a94 100644 --- a/encodings/onpair/src/canonical.rs +++ b/encodings/onpair/src/canonical.rs @@ -6,9 +6,11 @@ //! //! [`OnPairArray`]: crate::OnPairArray +use std::mem::MaybeUninit; use std::sync::Arc; use num_traits::AsPrimitive; +use onpair::CompactDictionaryView; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; @@ -44,68 +46,101 @@ pub(super) fn canonicalize_onpair( }) } -pub(crate) fn onpair_decode_bytes( - array: ArrayView<'_, OnPair>, - ctx: &mut ExecutionCtx, -) -> VortexResult<(ByteBufferMut, PrimitiveArray)> { - let lengths = array - .uncompressed_lengths() - .clone() - .execute::(ctx)?; +/// Everything needed to decode an OnPair array's values in one bulk `try_decode_into` call. +pub(crate) struct OnPairDecodePlan<'a> { + codes: Buffer, + dict: CompactDictionaryView<'a>, + /// Per-row uncompressed lengths, zero for null rows. + pub(crate) lengths: PrimitiveArray, + /// Total decoded size, i.e. the sum of `lengths`. + pub(crate) total_size: usize, +} + +impl<'a> OnPairDecodePlan<'a> { + pub(crate) fn new(array: ArrayView<'a, OnPair>, ctx: &mut ExecutionCtx) -> VortexResult { + let lengths = array + .uncompressed_lengths() + .clone() + .execute::(ctx)?; + + let total_size: usize = match_each_integer_ptype!(lengths.ptype(), |P| { + lengths + .as_slice::

() + .iter() + .map(|&l| AsPrimitive::::as_(l)) + .sum() + }); + + // `codes_offsets` holds the per-row code boundaries and may itself be a + // sliced or filtered view of the original. Its first and last entries + // bound the contiguous run of `codes` belonging to the rows present in + // this array: `slice` keeps the full `codes` child and only narrows + // `codes_offsets` (so `code_start > 0` and/or `code_end < codes.len()`), + // while `filter` rebuilds both children so the window is the whole stream. + // OnPair has no `TakeExecute`, so a reordering take is served from the + // canonical `VarBinView` and never reaches this path. We only need those + // two boundaries, so point-look them up rather than decoding every offset. + let codes_offsets = array.codes_offsets(); + let code_start = code_boundary_at(codes_offsets, 0, ctx)?; + let code_end = code_boundary_at(codes_offsets, array.len(), ctx)?; + vortex_ensure!( + code_start <= code_end, + "OnPair codes_offsets must be nondecreasing" + ); + vortex_ensure!( + code_end <= array.codes().len(), + "OnPair codes_offsets end {} exceeds codes len {}", + code_end, + array.codes().len() + ); - let total_size: usize = match_each_integer_ptype!(lengths.ptype(), |P| { - lengths - .as_slice::

() - .iter() - .map(|&l| AsPrimitive::::as_(l)) - .sum() - }); + // Slice the `codes` child to that window *before* unpacking it, so a sliced + // array materialises only its own codes rather than the whole column's. The + // contiguous decoder walks `codes` in order and never reads the per-row + // boundaries, so an empty boundary slice is sound. + let codes = collect_widened::(&array.codes().slice(code_start..code_end)?, ctx)?; + let dict = dict_view(array, ctx)?; - // `codes_offsets` holds the per-row code boundaries and may itself be a - // sliced or filtered view of the original. Its first and last entries - // bound the contiguous run of `codes` belonging to the rows present in - // this array: `slice` keeps the full `codes` child and only narrows - // `codes_offsets` (so `code_start > 0` and/or `code_end < codes.len()`), - // while `filter` rebuilds both children so the window is the whole stream. - // OnPair has no `TakeExecute`, so a reordering take is served from the - // canonical `VarBinView` and never reaches this path. We only need those - // two boundaries, so point-look them up rather than decoding every offset. - let codes_offsets = array.codes_offsets(); - let code_start = code_boundary_at(codes_offsets, 0, ctx)?; - let code_end = code_boundary_at(codes_offsets, array.len(), ctx)?; - vortex_ensure!( - code_start <= code_end, - "OnPair codes_offsets must be nondecreasing" - ); - vortex_ensure!( - code_end <= array.codes().len(), - "OnPair codes_offsets end {} exceeds codes len {}", - code_end, - array.codes().len() - ); + Ok(Self { + codes, + dict, + lengths, + total_size, + }) + } - // Slice the `codes` child to that window *before* unpacking it, so a sliced - // array materialises only its own codes rather than the whole column's. The - // contiguous decoder walks `codes` in order and never reads the per-row - // boundaries, so an empty boundary slice is sound. - let codes = collect_widened::(&array.codes().slice(code_start..code_end)?, ctx)?; - let dict = dict_view(array, ctx)?; - let mut out_bytes = ByteBufferMut::with_capacity(total_size); - let written = - match onpair::try_decode_into(codes.as_slice(), dict, out_bytes.spare_capacity_mut()) { + /// Bulk-decodes the whole code stream into `out`, which must hold at least `total_size` bytes. + /// + /// `try_decode_into` is generic over the dictionary view and only specializes its batched + /// copy loop once inlined into the caller, so leaving this wrapper out of line costs ~2.5x. + #[inline(always)] + pub(crate) fn decode_into(&self, out: &mut [MaybeUninit]) -> VortexResult { + let written = match onpair::try_decode_into(self.codes.as_slice(), self.dict, out) { Ok(written) => written, Err(_) => { vortex_panic!("OnPair codes decode to more bytes than uncompressed_lengths records") } }; - if written != total_size { - vortex_panic!( - "OnPair codes decoded to {written} bytes but uncompressed_lengths records {total_size}" - ); + if written != self.total_size { + vortex_panic!( + "OnPair codes decoded to {written} bytes but uncompressed_lengths records {}", + self.total_size + ); + } + Ok(written) } - // SAFETY: `try_decode_into` initialised exactly `written` bytes. +} + +pub(crate) fn onpair_decode_bytes( + array: ArrayView<'_, OnPair>, + ctx: &mut ExecutionCtx, +) -> VortexResult<(ByteBufferMut, PrimitiveArray)> { + let plan = OnPairDecodePlan::new(array, ctx)?; + let mut out_bytes = ByteBufferMut::with_capacity(plan.total_size); + let written = plan.decode_into(out_bytes.spare_capacity_mut())?; + // SAFETY: `decode_into` initialised exactly `written` bytes. unsafe { out_bytes.set_len(written) }; - Ok((out_bytes, lengths)) + Ok((out_bytes, plan.lengths)) } pub(crate) fn onpair_decode_views( diff --git a/encodings/onpair/src/tests.rs b/encodings/onpair/src/tests.rs index 61b5591f54d..e4504ad199f 100644 --- a/encodings/onpair/src/tests.rs +++ b/encodings/onpair/src/tests.rs @@ -16,7 +16,7 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::filter::FilterKernel; use vortex_array::assert_arrays_eq; use vortex_array::buffer::BufferHandle; -use vortex_array::builders::DynVarBinBuilder; +use vortex_array::builders::VarBinBuilder; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -66,7 +66,7 @@ fn test_direct_offset_builder() -> vortex_error::VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let input = sample_input(); let encoded = compress_onpair(input.as_ref(), &mut ctx)?; - let mut builder = DynVarBinBuilder::with_capacity(input.dtype().clone(), false, input.len()); + let mut builder = VarBinBuilder::::with_capacity(input.dtype().clone(), input.len()); encoded .into_array() .append_to_builder(&mut builder, &mut ctx)?; diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index 6a55554a6b9..9a92c153263 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -506,7 +506,7 @@ mod tests { use vortex_array::arrays::DictArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; - use vortex_array::builders::DynVarBinBuilder; + use vortex_array::builders::VarBinBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -564,7 +564,7 @@ mod tests { Some("c"), ]) .into_array(); - let mut builder = DynVarBinBuilder::with_capacity(arr.dtype().clone(), false, arr.len()); + let mut builder = VarBinBuilder::::with_capacity(arr.dtype().clone(), arr.len()); arr.append_to_builder(&mut builder, &mut ctx).unwrap(); assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx); assert_arrays_eq!(arr.into_array(), expected, &mut ctx); diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index b8976274d4c..733eb18da0d 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -702,7 +702,7 @@ mod test { use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; - use vortex_array::builders::DynVarBinBuilder; + use vortex_array::builders::VarBinBuilder; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -926,8 +926,7 @@ mod test { Some("last"), ]) .into_array(); - let mut builder = - DynVarBinBuilder::with_capacity(array.dtype().clone(), false, array.len()); + let mut builder = VarBinBuilder::::with_capacity(array.dtype().clone(), array.len()); array.append_to_builder(&mut builder, &mut ctx).unwrap(); assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx); } @@ -944,8 +943,7 @@ mod test { ) .unwrap(); let expected = VarBinViewArray::from_iter_str(["fill", "second"]).into_array(); - let mut builder = - DynVarBinBuilder::with_capacity(array.dtype().clone(), false, array.len()); + let mut builder = VarBinBuilder::::with_capacity(array.dtype().clone(), array.len()); array.append_to_builder(&mut builder, &mut ctx).unwrap(); assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx); } diff --git a/encodings/zstd/Cargo.toml b/encodings/zstd/Cargo.toml index 1b2cb4f2dd0..49f15619807 100644 --- a/encodings/zstd/Cargo.toml +++ b/encodings/zstd/Cargo.toml @@ -25,6 +25,7 @@ unstable_encodings = [] [dependencies] itertools = { workspace = true } +num-traits = { workspace = true } prost = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index 19d6fa8c736..945f9d0ef1d 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -6,9 +6,12 @@ use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hash; use std::hash::Hasher; +use std::mem::MaybeUninit; +use std::ops::Range; use std::sync::Arc; use itertools::Itertools as _; +use num_traits::AsPrimitive; use prost::Message as _; use vortex_array::Array; use vortex_array::ArrayEq; @@ -30,8 +33,11 @@ use vortex_array::arrays::varbinview::build_views::BinaryView; use vortex_array::arrays::varbinview::build_views::MAX_BUFFER_LEN; use vortex_array::buffer::BufferHandle; use vortex_array::builders::ArrayBuilder; -use vortex_array::builders::DynVarBinBuilder; +use vortex_array::builders::VarBinBuilder; +use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; +use vortex_array::dtype::IntegerPType; +use vortex_array::match_each_varbin_builder; use vortex_array::scalar::Scalar; use vortex_array::serde::ArrayChildren; use vortex_array::smallvec::smallvec; @@ -52,10 +58,10 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; -use vortex_error::vortex_panic; use vortex_mask::AllOr; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use zstd::zstd_safe::WriteBuf; use crate::ZstdFrameMetadata; use crate::ZstdMetadata; @@ -276,49 +282,20 @@ impl VTable for Zstd { builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - let Some(builder) = builder.as_any_mut().downcast_mut::() else { - return array - .array() - .clone() - .execute::(ctx)? - .into_array() - .append_to_builder(builder, ctx); - }; - - let unsliced_validity = - child_to_validity(array.slots()[0].as_ref(), array.dtype().nullability()); - let slice = array - .data() - .decompress_slice(array.dtype(), &unsliced_validity, ctx)?; - let value_start = slice.value_idx_start - slice.n_skipped_values; - let value_count = slice.value_idx_stop - slice.value_idx_start; - let mut values = zstd_values(slice.bytes.as_slice()) - .skip(value_start) - .take(value_count); - let mask = slice.validity.execute_mask(slice.n_rows, ctx)?; - match mask.indices() { - AllOr::All => { - for value in values { - builder.append_n_values(value, 1); - } - } - AllOr::None => builder.append_nulls(slice.n_rows), - AllOr::Some(valid_indices) => { - let mut row = 0; - for &valid_index in valid_indices { - builder.append_nulls(valid_index - row); - builder.append_n_values( - values - .next() - .vortex_expect("Zstd value count must match validity"), - 1, - ); - row = valid_index + 1; - } - builder.append_nulls(slice.n_rows - row); - } + if let Some(result) = + match_each_varbin_builder!(builder, |builder| append_to_varbin(array, builder, ctx)) + { + return result; } - Ok(()) + if let Some(builder) = builder.as_any_mut().downcast_mut::() { + return append_to_varbinview(array, builder, ctx); + } + array + .array() + .clone() + .execute::(ctx)? + .into_array() + .append_to_builder(builder, ctx) } fn reduce_parent( @@ -330,6 +307,83 @@ impl VTable for Zstd { } } +fn unsliced_validity(array: ArrayView<'_, Zstd>) -> Validity { + child_to_validity( + array.slots()[ZstdSlots::VALIDITY].as_ref(), + array.dtype().nullability(), + ) +} + +/// Copies the decompressed values straight into `builder`'s byte storage. +/// +/// The decompressed frames interleave a length prefix with each value, so the bytes have to be +/// compacted; sizing the offsets, byte storage and validity from the slice's own metadata keeps +/// that down to one offset store plus one `memcpy` per value. +fn append_to_varbin( + array: ArrayView<'_, Zstd>, + builder: &mut VarBinBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> +where + usize: AsPrimitive, +{ + let slice = array + .data() + .decompress_slice(array.dtype(), &unsliced_validity(array), ctx)?; + let mask = slice.validity.execute_mask(slice.n_rows, ctx)?; + let (values, num_bytes) = slice.value_bytes()?; + builder.append_value_slices(num_bytes, zstd_values(values), &mask) +} + +/// Hands the decompressed frames to `builder` as data buffers with views built over them. +/// +/// The frames already hold the values contiguously, so the views can reference them in place and +/// the only per-row work is building one view; going through the canonical array instead would +/// rewrite every view a second time to rebase its buffer index. +fn append_to_varbinview( + array: ArrayView<'_, Zstd>, + builder: &mut VarBinViewBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let slice = array + .data() + .decompress_slice(array.dtype(), &unsliced_validity(array), ctx)?; + let mask = slice.validity.execute_mask(slice.n_rows, ctx)?; + + let valid_indices = match mask.indices() { + // No values were stored, so there is nothing to reference and the frames can be dropped. + AllOr::None => { + builder.append_nulls(slice.n_rows); + return Ok(()); + } + AllOr::All => AllOr::All, + AllOr::Some(indices) => AllOr::Some(indices), + }; + + // Build the views against the index the pushed buffers 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 (buffers, all_views) = + try_reconstruct_views(&slice.bytes, next_buffer_index, MAX_BUFFER_LEN)?; + let valid_views = slice_views(&all_views, slice.value_range()?)?; + + let views = match valid_indices { + AllOr::All => valid_views, + AllOr::None => unreachable!("handled above"), + AllOr::Some(valid_indices) => { + // Null rows carry an empty view, so scatter the stored values into their rows. + let mut views = BufferMut::::zeroed(slice.n_rows); + for (view, index) in valid_views.into_iter().zip_eq(valid_indices) { + views[*index] = view; + } + views.freeze() + } + }; + + builder.push_buffer_and_adjusted_views(&buffers, &views, mask); + Ok(()) +} + #[derive(Clone, Debug)] /// Zstd array encoding marker. pub struct Zstd; @@ -528,24 +582,53 @@ fn collect_valid_vbv( /// /// Pass [`MAX_BUFFER_LEN`] for `max_buffer_len` in production; a smaller value can be used in /// tests to exercise the splitting path without allocating >2 GiB. +/// +/// The walk stops at the first length prefix that does not fit inside `buffer`, so a corrupt +/// buffer yields views for its decodable prefix only. pub fn reconstruct_views( buffer: &ByteBuffer, + start_buf_index: u32, max_buffer_len: usize, ) -> (Vec, Buffer) { + let (buffers, views, _) = walk_views(buffer, start_buf_index, max_buffer_len); + (buffers, views) +} + +/// [`reconstruct_views`], but rejecting a buffer that the length prefixes do not tile exactly. +fn try_reconstruct_views( + buffer: &ByteBuffer, + start_buf_index: u32, + max_buffer_len: usize, +) -> VortexResult<(Vec, Buffer)> { + match walk_views(buffer, start_buf_index, max_buffer_len) { + (buffers, views, None) => Ok((buffers, views)), + (_, _, Some(error)) => Err(error), + } +} + +/// Walks `buffer` until it is exhausted or a length prefix leaves it, returning what was decoded +/// along with the error that stopped the walk. +fn walk_views( + buffer: &ByteBuffer, + start_buf_index: u32, + max_buffer_len: usize, +) -> (Vec, Buffer, Option) { let mut views = BufferMut::::empty(); let mut buffers = Vec::new(); let mut segment_start: usize = 0; let mut offset = 0; + // Only a new segment changes the buffer index, so it is tracked instead of recomputed per view. + let mut buf_index = start_buf_index; + let mut error = None; while offset < buffer.len() { - let str_len = ViewLen::from_le_bytes( - buffer - .get(offset..offset + size_of::()) - .vortex_expect("corrupted zstd length") - .try_into() - .ok() - .vortex_expect("must fit ViewLen size"), - ) as usize; + let str_len = match zstd_value_len(buffer.as_slice(), offset) { + Ok(str_len) => str_len, + Err(err) => { + error = Some(err); + break; + } + }; let value_data_offset = offset + size_of::(); let local_offset = value_data_offset - segment_start; @@ -553,12 +636,28 @@ pub fn reconstruct_views( if local_offset + str_len > max_buffer_len && offset > segment_start { buffers.push(buffer.slice(segment_start..offset)); segment_start = offset; + let Some(next_index) = buf_index.checked_add(1) else { + error = Some(vortex_err!("Zstd values need more than u32::MAX buffers")); + break; + }; + buf_index = next_index; } - let local_offset = u32::try_from(value_data_offset - segment_start) - .vortex_expect("local offset within segment must fit in u32"); - let buf_index = u32::try_from(buffers.len()).vortex_expect("buffer index must fit in u32"); - let value = &buffer[value_data_offset..value_data_offset + str_len]; + let Ok(local_offset) = u32::try_from(value_data_offset - segment_start) else { + error = Some(vortex_err!( + "Zstd value offset {} does not fit in u32; max_buffer_len {max_buffer_len} is too large", + value_data_offset - segment_start + )); + break; + }; + let Some(value) = buffer.get(value_data_offset..value_data_offset + str_len) else { + error = Some(vortex_err!( + "Corrupt zstd value: {str_len} bytes at offset {value_data_offset} run past the \ + end of the {} byte frame buffer", + buffer.len() + )); + break; + }; views.push(BinaryView::make_view(value, buf_index, local_offset)); offset = value_data_offset + str_len; } @@ -567,7 +666,63 @@ pub fn reconstruct_views( buffers.push(buffer.slice(segment_start..buffer.len())); } - (buffers, views.freeze()) + (buffers, views.freeze(), error) +} + +/// Narrows the views decoded from the frames down to the values a slice requests. +fn slice_views( + views: &Buffer, + range: Range, +) -> VortexResult> { + vortex_ensure!( + range.end <= views.len(), + "Corrupt zstd metadata: values {}..{} are out of bounds of the {} values held by the \ + decompressed frames", + range.start, + range.end, + views.len() + ); + Ok(views.slice(range)) +} + +/// A zstd output buffer over uninitialized spare capacity. +/// +/// `decompress_to_buffer` writes through a raw pointer and reports how many bytes it produced, so +/// it never reads its destination — but handing it a `&mut [u8]` covering uninitialized memory +/// would be undefined behaviour regardless of what it does with it. [`WriteBuf`] is the interface +/// zstd provides for exactly this case, and it keeps the alternative (zeroing the whole buffer +/// before every decompression) off the hot path. +struct UninitDestination<'a> { + spare: &'a mut [MaybeUninit], + filled: usize, +} + +impl<'a> UninitDestination<'a> { + fn new(spare: &'a mut [MaybeUninit]) -> Self { + Self { spare, filled: 0 } + } +} + +// SAFETY: `as_mut_ptr` and `capacity` describe the whole spare region, so zstd only ever writes +// within it, and `filled_until` merely records the count it reports. `as_slice` is bounded by that +// count, so it never exposes a byte zstd did not write. +unsafe impl WriteBuf for UninitDestination<'_> { + fn as_slice(&self) -> &[u8] { + // SAFETY: zstd reported writing `filled` bytes from the start of `spare`. + unsafe { std::slice::from_raw_parts(self.spare.as_ptr().cast::(), self.filled) } + } + + fn capacity(&self) -> usize { + self.spare.len() + } + + fn as_mut_ptr(&mut self) -> *mut u8 { + self.spare.as_mut_ptr().cast::() + } + + unsafe fn filled_until(&mut self, n: usize) { + self.filled = n; + } } struct DecompressedSlice { @@ -578,23 +733,135 @@ struct DecompressedSlice { value_idx_start: usize, value_idx_stop: usize, n_skipped_values: usize, + /// Number of stored values held by `bytes`, which covers whole frames and so may extend past + /// the requested slice on either side. + n_buffered_values: usize, } +impl DecompressedSlice { + /// The range of stored values this slice requests, as an index into `bytes`. + /// + /// The frame metadata that drives `n_skipped_values` is untrusted, so a frame claiming to hold + /// values that precede the ones we decompressed is rejected instead of wrapping. + fn value_range(&self) -> VortexResult> { + let start = self + .value_idx_start + .checked_sub(self.n_skipped_values) + .ok_or_else(|| { + vortex_err!( + "Corrupt zstd metadata: skipped frames hold {} values, past the first \ + requested value {}", + self.n_skipped_values, + self.value_idx_start + ) + })?; + let end = self + .value_idx_stop + .checked_sub(self.n_skipped_values) + .ok_or_else(|| { + vortex_err!( + "Corrupt zstd metadata: skipped frames hold {} values, past the last \ + requested value {}", + self.n_skipped_values, + self.value_idx_stop + ) + })?; + vortex_ensure!( + start <= end, + "Corrupt zstd metadata: value range {start}..{end} is not ascending" + ); + Ok(start..end) + } + + /// The length-prefixed region of `bytes` holding exactly this slice's values, along with the + /// total size of those values once the length prefixes are dropped. + /// + /// Walking the length prefixes is a dependent load chain, so both ends are derived from the + /// slice metadata where possible: an unsliced array skips both walks. + fn value_bytes(&self) -> VortexResult<(&[u8], usize)> { + let Range { start, end } = self.value_range()?; + let buffer = self.bytes.as_slice(); + let from = zstd_value_offset(buffer, 0, start)?; + let to = if end == self.n_buffered_values { + buffer.len() + } else { + zstd_value_offset(buffer, from, end - start)? + }; + let bytes = buffer.get(from..to).ok_or_else(|| { + vortex_err!( + "Corrupt zstd metadata: values {from}..{to} are out of bounds of the {} byte \ + frame buffer", + buffer.len() + ) + })?; + // Every value carries a length prefix, so the region must be at least that large. + let prefix_bytes = (end - start) + .checked_mul(size_of::()) + .ok_or_else(|| { + vortex_err!( + "Corrupt zstd metadata: value count {} overflows a byte count", + end - start + ) + })?; + let num_bytes = bytes.len().checked_sub(prefix_bytes).ok_or_else(|| { + vortex_err!( + "Corrupt zstd metadata: {} values do not fit in the {} bytes holding them", + end - start, + bytes.len() + ) + })?; + Ok((bytes, num_bytes)) + } +} + +/// Returns the byte offset `count` length-prefixed values past `offset`. +/// +/// Each step is bounds-checked by the next prefix read, so only the offset the walk lands on needs +/// a check of its own. +fn zstd_value_offset(buffer: &[u8], mut offset: usize, count: usize) -> VortexResult { + for _ in 0..count { + offset += size_of::() + zstd_value_len(buffer, offset)?; + } + vortex_ensure!( + offset <= buffer.len(), + "Corrupt zstd values: walking {count} values ended at offset {offset}, past the end of \ + the {} byte frame buffer", + buffer.len() + ); + Ok(offset) +} + +/// Reads the length prefix of the value starting at `offset`. +#[inline] +fn zstd_value_len(buffer: &[u8], offset: usize) -> VortexResult { + let prefix = buffer + .get(offset..) + .and_then(|rest| rest.first_chunk::<{ size_of::() }>()) + .ok_or_else(|| { + vortex_err!( + "Corrupt zstd values: length prefix at offset {offset} runs past the end of the \ + {} byte frame buffer", + buffer.len() + ) + })?; + Ok(ViewLen::from_le_bytes(*prefix) as usize) +} + +/// Iterates the values of a length-prefixed region, stopping at the first prefix that leaves it. +/// +/// Stopping short leaves the value count and byte total below what the caller declared, which the +/// consuming builder rejects, so the walk does not need to report the error itself. fn zstd_values(buffer: &[u8]) -> impl Iterator { let mut offset = 0; std::iter::from_fn(move || { - if offset == buffer.len() { + if offset >= buffer.len() { return None; } - let len = ViewLen::from_le_bytes( - buffer[offset..offset + size_of::()] - .try_into() - .ok() - .vortex_expect("must fit ViewLen size"), - ) as usize; + let len = zstd_value_len(buffer, offset).ok()?; let value_start = offset + size_of::(); + let value = buffer.get(value_start..value_start + len)?; offset = value_start + len; - Some(&buffer[value_start..offset]) + Some(value) }) } @@ -991,36 +1258,66 @@ impl ZstdData { // what row offset into the first such frame. let byte_width = Self::byte_width(dtype); let slice_n_rows = self.slice_stop - self.slice_start; - let slice_value_indices = unsliced_validity - .execute_mask(self.unsliced_n_rows, ctx)? - .valid_counts_for_indices(&[self.slice_start, self.slice_stop]); + let unsliced_mask = unsliced_validity.execute_mask(self.unsliced_n_rows, ctx)?; + let slice_value_indices = + unsliced_mask.valid_counts_for_indices(&[self.slice_start, self.slice_stop]); let slice_value_idx_start = slice_value_indices[0]; let slice_value_idx_stop = slice_value_indices[1]; let mut frames_to_decompress = vec![]; let mut value_idx_start = 0; - let mut uncompressed_size_to_decompress = 0; + let mut uncompressed_size_to_decompress = 0usize; let mut n_skipped_values = 0; + let mut n_buffered_values = 0; for (frame, frame_meta) in self.frames.iter().zip(&self.metadata.frames) { if value_idx_start >= slice_value_idx_stop { break; } - let frame_uncompressed_size = usize::try_from(frame_meta.uncompressed_size) - .vortex_expect("Uncompressed size must fit in usize"); - let frame_n_values = if frame_meta.n_values == 0 { - // possibly older primitive-only metadata that just didn't store this + let frame_uncompressed_size = + usize::try_from(frame_meta.uncompressed_size).map_err(|_| { + vortex_err!( + "Zstd frame uncompressed size {} does not fit in a usize", + frame_meta.uncompressed_size + ) + })?; + let frame_n_values = if frame_meta.n_values != 0 { + usize::try_from(frame_meta.n_values).map_err(|_| { + vortex_err!( + "Zstd frame value count {} does not fit in a usize", + frame_meta.n_values + ) + })? + } else if dtype.is_primitive() { + // Possibly older primitive-only metadata that just didn't store this. Fixed-width + // values make the byte count an exact value count. frame_uncompressed_size / byte_width } else { - usize::try_from(frame_meta.n_values).vortex_expect("frame size must fit usize") + // The same fallback would read a byte count as a value count for variable-width + // values, which mis-attributes values to frames. A single frame holds every stored + // value, so that case is still recoverable; anything else is not. + vortex_ensure!( + self.frames.len() == 1, + "Zstd frame metadata for a variable-width array is missing its value count" + ); + unsliced_mask.true_count() }; - let value_idx_stop = value_idx_start + frame_n_values; + // Bounding the running total also bounds the two accumulators below, which partition + // it between the frames we keep and the ones we skip. + let value_idx_stop = value_idx_start.checked_add(frame_n_values).ok_or_else(|| { + vortex_err!("Corrupt zstd metadata: frame value counts overflow a usize") + })?; if value_idx_stop > slice_value_idx_start { // we need this frame frames_to_decompress.push(frame); - uncompressed_size_to_decompress += frame_uncompressed_size; + uncompressed_size_to_decompress = uncompressed_size_to_decompress + .checked_add(frame_uncompressed_size) + .ok_or_else(|| { + vortex_err!("Corrupt zstd metadata: frame sizes overflow a usize") + })?; + n_buffered_values += frame_n_values; } else { n_skipped_values += frame_n_values; } @@ -1037,24 +1334,28 @@ impl ZstdData { uncompressed_size_to_decompress, Alignment::new(byte_width), ); - unsafe { - // safety: we immediately fill all bytes in the following loop, - // assuming our metadata's uncompressed size is correct - decompressed.set_len(uncompressed_size_to_decompress); - } let mut uncompressed_start = 0; for frame in frames_to_decompress { - let uncompressed_written = decompressor - .decompress_to_buffer(frame.as_slice(), &mut decompressed[uncompressed_start..])?; - uncompressed_start += uncompressed_written; + // Decompress straight into the spare capacity. Each frame gets only the region after + // the ones before it, bounded by the size the metadata declared, so a frame that + // expands further than advertised is refused by zstd rather than overrunning. + let mut destination = UninitDestination::new( + &mut decompressed.spare_capacity_mut() + [uncompressed_start..uncompressed_size_to_decompress], + ); + uncompressed_start += + decompressor.decompress_to_buffer(frame.as_slice(), &mut destination)?; } if uncompressed_start != uncompressed_size_to_decompress { - vortex_panic!( + vortex_bail!( "Zstd metadata or frames were corrupt; expected {} bytes but decompressed {}", uncompressed_size_to_decompress, uncompressed_start ); } + // SAFETY: the loop above decompressed exactly `uncompressed_start` bytes into the front of + // the spare capacity, and the check above pins that to the requested length. + unsafe { decompressed.set_len(uncompressed_start) }; let decompressed = decompressed.freeze(); // Last, we slice the exact values requested out of the decompressed data. @@ -1069,7 +1370,7 @@ impl ZstdData { // We ensure that the validity of the decompressed array ALWAYS matches the validity // implied by the DType. if !dtype.is_nullable() && !matches!(slice_validity, Validity::NonNullable) { - assert!( + vortex_ensure!( matches!(slice_validity, Validity::AllValid), "ZSTD array expects to be non-nullable but there are nulls after decompression" ); @@ -1089,6 +1390,7 @@ impl ZstdData { value_idx_start: slice_value_idx_start, value_idx_stop: slice_value_idx_stop, n_skipped_values, + n_buffered_values, }) } @@ -1101,10 +1403,21 @@ impl ZstdData { let slice = self.decompress_slice(dtype, unsliced_validity, ctx)?; match dtype { DType::Primitive(..) => { - let slice_values_buffer = slice.bytes.slice( - (slice.value_idx_start - slice.n_skipped_values) * slice.byte_width - ..(slice.value_idx_stop - slice.n_skipped_values) * slice.byte_width, - ); + let Range { start, end } = slice.value_range()?; + let byte_range = start + .checked_mul(slice.byte_width) + .zip(end.checked_mul(slice.byte_width)) + .filter(|(_, byte_stop)| *byte_stop <= slice.bytes.len()) + .map(|(byte_start, byte_stop)| byte_start..byte_stop) + .ok_or_else(|| { + vortex_err!( + "Corrupt zstd metadata: values {start}..{end} of {} bytes each are \ + out of bounds of the {} byte frame buffer", + slice.byte_width, + slice.bytes.len() + ) + })?; + let slice_values_buffer = slice.bytes.slice(byte_range); let primitive = PrimitiveArray::from_values_byte_buffer( slice_values_buffer, dtype.as_ptype(), @@ -1118,11 +1431,9 @@ impl ZstdData { DType::Binary(_) | DType::Utf8(_) => { match slice.validity.execute_mask(slice.n_rows, ctx)?.indices() { AllOr::All => { - let (buffers, all_views) = reconstruct_views(&slice.bytes, MAX_BUFFER_LEN); - let valid_views = all_views.slice( - slice.value_idx_start - slice.n_skipped_values - ..slice.value_idx_stop - slice.n_skipped_values, - ); + let (buffers, all_views) = + try_reconstruct_views(&slice.bytes, 0, MAX_BUFFER_LEN)?; + let valid_views = slice_views(&all_views, slice.value_range()?)?; // SAFETY: we properly construct the views inside `reconstruct_views` Ok(unsafe { @@ -1141,11 +1452,9 @@ impl ZstdData { ) .into_array()), AllOr::Some(valid_indices) => { - let (buffers, all_views) = reconstruct_views(&slice.bytes, MAX_BUFFER_LEN); - let valid_views = all_views.slice( - slice.value_idx_start - slice.n_skipped_values - ..slice.value_idx_stop - slice.n_skipped_values, - ); + let (buffers, all_views) = + try_reconstruct_views(&slice.bytes, 0, MAX_BUFFER_LEN)?; + let valid_views = slice_views(&all_views, slice.value_range()?)?; let mut views = BufferMut::::zeroed(slice.n_rows); for (view, index) in valid_views.into_iter().zip_eq(valid_indices) { @@ -1165,7 +1474,7 @@ impl ZstdData { } } } - _ => vortex_panic!("Unsupported dtype for Zstd array: {}", dtype), + _ => vortex_bail!("Unsupported dtype for Zstd array: {}", dtype), } } @@ -1237,9 +1546,17 @@ impl OperationsVTable for Zstd { #[cfg(test)] #[expect(clippy::cast_possible_truncation)] mod tests { + use rstest::rstest; + use vortex_array::validity::Validity; use vortex_buffer::ByteBuffer; + use vortex_error::VortexResult; + use super::DecompressedSlice; + use super::ViewLen; use super::reconstruct_views; + use super::try_reconstruct_views; + use super::zstd_value_len; + use super::zstd_value_offset; use crate::array::BinaryView; /// Build a Zstd-style interleaved buffer: [u32-LE length][string bytes] repeated. @@ -1253,11 +1570,31 @@ mod tests { ByteBuffer::copy_from(buf.as_slice()) } + /// A slice over `bytes` that requests `value_idx_start..value_idx_stop`. + fn decompressed_slice( + bytes: ByteBuffer, + value_idx_start: usize, + value_idx_stop: usize, + n_skipped_values: usize, + n_buffered_values: usize, + ) -> DecompressedSlice { + DecompressedSlice { + bytes, + validity: Validity::NonNullable, + byte_width: 1, + n_rows: value_idx_stop - value_idx_start, + value_idx_start, + value_idx_stop, + n_skipped_values, + n_buffered_values, + } + } + #[test] fn test_reconstruct_views_no_split() { let strings: &[&[u8]] = &[b"hello", b"world"]; let buf = make_interleaved(strings); - let (buffers, views) = reconstruct_views(&buf, 1024); + let (buffers, views) = reconstruct_views(&buf, 0, 1024); assert_eq!(buffers.len(), 1); assert_eq!(views.len(), 2); @@ -1274,7 +1611,7 @@ mod tests { // so it rolls into a second segment. let strings: &[&[u8]] = &[b"aaaaaaaaaaaaa", b"bbbbbbbbbbbbb"]; let buf = make_interleaved(strings); - let (buffers, views) = reconstruct_views(&buf, 20); + let (buffers, views) = reconstruct_views(&buf, 0, 20); assert_eq!(buffers.len(), 2); assert_eq!(views.len(), 2); @@ -1282,4 +1619,84 @@ mod tests { // Second entry starts a new segment at byte 17 (the length prefix), so local offset = 4. assert_eq!(views[1], BinaryView::make_view(b"bbbbbbbbbbbbb", 1, 4)); } + + /// A buffer whose last entry claims more bytes than remain, as corrupt frame data would. + fn make_overrunning() -> ByteBuffer { + let mut buf = Vec::new(); + buf.extend_from_slice(&5u32.to_le_bytes()); + buf.extend_from_slice(b"hello"); + buf.extend_from_slice(&9u32.to_le_bytes()); + buf.extend_from_slice(b"ab"); + ByteBuffer::copy_from(buf.as_slice()) + } + + #[test] + fn test_reconstruct_views_rejects_overrunning_value() { + let buf = make_overrunning(); + assert!(try_reconstruct_views(&buf, 0, 1024).is_err()); + + // The lenient walk keeps the decodable prefix instead of panicking. + let (buffers, views) = reconstruct_views(&buf, 0, 1024); + assert_eq!(buffers.len(), 1); + assert_eq!(views.len(), 1); + assert_eq!(views[0], BinaryView::make_view(b"hello", 0, 4)); + } + + #[test] + fn test_reconstruct_views_rejects_truncated_length_prefix() { + // A trailing partial length prefix cannot start a value. + let buf = + ByteBuffer::copy_from([5u8, 0, 0, 0, b'h', b'e', b'l', b'l', b'o', 1, 0].as_ref()); + assert!(try_reconstruct_views(&buf, 0, 1024).is_err()); + assert_eq!(reconstruct_views(&buf, 0, 1024).1.len(), 1); + } + + #[rstest] + #[case::truncated_buffer(&[0u8, 0, 0], 0)] + #[case::truncated_tail(&[4u8, 0, 0, 0], 2)] + #[case::offset_at_end(&[4u8, 0, 0, 0], 4)] + #[case::offset_past_end(&[4u8, 0, 0, 0], 64)] + fn test_zstd_value_len_rejects_out_of_bounds(#[case] buffer: &[u8], #[case] offset: usize) { + assert!(zstd_value_len(buffer, offset).is_err()); + } + + #[test] + fn test_zstd_value_offset_rejects_walking_past_the_end() -> VortexResult<()> { + let buf = make_interleaved(&[b"hello", b"world"]); + assert_eq!(zstd_value_offset(buf.as_slice(), 0, 2)?, buf.len()); + // Only two values are stored, so the third step leaves the buffer. + assert!(zstd_value_offset(buf.as_slice(), 0, 3).is_err()); + Ok(()) + } + + #[test] + fn test_value_range_rejects_skipping_past_the_requested_values() { + // Frame metadata claiming more skipped values than the slice starts at would wrap. + let slice = decompressed_slice(make_interleaved(&[b"hello"]), 2, 3, 4, 1); + assert!(slice.value_range().is_err()); + assert!(slice.value_bytes().is_err()); + } + + #[rstest] + // The buffered value count matches, so both ends come from the metadata. + #[case::exact_metadata(2)] + // It does not, so the far end is walked instead. + #[case::walked_end(9)] + fn test_value_bytes_totals_the_stored_values( + #[case] n_buffered_values: usize, + ) -> VortexResult<()> { + let buf = make_interleaved(&[b"hello", b"world"]); + let slice = decompressed_slice(buf.clone(), 0, 2, 0, n_buffered_values); + let (bytes, num_bytes) = slice.value_bytes()?; + assert_eq!(bytes, buf.as_slice()); + assert_eq!(num_bytes, buf.len() - 2 * size_of::()); + Ok(()) + } + + #[test] + fn test_value_bytes_rejects_more_values_than_the_buffer_holds() { + // Frame metadata claims nine values but only two are stored. + let slice = decompressed_slice(make_interleaved(&[b"hello", b"world"]), 0, 5, 0, 9); + assert!(slice.value_bytes().is_err()); + } } diff --git a/encodings/zstd/src/test.rs b/encodings/zstd/src/test.rs index fac5fcdc631..7cee5d6a5d4 100644 --- a/encodings/zstd/src/test.rs +++ b/encodings/zstd/src/test.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors #![expect(clippy::cast_possible_truncation)] +use rstest::rstest; +use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; @@ -10,15 +12,19 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; use vortex_array::assert_nth_scalar; -use vortex_array::builders::DynVarBinBuilder; +use vortex_array::builders::VarBinBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::validity::Validity; use vortex_buffer::Alignment; use vortex_buffer::Buffer; +use vortex_error::VortexResult; use vortex_mask::Mask; use crate::Zstd; +use crate::ZstdArray; +use crate::ZstdData; +use crate::ZstdMetadata; #[test] fn test_zstd_compress_decompress() { @@ -224,7 +230,7 @@ fn test_zstd_append_to_offset_builder() { .slice(1..4) .unwrap(); let mut builder = - DynVarBinBuilder::with_capacity(compressed.dtype().clone(), false, compressed.len()); + VarBinBuilder::::with_capacity(compressed.dtype().clone(), compressed.len()); compressed .append_to_builder(&mut builder, &mut ctx) .unwrap(); @@ -276,6 +282,98 @@ fn test_sliced_array_children() { sliced.children(); } +/// Six rows, five of them stored, compressed into `values_per_frame`-sized frames with the frame +/// metadata a reader would have deserialized rewritten by `corrupt`. +fn corrupt_var_bin_view_metadata( + values_per_frame: usize, + corrupt: impl FnOnce(&mut ZstdMetadata), + ctx: &mut ExecutionCtx, +) -> VortexResult<(VarBinViewArray, ZstdArray)> { + let array = VarBinViewArray::from_iter( + [ + Some(b"foo".as_slice()), + Some(b"bar".as_slice()), + None, + Some(b"Lorem ipsum dolor sit amet".as_slice()), + Some(b"baz".as_slice()), + Some(b"quux".as_slice()), + ], + DType::Utf8(Nullability::Nullable), + ); + let mut data = ZstdData::from_var_bin_view(&array, 0, values_per_frame, ctx)?; + corrupt(&mut data.metadata); + let compressed = Zstd::try_new(array.dtype().clone(), data, array.validity()?)?; + Ok((array, compressed)) +} + +/// Frame metadata comes straight off disk, so an inconsistent value count has to surface as an +/// error from both read paths rather than a panic or a wrapped length. +#[rstest] +#[case::frame_holds_fewer_values_than_claimed(|metadata: &mut ZstdMetadata| { + metadata.frames[0].n_values = 1000; +})] +#[case::frame_value_counts_overflow(|metadata: &mut ZstdMetadata| { + metadata.frames[1].n_values = u64::MAX; +})] +#[case::missing_value_count_across_frames(|metadata: &mut ZstdMetadata| { + metadata.frames[0].n_values = 0; +})] +fn test_zstd_rejects_corrupt_frame_metadata( + #[case] corrupt: fn(&mut ZstdMetadata), +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let (_, compressed) = corrupt_var_bin_view_metadata(3, corrupt, &mut ctx)?; + + assert!(Zstd::decompress(&compressed, &mut ctx).is_err()); + + let mut builder = + VarBinBuilder::::with_capacity(compressed.dtype().clone(), compressed.len()); + assert!( + compressed + .append_to_builder(&mut builder, &mut ctx) + .is_err() + ); + Ok(()) +} + +/// Metadata written before frames recorded their value count leaves it at zero. A single frame +/// holds every stored value, so those arrays still read back. +#[test] +fn test_zstd_reads_legacy_single_frame_var_bin_metadata() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let (array, compressed) = + corrupt_var_bin_view_metadata(0, |metadata| metadata.frames[0].n_values = 0, &mut ctx)?; + + assert_arrays_eq!(compressed, array.clone().into_array(), &mut ctx); + assert_arrays_eq!( + compressed.slice(2..5)?, + array.into_array().slice(2..5)?, + &mut ctx + ); + Ok(()) +} + +/// The same legacy metadata for fixed-width values recovers the count from the frame size, which +/// stays exact across frames. +#[test] +fn test_zstd_reads_legacy_primitive_frame_metadata() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = PrimitiveArray::from_iter(0..200_i32); + let mut data = ZstdData::from_primitive(&array, 3, 30, &mut ctx)?; + for frame in &mut data.metadata.frames { + frame.n_values = 0; + } + let compressed = Zstd::try_new(array.dtype().clone(), data, array.validity()?)?; + + assert_arrays_eq!(compressed, PrimitiveArray::from_iter(0..200_i32), &mut ctx); + assert_arrays_eq!( + compressed.slice(100..105)?, + PrimitiveArray::from_iter(100..105_i32), + &mut ctx + ); + Ok(()) +} + /// Tests that each beginning of a frame in ZSTD matches /// the buffer alignment when compressing primitive arrays. #[test] diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index d9131e3a933..2e4c982a8ea 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -33,7 +33,6 @@ use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::BoolBuilder; use crate::builders::DecimalBuilder; -use crate::builders::DynVarBinBuilder; use crate::builders::NullBuilder; use crate::builders::PrimitiveBuilder; use crate::builders::VarBinViewBuilder; @@ -41,6 +40,7 @@ use crate::canonical::Canonical; use crate::dtype::DType; use crate::match_each_decimal_value; use crate::match_each_native_ptype; +use crate::match_each_varbin_builder; use crate::scalar::DecimalValue; use crate::scalar::Scalar; use crate::scalar::ScalarValue; @@ -223,8 +223,10 @@ impl VTable for Constant { }); } DType::Utf8(_) => { - if let Some(builder) = builder.as_any_mut().downcast_mut::() { - builder.append_scalar_repeated(scalar, n)?; + if let Some(result) = match_each_varbin_builder!(builder, |builder| { + builder.append_scalar_repeated(scalar, n) + }) { + result?; } else { append_value_or_nulls::(builder, scalar.is_null(), n, |b| { let value = scalar @@ -236,8 +238,10 @@ impl VTable for Constant { } } DType::Binary(_) => { - if let Some(builder) = builder.as_any_mut().downcast_mut::() { - builder.append_scalar_repeated(scalar, n)?; + if let Some(result) = match_each_varbin_builder!(builder, |builder| { + builder.append_scalar_repeated(scalar, n) + }) { + result?; } else { append_value_or_nulls::(builder, scalar.is_null(), n, |b| { let value = scalar diff --git a/vortex-array/src/arrays/dict/array.rs b/vortex-array/src/arrays/dict/array.rs index 10733402041..24832003052 100644 --- a/vortex-array/src/arrays/dict/array.rs +++ b/vortex-array/src/arrays/dict/array.rs @@ -298,7 +298,7 @@ mod test { use crate::arrays::PrimitiveArray; use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; - use crate::builders::DynVarBinBuilder; + use crate::builders::VarBinBuilder; use crate::builders::builder_with_capacity; use crate::dtype::DType; use crate::dtype::NativePType; @@ -445,11 +445,11 @@ mod test { } #[test] - fn test_dict_utf8_append_to_dyn_varbin_builder() -> VortexResult<()> { + fn test_dict_utf8_append_to_varbin_builder() -> VortexResult<()> { let values = VarBinViewArray::from_iter_str(["zero", "one", "two"]); let dict = DictArray::try_new(buffer![2u8, 0, 2, 1].into_array(), values.into_array())?; let expected = VarBinViewArray::from_iter_str(["two", "zero", "two", "one"]); - let mut builder = DynVarBinBuilder::with_capacity(dict.dtype().clone(), false, dict.len()); + let mut builder = VarBinBuilder::::with_capacity(dict.dtype().clone(), dict.len()); let mut ctx = array_session().create_execution_ctx(); dict.into_array() diff --git a/vortex-array/src/arrays/dict/vtable/mod.rs b/vortex-array/src/arrays/dict/vtable/mod.rs index d9ced9eeffa..897b7811179 100644 --- a/vortex-array/src/arrays/dict/vtable/mod.rs +++ b/vortex-array/src/arrays/dict/vtable/mod.rs @@ -3,13 +3,17 @@ use std::hash::Hasher; +use num_traits::AsPrimitive; use prost::Message; use smallvec::smallvec; +use vortex_buffer::BitBufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_error::vortex_panic; +use vortex_mask::AllOr; +use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -34,17 +38,23 @@ use crate::array::VTable; use crate::array::with_empty_buffers; use crate::arrays::ConstantArray; use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::arrays::VarBinViewArray; use crate::arrays::dict::DictArrayExt; use crate::arrays::dict::DictArraySlotsExt; use crate::arrays::dict::compute::rules::PARENT_RULES; use crate::arrays::dict::execute::take_canonical; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; +use crate::builders::VarBinBuilder; use crate::dtype::DType; +use crate::dtype::IntegerPType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; +use crate::match_each_integer_ptype; +use crate::match_each_varbin_builder; use crate::require_child; use crate::scalar::Scalar; use crate::serde::ArrayChildren; @@ -223,6 +233,14 @@ impl VTable for Dict { && !codes.validity()?.definitely_all_null() { let codes = codes.into_owned(); + if let Canonical::VarBinView(values) = Canonical::from(values) + && varbin_fast_path_is_host(&codes, &values) + && let Some(result) = match_each_varbin_builder!(builder, |builder| { + append_dict_to_varbin(&codes, &values, builder, ctx) + }) + { + return result; + } let canonical = take_canonical(values, &codes, ctx)?.into_array(); canonical.append_to_builder(builder, ctx)?; return Ok(()); @@ -245,3 +263,231 @@ impl VTable for Dict { PARENT_RULES.evaluate(array, parent, child_idx) } } + +/// Whether every buffer that [`append_dict_to_varbin`] dereferences directly is host-resident. +/// +/// The fast path reads the codes and the dictionary's views and data buffers as plain slices, +/// which is only possible for host allocations. Device-resident dictionaries fall back to the +/// canonical take, which goes through the compute stack and can therefore be served by a device +/// kernel. This mirrors [`DictArrayExt::validate_all_values_referenced`], which likewise skips +/// host-only work when the codes are not resident on the host. +/// +/// This is a handful of enum discriminant checks over the array's buffer handles, so it stays +/// outside the per-row loop and costs nothing measurable on the host-resident path. +fn varbin_fast_path_is_host(codes: &PrimitiveArray, values: &VarBinViewArray) -> bool { + codes.buffer_handle().is_on_host() + && values.views_handle().is_on_host() + && values.data_buffers().iter().all(BufferHandle::is_on_host) +} + +/// Gathers the dictionary values straight into `builder`. +/// +/// The canonical route first takes the values to full logical length, which allocates and then +/// re-reads a views buffer proportional to the row count. The dictionary is usually far smaller +/// than the column, so resolving each code against it in place skips that intermediate entirely +/// and leaves one `memcpy` per row as the only work. +/// +/// The caller must have checked [`varbin_fast_path_is_host`]; this function indexes the backing +/// buffers as host slices and panics otherwise. +fn append_dict_to_varbin( + codes: &PrimitiveArray, + values: &VarBinViewArray, + builder: &mut VarBinBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> +where + usize: AsPrimitive, +{ + let len = codes.as_ref().len(); + let codes_validity = codes.as_ref().validity()?.execute_mask(len, ctx)?; + let values_validity = values.validity()?.execute_mask(values.len(), ctx)?; + + // Resolve the dictionary's storage once so that looking up a code is an O(1) read. + let views = values.views(); + let buffers = values + .data_buffers() + .iter() + .map(|buffer| buffer.as_host().as_slice()) + .collect::>(); + + match_each_integer_ptype!(codes.ptype(), |C| { + let codes = codes.as_slice::(); + let validity = dict_validity(&codes_validity, codes, &values_validity); + let view = |row: usize| &views[AsPrimitive::::as_(codes[row])]; + match validity.bit_buffer() { + AllOr::All => { + let num_bytes = (0..len).map(|row| view(row).len() as usize).sum(); + builder.append_value_slices( + num_bytes, + (0..len).map(|row| view(row).bytes(&buffers)), + &validity, + ) + } + AllOr::None => { + builder.push_nulls(len); + Ok(()) + } + AllOr::Some(bits) => { + let mut num_bytes = 0; + bits.for_each_set_index(|row| num_bytes += view(row).len() as usize); + builder.append_value_slices( + num_bytes, + bits.set_indices().map(|row| view(row).bytes(&buffers)), + &validity, + ) + } + } + }) +} + +/// Combines the codes' validity with the dictionary values' validity resolved through the codes. +/// +/// A code pointing at a null dictionary entry produces a null row, matching what +/// [`take_canonical`] derives from [`Validity::take`](crate::validity::Validity::take). +fn dict_validity( + codes_validity: &Mask, + codes: &[C], + values_validity: &Mask, +) -> Mask { + if values_validity.all_true() { + return codes_validity.clone(); + } + + let mut valid = BitBufferMut::new_unset(codes.len()); + match codes_validity.bit_buffer() { + AllOr::All => { + for (row, &code) in codes.iter().enumerate() { + valid.set_to(row, values_validity.value(code.as_())); + } + } + AllOr::None => {} + AllOr::Some(bits) => bits.for_each_set_index(|row| { + valid.set_to(row, values_validity.value(codes[row].as_())); + }), + } + Mask::from_buffer(valid.freeze()) +} + +#[cfg(test)] +mod tests { + use std::any::Any; + use std::ops::Range; + use std::sync::Arc; + + use futures::FutureExt; + use futures::future::BoxFuture; + use vortex_buffer::Alignment; + use vortex_buffer::ByteBuffer; + use vortex_buffer::buffer; + + use super::*; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::VarBinViewArray; + use crate::arrays::dict::DictArray; + use crate::assert_arrays_eq; + use crate::buffer::DeviceBuffer; + use crate::dtype::Nullability::NonNullable; + use crate::dtype::Nullability::Nullable; + use crate::validity::Validity; + + const LONG: &str = "a string that is far too long to be inlined in a view"; + + /// A stand-in for a real device allocation so residency handling can be exercised without a + /// GPU. Only the host/device discriminant matters here; the payload is never read back. + #[derive(Debug, PartialEq, Eq, Hash)] + struct FakeDeviceBuffer(Vec); + + impl DeviceBuffer for FakeDeviceBuffer { + fn as_any(&self) -> &dyn Any { + self + } + + fn len(&self) -> usize { + self.0.len() + } + + fn alignment(&self) -> Alignment { + Alignment::of::() + } + + fn copy_to_host_sync(&self, alignment: Alignment) -> VortexResult { + Ok(ByteBuffer::copy_from_aligned(&self.0, alignment)) + } + + fn copy_to_host( + &self, + alignment: Alignment, + ) -> VortexResult>> { + let copied = self.copy_to_host_sync(alignment)?; + Ok(async move { Ok(copied) }.boxed()) + } + + fn slice(&self, range: Range) -> Arc { + Arc::new(Self(self.0[range].to_vec())) + } + + fn aligned(self: Arc, _alignment: Alignment) -> VortexResult> { + Ok(self) + } + } + + fn to_device(handle: &BufferHandle) -> BufferHandle { + BufferHandle::new_device(Arc::new(FakeDeviceBuffer(handle.as_host().to_vec()))) + } + + #[test] + fn append_to_builder_gathers_through_the_dictionary() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dict = DictArray::try_new( + PrimitiveArray::from_option_iter([Some(0u32), Some(2), None, Some(1), Some(0)]) + .into_array(), + VarBinViewArray::from_iter([Some(LONG), None, Some("short")], DType::Utf8(Nullable)) + .into_array(), + )?; + + let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); + dict.append_to_builder(&mut builder, &mut ctx)?; + + let expected = VarBinViewArray::from_iter( + [Some(LONG), Some("short"), None, None, Some(LONG)], + DType::Utf8(Nullable), + ); + assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx); + Ok(()) + } + + /// The fast path indexes the codes and the dictionary's buffers as host slices, so it must be + /// skipped whenever any of them lives on a device. Regression test for the fast path panicking + /// on device-resident dictionaries instead of deferring to the canonical take. + #[test] + fn fast_path_is_skipped_for_device_resident_buffers() { + let codes = PrimitiveArray::new(buffer![0u32, 1, 0], Validity::NonNullable); + let values = VarBinViewArray::from_iter_str([LONG, "short"]); + + assert!(varbin_fast_path_is_host(&codes, &values)); + + let device_codes = PrimitiveArray::from_buffer_handle( + to_device(codes.buffer_handle()), + PType::U32, + Validity::NonNullable, + ); + assert!(!varbin_fast_path_is_host(&device_codes, &values)); + + let device_views = VarBinViewArray::new_handle( + to_device(values.views_handle()), + Arc::clone(values.data_buffers()), + DType::Utf8(NonNullable), + Validity::NonNullable, + ); + assert!(!varbin_fast_path_is_host(&codes, &device_views)); + + let device_data = VarBinViewArray::new_handle( + values.views_handle().clone(), + values.data_buffers().iter().map(to_device).collect(), + DType::Utf8(NonNullable), + Validity::NonNullable, + ); + assert!(!varbin_fast_path_is_host(&codes, &device_data)); + } +} diff --git a/vortex-array/src/arrays/varbin/array.rs b/vortex-array/src/arrays/varbin/array.rs index 867e319bcb9..0b768d86571 100644 --- a/vortex-array/src/arrays/varbin/array.rs +++ b/vortex-array/src/arrays/varbin/array.rs @@ -383,11 +383,11 @@ impl Array { dtype: DType, ) -> Self { let iter = iter.into_iter(); - let mut builder = VarBinBuilder::::with_capacity(iter.size_hint().0); + let mut builder = VarBinBuilder::::with_capacity(dtype, iter.size_hint().0); for v in iter { builder.append(v.as_ref().map(|o| o.as_ref())); } - builder.finish(dtype) + builder.finish_into_varbin() } pub fn from_iter_nonnull, I: IntoIterator>( @@ -395,11 +395,11 @@ impl Array { dtype: DType, ) -> Self { let iter = iter.into_iter(); - let mut builder = VarBinBuilder::::with_capacity(iter.size_hint().0); + let mut builder = VarBinBuilder::::with_capacity(dtype, iter.size_hint().0); for v in iter { builder.append_value(v); } - builder.finish(dtype) + builder.finish_into_varbin() } fn from_vec_sized(vec: Vec, dtype: DType) -> Self @@ -407,11 +407,11 @@ impl Array { O: IntegerPType, T: AsRef<[u8]>, { - let mut builder = VarBinBuilder::::with_capacity(vec.len()); + let mut builder = VarBinBuilder::::with_capacity(dtype, vec.len()); for v in vec { builder.append_value(v.as_ref()); } - builder.finish(dtype) + builder.finish_into_varbin() } /// Create from a vector of string slices. diff --git a/vortex-array/src/arrays/varbin/builder.rs b/vortex-array/src/arrays/varbin/builder.rs index 9a0a2a8922f..6335013ce5e 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -2,19 +2,25 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::any::Any; +use std::mem::MaybeUninit; +use std::ops::Range; use num_traits::AsPrimitive; use vortex_buffer::BitBufferMut; +use vortex_buffer::BitIndexIterator; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; +use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; +use vortex_mask::AllOr; use vortex_mask::Mask; use crate::ArrayRef; +use crate::ArrayView; use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; @@ -26,6 +32,7 @@ use crate::arrays::VarBinArray; use crate::arrays::VarBinView; use crate::arrays::varbin::VarBinArrayExt; use crate::arrays::varbin::VarBinArraySlotsExt; +use crate::arrays::varbinview::BinaryView; use crate::arrays::varbinview::VarBinViewArrayExt; use crate::builders::ArrayBuilder; use crate::dtype::DType; @@ -34,71 +41,135 @@ use crate::expr::stats::Precision; use crate::expr::stats::Stat; #[cfg(debug_assertions)] use crate::legacy_session; +use crate::match_each_integer_ptype; use crate::scalar::Scalar; use crate::validity::Validity; + +/// Builder for [`VarBinArray`] values with `O`-typed offsets. +/// +/// This is the offset-based counterpart to +/// [`VarBinViewBuilder`](crate::builders::VarBinViewBuilder): values are laid out contiguously in +/// a single byte buffer described by a monotonically increasing offsets buffer. An `i32` or `i64` +/// builder can therefore be handed straight to Arrow as a `Utf8`/`LargeUtf8`/`Binary`/ +/// `LargeBinary` array without re-laying out the bytes. +/// +/// Encodings that can decode into this layout should specialize +/// [`append_to_builder`](crate::vtable::VTable::append_to_builder) with +/// [`match_each_varbin_builder!`](crate::match_each_varbin_builder), which recovers the concrete +/// offset type from a `&mut dyn ArrayBuilder` so the decode loop is monomorphized over it. pub struct VarBinBuilder { + dtype: DType, offsets: BufferMut, - data: BufferMut, + data: ByteBufferMut, validity: BitBufferMut, } -impl Default for VarBinBuilder { - fn default() -> Self { - Self::new() - } -} - impl VarBinBuilder { - pub fn new() -> Self { - Self::with_capacity(0) + /// Creates an empty builder for `dtype`. + pub fn new(dtype: DType) -> Self { + Self::with_capacity(dtype, 0) } - pub fn with_capacity(len: usize) -> Self { - let mut offsets = BufferMut::with_capacity(len + 1); + /// Creates a builder for `dtype` with room for `capacity` values. + pub fn with_capacity(dtype: DType, capacity: usize) -> Self { + assert!( + matches!(dtype, DType::Utf8(_) | DType::Binary(_)), + "VarBinBuilder dtype must be Utf8 or Binary, got {dtype}" + ); + let mut offsets = BufferMut::with_capacity(capacity + 1); offsets.push(O::zero()); Self { + dtype, offsets, data: BufferMut::empty(), - validity: BitBufferMut::with_capacity(len), + validity: BitBufferMut::with_capacity(capacity), } } + /// Creates a builder for `dtype` with room for `capacity` values totalling `bytes` bytes. + pub fn with_capacity_bytes(dtype: DType, capacity: usize, bytes: usize) -> Self { + let mut builder = Self::with_capacity(dtype, capacity); + builder.reserve_data(bytes); + builder + } + + /// Reserves room for `additional` value bytes. + /// + /// [`reserve_exact`](ArrayBuilder::reserve_exact) takes a row count and so cannot size the + /// value bytes; callers that know the byte total should use this to size the buffer once + /// instead of letting it grow. + pub fn reserve_data(&mut self, additional: usize) { + self.data.reserve(additional); + } + #[inline] pub fn append(&mut self, value: Option<&[u8]>) { match value { Some(v) => self.append_value(v), - None => self.append_null(), + None => self.push_null(), } } #[inline] pub fn append_value(&mut self, value: impl AsRef<[u8]>) { - let slice = value.as_ref(); - self.offsets - .push(O::from(self.data.len() + slice.len()).unwrap_or_else(|| { - vortex_panic!( - "Failed to convert sum of {} and {} to offset of type {}", - self.data.len(), - slice.len(), - std::any::type_name::() - ) - })); - self.data.extend_from_slice(slice); + self.push_value(value.as_ref()); self.validity.append_true(); } + /// Appends the same non-null value `n` times. + pub fn append_n_values(&mut self, value: impl AsRef<[u8]>, n: usize) { + let value = value.as_ref(); + self.offsets.reserve(n); + self.data.reserve(value.len().saturating_mul(n)); + for _ in 0..n { + self.push_value(value); + } + self.validity.append_n(true, n); + } + + /// Appends a null value. + /// + /// Unlike [`append_null`](ArrayBuilder::append_null) this does not check that the builder is + /// nullable; the offsets and validity stay consistent either way, and a non-nullable `dtype` + /// discards the validity bits on [`finish_into_varbin`](Self::finish_into_varbin). #[inline] - pub fn append_null(&mut self) { - self.offsets.push(self.offsets[self.offsets.len() - 1]); - self.validity.append_false(); + pub fn push_null(&mut self) { + self.push_nulls(1) } + /// Appends `n` null values. See [`push_null`](Self::push_null). #[inline] - pub fn append_n_nulls(&mut self, n: usize) { - self.offsets.push_n(self.offsets[self.offsets.len() - 1], n); + pub fn push_nulls(&mut self, n: usize) { + self.offsets.push_n(self.last_offset(), n); self.validity.append_n(false, n); } + /// Appends the same UTF-8 or binary scalar `n` times. + pub fn append_scalar_repeated(&mut self, scalar: &Scalar, n: usize) -> VortexResult<()> { + vortex_ensure!( + scalar.dtype() == &self.dtype, + "VarBinBuilder expected scalar with dtype {}, got {}", + self.dtype, + scalar.dtype() + ); + match &self.dtype { + DType::Utf8(_) => match scalar.as_utf8().value() { + Some(value) => self.append_n_values(value, n), + None => self.push_nulls(n), + }, + DType::Binary(_) => match scalar.as_binary().value() { + Some(value) => self.append_n_values(value, n), + None => self.push_nulls(n), + }, + dtype => vortex_bail!("VarBinBuilder cannot append scalar of dtype {dtype}"), + } + Ok(()) + } + + /// Appends values from one contiguous byte buffer described by relative end offsets. + /// + /// Each entry of `end_offsets` marks the end of one value relative to the start of `values`, + /// and there must be exactly one per entry in `validity`. #[inline] pub fn append_values

( &mut self, @@ -108,94 +179,191 @@ impl VarBinBuilder { ) -> VortexResult<()> where P: AsPrimitive, + usize: AsPrimitive, { - let offsets_start = self.offsets.len(); - let data_start = self.data.len(); - let mut previous_end = data_start; - let mut len = 0; - - for offset in end_offsets { - let relative_end = offset.as_(); - let Some(end) = data_start.checked_add(relative_end) else { - self.offsets.truncate(offsets_start); - vortex_bail!("Byte offset overflow: {data_start} + {relative_end}"); - }; - if end < previous_end { - self.offsets.truncate(offsets_start); - vortex_bail!("End offsets must be monotonically increasing"); - } - let Some(end_offset) = O::from(end) else { - self.offsets.truncate(offsets_start); - vortex_bail!( - "Byte offset {end} does not fit in {}", - std::any::type_name::() - ); - }; - self.offsets.push(end_offset); - previous_end = end; - len += 1; - } - - if len != validity.len() { - self.offsets.truncate(offsets_start); - vortex_bail!( - "End offset count {len} does not match validity length {}", - validity.len() - ); - } - if previous_end - data_start != values.len() { - self.offsets.truncate(offsets_start); - vortex_bail!( - "Final relative end offset {} does not match values length {}", - previous_end - data_start, - values.len() - ); - } - + // Offsets are committed first: they are the only fallible part, and failing before the + // bytes are appended leaves the builder untouched. + self.extend_offsets(values.len(), validity.len(), end_offsets)?; self.data.extend_from_slice(values); self.append_validity(validity); Ok(()) } - fn append_validity(&mut self, validity: &Mask) { - match validity { - Mask::AllTrue(len) => self.validity.append_n(true, *len), - Mask::AllFalse(len) => self.validity.append_n(false, *len), - Mask::Values(values) => self.validity.append_buffer(values.bit_buffer()), - } + /// Appends values whose bytes `decode` writes straight into the builder's byte storage. + /// + /// `decode` is handed at least `num_bytes + slack` bytes of uninitialized storage and returns + /// the number of bytes it initialized, which must be exactly `num_bytes`. Each entry of + /// `lengths` is the byte length of one value — including nulls, which are zero-length — so + /// there must be one per entry in `validity`. + /// + /// `slack` is spare headroom past the values themselves. It exists so that a decoder which + /// stores in wide fixed-size chunks can keep using them through the final value instead of + /// dropping into a byte-at-a-time tail; a decoder must still never write past the slice it is + /// handed, so `0` is correct for one that is already exact. + /// + /// Decoding in place saves staging the decoded heap in a temporary buffer and copying it in. + /// + /// `decode` is taken as a `dyn` reference so a caller that resolves the `lengths` type with + /// [`match_each_integer_ptype!`](crate::match_each_integer_ptype) can build the closure once + /// outside that match rather than inlining a whole decoder into each of its arms. + /// + /// # Safety + /// + /// `decode` must initialize the first `n` bytes of the slice it is passed, where `n` is the + /// value it returns. Those bytes are published as initialized without ever being read first, + /// so a `decode` that over-reports leaves the builder holding uninitialized memory. + pub unsafe fn append_decoded

( + &mut self, + num_bytes: usize, + slack: usize, + lengths: &[P], + validity: &Mask, + decode: &mut dyn FnMut(&mut [MaybeUninit]) -> VortexResult, + ) -> VortexResult<()> + where + P: AsPrimitive, + usize: AsPrimitive, + { + let Some(capacity) = num_bytes.checked_add(slack) else { + vortex_bail!("Decoded size overflow: {num_bytes} + {slack}"); + }; + self.data.reserve(capacity); + + let data_len = self.data.len(); + let written = decode(self.data.spare_capacity_mut())?; + vortex_ensure!( + written == num_bytes, + "Decoded {written} bytes, expected {num_bytes}" + ); + + // The decoded bytes live in spare capacity until `set_len` below, so an invalid `lengths` + // still leaves the builder unchanged. + self.extend_offsets(num_bytes, validity.len(), prefix_sums(lengths))?; + + // SAFETY: `decode` reported initializing `written` spare bytes, and the caller guarantees + // that report is accurate; `written == num_bytes` was checked above. + unsafe { self.data.set_len(data_len + num_bytes) }; + self.append_validity(validity); + Ok(()) } - fn len(&self) -> usize { - self.validity.len() + /// Appends `validity.len()` values by copying the slices yielded by `values`. + /// + /// `values` must yield exactly `validity.true_count()` slices totalling `num_bytes` bytes. + /// Null entries consume no bytes and repeat the preceding offset. + /// + /// Both buffers are sized once up front and the validity bits are appended in bulk, so the + /// copy loop costs one offset store plus one `memcpy` per value. + pub fn append_value_slices<'a, I>( + &mut self, + num_bytes: usize, + values: I, + validity: &Mask, + ) -> VortexResult<()> + where + I: Iterator, + usize: AsPrimitive, + { + let data_start = self.data.len(); + match self.gather_value_slices(num_bytes, values, validity) { + Ok(()) => { + self.append_validity(validity); + Ok(()) + } + Err(error) => { + // The offsets are never committed on failure, so only the copied bytes need to go. + self.data.truncate(data_start); + Err(error) + } + } } - fn reserve_exact(&mut self, additional: usize) { - self.offsets.reserve(additional); - self.validity.reserve(additional); + /// Appends a [`VarBinArray`], reusing its offsets instead of walking its values. + pub fn append_varbin( + &mut self, + array: ArrayView<'_, VarBin>, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> + where + usize: AsPrimitive, + { + let offsets = array.offsets().clone().execute::(ctx)?; + let bytes: ByteBuffer = array.sliced_bytes(); + let validity = array + .varbin_validity() + .execute_mask(array.as_ref().len(), ctx)?; + match_each_integer_ptype!(offsets.ptype(), |P| { + let offsets = offsets.as_slice::

(); + let first: usize = offsets[0].as_(); + self.append_values( + bytes.as_slice(), + // Wrapping keeps a corrupt offsets child from panicking here; `append_values` + // rejects the resulting non-monotonic offsets. + offsets[1..] + .iter() + .map(|offset| AsPrimitive::::as_(*offset).wrapping_sub(first)), + &validity, + ) + }) } - fn set_validity(&mut self, validity: Mask) { - self.validity = match validity { - Mask::AllTrue(len) => BitBufferMut::new_set(len), - Mask::AllFalse(len) => BitBufferMut::new_unset(len), - values @ Mask::Values(_) => values - .into_bit_buffer() - .try_into_mut() - .unwrap_or_else(|buffer| BitBufferMut::copy_from(&buffer)), + /// Appends a [`VarBinViewArray`](crate::arrays::VarBinViewArray), compacting its values. + pub fn append_varbinview( + &mut self, + array: ArrayView<'_, VarBinView>, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> + where + usize: AsPrimitive, + { + let len = array.as_ref().len(); + let validity = array.varbinview_validity().execute_mask(len, ctx)?; + + // Resolve the views slice and the data buffers once. Reading them per row costs a buffer + // handle clone, and the byte total needs only the fixed-width view headers. + let views = array.views(); + let buffers = array + .data_buffers() + .iter() + .map(|buffer| buffer.as_host().as_slice()) + .collect::>(); + + let num_bytes = match validity.bit_buffer() { + AllOr::All => views.iter().map(|view| view.len() as usize).sum(), + AllOr::None => 0, + AllOr::Some(bits) => { + let mut total = 0; + bits.for_each_set_index(|index| total += views[index].len() as usize); + total + } }; + + self.append_value_slices( + num_bytes, + ValidValues::new(&validity, views, &buffers), + &validity, + ) } + /// Finishes the appended values into a [`VarBinArray`] and resets the builder. #[allow(clippy::disallowed_methods)] - pub fn finish(self, dtype: DType) -> VarBinArray { + pub fn finish_into_varbin(&mut self) -> VarBinArray { assert_eq!( self.offsets.len() - 1, self.validity.len(), "The offset count must be one more than the validity length" ); - let offsets = PrimitiveArray::new(self.offsets.freeze(), Validity::NonNullable); - let nulls = self.validity.freeze(); - let validity = Validity::from_bit_buffer(nulls, dtype.nullability()); + let mut fresh_offsets = BufferMut::with_capacity(1); + fresh_offsets.push(O::zero()); + let offsets = PrimitiveArray::new( + std::mem::replace(&mut self.offsets, fresh_offsets).freeze(), + Validity::NonNullable, + ); + let data = std::mem::replace(&mut self.data, BufferMut::empty()); + let nulls = std::mem::replace(&mut self.validity, BitBufferMut::empty()).freeze(); + + let validity = Validity::from_bit_buffer(nulls, self.dtype.nullability()); // The builder adds offsets in monotonically increasing order. Store this statistic to // prevent VarBinArray::validate from recomputing it after deserialization. @@ -217,175 +385,200 @@ impl VarBinBuilder { // - Validity matches the dtype nullability. // - UTF-8 validity is ensured by the caller when using DType::Utf8. unsafe { - VarBinArray::new_unchecked(offsets.into_array(), self.data.freeze(), dtype, validity) + VarBinArray::new_unchecked( + offsets.into_array(), + data.freeze(), + self.dtype.clone(), + validity, + ) } } -} -/// Builder for UTF-8 and binary [`VarBinArray`] values. -/// -/// Unlike [`crate::builders::VarBinViewBuilder`], this builder stores 32-bit or 64-bit offsets. -/// Encodings can decode into this builder without first creating a -/// [`VarBinViewArray`](crate::arrays::VarBinViewArray). -pub struct DynVarBinBuilder { - dtype: DType, - storage: DynOffsets, -} + #[inline] + fn last_offset(&self) -> O { + self.offsets[self.offsets.len() - 1] + } -enum DynOffsets { - I32(VarBinBuilder), - I64(VarBinBuilder), -} + /// Appends `value`'s bytes and its end offset, leaving the validity bits to the caller. + #[inline] + fn push_value(&mut self, value: &[u8]) { + self.offsets + .push(O::from(self.data.len() + value.len()).unwrap_or_else(|| { + vortex_panic!( + "Failed to convert sum of {} and {} to offset of type {}", + self.data.len(), + value.len(), + std::any::type_name::() + ) + })); + self.data.extend_from_slice(value); + } -impl DynVarBinBuilder { - /// Creates a builder with 32-bit offsets, or 64-bit offsets when `large_offsets` is true. - pub fn with_capacity(dtype: DType, large_offsets: bool, capacity: usize) -> Self { - assert!( - matches!(dtype, DType::Utf8(_) | DType::Binary(_)), - "DynVarBinBuilder dtype must be Utf8 or Binary" - ); - let storage = if large_offsets { - DynOffsets::I64(VarBinBuilder::with_capacity(capacity)) - } else { - DynOffsets::I32(VarBinBuilder::with_capacity(capacity)) + fn append_validity(&mut self, validity: &Mask) { + match validity { + Mask::AllTrue(len) => self.validity.append_n(true, *len), + Mask::AllFalse(len) => self.validity.append_n(false, *len), + Mask::Values(values) => self.validity.append_buffer(values.bit_buffer()), + } + } + + fn replace_validity(&mut self, validity: Mask) { + self.validity = match validity { + Mask::AllTrue(len) => BitBufferMut::new_set(len), + Mask::AllFalse(len) => BitBufferMut::new_unset(len), + values @ Mask::Values(_) => values + .into_bit_buffer() + .try_into_mut() + .unwrap_or_else(|buffer| BitBufferMut::copy_from(&buffer)), }; - Self { dtype, storage } } - /// Appends decompressed values from one contiguous byte buffer. + /// Appends `count` end offsets derived from `end_offsets`, shifted past the current data end. /// - /// Each offset in `end_offsets` marks the end of one value relative to the start of `values`. - pub fn append_values

( + /// `end_offsets` must be monotonically non-decreasing and end at exactly `num_bytes`. Offsets + /// are written through the reserved spare capacity and committed with a single `set_len`, so a + /// rejected input leaves the buffer untouched and the caller can propagate the error. + fn extend_offsets

( &mut self, - values: &[u8], + num_bytes: usize, + count: usize, end_offsets: impl Iterator, - validity: &Mask, ) -> VortexResult<()> where P: AsPrimitive, + usize: AsPrimitive, { - match &mut self.storage { - DynOffsets::I32(builder) => builder.append_values(values, end_offsets, validity), - DynOffsets::I64(builder) => builder.append_values(values, end_offsets, validity), - } - } - - /// Appends the same non-null value `n` times. - pub fn append_n_values(&mut self, value: impl AsRef<[u8]>, n: usize) { - let value = value.as_ref(); - for _ in 0..n { - self.append_value(value); + let data_start = self.data.len(); + let offsets_len = self.offsets.len(); + self.check_offset_limit(data_start, num_bytes)?; + self.offsets.reserve(count); + + // Writing into the spare capacity keeps the output cursor in a register: `push` rewrites + // the buffer length on every value, which the optimizer cannot hoist out of the loop. + let spare = &mut self.offsets.spare_capacity_mut()[..count]; + let mut end_offsets = end_offsets; + let mut previous = 0usize; + for slot in spare.iter_mut() { + let Some(end) = end_offsets.next() else { + vortex_bail!("End offset count is less than the validity length {count}"); + }; + let end = end.as_(); + vortex_ensure!( + end >= previous && end <= num_bytes, + "End offsets must be monotonically increasing within {num_bytes} bytes, \ + got {end} after {previous}" + ); + slot.write((data_start + end).as_()); + previous = end; } - } - - /// Appends the same UTF-8 or binary scalar `n` times. - pub fn append_scalar_repeated(&mut self, scalar: &Scalar, n: usize) -> VortexResult<()> { vortex_ensure!( - scalar.dtype() == self.dtype(), - "DynVarBinBuilder expected scalar with dtype {}, got {}", - self.dtype(), - scalar.dtype() + end_offsets.next().is_none(), + "End offset count exceeds the validity length {count}" + ); + vortex_ensure!( + previous == num_bytes, + "Final end offset {previous} does not match the value byte count {num_bytes}" ); - match self.dtype() { - DType::Utf8(_) => match scalar.as_utf8().value() { - Some(value) => self.append_n_values(value, n), - None => self.push_nulls(n), - }, - DType::Binary(_) => match scalar.as_binary().value() { - Some(value) => self.append_n_values(value, n), - None => self.push_nulls(n), - }, - dtype => vortex_bail!("DynVarBinBuilder cannot append scalar of dtype {dtype}"), - } - Ok(()) - } - /// Appends an existing [`VarBinArray`] without constructing views. - pub fn append_varbin( - &mut self, - array: crate::ArrayView<'_, VarBin>, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - let offsets = array.offsets().clone().execute::(ctx)?; - let bytes: ByteBuffer = array.sliced_bytes(); - let validity = array - .varbin_validity() - .execute_mask(array.as_ref().len(), ctx)?; - crate::match_each_integer_ptype!(offsets.ptype(), |P| { - let offsets = offsets.as_slice::

(); - let first_offset: usize = offsets[0].as_(); - self.append_values( - bytes.as_slice(), - offsets - .iter() - .skip(1) - .map(|offset| AsPrimitive::::as_(*offset) - first_offset), - &validity, - ) - })?; + // SAFETY: the loop initialized the first `count` spare slots. + unsafe { self.offsets.set_len(offsets_len + count) }; Ok(()) } - /// Appends a [`VarBinViewArray`](crate::arrays::VarBinViewArray). - pub fn append_varbinview( + /// Copies `values` into the byte storage and records their offsets. See + /// [`append_value_slices`](Self::append_value_slices); the validity bits are the caller's job + /// so that a failure here can be unwound. + fn gather_value_slices<'a, I>( &mut self, - array: crate::ArrayView<'_, VarBinView>, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - let validity = array - .varbinview_validity() - .execute_mask(array.as_ref().len(), ctx)?; - match &validity { - Mask::AllTrue(_) => { - for index in 0..array.as_ref().len() { - self.append_value(array.bytes_at(index)); + num_bytes: usize, + values: I, + validity: &Mask, + ) -> VortexResult<()> + where + I: Iterator, + usize: AsPrimitive, + { + let count = validity.len(); + let data_start = self.data.len(); + let offsets_len = self.offsets.len(); + self.check_offset_limit(data_start, num_bytes)?; + self.offsets.reserve(count); + self.data.reserve(num_bytes); + + // Disjoint field borrows: the offsets spare capacity stays valid while the byte buffer + // grows, because the reserve above means it never reallocates. + let Self { offsets, data, .. } = self; + let spare = &mut offsets.spare_capacity_mut()[..count]; + let mut values = values; + let limit = data_start + num_bytes; + + match validity.bit_buffer() { + AllOr::All => { + for slot in spare.iter_mut() { + let Some(value) = values.next() else { + vortex_bail!("Value slice count is less than the row count {count}"); + }; + vortex_ensure!( + data.len() + value.len() <= limit, + "Value slices exceed the declared byte count {num_bytes}" + ); + data.extend_from_slice(value); + slot.write(data.len().as_()); } } - Mask::AllFalse(len) => self.push_nulls(*len), - Mask::Values(values) => { - for (index, is_valid) in values.bit_buffer().iter().enumerate() { - if is_valid { - self.append_value(array.bytes_at(index)); - } else { - self.push_null(); - } + AllOr::None => { + spare.fill(MaybeUninit::new(data_start.as_())); + } + AllOr::Some(bits) => { + let mut row = 0; + for index in bits.set_indices() { + let Some(value) = values.next() else { + vortex_bail!("Value slice count is less than the valid row count"); + }; + vortex_ensure!( + data.len() + value.len() <= limit, + "Value slices exceed the declared byte count {num_bytes}" + ); + // Null rows between the previous valid row and this one repeat its end offset. + spare[row..index].fill(MaybeUninit::new(data.len().as_())); + data.extend_from_slice(value); + spare[index].write(data.len().as_()); + row = index + 1; } + spare[row..].fill(MaybeUninit::new(data.len().as_())); } } - Ok(()) - } - - /// Finishes the current values as a [`VarBinArray`] and resets the builder. - pub fn finish_into_varbin(&mut self) -> VarBinArray { - match &mut self.storage { - DynOffsets::I32(builder) => std::mem::take(builder).finish(self.dtype.clone()), - DynOffsets::I64(builder) => std::mem::take(builder).finish(self.dtype.clone()), - } - } - fn append_value(&mut self, value: impl AsRef<[u8]>) { - match &mut self.storage { - DynOffsets::I32(builder) => builder.append_value(value), - DynOffsets::I64(builder) => builder.append_value(value), - } - } + vortex_ensure!( + values.next().is_none(), + "Value slice count exceeds the valid row count" + ); + vortex_ensure!( + data.len() == limit, + "Value slices total {} bytes, expected {num_bytes}", + data.len() - data_start + ); - fn push_null(&mut self) { - match &mut self.storage { - DynOffsets::I32(builder) => builder.append_null(), - DynOffsets::I64(builder) => builder.append_null(), - } + // SAFETY: every branch above initialized all `count` spare slots. + unsafe { self.offsets.set_len(offsets_len + count) }; + Ok(()) } - fn push_nulls(&mut self, n: usize) { - match &mut self.storage { - DynOffsets::I32(builder) => builder.append_n_nulls(n), - DynOffsets::I64(builder) => builder.append_n_nulls(n), - } + /// Checks that an offset past `num_bytes` more bytes is representable as an `O`. + fn check_offset_limit(&self, data_start: usize, num_bytes: usize) -> VortexResult<()> { + let Some(limit) = data_start.checked_add(num_bytes) else { + vortex_bail!("Byte offset overflow: {data_start} + {num_bytes}"); + }; + vortex_ensure!( + u64::try_from(limit).is_ok_and(|limit| limit <= O::max_value_as_u64()), + "Byte offset {limit} does not fit in {}", + std::any::type_name::() + ); + Ok(()) } } -impl ArrayBuilder for DynVarBinBuilder { +impl ArrayBuilder for VarBinBuilder { fn as_any(&self) -> &dyn Any { self } @@ -399,16 +592,12 @@ impl ArrayBuilder for DynVarBinBuilder { } fn len(&self) -> usize { - match &self.storage { - DynOffsets::I32(builder) => builder.len(), - DynOffsets::I64(builder) => builder.len(), - } + self.validity.len() } fn append_zeros(&mut self, n: usize) { - for _ in 0..n { - self.append_value([]); - } + self.offsets.push_n(self.last_offset(), n); + self.validity.append_n(true, n); } unsafe fn append_nulls_unchecked(&mut self, n: usize) { @@ -420,17 +609,12 @@ impl ArrayBuilder for DynVarBinBuilder { } fn reserve_exact(&mut self, additional: usize) { - match &mut self.storage { - DynOffsets::I32(builder) => builder.reserve_exact(additional), - DynOffsets::I64(builder) => builder.reserve_exact(additional), - } + self.offsets.reserve(additional); + self.validity.reserve(additional); } unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - match &mut self.storage { - DynOffsets::I32(builder) => builder.set_validity(validity), - DynOffsets::I64(builder) => builder.set_validity(validity), - } + self.replace_validity(validity) } fn finish(&mut self) -> ArrayRef { @@ -440,23 +624,162 @@ impl ArrayBuilder for DynVarBinBuilder { fn finish_into_canonical(&mut self, ctx: &mut ExecutionCtx) -> Canonical { self.finish() .execute::(ctx) - .vortex_expect("varbin buffer builder should canonicalize") + .vortex_expect("varbin builder should canonicalize") + } +} + +/// Recovers a concrete [`VarBinBuilder`] from a `&mut dyn ArrayBuilder`. +/// +/// Evaluates to `Some($body)` with `$typed` bound to the `&mut VarBinBuilder` when the builder +/// is one, and `None` otherwise, so an encoding can fall through to its other output paths: +/// +/// ```ignore +/// if let Some(result) = match_each_varbin_builder!(builder, |builder| { +/// append_my_encoding(array, builder, ctx) +/// }) { +/// return result; +/// } +/// ``` +/// +/// The body is expanded once per offset width, which is the point: the decode loop specializes on +/// the width instead of dispatching through `dyn`. That is also why only `i32` and `i64` are +/// matched — the two widths Arrow's byte arrays use. Widening it would duplicate every caller's +/// decode loop again for widths no hot path uses. An encoding that falls through still reaches a +/// builder of any width by canonicalizing first, because the canonical `VarBin` and `VarBinView` +/// appends use [`match_each_any_varbin_builder!`](crate::match_each_any_varbin_builder), whose +/// bodies are small enough to expand for every width. +#[macro_export] +macro_rules! match_each_varbin_builder { + ($builder:expr, | $typed:ident | $body:expr) => { + $crate::__match_varbin_builder_widths!($builder, |$typed| $body, [i32, i64]) + }; +} + +/// Recovers a concrete [`VarBinBuilder`] of *any* offset width from a `&mut dyn ArrayBuilder`. +/// +/// Same shape as [`match_each_varbin_builder!`], but covering every width `VarBinBuilder` can be +/// instantiated with. Reserve it for small bodies — it expands them eight times. It exists so the +/// canonical `VarBin` and `VarBinView` appends accept a builder of any width, which is what makes +/// an unusual width usable at all: every encoding without a specialization for it canonicalizes +/// and lands here. +#[macro_export] +macro_rules! match_each_any_varbin_builder { + ($builder:expr, | $typed:ident | $body:expr) => { + $crate::__match_varbin_builder_widths!( + $builder, + |$typed| $body, + [u8, u16, u32, u64, i8, i16, i32, i64] + ) + }; +} + +/// Expands `$body` once per listed offset width, guarded by a downcast. See +/// [`match_each_varbin_builder!`]. +#[doc(hidden)] +#[macro_export] +macro_rules! __match_varbin_builder_widths { + ($builder:expr, | $typed:ident | $body:expr, [$($width:ty),+ $(,)?]) => {{ + let __varbin_builder: &mut dyn $crate::builders::ArrayBuilder = $builder; + $crate::__match_varbin_builder_arms!(__varbin_builder, |$typed| $body, [$($width),+]) + }}; +} + +/// The `if`/`else if` chain behind [`__match_varbin_builder_widths!`], one arm per width. +#[doc(hidden)] +#[macro_export] +macro_rules! __match_varbin_builder_arms { + ($builder:expr, | $typed:ident | $body:expr, []) => { + None + }; + ($builder:expr, | $typed:ident | $body:expr, [$head:ty $(, $tail:ty)*]) => { + if $builder + .as_any() + .is::<$crate::builders::VarBinBuilder<$head>>() + { + let $typed = match $builder + .as_any_mut() + .downcast_mut::<$crate::builders::VarBinBuilder<$head>>() + { + Some(typed) => typed, + None => unreachable!("builder type checked above"), + }; + Some($body) + } else { + $crate::__match_varbin_builder_arms!($builder, |$typed| $body, [$($tail),*]) + } + }; +} + +/// Running totals of `lengths`, wrapping so a corrupt lengths child is rejected rather than +/// panicking; [`VarBinBuilder::extend_offsets`] catches the resulting non-monotonic offsets. +#[inline] +fn prefix_sums>(lengths: &[P]) -> impl Iterator { + lengths.iter().scan(0usize, |end, length| { + *end = end.wrapping_add(length.as_()); + Some(*end) + }) +} + +/// Iterator over the bytes of the valid rows of a `VarBinViewArray`. +struct ValidValues<'a> { + views: &'a [BinaryView], + buffers: &'a [&'a [u8]], + indices: ValidIndices<'a>, +} + +enum ValidIndices<'a> { + All(Range), + None, + Some(BitIndexIterator<'a>), +} + +impl<'a> ValidValues<'a> { + fn new(validity: &'a Mask, views: &'a [BinaryView], buffers: &'a [&'a [u8]]) -> Self { + let indices = match validity.bit_buffer() { + AllOr::All => ValidIndices::All(0..views.len()), + AllOr::None => ValidIndices::None, + AllOr::Some(bits) => ValidIndices::Some(bits.set_indices()), + }; + Self { + views, + buffers, + indices, + } + } +} + +impl<'a> Iterator for ValidValues<'a> { + type Item = &'a [u8]; + + #[inline] + fn next(&mut self) -> Option { + let index = match &mut self.indices { + ValidIndices::All(range) => range.next()?, + ValidIndices::None => return None, + ValidIndices::Some(indices) => indices.next()?, + }; + Some(self.views[index].bytes(self.buffers)) } } #[cfg(test)] mod tests { + use std::mem::MaybeUninit; + use rstest::rstest; + use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_mask::Mask; + use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; + use crate::arrays::DictArray; + use crate::arrays::PrimitiveArray; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::arrays::varbin::VarBinArraySlotsExt; - use crate::arrays::varbin::builder::DynVarBinBuilder; use crate::arrays::varbin::builder::VarBinBuilder; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; @@ -466,14 +789,15 @@ mod tests { use crate::expr::stats::Stat; use crate::expr::stats::StatsProviderExt; use crate::scalar::Scalar; + use crate::validity::Validity; #[test] fn test_builder() { - let mut builder = VarBinBuilder::::with_capacity(0); + let mut builder = VarBinBuilder::::with_capacity(DType::Utf8(Nullable), 0); builder.append(Some(b"hello")); builder.append(None); builder.append(Some(b"world")); - let array = builder.finish(DType::Utf8(Nullable)); + let array = builder.finish_into_varbin(); assert_eq!(array.len(), 3); assert_eq!(array.dtype().nullability(), Nullable); @@ -507,19 +831,19 @@ mod tests { ) .into_array() .slice(1..4)?; - let mut builder = - DynVarBinBuilder::with_capacity(source.dtype().clone(), large_offsets, source.len()); let mut ctx = array_session().create_execution_ctx(); - source.append_to_builder(&mut builder, &mut ctx)?; + let actual = with_offsets(large_offsets, source.dtype().clone(), |builder| { + source.append_to_builder(builder, &mut ctx) + })?; - assert_arrays_eq!(builder.finish_into_varbin(), source, &mut ctx); + assert_arrays_eq!(actual, source, &mut ctx); Ok(()) } #[test] fn append_values_offset_overflow_returns_error() { - let mut builder = VarBinBuilder::::new(); + let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); let values = [0u8; 128]; let result = builder.append_values(&values, [values.len()].into_iter(), &Mask::new_true(1)); @@ -530,12 +854,94 @@ mod tests { assert_eq!(builder.validity.len(), 0); } + #[test] + fn append_values_rejects_a_short_offset_count() { + let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); + + let result = builder.append_values(b"ab", [1usize, 2].into_iter(), &Mask::new_true(3)); + + assert!(result.is_err()); + assert_eq!(builder.offsets.len(), 1); + assert!(builder.data.is_empty()); + } + + #[test] + fn append_values_rejects_non_monotonic_offsets() { + let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); + + let result = builder.append_values(b"ab", [2usize, 1].into_iter(), &Mask::new_true(2)); + + assert!(result.is_err()); + assert_eq!(builder.offsets.len(), 1); + } + + #[test] + fn append_decoded_writes_into_the_builder_storage() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); + + // SAFETY: the closure initializes exactly the 6 bytes it reports. + unsafe { + builder.append_decoded( + 6, + 4, + &[3usize, 0, 3], + &Mask::from_iter([true, false, true]), + &mut |spare: &mut [MaybeUninit]| { + for (slot, byte) in spare.iter_mut().zip(b"foobar") { + slot.write(*byte); + } + Ok(6) + }, + )?; + } + + let expected = + VarBinViewArray::from_iter([Some("foo"), None, Some("bar")], DType::Utf8(Nullable)); + assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx); + Ok(()) + } + + #[test] + fn append_decoded_rejects_a_short_decode() { + let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); + + // SAFETY: the closure initializes the 3 bytes it reports (none, and it reports 3 — but the + // length mismatch is rejected before anything is published). + let result = unsafe { + builder.append_decoded(6, 0, &[3usize, 3], &Mask::new_true(2), &mut |spare| { + spare[..3].fill(MaybeUninit::new(b'x')); + Ok(3) + }) + }; + + assert!(result.is_err()); + assert_eq!(builder.offsets.len(), 1); + assert_eq!(builder.validity.len(), 0); + } + + #[test] + fn append_value_slices_rejects_a_byte_count_mismatch() { + let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); + + let result = builder.append_value_slices( + 6, + [b"foo".as_slice(), b"quux".as_slice()].into_iter(), + &Mask::new_true(2), + ); + + assert!(result.is_err()); + assert_eq!(builder.offsets.len(), 1); + assert!(builder.data.is_empty()); + assert_eq!(builder.validity.len(), 0); + } + #[test] #[should_panic(expected = "The offset count must be one more than the validity length")] fn finish_rejects_mismatched_validity() { - let mut builder = VarBinBuilder::::new(); + let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); builder.validity.append_true(); - drop(builder.finish(DType::Utf8(Nullable))); + drop(builder.finish_into_varbin()); } #[rstest] @@ -548,53 +954,96 @@ mod tests { Mask::new_false(3), Mask::from_iter([true, false, true]), ] { - let mut builder = - DynVarBinBuilder::with_capacity(DType::Utf8(Nullable), large_offsets, 0); - assert!(builder.as_any().is::()); - builder.reserve_exact(3); - builder.append_zero(); - builder.append_scalar(&Scalar::utf8("hello", Nullable))?; - builder.append_null(); - assert_eq!(builder.len(), 3); - builder.set_validity(validity.clone()); - - let result = builder.finish_into_canonical(&mut ctx).into_array(); + let result = with_offsets(large_offsets, DType::Utf8(Nullable), |builder| { + builder.reserve_exact(3); + builder.append_zero(); + builder.append_scalar(&Scalar::utf8("hello", Nullable))?; + builder.append_null(); + assert_eq!(builder.len(), 3); + builder.set_validity(validity.clone()); + Ok(()) + })?; assert_eq!(result.validity()?.execute_mask(3, &mut ctx)?, validity); } Ok(()) } - #[test] - fn test_append_varbinview_validity_to_dyn_builder() -> VortexResult<()> { + #[rstest] + #[case(false)] + #[case(true)] + fn test_append_varbinview_validity_to_builder(#[case] large_offsets: bool) -> VortexResult<()> { + let long = "a value that does not fit inline"; let all_null = VarBinViewArray::from_iter([None::<&str>, None], DType::Utf8(Nullable)); let mixed = - VarBinViewArray::from_iter([Some("hello"), None, Some("world")], DType::Utf8(Nullable)); + VarBinViewArray::from_iter([Some("hello"), None, Some(long)], DType::Utf8(Nullable)); let expected = VarBinViewArray::from_iter( - [None, None, Some("hello"), None, Some("world")], + [None, None, Some("hello"), None, Some(long)], DType::Utf8(Nullable), ); - let mut builder = - DynVarBinBuilder::with_capacity(expected.dtype().clone(), false, expected.len()); let mut ctx = array_session().create_execution_ctx(); - all_null - .into_array() - .append_to_builder(&mut builder, &mut ctx)?; - mixed + let actual = with_offsets(large_offsets, expected.dtype().clone(), |builder| { + all_null + .clone() + .into_array() + .append_to_builder(builder, &mut ctx)?; + mixed + .clone() + .into_array() + .append_to_builder(builder, &mut ctx) + })?; + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + /// `match_each_varbin_builder!` covers only `i32` and `i64`, but a builder of any other offset + /// width must still be fillable — it reaches the same code through `AnyVarBinBuilder`. + #[rstest] + #[case::u32(VarBinBuilder::::new(DType::Utf8(Nullable)))] + #[case::u64(VarBinBuilder::::new(DType::Utf8(Nullable)))] + #[case::i16(VarBinBuilder::::new(DType::Utf8(Nullable)))] + fn append_to_an_unmatched_offset_width_still_works( + #[case] mut builder: impl ArrayBuilder, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let expected = + VarBinViewArray::from_iter([Some("hello"), None, Some("world")], DType::Utf8(Nullable)); + + expected + .clone() .into_array() .append_to_builder(&mut builder, &mut ctx)?; + assert_arrays_eq!(builder.finish(), expected, &mut ctx); + Ok(()) + } + + /// Reaching a `VarBinBuilder` of an unmatched width from a compressed encoding goes the long + /// way round: the encoding's specialization declines, it canonicalizes, and the canonical + /// `VarBinView` append dispatches through `AnyVarBinBuilder`. + #[test] + fn append_dict_to_an_unmatched_offset_width_still_works() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let values = VarBinViewArray::from_iter_str(["hello", "world"]).into_array(); + let codes = PrimitiveArray::new(buffer![0u8, 1, 1, 0], Validity::NonNullable); + let dict = DictArray::try_new(codes.into_array(), values)?.into_array(); + let expected = dict.clone().execute::(&mut ctx)?.into_array(); + + let mut builder = VarBinBuilder::::new(dict.dtype().clone()); + dict.append_to_builder(&mut builder, &mut ctx)?; + assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx); Ok(()) } #[test] fn offsets_have_is_sorted_stat() -> VortexResult<()> { - let mut builder = VarBinBuilder::::with_capacity(0); + let mut builder = VarBinBuilder::::with_capacity(DType::Utf8(Nullable), 0); builder.append_value(b"aaa"); - builder.append_null(); + builder.push_null(); builder.append_value(b"bbb"); - let array = builder.finish(DType::Utf8(Nullable)); + let array = builder.finish_into_varbin(); let is_sorted = array .offsets() @@ -606,8 +1055,8 @@ mod tests { #[test] fn empty_builder_offsets_have_is_sorted_stat() -> VortexResult<()> { - let builder = VarBinBuilder::::new(); - let array = builder.finish(DType::Utf8(Nullable)); + let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); + let array = builder.finish_into_varbin(); let is_sorted = array .offsets() @@ -616,4 +1065,21 @@ mod tests { assert_eq!(is_sorted, Precision::Exact(true)); Ok(()) } + + /// Runs `f` against an `i32` or `i64` builder and returns the finished array. + fn with_offsets( + large_offsets: bool, + dtype: DType, + f: impl FnOnce(&mut dyn ArrayBuilder) -> VortexResult<()>, + ) -> VortexResult { + if large_offsets { + let mut builder = VarBinBuilder::::with_capacity(dtype, 8); + f(&mut builder)?; + Ok(builder.finish_into_varbin()) + } else { + let mut builder = VarBinBuilder::::with_capacity(dtype, 8); + f(&mut builder)?; + Ok(builder.finish_into_varbin()) + } + } } diff --git a/vortex-array/src/arrays/varbin/compute/compare.rs b/vortex-array/src/arrays/varbin/compute/compare.rs index 3f5b2b0b01f..3fff845b361 100644 --- a/vortex-array/src/arrays/varbin/compute/compare.rs +++ b/vortex-array/src/arrays/varbin/compute/compare.rs @@ -301,11 +301,12 @@ mod tests { #[test] fn varbin_i64_offsets_compare_constant() { let mut ctx = array_session().create_execution_ctx(); - let mut builder = VarBinBuilder::::with_capacity(3); + let mut builder = + VarBinBuilder::::with_capacity(DType::Utf8(Nullability::NonNullable), 3); builder.append_value(b"abc"); builder.append_value(b"xyz"); builder.append_value(b"abc"); - let array = builder.finish(DType::Utf8(Nullability::NonNullable)); + let array = builder.finish_into_varbin(); let result = array .into_array() @@ -322,11 +323,12 @@ mod tests { #[test] fn varbin_i64_offsets_compare_constant_binary() { let mut ctx = array_session().create_execution_ctx(); - let mut builder = VarBinBuilder::::with_capacity(3); + let mut builder = + VarBinBuilder::::with_capacity(DType::Binary(Nullability::NonNullable), 3); builder.append_value(b"abc"); builder.append_value(b"xyz"); builder.append_value(b"abc"); - let array = builder.finish(DType::Binary(Nullability::NonNullable)); + let array = builder.finish_into_varbin(); let result = array .into_array() diff --git a/vortex-array/src/arrays/varbin/compute/filter.rs b/vortex-array/src/arrays/varbin/compute/filter.rs index a90bc87ef07..00b3d1b2fd9 100644 --- a/vortex-array/src/arrays/varbin/compute/filter.rs +++ b/vortex-array/src/arrays/varbin/compute/filter.rs @@ -88,7 +88,7 @@ where O: IntegerPType, usize: AsPrimitive, { - let mut builder = VarBinBuilder::::with_capacity(selection_count); + let mut builder = VarBinBuilder::::with_capacity(dtype, selection_count); match logical_validity.bit_buffer() { AllOr::All => { for &(start, end) in mask_slices { @@ -96,7 +96,7 @@ where } } AllOr::None => { - builder.append_n_nulls(selection_count); + builder.push_nulls(selection_count); } AllOr::Some(validity) => { for (start, end) in mask_slices.iter().copied() { @@ -120,14 +120,14 @@ where })?; builder.append_value(&data[s..e]) } else { - builder.append_null() + builder.push_null() } } } } } } - Ok(builder.finish(dtype)) + Ok(builder.finish_into_varbin()) } fn update_non_nullable_slice( @@ -195,7 +195,7 @@ fn filter_select_var_bin_by_index_primitive_offset( Ok(&data[start..end]) }; - let mut builder = VarBinBuilder::::with_capacity(selection_count); + let mut builder = VarBinBuilder::::with_capacity(dtype, selection_count); match mask.bit_buffer() { AllOr::All => { for idx in mask_indices.iter().copied() { @@ -203,19 +203,19 @@ fn filter_select_var_bin_by_index_primitive_offset( } } AllOr::None => { - builder.append_n_nulls(selection_count); + builder.push_nulls(selection_count); } AllOr::Some(validity) => { for idx in mask_indices.iter().copied() { if validity.value(idx) { builder.append_value(value_at(idx)?) } else { - builder.append_null() + builder.push_null() } } } } - Ok(builder.finish(dtype)) + Ok(builder.finish_into_varbin()) } #[cfg(test)] diff --git a/vortex-array/src/arrays/varbin/vtable/canonical.rs b/vortex-array/src/arrays/varbin/vtable/canonical.rs index fca9671e3b4..acea7bd92f4 100644 --- a/vortex-array/src/arrays/varbin/vtable/canonical.rs +++ b/vortex-array/src/arrays/varbin/vtable/canonical.rs @@ -60,14 +60,14 @@ mod tests { #[case(DType::Utf8(Nullability::Nullable))] #[case(DType::Binary(Nullability::Nullable))] fn test_canonical_varbin_sliced(#[case] dtype: DType) { - let mut varbin = VarBinBuilder::::with_capacity(10); - varbin.append_null(); - varbin.append_null(); + let mut varbin = VarBinBuilder::::with_capacity(dtype.clone(), 10); + varbin.push_null(); + varbin.push_null(); // inlined value varbin.append_value("123456789012".as_bytes()); // non-inlinable value varbin.append_value("1234567890123".as_bytes()); - let varbin = varbin.finish(dtype.clone()); + let varbin = varbin.finish_into_varbin(); let varbin = varbin.slice(1..4).unwrap(); diff --git a/vortex-array/src/arrays/varbin/vtable/mod.rs b/vortex-array/src/arrays/varbin/vtable/mod.rs index 6a51e61e3d3..5a11aee863a 100644 --- a/vortex-array/src/arrays/varbin/vtable/mod.rs +++ b/vortex-array/src/arrays/varbin/vtable/mod.rs @@ -25,10 +25,10 @@ use crate::arrays::varbin::VarBinData; use crate::arrays::varbin::VarBinSlots; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; -use crate::builders::DynVarBinBuilder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; +use crate::match_each_any_varbin_builder; use crate::serde::ArrayChildren; use crate::validity::Validity; mod canonical; @@ -209,8 +209,10 @@ impl VTable for VarBin { builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - if let Some(builder) = builder.as_any_mut().downcast_mut::() { - return builder.append_varbin(array, ctx); + if let Some(result) = + match_each_any_varbin_builder!(builder, |builder| builder.append_varbin(array, ctx)) + { + return result; } varbin_to_canonical(array, ctx)? .into_array() diff --git a/vortex-array/src/arrays/varbinview/view.rs b/vortex-array/src/arrays/varbinview/view.rs index e8e6d37d502..94392529715 100644 --- a/vortex-array/src/arrays/varbinview/view.rs +++ b/vortex-array/src/arrays/varbinview/view.rs @@ -246,6 +246,26 @@ impl BinaryView { unsafe { &mut self._ref } } + /// Returns the bytes of this value, reading out of `buffers` when it is not inlined. + /// + /// Unlike [`VarBinViewData::bytes_at`](crate::arrays::varbinview::VarBinViewData::bytes_at) + /// this borrows + /// from slices the caller has already resolved instead of cloning a buffer handle, so it is + /// safe to call once per row in a loop. + /// + /// # Panics + /// + /// Panics if the view references a buffer or range not covered by `buffers`. + #[inline] + pub fn bytes<'a>(&'a self, buffers: &[&'a [u8]]) -> &'a [u8] { + if self.is_inlined() { + self.as_inlined().value() + } else { + let view = self.as_view(); + &buffers[view.buffer_index as usize][view.as_range()] + } + } + /// Returns the binary view as u128 representation. pub fn as_u128(&self) -> u128 { // SAFETY: binary view always safe to read as u128 LE bytes diff --git a/vortex-array/src/arrays/varbinview/vtable/mod.rs b/vortex-array/src/arrays/varbinview/vtable/mod.rs index b3ad6767f0e..48bc4abeb0e 100644 --- a/vortex-array/src/arrays/varbinview/vtable/mod.rs +++ b/vortex-array/src/arrays/varbinview/vtable/mod.rs @@ -29,11 +29,11 @@ use crate::arrays::varbinview::array::VarBinViewSlots; use crate::arrays::varbinview::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; -use crate::builders::DynVarBinBuilder; use crate::builders::VarBinViewBuilder; use crate::dtype::DType; use crate::hash::ArrayEq; use crate::hash::ArrayHash; +use crate::match_each_any_varbin_builder; use crate::serde::ArrayChildren; use crate::validity::Validity; mod kernel; @@ -249,8 +249,10 @@ impl VTable for VarBinView { if let Some(builder) = builder.as_any_mut().downcast_mut::() { return builder.append_varbinview_array(&array.into_owned(), ctx); } - if let Some(builder) = builder.as_any_mut().downcast_mut::() { - return builder.append_varbinview(array, ctx); + if let Some(result) = + match_each_any_varbin_builder!(builder, |builder| builder.append_varbinview(array, ctx)) + { + return result; } vortex_bail!("append_to_builder for VarBinView requires a variable-binary builder") } diff --git a/vortex-array/src/builders/mod.rs b/vortex-array/src/builders/mod.rs index ac1183177bc..80c450a2c0f 100644 --- a/vortex-array/src/builders/mod.rs +++ b/vortex-array/src/builders/mod.rs @@ -75,7 +75,7 @@ pub use primitive::*; pub use struct_::*; pub use varbinview::*; -pub use crate::arrays::varbin::builder::DynVarBinBuilder; +pub use crate::arrays::varbin::builder::VarBinBuilder; #[cfg(test)] mod tests; diff --git a/vortex-arrow/Cargo.toml b/vortex-arrow/Cargo.toml index 540deb58a2f..11f2308451c 100644 --- a/vortex-arrow/Cargo.toml +++ b/vortex-arrow/Cargo.toml @@ -28,6 +28,7 @@ arrow-schema = { workspace = true, features = ["canonical_extension_types"] } arrow-select = { workspace = true } itertools = { workspace = true } num-traits = { workspace = true } +simdutf8 = { workspace = true } tracing = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true, features = ["arrow"] } @@ -41,6 +42,7 @@ divan = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } vortex-fsst = { workspace = true } +vortex-onpair = { workspace = true } vortex-zstd = { workspace = true } [[bench]] diff --git a/vortex-arrow/benches/to_arrow.rs b/vortex-arrow/benches/to_arrow.rs index 4a03cf3c6a9..08442a9f659 100644 --- a/vortex-arrow/benches/to_arrow.rs +++ b/vortex-arrow/benches/to_arrow.rs @@ -11,6 +11,7 @@ use arrow_schema::Field; use divan::Bencher; use divan::counter::ItemsCount; use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; @@ -22,6 +23,8 @@ use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; +use vortex_array::builders::VarBinBuilder; +use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; @@ -39,6 +42,8 @@ use vortex_arrow::dtype::ToArrowType as _; use vortex_fsst::fsst_compress; use vortex_fsst::fsst_train_compressor; use vortex_mask::Mask; +use vortex_onpair::DEFAULT_DICT12_CONFIG; +use vortex_onpair::onpair_compress; use vortex_session::VortexSession; use vortex_zstd::Zstd; @@ -50,6 +55,7 @@ fn main() { static SESSION: LazyLock = LazyLock::new(|| { let session = array_session(); vortex_fsst::initialize(&session); + vortex_onpair::initialize(&session); session.arrays().register(Zstd); session }); @@ -116,9 +122,10 @@ fn ArrowExportVTable_to_arrow_field(bencher: Bencher) { } #[derive(Clone, Copy, Debug)] -enum OffsetStringEncoding { +enum StringEncoding { View, Fsst, + OnPair, Zstd, Dict, DictFsst, @@ -127,20 +134,48 @@ enum OffsetStringEncoding { FilterZstd, FilterDictFsst, ChunkedFsst, + /// Every third row null, so the export walks a partial validity mask rather than an all-valid + /// one and has to interleave nulls with the decoded values. + NullableFsst, + NullableZstd, + NullableDict, } -const OFFSET_STRING_ENCODINGS: &[OffsetStringEncoding] = &[ - OffsetStringEncoding::View, - OffsetStringEncoding::Fsst, - OffsetStringEncoding::Zstd, - OffsetStringEncoding::Dict, - OffsetStringEncoding::DictFsst, - OffsetStringEncoding::DictZstd, - OffsetStringEncoding::FilterFsst, - OffsetStringEncoding::FilterZstd, - OffsetStringEncoding::FilterDictFsst, - OffsetStringEncoding::ChunkedFsst, +const STRING_ENCODINGS: &[StringEncoding] = &[ + StringEncoding::View, + StringEncoding::Fsst, + StringEncoding::OnPair, + StringEncoding::Zstd, + StringEncoding::Dict, + StringEncoding::DictFsst, + StringEncoding::DictZstd, + StringEncoding::FilterFsst, + StringEncoding::FilterZstd, + StringEncoding::FilterDictFsst, + StringEncoding::ChunkedFsst, + StringEncoding::NullableFsst, + StringEncoding::NullableZstd, + StringEncoding::NullableDict, ]; + +/// Encodings whose `append_to_builder` the builder benchmarks reach directly. +/// +/// The Arrow export cannot stand in for these: `execute_until` stops at the first canonical array, +/// so a bare FSST/OnPair/Zstd root is canonicalized to `VarBinView` before any builder sees it. +/// Only `Chunked`, `Constant` and `VarBin` roots reach an encoding's own `append_to_builder` that +/// way, whereas the scan machinery appends encoded arrays into a builder directly. +const BUILDER_STRING_ENCODINGS: &[StringEncoding] = &[ + StringEncoding::View, + StringEncoding::Fsst, + StringEncoding::OnPair, + StringEncoding::Zstd, + StringEncoding::Dict, + StringEncoding::ChunkedFsst, + StringEncoding::NullableFsst, + StringEncoding::NullableZstd, + StringEncoding::NullableDict, +]; + const OFFSET_STRING_ROWS: usize = 100_000; const OFFSET_STRING_CHUNKS: usize = 4; const DICTIONARY_SIZE: usize = 2_048; @@ -152,6 +187,19 @@ fn structured_strings(len: usize) -> VarBinViewArray { VarBinViewArray::from_iter_str(values.iter().map(String::as_str)) } +fn nullable_structured_strings(len: usize) -> VarBinViewArray { + let values = (0..len) + .map(|index| { + (!index.is_multiple_of(3)) + .then(|| format!("https://example.com/common/path/{index:06}/shared-suffix")) + }) + .collect::>(); + VarBinViewArray::from_iter( + values.iter().map(|value| value.as_deref()), + DType::Utf8(Nullability::Nullable), + ) +} + fn dictionary_values() -> VarBinViewArray { structured_strings(DICTIONARY_SIZE) } @@ -172,7 +220,7 @@ fn filtered(array: ArrayRef) -> ArrayRef { FilterArray::new(array, half_rows_mask()).into_array() } -fn chunked_fsst(ctx: &mut vortex_array::ExecutionCtx) -> ArrayRef { +fn chunked_fsst(ctx: &mut ExecutionCtx) -> ArrayRef { let source = structured_strings(OFFSET_STRING_ROWS).into_array(); let compressor = fsst_train_compressor(&source, ctx).unwrap(); let chunk_size = OFFSET_STRING_ROWS / OFFSET_STRING_CHUNKS; @@ -195,39 +243,45 @@ fn chunked_fsst(ctx: &mut vortex_array::ExecutionCtx) -> ArrayRef { .into_array() } -fn offset_string_array(encoding: OffsetStringEncoding) -> ArrayRef { +fn fsst(source: ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { + let compressor = fsst_train_compressor(&source, ctx).unwrap(); + fsst_compress(&source, &compressor, ctx) + .unwrap() + .into_array() +} + +fn string_array(encoding: StringEncoding) -> ArrayRef { let mut ctx = SESSION.create_execution_ctx(); match encoding { - OffsetStringEncoding::View => structured_strings(OFFSET_STRING_ROWS).into_array(), - OffsetStringEncoding::Fsst => { - let source = structured_strings(OFFSET_STRING_ROWS).into_array(); - let compressor = fsst_train_compressor(&source, &mut ctx).unwrap(); - fsst_compress(&source, &compressor, &mut ctx) - .unwrap() - .into_array() - } - OffsetStringEncoding::Zstd => { + StringEncoding::View => structured_strings(OFFSET_STRING_ROWS).into_array(), + StringEncoding::Fsst => fsst( + structured_strings(OFFSET_STRING_ROWS).into_array(), + &mut ctx, + ), + StringEncoding::OnPair => onpair_compress( + &structured_strings(OFFSET_STRING_ROWS).into_array(), + DEFAULT_DICT12_CONFIG, + &mut ctx, + ) + .unwrap(), + StringEncoding::Zstd => { let source = structured_strings(OFFSET_STRING_ROWS); Zstd::from_var_bin_view_without_dict(&source, 3, 8_192, &mut ctx) .unwrap() .into_array() } - OffsetStringEncoding::Dict => { + StringEncoding::Dict => { DictArray::try_new(dictionary_codes(), dictionary_values().into_array()) .unwrap() .into_array() } - OffsetStringEncoding::DictFsst => { - let values = dictionary_values().into_array(); - let compressor = fsst_train_compressor(&values, &mut ctx).unwrap(); - let compressed_values = fsst_compress(&values, &compressor, &mut ctx) - .unwrap() - .into_array(); - DictArray::try_new(dictionary_codes(), compressed_values) + StringEncoding::DictFsst => { + let values = fsst(dictionary_values().into_array(), &mut ctx); + DictArray::try_new(dictionary_codes(), values) .unwrap() .into_array() } - OffsetStringEncoding::DictZstd => { + StringEncoding::DictZstd => { let values = dictionary_values(); let compressed_values = Zstd::from_var_bin_view_without_dict(&values, 3, 8_192, &mut ctx) @@ -237,22 +291,35 @@ fn offset_string_array(encoding: OffsetStringEncoding) -> ArrayRef { .unwrap() .into_array() } - OffsetStringEncoding::FilterFsst => { - filtered(offset_string_array(OffsetStringEncoding::Fsst)) - } - OffsetStringEncoding::FilterZstd => { - filtered(offset_string_array(OffsetStringEncoding::Zstd)) + StringEncoding::FilterFsst => filtered(string_array(StringEncoding::Fsst)), + StringEncoding::FilterZstd => filtered(string_array(StringEncoding::Zstd)), + StringEncoding::FilterDictFsst => filtered(string_array(StringEncoding::DictFsst)), + StringEncoding::ChunkedFsst => chunked_fsst(&mut ctx), + StringEncoding::NullableFsst => fsst( + nullable_structured_strings(OFFSET_STRING_ROWS).into_array(), + &mut ctx, + ), + StringEncoding::NullableZstd => { + let source = nullable_structured_strings(OFFSET_STRING_ROWS); + Zstd::from_var_bin_view_without_dict(&source, 3, 8_192, &mut ctx) + .unwrap() + .into_array() } - OffsetStringEncoding::FilterDictFsst => { - filtered(offset_string_array(OffsetStringEncoding::DictFsst)) + StringEncoding::NullableDict => { + // Nulls live in the dictionary rather than the codes, so the export has to combine + // the two validities. + let values = nullable_structured_strings(DICTIONARY_SIZE); + DictArray::try_new(dictionary_codes(), values.into_array()) + .unwrap() + .into_array() } - OffsetStringEncoding::ChunkedFsst => chunked_fsst(&mut ctx), } } -#[divan::bench(args = OFFSET_STRING_ENCODINGS)] -fn offset_string_export(bencher: Bencher, encoding: OffsetStringEncoding) { - let array = offset_string_array(encoding); +/// End-to-end export to Arrow `Utf8`, which is served through a `VarBinBuilder`. +#[divan::bench(args = STRING_ENCODINGS)] +fn offset_string_export(bencher: Bencher, encoding: StringEncoding) { + let array = string_array(encoding); let field = Field::new("value", DataType::Utf8, array.dtype().is_nullable()); bencher @@ -266,6 +333,37 @@ fn offset_string_export(bencher: Bencher, encoding: OffsetStringEncoding) { }); } +/// Appends an encoded array straight into an offset builder. +#[divan::bench(args = BUILDER_STRING_ENCODINGS)] +fn append_to_varbin_builder(bencher: Bencher, encoding: StringEncoding) { + let array = string_array(encoding); + + bencher + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .input_counter(|(array, _)| ItemsCount::new(array.len())) + .bench_values(|(array, mut ctx)| { + let mut builder = + VarBinBuilder::::with_capacity(array.dtype().clone(), array.len()); + array.append_to_builder(&mut builder, &mut ctx).unwrap(); + builder.finish_into_varbin() + }); +} + +/// Appends an encoded array straight into a view builder. +#[divan::bench(args = BUILDER_STRING_ENCODINGS)] +fn append_to_view_builder(bencher: Bencher, encoding: StringEncoding) { + let array = string_array(encoding); + + bencher + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .input_counter(|(array, _)| ItemsCount::new(array.len())) + .bench_values(|(array, mut ctx)| { + let mut builder = VarBinViewBuilder::with_capacity(array.dtype().clone(), array.len()); + array.append_to_builder(&mut builder, &mut ctx).unwrap(); + builder.finish_into_varbinview() + }); +} + #[divan::bench] fn to_arrow_array(bencher: Bencher) { bencher diff --git a/vortex-arrow/src/executor/byte.rs b/vortex-arrow/src/executor/byte.rs index 7df3aae5701..f3558fb77cc 100644 --- a/vortex-arrow/src/executor/byte.rs +++ b/vortex-arrow/src/executor/byte.rs @@ -5,10 +5,13 @@ use std::sync::Arc; use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::GenericByteArray; -use arrow_array::types::BinaryViewType; use arrow_array::types::ByteArrayType; -use arrow_array::types::StringViewType; +use arrow_buffer::ArrowNativeType; +use arrow_buffer::Buffer as ArrowBuffer; +use arrow_buffer::NullBuffer; +use arrow_buffer::OffsetBuffer; use arrow_schema::DataType; +use num_traits::AsPrimitive; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::Canonical; @@ -16,20 +19,19 @@ use vortex_array::ExecutionCtx; use vortex_array::arrays::Chunked; use vortex_array::arrays::Constant; use vortex_array::arrays::VarBin; -use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::varbin::VarBinArraySlotsExt; -use vortex_array::builders::DynVarBinBuilder; +use vortex_array::builders::VarBinBuilder; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; +use vortex_array::dtype::IntegerPType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; -use vortex_array::dtype::PType; use vortex_array::matcher::Matcher; -use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; -use crate::byte_view::execute_varbinview_to_arrow; use crate::executor::validity::to_arrow_null_buffer; /// Matches the encodings [`to_arrow_byte_array`] requires for export. @@ -53,7 +55,8 @@ pub(super) fn to_arrow_byte_array( ctx: &mut ExecutionCtx, ) -> VortexResult where - T::Offset: NativePType, + T::Offset: IntegerPType, + usize: AsPrimitive, { if !matches!(array.dtype(), DType::Utf8(_) | DType::Binary(_)) { vortex_bail!( @@ -63,37 +66,33 @@ where ); } + // A logical dtype mismatch changes nothing about the physical export except that a `Binary` + // source exported to `Utf8` has to have its bytes validated. let source_is_utf8 = matches!(array.dtype(), DType::Utf8(_)); let target_is_utf8 = matches!(T::DATA_TYPE, DataType::Utf8 | DataType::LargeUtf8); - if source_is_utf8 != target_is_utf8 { - let varbinview = array.execute::(ctx)?; - let binary_view = match varbinview.dtype() { - DType::Utf8(_) => execute_varbinview_to_arrow::(&varbinview, ctx), - DType::Binary(_) => execute_varbinview_to_arrow::(&varbinview, ctx), - _ => unreachable!("VarBinViewArray must have Utf8 or Binary dtype"), - }?; - return arrow_cast::cast(&binary_view, &T::DATA_TYPE).map_err(VortexError::from); - } + let validate_utf8 = target_is_utf8 && !source_is_utf8; let array = array.execute_until::(ctx)?; // If the Vortex array is in VarBin format, we can directly convert it. if let Some(array) = array.as_opt::() { - return varbin_to_byte_array::(array, ctx); + return varbin_to_byte_array::(array, validate_utf8, ctx); } - let mut builder = DynVarBinBuilder::with_capacity( - array.dtype().clone(), - T::Offset::PTYPE == PType::I64, - array.len(), - ); + // The builder's offset type matches the Arrow target, so `varbin_to_byte_array` hands the + // offsets buffer straight to Arrow without a cast. + let mut builder = VarBinBuilder::::with_capacity(array.dtype().clone(), array.len()); array.append_to_builder(&mut builder, ctx)?; - varbin_to_byte_array::(builder.finish_into_varbin().as_view(), ctx) + varbin_to_byte_array::(builder.finish_into_varbin().as_view(), validate_utf8, ctx) } /// Convert a Vortex VarBinArray into an Arrow GenericBinaryArray. +/// +/// `validate_utf8` must be set when the Vortex dtype does not already guarantee the bytes are +/// valid for `T`, i.e. a `Binary` source exported to a `Utf8` target. fn varbin_to_byte_array( array: ArrayView<'_, VarBin>, + validate_utf8: bool, ctx: &mut ExecutionCtx, ) -> VortexResult where @@ -111,11 +110,78 @@ where let data = array.bytes().clone().into_arrow_buffer(); let null_buffer = to_arrow_null_buffer(array.validity()?, array.len(), ctx)?; + if validate_utf8 { + validate_live_values_utf8::(&offsets, &data, null_buffer.as_ref())?; + } + // SAFETY: `T::DATA_TYPE` is a `Utf8` variant only when the source dtype already guarantees the + // bytes are UTF-8, or when `validate_utf8` made us prove it above. Ok(Arc::new(unsafe { GenericByteArray::::new_unchecked(offsets, data, null_buffer) })) } +/// Checks that every non-null value spanned by `offsets` is valid UTF-8. +/// +/// [`GenericByteArray::try_new`] would instead validate the whole byte range the offsets span, +/// which also covers the extents of null slots. Those bytes are never exposed as a value, and +/// Vortex only requires non-null values to be UTF-8 (see `VarBinArray::validate_utf8`), so +/// checking them would reject perfectly well-formed exports. +fn validate_live_values_utf8( + offsets: &OffsetBuffer, + values: &ArrowBuffer, + nulls: Option<&NullBuffer>, +) -> VortexResult<()> { + let value_at = |index: usize, start: usize, end: usize| -> VortexResult<&[u8]> { + values + .get(start..end) + .ok_or_else(|| vortex_err!("Offsets {start}..{end} at index {index} are out of bounds")) + }; + + let Some(nulls) = nulls.filter(|nulls| nulls.null_count() > 0) else { + // With no nulls every byte the offsets span belongs to a value, so a single pass validates + // all of them at once. The offsets then only have to land on character boundaries, which + // is what stops `value(i)` from slicing a multi-byte character in half. + let (Some(first), Some(last)) = (offsets.first(), offsets.last()) else { + return Ok(()); + }; + let start = first.as_usize(); + let bytes = value_at(0, start, last.as_usize())?; + let validated = utf8_from_bytes(bytes) + .map_err(|err| vortex_err!("Encountered non UTF-8 data: {err}"))?; + for (index, offset) in offsets.iter().enumerate() { + let boundary = offset + .as_usize() + .checked_sub(start) + .filter(|boundary| validated.is_char_boundary(*boundary)); + vortex_ensure!( + boundary.is_some(), + "Offset {} at index {index} does not fall on a UTF-8 character boundary", + offset.as_usize() + ); + } + return Ok(()); + }; + + for (index, window) in offsets.windows(2).enumerate() { + if nulls.is_null(index) { + continue; + } + let bytes = value_at(index, window[0].as_usize(), window[1].as_usize())?; + utf8_from_bytes(bytes) + .map_err(|err| vortex_err!("Encountered non UTF-8 data at index {index}: {err}"))?; + } + Ok(()) +} + +/// Validates `bytes` as UTF-8, re-running the slower checker only to describe a failure. +fn utf8_from_bytes(bytes: &[u8]) -> Result<&str, simdutf8::compat::Utf8Error> { + simdutf8::basic::from_utf8(bytes).map_err(|_| { + #[expect(clippy::unwrap_used)] + // Run validation again using the `compat` module to get a more detailed error message. + simdutf8::compat::from_utf8(bytes).unwrap_err() + }) +} + #[cfg(test)] mod tests { use arrow_array::Array; @@ -135,6 +201,9 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::fns::mask::Mask as MaskFn; + use vortex_array::validity::Validity; + use vortex_buffer::ByteBuffer; + use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -180,13 +249,13 @@ mod tests { } #[rstest] - // Utf8 source can convert to all string types and binary types (via arrow_cast) + // Utf8 source can convert to all string types and binary types #[case::utf8_to_binary(make_utf8_array(), DataType::Binary)] #[case::utf8_to_large_binary(make_utf8_array(), DataType::LargeBinary)] #[case::utf8_to_utf8(make_utf8_array(), DataType::Utf8)] #[case::utf8_to_large_utf8(make_utf8_array(), DataType::LargeUtf8)] #[case::utf8_to_utf8_view(make_utf8_array(), DataType::Utf8View)] - // Binary source can convert to all binary types and string types (via arrow_cast) + // Binary source can convert to all binary types and string types #[case::binary_to_binary(make_binary_array(), DataType::Binary)] #[case::binary_to_large_binary(make_binary_array(), DataType::LargeBinary)] #[case::binary_to_utf8(make_binary_array(), DataType::Utf8)] @@ -262,6 +331,107 @@ mod tests { assert!(!arrow.is_null(2)); } + #[rstest] + #[case(DataType::Utf8)] + #[case(DataType::LargeUtf8)] + fn binary_with_invalid_utf8_to_string_returns_error(#[case] target_dtype: DataType) { + let vortex_array = VarBinViewArray::from_iter_bin([ + b"hello".as_slice(), + b"\xff\xfe invalid utf8".as_slice(), + ]); + + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + let field = Field::new("test_field", target_dtype, true); + let result = + session + .arrow() + .execute_arrow(vortex_array.into_array(), Some(&field), &mut ctx); + + assert!(result.is_err()); + } + + /// The whole values buffer is valid UTF-8, but each individual value slices a two-byte + /// character in half. Arrow's `value(i)` reinterprets the bytes unchecked, so the export has + /// to reject this rather than hand out a malformed `&str`. + #[rstest] + #[case(DataType::Utf8)] + #[case(DataType::LargeUtf8)] + fn binary_to_string_rejects_offsets_splitting_a_character(#[case] target_dtype: DataType) { + let array = VarBinArray::new( + buffer![0i32, 1, 2].into_array(), + ByteBuffer::copy_from("é".as_bytes()), + DType::Binary(Nullability::NonNullable), + Validity::NonNullable, + ); + + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + let field = Field::new("test_field", target_dtype, false); + let result = session + .arrow() + .execute_arrow(array.into_array(), Some(&field), &mut ctx); + + assert!(result.is_err()); + } + + /// Multi-byte characters that do line up with the offsets must still export. + #[rstest] + #[case(DataType::Utf8)] + #[case(DataType::LargeUtf8)] + fn binary_to_string_accepts_multibyte_characters( + #[case] target_dtype: DataType, + ) -> VortexResult<()> { + let array = VarBinArray::new( + buffer![0i32, 2, 5].into_array(), + ByteBuffer::copy_from("é☃".as_bytes()), + DType::Binary(Nullability::NonNullable), + Validity::NonNullable, + ); + + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + let field = Field::new("test_field", target_dtype.clone(), false); + let arrow = session + .arrow() + .execute_arrow(array.into_array(), Some(&field), &mut ctx)?; + + let values: Vec<&str> = match target_dtype { + DataType::Utf8 => (0..2).map(|i| arrow.as_string::().value(i)).collect(), + _ => (0..2).map(|i| arrow.as_string::().value(i)).collect(), + }; + assert_eq!(values, ["é", "☃"]); + Ok(()) + } + + /// The `VarBin` fast path hands Arrow the untrimmed bytes buffer, so a `Binary` to `Utf8` + /// export must not validate bytes that no live value points at. + #[rstest] + #[case(DataType::Utf8)] + #[case(DataType::LargeUtf8)] + fn varbin_binary_to_string_ignores_bytes_outside_values( + #[case] target_dtype: DataType, + ) -> VortexResult<()> { + // Slot 1 is null and its offsets still span two bytes that are not valid UTF-8. + let array = VarBinArray::new( + buffer![0i32, 5, 7, 12].into_array(), + ByteBuffer::copy_from(b"hello\xff\xfeworld".as_slice()), + DType::Binary(Nullability::Nullable), + Validity::from_mask(Mask::from_iter([true, false, true]), Nullability::Nullable), + ); + + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + let field = Field::new("test_field", target_dtype, true); + let arrow = session + .arrow() + .execute_arrow(array.into_array(), Some(&field), &mut ctx)?; + + assert_eq!(arrow.len(), 3); + assert!(arrow.is_null(1)); + Ok(()) + } + #[rstest] #[case(DataType::Utf8)] #[case(DataType::LargeUtf8)] diff --git a/vortex-cuda/src/kernel/encodings/zstd.rs b/vortex-cuda/src/kernel/encodings/zstd.rs index 6d14cc46e86..01f8f964615 100644 --- a/vortex-cuda/src/kernel/encodings/zstd.rs +++ b/vortex-cuda/src/kernel/encodings/zstd.rs @@ -331,7 +331,7 @@ async fn decode_zstd(array: ZstdArray, ctx: &mut CudaExecutionCtx) -> VortexResu .indices() { AllOr::All => { - let (buffers, all_views) = reconstruct_views(&host_buffer, MAX_BUFFER_LEN); + let (buffers, all_views) = reconstruct_views(&host_buffer, 0, MAX_BUFFER_LEN); let sliced_views = all_views.slice(slice_value_idx_start..slice_value_idx_stop); Ok(Canonical::VarBinView(unsafe { diff --git a/vortex-duckdb/src/e2e_test/vortex_scan_test.rs b/vortex-duckdb/src/e2e_test/vortex_scan_test.rs index 2a046b1627d..246317b28bd 100644 --- a/vortex-duckdb/src/e2e_test/vortex_scan_test.rs +++ b/vortex-duckdb/src/e2e_test/vortex_scan_test.rs @@ -983,11 +983,12 @@ fn test_geometry() { let mut wkb_binary: Vec = Vec::new(); wkb::writer::write_polygon(&mut wkb_binary, &rect10, &WriteOptions::default()) .expect("serializing WKB"); - let mut geometry = VarBinBuilder::::with_capacity(10); + let mut geometry = + VarBinBuilder::::with_capacity(DType::Binary(Nullability::NonNullable), 10); for _ in 0..10 { geometry.append_value(wkb_binary.as_slice()); } - let geometry = geometry.finish(DType::Binary(Nullability::NonNullable)); + let geometry = geometry.finish_into_varbin(); let geometry = ExtensionArray::new( ExtDType::::try_new( diff --git a/vortex-geo/src/tests/wkb.rs b/vortex-geo/src/tests/wkb.rs index d61c48bdc32..9be8871080f 100644 --- a/vortex-geo/src/tests/wkb.rs +++ b/vortex-geo/src/tests/wkb.rs @@ -62,7 +62,8 @@ fn wkb_extension_array() -> VortexResult<(Vec, vortex_array::ArrayRef)> { wkb::writer::write_geometry(&mut buf, &test_polygon(), &WriteOptions::default()) .map_err(|e| vortex_err!("writing WKB failed: {e}"))?; - let mut builder = VarBinBuilder::::with_capacity(3); + let mut builder = + VarBinBuilder::::with_capacity(DType::Binary(Nullability::NonNullable), 3); builder.append_value(&buf); builder.append_value(&buf); builder.append_value(&buf); @@ -73,7 +74,7 @@ fn wkb_extension_array() -> VortexResult<(Vec, vortex_array::ArrayRef)> { }, DType::Binary(Nullability::NonNullable), )?; - let storage = builder.finish(DType::Binary(Nullability::NonNullable)); + let storage = builder.finish_into_varbin(); let array = ExtensionArray::new(dtype.erased(), storage.into_array()).into_array(); Ok((buf, array)) }