From b980c097b02e9dfdcfb1b2efb0aee27cf095bdda Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 30 Jul 2026 22:15:04 +0100 Subject: [PATCH 1/2] Replace DynVarBinBuilder with a generic VarBinBuilder `VarBinBuilder` becomes an actual `ArrayBuilder`, carrying its own `DType` and offset width, so `DynVarBinBuilder` is no longer needed as a separate type-erasing wrapper. Encodings recover the concrete offset width from a `&mut dyn ArrayBuilder` with the new `match_each_varbin_builder!` (i32/i64) and `match_each_any_varbin_builder!` (every width) macros, which monomorphize the append body over the width instead of dispatching through `dyn`. The existing `append_to_builder` specializations for FSST, OnPair and Zstd are ported over mechanically; their optimizations follow in later commits. Because the builder's offset type now matches the Arrow target, the Arrow byte export hands the offsets buffer straight to Arrow with no cast, which also removes the `VarBinView` + `arrow_cast` round trip a `Utf8`/`Binary` dtype mismatch used to take. A `Binary` source exported to `Utf8` instead validates the bytes of its live values, skipping the extents of null slots that `GenericByteArray::try_new` would have rejected. Also adds `BinaryView::bytes`, which borrows from already-resolved buffer slices instead of cloning a buffer handle per row. Signed-off-by: Robert Kruszewski --- Cargo.lock | 1 + encodings/fsst/src/array.rs | 51 +- encodings/fsst/src/canonical.rs | 4 +- encodings/fsst/src/compress.rs | 22 +- encodings/fsst/src/kernel.rs | 28 +- encodings/fsst/src/tests.rs | 13 +- encodings/onpair/src/array.rs | 51 +- encodings/onpair/src/tests.rs | 4 +- encodings/runend/src/array.rs | 4 +- encodings/sparse/src/lib.rs | 8 +- encodings/zstd/src/array.rs | 100 +- encodings/zstd/src/test.rs | 4 +- .../src/arrays/constant/vtable/mod.rs | 14 +- vortex-array/src/arrays/dict/array.rs | 6 +- vortex-array/src/arrays/varbin/array.rs | 12 +- vortex-array/src/arrays/varbin/builder.rs | 1102 ++++++++++++----- .../src/arrays/varbin/compute/compare.rs | 10 +- .../src/arrays/varbin/compute/filter.rs | 16 +- .../src/arrays/varbin/vtable/canonical.rs | 8 +- vortex-array/src/arrays/varbin/vtable/mod.rs | 8 +- vortex-array/src/arrays/varbinview/view.rs | 20 + .../src/arrays/varbinview/vtable/mod.rs | 8 +- vortex-array/src/builders/mod.rs | 2 +- vortex-arrow/Cargo.toml | 1 + vortex-arrow/benches/to_arrow.rs | 172 ++- vortex-arrow/src/executor/byte.rs | 222 +++- .../src/e2e_test/vortex_scan_test.rs | 5 +- vortex-geo/src/tests/wkb.rs | 5 +- 28 files changed, 1390 insertions(+), 511 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 25fb98611d0..d72cb6af8b9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9443,6 +9443,7 @@ dependencies = [ "itertools 0.14.0", "num-traits", "rstest", + "simdutf8", "tracing", "vortex-array", "vortex-buffer", diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index 4a740417f7e..5e60b83bdfb 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -34,13 +34,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; @@ -307,23 +309,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 +350,32 @@ impl VTable for FSST { } } +/// Decompresses the values and appends them to `builder`. +fn append_to_varbin( + array: ArrayView<'_, FSST>, + builder: &mut VarBinBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> +where + usize: AsPrimitive, +{ + 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, + ) + }) +} + #[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..f248e123fe9 100644 --- a/encodings/fsst/src/canonical.rs +++ b/encodings/fsst/src/canonical.rs @@ -106,7 +106,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 +212,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..b785dd8fce2 100644 --- a/encodings/onpair/src/array.rs +++ b/encodings/onpair/src/array.rs @@ -26,12 +26,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; @@ -553,23 +555,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 +592,32 @@ impl VTable for OnPair { } } +/// Decodes the values and appends them to `builder`. +fn append_to_varbin( + array: ArrayView<'_, OnPair>, + builder: &mut VarBinBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> +where + usize: AsPrimitive, +{ + 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, + ) + }) +} + impl ValidityVTable for OnPair { fn validity(array: ArrayView<'_, OnPair>) -> VortexResult { Ok(child_to_validity( 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/src/array.rs b/encodings/zstd/src/array.rs index 19d6fa8c736..eb6a1cc27ca 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -30,8 +30,10 @@ 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::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; @@ -276,49 +278,17 @@ 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(()) + array + .array() + .clone() + .execute::(ctx)? + .into_array() + .append_to_builder(builder, ctx) } fn reduce_parent( @@ -330,6 +300,50 @@ impl VTable for Zstd { } } +/// Copies the decompressed values into `builder`. +fn append_to_varbin( + array: ArrayView<'_, Zstd>, + builder: &mut VarBinBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let unsliced_validity = child_to_validity( + array.slots()[ZstdSlots::VALIDITY].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.push_nulls(slice.n_rows), + AllOr::Some(valid_indices) => { + let mut row = 0; + for &valid_index in valid_indices { + builder.push_nulls(valid_index - row); + builder.append_n_values( + values + .next() + .vortex_expect("Zstd value count must match validity"), + 1, + )?; + row = valid_index + 1; + } + builder.push_nulls(slice.n_rows - row); + } + } + Ok(()) +} + #[derive(Clone, Debug)] /// Zstd array encoding marker. pub struct Zstd; diff --git a/encodings/zstd/src/test.rs b/encodings/zstd/src/test.rs index fac5fcdc631..eea353a2a34 100644 --- a/encodings/zstd/src/test.rs +++ b/encodings/zstd/src/test.rs @@ -10,7 +10,7 @@ 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; @@ -224,7 +224,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(); 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/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..20ec7200bfe 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,164 @@ 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); + } + + /// Appends one value, or a null when it is `None`. + /// + /// # Panics + /// + /// Panics if the resulting end offset does not fit in `O`. See + /// [`append_value`](Self::append_value). #[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(), } } + /// Appends one non-null value. + /// + /// # Panics + /// + /// Panics if the resulting end offset does not fit in `O`. The bulk appends report that as an + /// error instead; use [`append_n_values`](Self::append_n_values) for a fallible single append, + /// or an `i64` builder for byte totals past `i32::MAX`. #[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. + /// + /// # Errors + /// + /// Returns an error, leaving the builder unchanged, if the resulting end offsets do not fit + /// in `O`. + pub fn append_n_values(&mut self, value: impl AsRef<[u8]>, n: usize) -> VortexResult<()> { + let value = value.as_ref(); + let Some(num_bytes) = value.len().checked_mul(n) else { + vortex_bail!("Byte count overflow: {} values of {} bytes", n, value.len()); + }; + // Checking the total up front is what keeps `push_value` below from panicking. + self.check_offset_limit(self.data.len(), num_bytes)?; + self.offsets.reserve(n); + self.data.reserve(num_bytes); + for _ in 0..n { + self.push_value(value); + } + self.validity.append_n(true, n); + Ok(()) + } + + /// 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. + /// + /// # Errors + /// + /// Returns an error if `scalar` has a different dtype than the builder, or if the resulting + /// end offsets do not fit in `O`. + 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 +208,200 @@ 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. + /// + /// # Errors + /// + /// Returns an error, leaving the builder unchanged, if `num_bytes` would push an offset past + /// what `O` can hold — checked before `decode` runs — or if `decode` reports a different byte + /// count than `num_bytes`, or if `lengths` does not describe those bytes. + /// + /// # 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}"); + }; + // Checked before decoding: an `i32` builder that the decoded bytes would overflow should + // not pay for the decompression first. + self.check_offset_limit(self.data.len(), num_bytes)?; + 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 +423,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 +630,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 +647,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 +662,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 +827,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 +869,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 +892,130 @@ 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(()) + } + + /// The offset width has to be checked before `decode` runs: an overflowing `i32` builder + /// should not pay for a full decompression first. + #[test] + fn append_decoded_rejects_an_offset_overflow_without_decoding() { + let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); + let mut decoded = false; + + // SAFETY: the closure is never called, and reports only what it initializes if it were. + let result = unsafe { + builder.append_decoded(128, 0, &[128usize], &Mask::new_true(1), &mut |spare| { + decoded = true; + spare[..128].fill(MaybeUninit::new(b'x')); + Ok(128) + }) + }; + + assert!(result.is_err()); + assert!(!decoded, "decode ran before the offset width was checked"); + assert_eq!(builder.offsets.len(), 1); + assert_eq!(builder.validity.len(), 0); + } + + /// `append_scalar_repeated` reports an offset overflow rather than panicking, so a constant + /// column wider than the offset type errors like the bulk appends do. + #[test] + fn append_scalar_repeated_rejects_an_offset_overflow() { + let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); + + let result = builder.append_scalar_repeated(&Scalar::utf8("hello", Nullable), 64); + + assert!(result.is_err()); + assert_eq!(builder.offsets.len(), 1); + assert!(builder.data.is_empty()); + assert_eq!(builder.validity.len(), 0); + } + + #[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 +1028,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 +1129,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 +1139,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..e69e6b656fa 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"] } diff --git a/vortex-arrow/benches/to_arrow.rs b/vortex-arrow/benches/to_arrow.rs index b1dc190da17..3d38b3813a0 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; @@ -116,7 +119,7 @@ fn ArrowExportVTable_to_arrow_field(bencher: Bencher) { } #[derive(Clone, Copy, Debug)] -enum OffsetStringEncoding { +enum StringEncoding { View, Fsst, Zstd, @@ -127,20 +130,46 @@ 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::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/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::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 +181,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 +214,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; @@ -193,39 +235,39 @@ 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::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) @@ -235,22 +277,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 @@ -264,6 +319,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-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)) } From 2de8a74c3de65c3151fdd7ae59d5280880a4a9bd Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 30 Jul 2026 22:15:39 +0100 Subject: [PATCH 2/2] FSST: decompress straight into the builder's byte storage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `append_to_builder` staged the decompressed heap in a temporary `ByteBufferMut` and then copied it into the builder. `FsstDecodePlan` hoists the parts of `fsst_decode_bytes` that describe the decode — the code stream, the per-row lengths and the total decoded size — so the decompression can target the builder's spare capacity directly via `append_decoded`, leaving one prefix sum over the lengths as the only work beyond the bulk `decompress_into`. `FSST_DECODE_SLACK` names the 7 spare bytes that keep `decompress_into` on its wide-store path through the final value; it is a performance knob, not a safety requirement. Signed-off-by: Robert Kruszewski --- encodings/fsst/src/array.rs | 35 ++++++++---- encodings/fsst/src/canonical.rs | 98 ++++++++++++++++++++++++--------- 2 files changed, 96 insertions(+), 37 deletions(-) diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index 5e60b83bdfb..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; @@ -60,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; @@ -350,7 +352,10 @@ impl VTable for FSST { } } -/// Decompresses the values and appends them to `builder`. +/// 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, @@ -359,20 +364,26 @@ fn append_to_varbin( where usize: AsPrimitive, { - let (bytes, lengths) = fsst_decode_bytes(array, ctx)?; + let plan = FsstDecodePlan::new(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, - ) + 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, + ) + } }) } diff --git a/encodings/fsst/src/canonical.rs b/encodings/fsst/src/canonical.rs index f248e123fe9..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(