diff --git a/Cargo.lock b/Cargo.lock index 5be1371fac3..a2281e5728b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9465,6 +9465,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..54e88e109ef 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::Nullability; +use vortex_array::dtype::OffsetBuilderPType; 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..8133dc6c6d9 100644 --- a/encodings/fsst/src/compress.rs +++ b/encodings/fsst/src/compress.rs @@ -24,6 +24,7 @@ use vortex_array::arrays::varbinview::BinaryView; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::OffsetBuilderPType; use vortex_array::match_each_integer_ptype; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -197,9 +198,13 @@ fn compress_views( ctx: &mut ExecutionCtx, ) -> VortexResult where - O: IntegerPType + 'static, + O: OffsetBuilderPType + '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() { @@ -230,9 +235,13 @@ fn compress_varbin( ctx: &mut ExecutionCtx, ) -> VortexResult where - O: IntegerPType + 'static, + O: OffsetBuilderPType + '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::(); @@ -269,18 +278,18 @@ where } /// Per-row output state for an FSST compression pass. -struct FsstSink<'c, O: IntegerPType + 'static> { +struct FsstSink<'c, O: OffsetBuilderPType + 'static> { buffer: Vec, builder: VarBinBuilder, uncompressed_lengths: BufferMut, compressor: &'c Compressor, } -impl<'c, O: IntegerPType + 'static> FsstSink<'c, O> { - fn with_capacity(len: usize, compressor: &'c Compressor) -> Self { +impl<'c, O: OffsetBuilderPType + 'static> FsstSink<'c, O> { + 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 +298,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 +319,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..9be2d3075b3 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::Nullability; +use vortex_array::dtype::OffsetBuilderPType; 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..41ea51b250e 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::OffsetBuilderPType; +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..7cabbefa0bd 100644 --- a/vortex-array/src/arrays/varbin/array.rs +++ b/vortex-array/src/arrays/varbin/array.rs @@ -25,8 +25,8 @@ use crate::arrays::VarBin; use crate::arrays::varbin::builder::VarBinBuilder; use crate::buffer::BufferHandle; use crate::dtype::DType; -use crate::dtype::IntegerPType; use crate::dtype::Nullability; +use crate::dtype::OffsetBuilderPType; use crate::legacy_session; use crate::match_each_integer_ptype; use crate::validity::Validity; @@ -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,23 +395,23 @@ 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 where - O: IntegerPType, + O: OffsetBuilderPType, 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..c798d7ba02c 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -2,19 +2,23 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::any::Any; +use std::mem::MaybeUninit; use num_traits::AsPrimitive; use vortex_buffer::BitBufferMut; 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; @@ -29,76 +33,169 @@ use crate::arrays::varbin::VarBinArraySlotsExt; use crate::arrays::varbinview::VarBinViewArrayExt; use crate::builders::ArrayBuilder; use crate::dtype::DType; -use crate::dtype::IntegerPType; +use crate::dtype::OffsetBuilderPType; 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; -pub struct VarBinBuilder { + +/// 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) +impl VarBinBuilder { + /// 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 +205,199 @@ 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, taking each valid row's bytes from `value`. + /// + /// `value` is called once per valid row in ascending row order and the slices it returns must + /// total `num_bytes` bytes. Null rows consume no bytes and repeat the preceding offset. + /// + /// The valid rows are walked a `u64` word at a time rather than pulled one at a time through + /// an iterator, and both buffers are sized once up front, so the copy loop costs one offset + /// store plus one `memcpy` per value. A caller whose values only exist as a stream can still + /// use this by pulling from the stream inside `value`: the rows are visited in ascending + /// order. + pub fn append_valid_slices<'a, F>( + &mut self, + num_bytes: usize, + validity: &Mask, + value: F, + ) -> VortexResult<()> + where + F: FnMut(usize) -> &'a [u8], + usize: AsPrimitive, + { + let data_start = self.data.len(); + match self.gather_valid_slices(num_bytes, validity, value) { + 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_valid_slices(num_bytes, &validity, |index| views[index].bytes(&buffers)) } + /// 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 +419,183 @@ 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 each valid row's bytes from `value` into the byte storage and records the offsets. + /// See [`append_valid_slices`](Self::append_valid_slices); the validity bits are the caller's + /// job so that a failure here can be unwound. + fn gather_valid_slices<'a, F>( &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, + validity: &Mask, + mut value: F, + ) -> VortexResult<()> + where + F: FnMut(usize) -> &'a [u8], + 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, since the two are separate allocations. + let Self { offsets, data, .. } = self; + let spare = &mut offsets.spare_capacity_mut()[..count]; + + match validity.bit_buffer() { + AllOr::All => { + for (row, slot) in spare.iter_mut().enumerate() { + data.extend_from_slice(value(row)); + 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; + bits.for_each_set_index(|index| { + // 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(index)); + 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), - } - } + // A caller whose slices overrun `num_bytes` only grows the byte buffer past the reservation, + // so the overrun is caught here rather than per value. The offsets are still uncommitted, so + // rejecting it now leaves them untouched. + vortex_ensure!( + data.len() == data_start + num_bytes, + "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 +609,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 +626,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,12 +641,80 @@ 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; +/// } +/// ``` +#[macro_export] +macro_rules! match_each_varbin_builder { + ($builder:expr, | $typed:ident | $body:expr) => { + $crate::__match_varbin_builder_widths!($builder, |$typed| $body, [u32, u64, 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) + }) +} + #[cfg(test)] mod tests { + use std::mem::MaybeUninit; + use rstest::rstest; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -453,10 +722,11 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; + use crate::arrays::ChunkedArray; + use crate::arrays::ConstantArray; 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; @@ -469,11 +739,11 @@ mod tests { #[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,22 +777,147 @@ 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 values = [0u8; 128]; + fn append_n_values_offset_overflow_returns_error() { + let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); - let result = builder.append_values(&values, [values.len()].into_iter(), &Mask::new_true(1)); + // The limit is checked before anything is reserved, so the huge byte total is free. + let result = builder.append_n_values(b"hello", i32::MAX as usize / 5 + 1); + + assert!(result.is_err()); + assert_eq!(builder.offsets.len(), 1); + assert!(builder.data.is_empty()); + 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 num_bytes = i32::MAX as usize + 1; + 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(num_bytes, 0, &[num_bytes], &Mask::new_true(1), &mut |_| { + decoded = true; + Ok(0) + }) + }; + + 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), i32::MAX as usize / 5 + 1); + + 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); + } + + /// Slices that overrun the declared byte count are rejected, and the builder is left exactly + /// as it was — the overrun bytes are copied before the total can be checked, so the unwind has + /// to put them back. + #[rstest] + #[case::all_valid(Mask::new_true(2))] + #[case::some_valid(Mask::from_iter([true, false, true]))] + fn append_valid_slices_rejects_a_byte_count_mismatch(#[case] validity: Mask) { + let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); + let values = [b"foo".as_slice(), b"quux".as_slice()]; + + let mut next = 0; + let result = builder.append_valid_slices(6, &validity, |_| { + next += 1; + values[next - 1] + }); assert!(result.is_err()); assert_eq!(builder.offsets.len(), 1); @@ -533,9 +928,9 @@ mod tests { #[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 +943,83 @@ 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 - .into_array() - .append_to_builder(&mut builder, &mut ctx)?; + 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!(builder.finish_into_varbin(), expected, &mut ctx); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + /// `match_each_varbin_builder!` covers every offset width a `VarBinBuilder` can be built with, + /// so both canonical appends must reach their specialization for all of them. A macro that + /// only matched the signed pair would send an unsigned builder down a downcast that assumes + /// `VarBinViewBuilder` and panic. + #[rstest] + #[case::u32(VarBinBuilder::::new(DType::Utf8(Nullable)))] + #[case::u64(VarBinBuilder::::new(DType::Utf8(Nullable)))] + #[case::i32(VarBinBuilder::::new(DType::Utf8(Nullable)))] + #[case::i64(VarBinBuilder::::new(DType::Utf8(Nullable)))] + fn append_to_every_offset_width(#[case] mut builder: impl ArrayBuilder) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let long = "a string that is far too long to be inlined in a view"; + let values = [Some("hello"), None, Some(long), Some("")]; + + let view = VarBinViewArray::from_iter(values, DType::Utf8(Nullable)).into_array(); + let varbin = VarBinArray::from_iter(values, DType::Utf8(Nullable)).into_array(); + let constant = ConstantArray::new(Scalar::from("hello").into_nullable(), 2).into_array(); + + for chunk in [&view, &varbin, &constant] { + chunk.append_to_builder(&mut builder, &mut ctx)?; + } + + let expected = ChunkedArray::try_new(vec![view, varbin, constant], DType::Utf8(Nullable))?; + assert_arrays_eq!(builder.finish(), 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 +1031,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 +1041,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..e6ff4b9734c 100644 --- a/vortex-array/src/arrays/varbin/compute/filter.rs +++ b/vortex-array/src/arrays/varbin/compute/filter.rs @@ -23,7 +23,9 @@ use crate::arrays::varbin::VarBinArraySlotsExt; use crate::arrays::varbin::builder::VarBinBuilder; use crate::dtype::DType; use crate::dtype::IntegerPType; +use crate::dtype::OffsetBuilderPType; use crate::match_each_integer_ptype; +use crate::match_smallest_list_offset_type; impl FilterKernel for VarBin { fn filter( @@ -62,21 +64,23 @@ fn filter_select_var_bin_by_slice( ) -> VortexResult { let offsets = values.offsets().clone().execute::(ctx)?; match_each_integer_ptype!(offsets.ptype(), |O| { - filter_select_var_bin_by_slice_primitive_offset( - values.dtype().clone(), - offsets.as_slice::(), - values.bytes().as_slice(), - mask_slices, - values - .varbin_validity() - .execute_mask(values.as_ref().len(), ctx) - .vortex_expect("Failed to compute validity mask"), - selection_count, - ) + match_smallest_list_offset_type!(values.bytes().len(), |B| { + filter_select_var_bin_by_slice_primitive_offset::( + values.dtype().clone(), + offsets.as_slice::(), + values.bytes().as_slice(), + mask_slices, + values + .varbin_validity() + .execute_mask(values.as_ref().len(), ctx) + .vortex_expect("Failed to compute validity mask"), + selection_count, + ) + }) }) } -fn filter_select_var_bin_by_slice_primitive_offset( +fn filter_select_var_bin_by_slice_primitive_offset( dtype: DType, offsets: &[O], data: &[u8], @@ -86,9 +90,10 @@ fn filter_select_var_bin_by_slice_primitive_offset( ) -> VortexResult where O: IntegerPType, - usize: AsPrimitive, + B: OffsetBuilderPType, + 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 +101,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,26 +125,27 @@ 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( +fn update_non_nullable_slice( data: &[u8], offsets: &[O], - builder: &mut VarBinBuilder, + builder: &mut VarBinBuilder, start: usize, end: usize, ) -> VortexResult<()> where O: IntegerPType, - usize: AsPrimitive, + B: OffsetBuilderPType, + usize: AsPrimitive, { let offset_start = offsets[start] .to_usize() @@ -163,21 +169,23 @@ fn filter_select_var_bin_by_index( ) -> VortexResult { let offsets = values.offsets().clone().execute::(ctx)?; match_each_integer_ptype!(offsets.ptype(), |O| { - filter_select_var_bin_by_index_primitive_offset( - values.dtype().clone(), - offsets.as_slice::(), - values.bytes().as_slice(), - mask_indices, - values - .varbin_validity() - .execute_mask(values.as_ref().len(), ctx) - .vortex_expect("Failed to compute validity mask"), - selection_count, - ) + match_smallest_list_offset_type!(values.bytes().len(), |B| { + filter_select_var_bin_by_index_primitive_offset::( + values.dtype().clone(), + offsets.as_slice::(), + values.bytes().as_slice(), + mask_indices, + values + .varbin_validity() + .execute_mask(values.as_ref().len(), ctx) + .vortex_expect("Failed to compute validity mask"), + selection_count, + ) + }) }) } -fn filter_select_var_bin_by_index_primitive_offset( +fn filter_select_var_bin_by_index_primitive_offset( dtype: DType, offsets: &[O], data: &[u8], @@ -195,7 +203,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 +211,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..22cace7ac2b 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_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_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..33e9ce289a5 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_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_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 b9c05cec3e1..85dafb69cb8 100644 --- a/vortex-array/src/builders/mod.rs +++ b/vortex-array/src/builders/mod.rs @@ -71,7 +71,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-array/src/dtype/ptype.rs b/vortex-array/src/dtype/ptype.rs index 9948a74a129..928a068f103 100644 --- a/vortex-array/src/dtype/ptype.rs +++ b/vortex-array/src/dtype/ptype.rs @@ -91,13 +91,15 @@ pub trait UnsignedPType: IntegerPType + Unsigned {} /// Implements [`UnsignedPType`] for all possible `T` that have the correct bounds. impl UnsignedPType for T where T: IntegerPType + Unsigned {} -/// Trait for the integer types that can back a list builder's `offsets` and `sizes`: `u32`, -/// `i32`, `u64`, and `i64`. +/// Trait for the integer types that can back a builder's `offsets` and `sizes`: `u32`, `i32`, +/// `u64`, and `i64`. /// /// This trait is sealed and cannot be implemented for any other type. Restricting the set of -/// widths keeps the concrete [`ListBuilder`](crate::builders::ListBuilder) and -/// [`ListViewBuilder`](crate::builders::ListViewBuilder) instantiations small enough for -/// [`match_each_list_builder!`](crate::match_each_list_builder) to enumerate them. +/// widths keeps the concrete [`ListBuilder`](crate::builders::ListBuilder), +/// [`ListViewBuilder`](crate::builders::ListViewBuilder), and +/// [`VarBinBuilder`](crate::builders::VarBinBuilder) instantiations small enough for +/// [`match_each_list_builder!`](crate::match_each_list_builder) and +/// [`match_each_varbin_builder!`](crate::match_each_varbin_builder) to enumerate them. pub trait OffsetBuilderPType: IntegerPType + offset_builder_sealed::Sealed {} mod offset_builder_sealed { @@ -718,7 +720,7 @@ macro_rules! match_smallest_offset_type { /// Macro to match the smallest [`OffsetBuilderPType`] able to index a given number of elements. /// /// Like [`match_smallest_offset_type!`](crate::match_smallest_offset_type), but restricted to the -/// unsigned types valid as list builder offsets (`u32` and `u64`). +/// unsigned types valid as builder offsets (`u32` and `u64`). #[macro_export] macro_rules! match_smallest_list_offset_type { ($n_elements:expr, | $offset_type:ident | $body:block) => {{ 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..434269b04e7 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::NativePType; use vortex_array::dtype::Nullability; -use vortex_array::dtype::PType; +use vortex_array::dtype::OffsetBuilderPType; 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: OffsetBuilderPType, + 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)) }