From a377a9bf813332725b149588909e4821d3f1b3aa Mon Sep 17 00:00:00 2001 From: cl Date: Thu, 23 Jul 2026 01:00:00 +0800 Subject: [PATCH 01/10] feat(arrow): export strings and binary arrays with offsets Problem Vortex exports UTF-8 and binary data as Arrow view arrays by default. Some consumers require Arrow arrays that use 32-bit or 64-bit offsets. These consumers must cast each view array and rebuild its buffers. Implementation Add VarBinBufferBuilder for 32-bit and 64-bit offsets. Write Arrow value, offset, and validity buffers in this builder. Decode FSST, OnPair, Zstd, Dict, RunEnd, and Sparse arrays directly into this builder. Do not create an intermediate VarBinViewArray for these paths. Interface Add an optional Arrow type to the Rust and Python array APIs. Add an optional Arrow schema to the Python file API. Keep Arrow view arrays as the default output. Result Arrow uses the builder buffers without another copy. Callers can select StringArray, BinaryArray, LargeStringArray, or LargeBinaryArray. Signed-off-by: cl --- encodings/fsst/src/array.rs | 23 ++ encodings/fsst/src/canonical.rs | 66 +++- encodings/onpair/src/array.rs | 23 ++ encodings/onpair/src/canonical.rs | 28 +- encodings/onpair/src/tests.rs | 14 + encodings/runend/src/array.rs | 82 ++++- encodings/sparse/src/lib.rs | 133 ++++++++ encodings/zstd/src/array.rs | 146 +++++++-- encodings/zstd/src/test.rs | 30 ++ .../src/arrays/constant/vtable/mod.rs | 37 ++- vortex-array/src/arrays/dict/vtable/mod.rs | 22 ++ vortex-array/src/arrays/varbin/builder.rs | 289 +++++++++++++++++- vortex-array/src/arrays/varbin/vtable/mod.rs | 15 + .../src/arrays/varbinview/vtable/mod.rs | 12 +- vortex-array/src/builders/mod.rs | 2 + vortex-arrow/src/executor/byte.rs | 33 +- vortex-python/python/vortex/_lib/arrays.pyi | 4 +- vortex-python/python/vortex/_lib/file.pyi | 1 + vortex-python/python/vortex/file.py | 8 +- vortex-python/src/arrays/mod.rs | 51 +++- vortex-python/src/file.rs | 14 +- vortex-python/test/test_array.py | 18 ++ vortex-python/test/test_file.py | 10 + 23 files changed, 967 insertions(+), 94 deletions(-) diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index de0cf3157ba..363180a4fec 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -11,6 +11,7 @@ use std::sync::OnceLock; use fsst::Compressor; use fsst::Decompressor; use fsst::Symbol; +use num_traits::AsPrimitive; use prost::Message as _; use vortex_array::Array; use vortex_array::ArrayEq; @@ -32,11 +33,13 @@ use vortex_array::arrays::VarBinArray; use vortex_array::arrays::varbin::VarBinArrayExt; use vortex_array::buffer::BufferHandle; use vortex_array::builders::ArrayBuilder; +use vortex_array::builders::VarBinBufferBuilder; use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::legacy_session; +use vortex_array::match_each_integer_ptype; use vortex_array::serde::ArrayChildren; use vortex_array::smallvec::smallvec; use vortex_array::validity::Validity; @@ -56,6 +59,7 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::canonical::canonicalize_fsst; +use crate::canonical::fsst_decode_bytes; use crate::canonical::fsst_decode_views; use crate::rules::RULES; @@ -304,6 +308,25 @@ 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() + .map(|length| AsPrimitive::::as_(*length)), + &validity, + ); + }); + return Ok(()); + } + let Some(builder) = builder.as_any_mut().downcast_mut::() else { return array .array() diff --git a/encodings/fsst/src/canonical.rs b/encodings/fsst/src/canonical.rs index 5c58dbbd501..dd14090541b 100644 --- a/encodings/fsst/src/canonical.rs +++ b/encodings/fsst/src/canonical.rs @@ -18,6 +18,7 @@ use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use crate::FSST; use crate::FSSTArrayExt; @@ -40,20 +41,11 @@ pub(super) fn canonicalize_fsst( }) } -pub(crate) fn fsst_decode_views( +pub(crate) fn fsst_decode_bytes( fsst_array: ArrayView<'_, FSST>, - start_buf_index: u32, ctx: &mut ExecutionCtx, -) -> VortexResult<(Vec, Buffer)> { - // FSSTArray has two child arrays: - // 1. A VarBinArray, which holds the string heap of the compressed codes. - // 2. An uncompressed_lengths primitive array, storing the length of each original - // string element. - // To speed up canonicalization, we can decompress the entire string-heap in a single - // call. We then turn our uncompressed_lengths into an offsets buffer - // necessary for a VarBinViewArray and construct the canonical array. +) -> VortexResult<(ByteBufferMut, PrimitiveArray)> { let bytes = fsst_array.codes().sliced_bytes(); - let uncompressed_lens_array = fsst_array .uncompressed_lengths() .clone() @@ -68,14 +60,24 @@ pub(crate) fn fsst_decode_views( .sum() }); - // Bulk-decompress the entire array. let decompressor = fsst_array.decompressor(); let mut uncompressed_bytes = ByteBufferMut::with_capacity(total_size + 7); let len = decompressor.decompress_into(bytes.as_slice(), uncompressed_bytes.spare_capacity_mut()); + vortex_ensure!( + len == total_size, + "FSST decoded {len} bytes, expected {total_size}" + ); unsafe { uncompressed_bytes.set_len(len) }; + Ok((uncompressed_bytes, uncompressed_lens_array)) +} - // Directly create the binary views. +pub(crate) fn fsst_decode_views( + fsst_array: ArrayView<'_, FSST>, + start_buf_index: u32, + ctx: &mut ExecutionCtx, +) -> VortexResult<(Vec, Buffer)> { + let (uncompressed_bytes, uncompressed_lens_array) = fsst_decode_bytes(fsst_array, ctx)?; match_each_integer_ptype!(uncompressed_lens_array.ptype(), |P| { Ok(build_views( start_buf_index, @@ -97,15 +99,21 @@ mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; + use vortex_array::arrays::varbin::VarBinArrayExt; use vortex_array::builders::ArrayBuilder; + use vortex_array::builders::VarBinBufferBuilder; use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_error::VortexResult; use vortex_session::VortexSession; + use super::fsst_decode_bytes; + use crate::FSST; + use crate::FSSTArrayExt; use crate::fsst_compress; use crate::fsst_train_compressor; @@ -199,6 +207,20 @@ mod tests { .collect::>(); assert_eq!(data, res2) }; + + { + let mut builder = + VarBinBufferBuilder::with_capacity(chunked_arr.dtype().clone(), false, data.len()); + chunked_arr + .into_array() + .append_to_builder(&mut builder, &mut ctx)?; + let arr = builder.finish_into_varbin(); + let mask = arr.validity()?.execute_mask(arr.len(), &mut ctx)?; + let actual = (0..arr.len()) + .map(|i| mask.value(i).then(|| arr.bytes_at(i).to_vec())) + .collect::>(); + assert_eq!(data, actual); + } Ok(()) } @@ -225,4 +247,22 @@ mod tests { let _result = builder.finish_into_varbinview(); Ok(()) } + + #[test] + fn rejects_incorrect_uncompressed_lengths() -> VortexResult<()> { + let input = VarBinViewArray::from_iter_str(["hello"]).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let encoded = fsst_compress(&input, &fsst_train_compressor(&input, &mut ctx)?, &mut ctx)?; + let invalid = FSST::try_new( + encoded.dtype().clone(), + encoded.symbols().clone(), + encoded.symbol_lengths().clone(), + encoded.codes(), + PrimitiveArray::from_iter([4i32]).into_array(), + &mut ctx, + )?; + + assert!(fsst_decode_bytes(invalid.as_view(), &mut ctx).is_err()); + Ok(()) + } } diff --git a/encodings/onpair/src/array.rs b/encodings/onpair/src/array.rs index 111269ac484..b447a268781 100644 --- a/encodings/onpair/src/array.rs +++ b/encodings/onpair/src/array.rs @@ -8,6 +8,7 @@ use std::hash::Hasher; use std::sync::Arc; use std::sync::OnceLock; +use num_traits::AsPrimitive; use onpair::CompactDictionaryView; use prost::Message as _; use vortex_array::Array; @@ -25,10 +26,12 @@ use vortex_array::IntoArray; use vortex_array::array_slots; use vortex_array::buffer::BufferHandle; use vortex_array::builders::ArrayBuilder; +use vortex_array::builders::VarBinBufferBuilder; use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; +use vortex_array::match_each_integer_ptype; use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::VTable; @@ -46,6 +49,7 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::canonical::canonicalize_onpair; +use crate::canonical::onpair_decode_bytes; use crate::canonical::onpair_decode_views; use crate::decode::collect_widened; use crate::rules::RULES; @@ -544,6 +548,25 @@ 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() + .map(|length| AsPrimitive::::as_(*length)), + &validity, + ); + }); + return Ok(()); + } + let Some(builder) = builder.as_any_mut().downcast_mut::() else { return array .array() diff --git a/encodings/onpair/src/canonical.rs b/encodings/onpair/src/canonical.rs index f5e55467eff..e955bd3219b 100644 --- a/encodings/onpair/src/canonical.rs +++ b/encodings/onpair/src/canonical.rs @@ -44,11 +44,10 @@ pub(super) fn canonicalize_onpair( }) } -pub(crate) fn onpair_decode_views( +pub(crate) fn onpair_decode_bytes( array: ArrayView<'_, OnPair>, - start_buf_index: u32, ctx: &mut ExecutionCtx, -) -> VortexResult<(Vec, Buffer)> { +) -> VortexResult<(ByteBufferMut, PrimitiveArray)> { let lengths = array .uncompressed_lengths() .clone() @@ -62,15 +61,6 @@ pub(crate) fn onpair_decode_views( .sum() }); - // `codes_offsets` holds the per-row code boundaries and may itself be a - // sliced or filtered view of the original. Its first and last entries - // bound the contiguous run of `codes` belonging to the rows present in - // this array: `slice` keeps the full `codes` child and only narrows - // `codes_offsets` (so `code_start > 0` and/or `code_end < codes.len()`), - // while `filter` rebuilds both children so the window is the whole stream. - // OnPair has no `TakeExecute`, so a reordering take is served from the - // canonical `VarBinView` and never reaches this path. We only need those - // two boundaries, so point-look them up rather than decoding every offset. let codes_offsets = array.codes_offsets(); let code_start = code_boundary_at(codes_offsets, 0, ctx)?; let code_end = code_boundary_at(codes_offsets, array.len(), ctx)?; @@ -85,13 +75,8 @@ pub(crate) fn onpair_decode_views( array.codes().len() ); - // Slice the `codes` child to that window *before* unpacking it, so a sliced - // array materialises only its own codes rather than the whole column's. The - // contiguous decoder walks `codes` in order and never reads the per-row - // boundaries, so an empty boundary slice is sound. let codes = collect_widened::(&array.codes().slice(code_start..code_end)?, ctx)?; let dict = dict_view(array, ctx)?; - let mut out_bytes = ByteBufferMut::with_capacity(total_size); let written = match onpair::try_decode_into(codes.as_slice(), dict, out_bytes.spare_capacity_mut()) { @@ -105,9 +90,16 @@ pub(crate) fn onpair_decode_views( "OnPair codes decoded to {written} bytes but uncompressed_lengths records {total_size}" ); } - // SAFETY: `try_decode_into` initialised exactly `written` bytes. unsafe { out_bytes.set_len(written) }; + Ok((out_bytes, lengths)) +} +pub(crate) fn onpair_decode_views( + array: ArrayView<'_, OnPair>, + start_buf_index: u32, + ctx: &mut ExecutionCtx, +) -> VortexResult<(Vec, Buffer)> { + let (out_bytes, lengths) = onpair_decode_bytes(array, ctx)?; match_each_integer_ptype!(lengths.ptype(), |P| { Ok(build_views( start_buf_index, diff --git a/encodings/onpair/src/tests.rs b/encodings/onpair/src/tests.rs index 20aa7b4e061..41c3ee73bce 100644 --- a/encodings/onpair/src/tests.rs +++ b/encodings/onpair/src/tests.rs @@ -16,6 +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::VarBinBufferBuilder; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -60,6 +61,19 @@ fn sample_input() -> VarBinArray { ) } +#[test] +fn 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 = VarBinBufferBuilder::with_capacity(input.dtype().clone(), false, input.len()); + encoded + .into_array() + .append_to_builder(&mut builder, &mut ctx)?; + assert_arrays_eq!(builder.finish_into_varbin(), input, &mut ctx); + Ok(()) +} + #[test] fn test_onpair_rejects_100k_token_dictionary() -> vortex_error::VortexResult<()> { let num_tokens = 100_000usize; diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index 2a9659efdd8..b28a81b5375 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -15,6 +15,7 @@ use vortex_array::ArrayId; use vortex_array::ArrayParts; use vortex_array::ArrayRef; use vortex_array::ArrayView; +use vortex_array::Canonical; use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; @@ -22,12 +23,17 @@ use vortex_array::IntoArray; use vortex_array::TypedArrayRef; use vortex_array::VortexSessionExecute; use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::varbin::VarBinArrayExt; use vortex_array::buffer::BufferHandle; +use vortex_array::builders::ArrayBuilder; +use vortex_array::builders::VarBinBufferBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::legacy_session; +use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::serde::ArrayChildren; use vortex_array::smallvec::smallvec; use vortex_array::validity::Validity; @@ -45,6 +51,7 @@ use crate::compress::runend_decode_primitive; use crate::compress::runend_decode_varbinview; use crate::compress::runend_encode; use crate::decompress_bool::runend_decode_bools; +use crate::iter::trimmed_ends_iter; use crate::ops::find_physical_index; use crate::ops::find_slice_end_index; use crate::rules::RULES; @@ -182,6 +189,54 @@ impl VTable for RunEnd { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { run_end_canonicalize(&array, ctx).map(ExecutionResult::done) } + + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if matches!(array.dtype(), DType::Utf8(_) | DType::Binary(_)) + && let Some(builder) = builder.as_any_mut().downcast_mut::() + { + let ends = array + .ends() + .clone() + .execute_as::("ends", ctx)?; + let mut values_builder = VarBinBufferBuilder::with_capacity( + array.dtype().clone(), + builder.has_large_offsets(), + array.values().len(), + ); + array.values().append_to_builder(&mut values_builder, ctx)?; + let values = values_builder.finish_into_varbin(); + let validity = values + .as_ref() + .validity()? + .execute_mask(values.len(), ctx)?; + + match_each_unsigned_integer_ptype!(ends.ptype(), |E| { + let run_ends = trimmed_ends_iter(ends.as_slice::(), array.offset(), array.len()); + let mut previous_end = 0; + for (index, end) in run_ends.enumerate() { + let run_length = end - previous_end; + if validity.value(index) { + builder.append_n_values(values.bytes_at(index), run_length); + } else { + builder.append_nulls(run_length); + } + previous_end = end; + } + }); + return Ok(()); + } + + array + .array() + .clone() + .execute::(ctx)? + .into_array() + .append_to_builder(builder, ctx) + } } /// The run-end positions marking where each run terminates. @@ -518,6 +573,7 @@ mod tests { use vortex_array::arrays::DictArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; + use vortex_array::builders::VarBinBufferBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -556,14 +612,28 @@ mod tests { #[test] fn test_runend_utf8() { let mut ctx = SESSION.create_execution_ctx(); - let values = VarBinViewArray::from_iter_str(["a", "b", "c"]).into_array(); + let values = + VarBinViewArray::from_iter_nullable_str([Some("a"), None, Some("c")]).into_array(); let arr = RunEnd::new(buffer![2u32, 5, 10].into_array(), values, &mut ctx); assert_eq!(arr.len(), 10); - assert_eq!(arr.dtype(), &DType::Utf8(Nullability::NonNullable)); - - let expected = - VarBinViewArray::from_iter_str(["a", "a", "b", "b", "b", "c", "c", "c", "c", "c"]) - .into_array(); + assert_eq!(arr.dtype(), &DType::Utf8(Nullability::Nullable)); + + let expected = VarBinViewArray::from_iter_nullable_str([ + Some("a"), + Some("a"), + None, + None, + None, + Some("c"), + Some("c"), + Some("c"), + Some("c"), + Some("c"), + ]) + .into_array(); + let mut builder = VarBinBufferBuilder::with_capacity(arr.dtype().clone(), false, 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 4d4ec2623de..282ac3e402d 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -26,11 +26,17 @@ use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::VarBinArray; use vortex_array::arrays::bool::BoolArrayExt; +use vortex_array::arrays::varbin::VarBinArrayExt; use vortex_array::buffer::BufferHandle; +use vortex_array::builders::ArrayBuilder; +use vortex_array::builders::VarBinBufferBuilder; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; +use vortex_array::dtype::IntegerPType; use vortex_array::dtype::Nullability; +use vortex_array::match_each_integer_ptype; use vortex_array::patches::PatchSlotIndices; use vortex_array::patches::Patches; use vortex_array::patches::PatchesData; @@ -350,6 +356,91 @@ impl VTable for Sparse { // TODO(joe): remove ctx from execute_sparse since all slots should be canonical. execute_sparse(parts, ctx).map(ExecutionResult::done) } + + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if matches!(array.dtype(), DType::Utf8(_) | DType::Binary(_)) + && let Some(builder) = builder.as_any_mut().downcast_mut::() + { + return append_sparse_varbin(array, builder, ctx); + } + + array + .array() + .clone() + .execute::(ctx)? + .into_array() + .append_to_builder(builder, ctx) + } +} + +fn append_sparse_varbin( + array: ArrayView<'_, Sparse>, + builder: &mut VarBinBufferBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let patches = array.resolved_patches()?; + let indices = patches.indices().clone().execute::(ctx)?; + let mut values_builder = VarBinBufferBuilder::with_capacity( + array.dtype().clone(), + builder.has_large_offsets(), + patches.values().len(), + ); + patches + .values() + .append_to_builder(&mut values_builder, ctx)?; + let values = values_builder.finish_into_varbin(); + let validity = values + .as_ref() + .validity()? + .execute_mask(values.len(), ctx)?; + + match_each_integer_ptype!(indices.ptype(), |I| { + append_varbin_patches( + indices.as_slice::(), + &values, + &validity, + array.fill_scalar(), + array.len(), + builder, + ) + }) +} + +fn append_varbin_patches( + indices: &[I], + values: &VarBinArray, + validity: &Mask, + fill: &Scalar, + len: usize, + builder: &mut VarBinBufferBuilder, +) -> VortexResult<()> { + let mut row = 0; + let mut patch_index = 0; + while patch_index < indices.len() { + let index: usize = indices[patch_index].as_(); + let mut last_patch = patch_index; + while last_patch + 1 < indices.len() { + let next_index: usize = indices[last_patch + 1].as_(); + if next_index != index { + break; + } + last_patch += 1; + } + + builder.append_scalar_repeated(fill, index - row)?; + if validity.value(last_patch) { + builder.append_n_values(values.bytes_at(last_patch), 1); + } else { + builder.append_nulls(1); + } + row = index + 1; + patch_index = last_patch + 1; + } + builder.append_scalar_repeated(fill, len - row) } const PATCH_SLOTS: PatchSlotIndices = PatchSlotIndices { @@ -697,7 +788,9 @@ mod test { use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; + use vortex_array::builders::VarBinBufferBuilder; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -905,6 +998,46 @@ mod test { ); } + #[test] + fn append_utf8_to_varbin_buffer_builder() { + let mut ctx = SESSION.create_execution_ctx(); + let values = VarBinViewArray::from_iter_nullable_str([Some("patched"), None, Some("last")]) + .into_array(); + let fill = Scalar::null(values.dtype().clone()); + let array = Sparse::try_new(buffer![1u8, 3, 5].into_array(), values, 6, fill).unwrap(); + let expected = VarBinViewArray::from_iter_nullable_str([ + None, + Some("patched"), + None, + None, + None, + Some("last"), + ]) + .into_array(); + let mut builder = + VarBinBufferBuilder::with_capacity(array.dtype().clone(), false, array.len()); + array.append_to_builder(&mut builder, &mut ctx).unwrap(); + assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx); + } + + #[test] + fn append_utf8_with_duplicate_patch_indices() { + let mut ctx = SESSION.create_execution_ctx(); + let values = VarBinViewArray::from_iter_str(["first", "second"]).into_array(); + let array = Sparse::try_new( + buffer![1u8, 1].into_array(), + values, + 2, + Scalar::utf8("fill".to_string(), Nullability::NonNullable), + ) + .unwrap(); + let expected = VarBinViewArray::from_iter_str(["fill", "second"]).into_array(); + let mut builder = + VarBinBufferBuilder::with_capacity(array.dtype().clone(), false, array.len()); + array.append_to_builder(&mut builder, &mut ctx).unwrap(); + assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx); + } + #[test] #[should_panic] fn test_invalid_length() { diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index dc9f075e44e..53545b10955 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -28,6 +28,8 @@ use vortex_array::arrays::VarBinViewArray; 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::VarBinBufferBuilder; use vortex_array::dtype::DType; use vortex_array::scalar::Scalar; use vortex_array::serde::ArrayChildren; @@ -268,6 +270,29 @@ impl VTable for Zstd { .map(ExecutionResult::done) } + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if matches!(array.dtype(), DType::Binary(_) | DType::Utf8(_)) + && let Some(builder) = builder.as_any_mut().downcast_mut::() + { + let unsliced_validity = + child_to_validity(array.slots()[0].as_ref(), array.dtype().nullability()); + return array + .data() + .append_varbin(array.dtype(), &unsliced_validity, builder, ctx); + } + + array + .array() + .clone() + .execute::(ctx)? + .into_array() + .append_to_builder(builder, ctx) + } + fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, @@ -514,6 +539,34 @@ pub fn reconstruct_views( (buffers, views.freeze()) } +struct DecompressedSlice { + bytes: ByteBuffer, + validity: Validity, + byte_width: usize, + n_rows: usize, + value_idx_start: usize, + value_idx_stop: usize, + n_skipped_values: usize, +} + +fn zstd_values(buffer: &[u8]) -> impl Iterator { + let mut offset = 0; + std::iter::from_fn(move || { + if offset == buffer.len() { + return None; + } + let len = ViewLen::from_le_bytes( + buffer[offset..offset + size_of::()] + .try_into() + .ok() + .vortex_expect("must fit ViewLen size"), + ) as usize; + let value_start = offset + size_of::(); + offset = value_start + len; + Some(&buffer[value_start..offset]) + }) +} + impl ZstdData { /// Construct unsliced zstd data from raw frames and metadata. pub fn new( @@ -897,12 +950,12 @@ impl ZstdData { } } - fn decompress( + fn decompress_slice( &self, dtype: &DType, unsliced_validity: &Validity, ctx: &mut ExecutionCtx, - ) -> VortexResult { + ) -> VortexResult { // To start, we figure out which frames we need to decompress, and with // what row offset into the first such frame. let byte_width = Self::byte_width(dtype); @@ -997,29 +1050,47 @@ impl ZstdData { // END OF IMPORTANT BLOCK // + Ok(DecompressedSlice { + bytes: decompressed, + validity: slice_validity, + byte_width, + n_rows: slice_n_rows, + value_idx_start: slice_value_idx_start, + value_idx_stop: slice_value_idx_stop, + n_skipped_values, + }) + } + + fn decompress( + &self, + dtype: &DType, + unsliced_validity: &Validity, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let slice = self.decompress_slice(dtype, unsliced_validity, ctx)?; match dtype { DType::Primitive(..) => { - let slice_values_buffer = decompressed.slice( - (slice_value_idx_start - n_skipped_values) * byte_width - ..(slice_value_idx_stop - n_skipped_values) * byte_width, + let slice_values_buffer = slice.bytes.slice( + (slice.value_idx_start - slice.n_skipped_values) * slice.byte_width + ..(slice.value_idx_stop - slice.n_skipped_values) * slice.byte_width, ); let primitive = PrimitiveArray::from_values_byte_buffer( slice_values_buffer, dtype.as_ptype(), - slice_validity, - slice_n_rows, + slice.validity, + slice.n_rows, ctx, ); Ok(primitive.into_array()) } DType::Binary(_) | DType::Utf8(_) => { - match slice_validity.execute_mask(slice_n_rows, ctx)?.indices() { + match slice.validity.execute_mask(slice.n_rows, ctx)?.indices() { AllOr::All => { - let (buffers, all_views) = reconstruct_views(&decompressed, MAX_BUFFER_LEN); + let (buffers, all_views) = reconstruct_views(&slice.bytes, MAX_BUFFER_LEN); let valid_views = all_views.slice( - slice_value_idx_start - n_skipped_values - ..slice_value_idx_stop - n_skipped_values, + slice.value_idx_start - slice.n_skipped_values + ..slice.value_idx_stop - slice.n_skipped_values, ); // SAFETY: we properly construct the views inside `reconstruct_views` @@ -1028,24 +1099,24 @@ impl ZstdData { valid_views, Arc::from(buffers), dtype.clone(), - slice_validity, + slice.validity, ) } .into_array()) } AllOr::None => Ok(ConstantArray::new( Scalar::null(dtype.clone()), - slice_n_rows, + slice.n_rows, ) .into_array()), AllOr::Some(valid_indices) => { - let (buffers, all_views) = reconstruct_views(&decompressed, MAX_BUFFER_LEN); + let (buffers, all_views) = reconstruct_views(&slice.bytes, MAX_BUFFER_LEN); let valid_views = all_views.slice( - slice_value_idx_start - n_skipped_values - ..slice_value_idx_stop - n_skipped_values, + slice.value_idx_start - slice.n_skipped_values + ..slice.value_idx_stop - slice.n_skipped_values, ); - let mut views = BufferMut::::zeroed(slice_n_rows); + let mut views = BufferMut::::zeroed(slice.n_rows); for (view, index) in valid_views.into_iter().zip_eq(valid_indices) { views[*index] = view } @@ -1056,7 +1127,7 @@ impl ZstdData { views.freeze(), Arc::from(buffers), dtype.clone(), - slice_validity, + slice.validity, ) } .into_array()) @@ -1067,6 +1138,45 @@ impl ZstdData { } } + fn append_varbin( + &self, + dtype: &DType, + unsliced_validity: &Validity, + builder: &mut VarBinBufferBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let slice = self.decompress_slice(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); + } + } + Ok(()) + } + /// Returns the length of the array. #[inline] pub fn len(&self) -> usize { diff --git a/encodings/zstd/src/test.rs b/encodings/zstd/src/test.rs index cf45eb97c04..9552475fc17 100644 --- a/encodings/zstd/src/test.rs +++ b/encodings/zstd/src/test.rs @@ -10,6 +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::VarBinBufferBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::validity::Validity; @@ -205,6 +206,35 @@ fn test_zstd_var_bin_view() { assert_nth_scalar!(sliced, 2, "Lorem ipsum dolor sit amet", &mut ctx); } +#[test] +fn test_zstd_append_to_offset_builder() { + let mut ctx = array_session().create_execution_ctx(); + let array = VarBinViewArray::from_iter( + [ + Some(b"foo".as_slice()), + Some(b"bar".as_slice()), + None, + Some(b"Lorem ipsum dolor sit amet".as_slice()), + Some(b"baz".as_slice()), + ], + DType::Utf8(Nullability::Nullable), + ); + let compressed = Zstd::from_var_bin_view(&array, 0, 3, &mut ctx) + .unwrap() + .slice(1..4) + .unwrap(); + let mut builder = + VarBinBufferBuilder::with_capacity(compressed.dtype().clone(), false, compressed.len()); + compressed + .append_to_builder(&mut builder, &mut ctx) + .unwrap(); + assert_arrays_eq!( + builder.finish_into_varbin(), + array.into_array().slice(1..4).unwrap(), + &mut ctx + ); +} + #[test] fn test_zstd_decompress_var_bin_view() { let mut ctx = array_session().create_execution_ctx(); diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index a2057077335..e3255693d42 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -35,6 +35,7 @@ use crate::builders::BoolBuilder; use crate::builders::DecimalBuilder; use crate::builders::NullBuilder; use crate::builders::PrimitiveBuilder; +use crate::builders::VarBinBufferBuilder; use crate::builders::VarBinViewBuilder; use crate::canonical::Canonical; use crate::dtype::DType; @@ -222,22 +223,30 @@ impl VTable for Constant { }); } DType::Utf8(_) => { - append_value_or_nulls::(builder, scalar.is_null(), n, |b| { - let typed = scalar.as_utf8(); - let value = typed - .value() - .vortex_expect("non-null utf8 scalar must have a value"); - b.append_n_values(value.as_bytes(), n); - }); + if let Some(builder) = builder.as_any_mut().downcast_mut::() { + builder.append_scalar_repeated(scalar, n)?; + } else { + append_value_or_nulls::(builder, scalar.is_null(), n, |b| { + let value = scalar + .as_utf8() + .value() + .vortex_expect("non-null utf8 scalar must have a value"); + b.append_n_values(value.as_bytes(), n); + }); + } } DType::Binary(_) => { - append_value_or_nulls::(builder, scalar.is_null(), n, |b| { - let typed = scalar.as_binary(); - let value = typed - .value() - .vortex_expect("non-null binary scalar must have a value"); - b.append_n_values(value, n); - }); + if let Some(builder) = builder.as_any_mut().downcast_mut::() { + builder.append_scalar_repeated(scalar, n)?; + } else { + append_value_or_nulls::(builder, scalar.is_null(), n, |b| { + let value = scalar + .as_binary() + .value() + .vortex_expect("non-null binary scalar must have a value"); + b.append_n_values(value, n); + }); + } } // TODO: add fast paths for DType::Struct, DType::List, DType::FixedSizeList, DType::Extension. _ => { diff --git a/vortex-array/src/arrays/dict/vtable/mod.rs b/vortex-array/src/arrays/dict/vtable/mod.rs index d9ced9eeffa..36adc11a37a 100644 --- a/vortex-array/src/arrays/dict/vtable/mod.rs +++ b/vortex-array/src/arrays/dict/vtable/mod.rs @@ -5,6 +5,7 @@ use std::hash::Hasher; use prost::Message; use smallvec::smallvec; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -34,12 +35,15 @@ use crate::array::VTable; use crate::array::with_empty_buffers; use crate::arrays::ConstantArray; use crate::arrays::Primitive; +use crate::arrays::VarBin; use crate::arrays::dict::DictArrayExt; use crate::arrays::dict::DictArraySlotsExt; +use crate::arrays::dict::TakeExecute; use crate::arrays::dict::compute::rules::PARENT_RULES; use crate::arrays::dict::execute::take_canonical; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; +use crate::builders::VarBinBufferBuilder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -215,6 +219,24 @@ impl VTable for Dict { builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { + if let Some(builder) = builder.as_any_mut().downcast_mut::() { + // Dictionary values may exceed the target's offset width even when the selected + // output does not. Keep the intermediate wide and narrow only after applying codes. + let mut values_builder = VarBinBufferBuilder::with_capacity( + array.values().dtype().clone(), + true, + array.values().len(), + ); + array + .values() + .clone() + .append_to_builder(&mut values_builder, ctx)?; + let values = values_builder.finish_into_varbin(); + let taken = ::take(values.as_view(), array.codes(), ctx)? + .vortex_expect("taking dictionary values should produce an array"); + return builder.append_varbin(taken.as_::(), ctx); + } + if !array.is_empty() && let (Some(codes), Some(values)) = ( array.codes().as_opt::(), diff --git a/vortex-array/src/arrays/varbin/builder.rs b/vortex-array/src/arrays/varbin/builder.rs index 422139cc851..947d619769c 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -1,24 +1,40 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::Any; + use num_traits::AsPrimitive; use vortex_buffer::BitBufferMut; use vortex_buffer::BufferMut; +use vortex_buffer::ByteBuffer; +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::Mask; +use crate::ArrayRef; +use crate::Canonical; +use crate::ExecutionCtx; use crate::IntoArray; #[cfg(debug_assertions)] use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; +use crate::arrays::VarBin; use crate::arrays::VarBinArray; +use crate::arrays::VarBinView; +use crate::arrays::varbin::VarBinArrayExt; +use crate::arrays::varbinview::VarBinViewArrayExt; +use crate::builders::ArrayBuilder; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::expr::stats::Precision; use crate::expr::stats::Stat; #[cfg(debug_assertions)] use crate::legacy_session; +use crate::scalar::Scalar; use crate::validity::Validity; - pub struct VarBinBuilder { offsets: BufferMut, data: BufferMut, @@ -94,6 +110,54 @@ impl VarBinBuilder { self.validity.append_n(true, num); } + fn append_values_with_lengths( + &mut self, + values: &[u8], + lengths: impl Iterator, + validity: &Mask, + ) { + let mut end = self.data.len(); + let mut len = 0; + for value_len in lengths { + end += value_len; + self.offsets.push(O::from(end).unwrap_or_else(|| { + vortex_panic!( + "Failed to convert byte offset {end} to {}", + std::any::type_name::() + ) + })); + len += 1; + } + debug_assert_eq!(len, validity.len()); + debug_assert_eq!(end - self.data.len(), values.len()); + self.data.extend_from_slice(values); + 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 len(&self) -> usize { + self.validity.len() + } + + fn reserve_exact(&mut self, additional: usize) { + self.offsets.reserve(additional); + self.validity.reserve(additional); + } + + 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)), + }; + } + #[allow(clippy::disallowed_methods)] pub fn finish(self, dtype: DType) -> VarBinArray { let offsets = PrimitiveArray::new(self.offsets.freeze(), Validity::NonNullable); @@ -127,6 +191,229 @@ impl VarBinBuilder { } } +/// Builder for offset-based UTF-8 and binary arrays with either 32-bit or 64-bit offsets. +/// +/// Unlike [`crate::builders::VarBinViewBuilder`], this builder produces a [`VarBinArray`]. +/// Encodings can downcast an [`ArrayBuilder`] to this type and decode directly into Arrow's +/// offset-based physical representation without first materializing a view array. +pub struct VarBinBufferBuilder { + dtype: DType, + storage: TypedBuilder, +} + +enum TypedBuilder { + I32(VarBinBuilder), + I64(VarBinBuilder), +} + +impl VarBinBufferBuilder { + /// 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(_))); + let storage = if large_offsets { + TypedBuilder::I64(VarBinBuilder::with_capacity(capacity)) + } else { + TypedBuilder::I32(VarBinBuilder::with_capacity(capacity)) + }; + Self { dtype, storage } + } + + /// Returns whether this builder uses 64-bit offsets. + pub fn has_large_offsets(&self) -> bool { + matches!(self.storage, TypedBuilder::I64(_)) + } + + /// Appends decompressed values represented by one contiguous byte buffer and per-row lengths. + pub fn append_values( + &mut self, + values: &[u8], + lengths: impl Iterator, + validity: &Mask, + ) { + match &mut self.storage { + TypedBuilder::I32(builder) => { + builder.append_values_with_lengths(values, lengths, validity) + } + TypedBuilder::I64(builder) => { + builder.append_values_with_lengths(values, lengths, 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); + } + } + + /// 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(), + "VarBinBufferBuilder 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!("VarBinBufferBuilder cannot append scalar of dtype {dtype}"), + } + Ok(()) + } + + /// Appends an existing offset-based Vortex array 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::

(); + self.append_values( + bytes.as_slice(), + offsets.windows(2).map(|window| { + let start: usize = window[0].as_(); + let end: usize = window[1].as_(); + end - start + }), + &validity, + ); + }); + Ok(()) + } + + /// Appends a Vortex view array into the offset-based output. + pub fn append_varbinview( + &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)); + } + } + 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(); + } + } + } + } + Ok(()) + } + + /// Finishes the current values as a [`VarBinArray`] and resets the builder. + pub fn finish_into_varbin(&mut self) -> VarBinArray { + match &mut self.storage { + TypedBuilder::I32(builder) => std::mem::take(builder).finish(self.dtype.clone()), + TypedBuilder::I64(builder) => std::mem::take(builder).finish(self.dtype.clone()), + } + } + + fn append_value(&mut self, value: impl AsRef<[u8]>) { + match &mut self.storage { + TypedBuilder::I32(builder) => builder.append_value(value), + TypedBuilder::I64(builder) => builder.append_value(value), + } + } + + fn push_null(&mut self) { + match &mut self.storage { + TypedBuilder::I32(builder) => builder.append_null(), + TypedBuilder::I64(builder) => builder.append_null(), + } + } + + fn push_nulls(&mut self, n: usize) { + match &mut self.storage { + TypedBuilder::I32(builder) => builder.append_n_nulls(n), + TypedBuilder::I64(builder) => builder.append_n_nulls(n), + } + } +} + +impl ArrayBuilder for VarBinBufferBuilder { + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn len(&self) -> usize { + match &self.storage { + TypedBuilder::I32(builder) => builder.len(), + TypedBuilder::I64(builder) => builder.len(), + } + } + + fn append_zeros(&mut self, n: usize) { + for _ in 0..n { + self.append_value([]); + } + } + + unsafe fn append_nulls_unchecked(&mut self, n: usize) { + self.push_nulls(n); + } + + fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> { + self.append_scalar_repeated(scalar, 1) + } + + fn reserve_exact(&mut self, additional: usize) { + match &mut self.storage { + TypedBuilder::I32(builder) => builder.reserve_exact(additional), + TypedBuilder::I64(builder) => builder.reserve_exact(additional), + } + } + + unsafe fn set_validity_unchecked(&mut self, validity: Mask) { + match &mut self.storage { + TypedBuilder::I32(builder) => builder.set_validity(validity), + TypedBuilder::I64(builder) => builder.set_validity(validity), + } + } + + fn finish(&mut self) -> ArrayRef { + self.finish_into_varbin().into_array() + } + + fn finish_into_canonical(&mut self, ctx: &mut ExecutionCtx) -> Canonical { + self.finish() + .execute::(ctx) + .vortex_expect("varbin buffer builder should canonicalize") + } +} + #[cfg(test)] mod tests { use vortex_error::VortexResult; diff --git a/vortex-array/src/arrays/varbin/vtable/mod.rs b/vortex-array/src/arrays/varbin/vtable/mod.rs index 301fdebc084..48314665292 100644 --- a/vortex-array/src/arrays/varbin/vtable/mod.rs +++ b/vortex-array/src/arrays/varbin/vtable/mod.rs @@ -26,6 +26,8 @@ use crate::arrays::varbin::array::NUM_SLOTS; use crate::arrays::varbin::array::OFFSETS_SLOT; use crate::arrays::varbin::array::SLOT_NAMES; use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::builders::VarBinBufferBuilder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -203,6 +205,19 @@ impl VTable for VarBin { PARENT_RULES.evaluate(array, parent, child_idx) } + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if let Some(builder) = builder.as_any_mut().downcast_mut::() { + return builder.append_varbin(array, ctx); + } + varbin_to_canonical(array, ctx)? + .into_array() + .append_to_builder(builder, ctx) + } + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { Ok(ExecutionResult::done( varbin_to_canonical(array.as_view(), ctx)?.into_array(), diff --git a/vortex-array/src/arrays/varbinview/vtable/mod.rs b/vortex-array/src/arrays/varbinview/vtable/mod.rs index f730282ddb4..cb813b54658 100644 --- a/vortex-array/src/arrays/varbinview/vtable/mod.rs +++ b/vortex-array/src/arrays/varbinview/vtable/mod.rs @@ -30,6 +30,7 @@ use crate::arrays::varbinview::array::SLOT_NAMES; use crate::arrays::varbinview::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; +use crate::builders::VarBinBufferBuilder; use crate::builders::VarBinViewBuilder; use crate::dtype::DType; use crate::hash::ArrayEq; @@ -245,10 +246,13 @@ impl VTable for VarBinView { builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - let Some(builder) = builder.as_any_mut().downcast_mut::() else { - vortex_bail!("append_to_builder for VarBinView requires a VarBinViewBuilder"); - }; - builder.append_varbinview_array(&array.into_owned(), ctx) + 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); + } + vortex_bail!("append_to_builder for VarBinView requires a variable-binary builder") } fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { diff --git a/vortex-array/src/builders/mod.rs b/vortex-array/src/builders/mod.rs index 561efc2711e..343edd2aba0 100644 --- a/vortex-array/src/builders/mod.rs +++ b/vortex-array/src/builders/mod.rs @@ -75,6 +75,8 @@ pub use primitive::*; pub use struct_::*; pub use varbinview::*; +pub use crate::arrays::varbin::builder::VarBinBufferBuilder; + #[cfg(test)] mod tests; diff --git a/vortex-arrow/src/executor/byte.rs b/vortex-arrow/src/executor/byte.rs index ae40ee01c26..602a87dcd6b 100644 --- a/vortex-arrow/src/executor/byte.rs +++ b/vortex-arrow/src/executor/byte.rs @@ -8,6 +8,7 @@ use arrow_array::GenericByteArray; use arrow_array::types::BinaryViewType; use arrow_array::types::ByteArrayType; use arrow_array::types::StringViewType; +use arrow_schema::DataType; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::Canonical; @@ -15,10 +16,12 @@ use vortex_array::ExecutionCtx; use vortex_array::arrays::VarBin; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::varbin::VarBinArrayExt; +use vortex_array::builders::VarBinBufferBuilder; 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_error::VortexError; use vortex_error::VortexResult; @@ -33,20 +36,30 @@ pub(super) fn to_arrow_byte_array( where T::Offset: NativePType, { + 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); + } + // If the Vortex array is already in VarBin format, we can directly convert it. if let Some(array) = array.as_opt::() { return varbin_to_byte_array::(array, ctx); } - // Otherwise, we execute the array to a VarBinViewArray and convert to Arrow ByteView, - // then cast to the target byte array type. - 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"), - }?; - arrow_cast::cast(&binary_view, &T::DATA_TYPE).map_err(VortexError::from) + let mut builder = VarBinBufferBuilder::with_capacity( + array.dtype().clone(), + T::Offset::PTYPE == PType::I64, + array.len(), + ); + array.append_to_builder(&mut builder, ctx)?; + varbin_to_byte_array::(builder.finish_into_varbin().as_view(), ctx) } /// Convert a Vortex VarBinArray into an Arrow GenericBinaryArray. @@ -84,13 +97,13 @@ mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; + use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_error::VortexResult; use vortex_mask::Mask; use crate::ArrowSessionExt; - use crate::executor::byte::VarBinViewArray; fn make_utf8_array() -> VarBinViewArray { VarBinViewArray::from_iter_str(["hello", "world", "this is a longer string for testing"]) diff --git a/vortex-python/python/vortex/_lib/arrays.pyi b/vortex-python/python/vortex/_lib/arrays.pyi index 8da62190291..44c7c824ac9 100644 --- a/vortex-python/python/vortex/_lib/arrays.pyi +++ b/vortex-python/python/vortex/_lib/arrays.pyi @@ -25,7 +25,9 @@ class Array: ) -> Array: ... @staticmethod def from_range(obj: range, *, dtype: DType | None = None) -> Array: ... - def to_arrow_array(self) -> pa.Array[pa.Scalar[pa.DataType]]: ... + def to_arrow_array( + self, *, arrow_type: pa.DataType | None = None + ) -> pa.Array[pa.Scalar[pa.DataType]]: ... def __vortex_array_metadata__(self) -> tuple[str, bytes, int, bytes, list[object], list[object]]: ... @property def id(self) -> str: ... diff --git a/vortex-python/python/vortex/_lib/file.pyi b/vortex-python/python/vortex/_lib/file.pyi index 4a2045bc0b8..ed8109210c4 100644 --- a/vortex-python/python/vortex/_lib/file.pyi +++ b/vortex-python/python/vortex/_lib/file.pyi @@ -46,6 +46,7 @@ class VortexFile: expr: Expr | None = None, limit: int | None = None, batch_size: int | None = None, + schema: pa.Schema | None = None, ) -> pa.RecordBatchReader: ... def to_dataset(self) -> VortexDataset: ... def to_polars(self) -> pl.LazyFrame: ... diff --git a/vortex-python/python/vortex/file.py b/vortex-python/python/vortex/file.py index 0167aae7318..e077fdf852e 100644 --- a/vortex-python/python/vortex/file.py +++ b/vortex-python/python/vortex/file.py @@ -202,6 +202,7 @@ def to_arrow( limit: int | None = None, expr: Expr | None = None, batch_size: int | None = None, + schema: pa.Schema | None = None, ) -> RecordBatchReader: """Scan the Vortex file as a :class:`pyarrow.RecordBatchReader`. @@ -214,9 +215,14 @@ def to_arrow( The predicate used to filter rows. The filter columns need not appear in the projection. batch_size : :class:`int` | None The number of rows to read per chunk. + schema : :class:`pyarrow.Schema` | None + The Arrow physical schema to return. Use ``pyarrow.string()`` and + ``pyarrow.binary()`` fields to request offset-based arrays instead of view arrays. """ - return self._file.to_arrow(projection, expr=expr, limit=limit, batch_size=batch_size) + return self._file.to_arrow( + projection, expr=expr, limit=limit, batch_size=batch_size, schema=schema + ) def to_dataset(self) -> VortexDataset: """Scan the Vortex file using the :class:`pyarrow.dataset.Dataset` API.""" diff --git a/vortex-python/src/arrays/mod.rs b/vortex-python/src/arrays/mod.rs index 4ea6745f707..51c8dde7033 100644 --- a/vortex-python/src/arrays/mod.rs +++ b/vortex-python/src/arrays/mod.rs @@ -15,6 +15,8 @@ use std::ptr::NonNull; use arrow_array::Array as ArrowArray; use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_schema::DataType; +use arrow_schema::Field; use pyo3::IntoPyObjectExt; use pyo3::exceptions::PyIndexError; use pyo3::exceptions::PyNotImplementedError; @@ -62,6 +64,7 @@ use crate::arrays::native::PyNativeArray; use crate::arrays::py::PyPythonArray; use crate::arrays::py::PythonArray; use crate::arrays::py::PythonVTable; +use crate::arrow::FromPyArrow; use crate::arrow::ToPyArrow; use crate::dtype::PyDType; use crate::error::PyVortexError; @@ -429,6 +432,12 @@ impl PyArray { /// .. seealso:: /// :meth:`.to_arrow_table` /// + /// Parameters + /// ---------- + /// arrow_type : :class:`pyarrow.DataType`, optional + /// The Arrow physical type to return. When omitted, Vortex chooses its preferred type, + /// including ``string_view`` and ``binary_view`` for UTF-8 and binary arrays. + /// /// Returns /// ------- /// :class:`pyarrow.Array` @@ -448,24 +457,50 @@ impl PyArray { /// 3 /// ] /// ``` - fn to_arrow_array<'py>(self_: &'py Bound<'py, Self>) -> PyVortexResult> { + /// Export an offset-based Arrow string array instead of the default string-view array: + /// + /// ```python + /// >>> import pyarrow + /// >>> import vortex as vx + /// >>> vx.array(["hello", "world"]).to_arrow_array(arrow_type=pyarrow.string()) + /// + /// [ + /// "hello", + /// "world" + /// ] + /// ``` + #[pyo3(signature = (*, arrow_type = None))] + fn to_arrow_array<'py>( + self_: &'py Bound<'py, Self>, + arrow_type: Option<&Bound<'py, PyAny>>, + ) -> PyVortexResult> { // NOTE(ngates): for struct arrays, we could also return a RecordBatchStreamReader. let array = PyArrayRef::extract(self_.as_any().as_borrowed())?.into_inner(); let py = self_.py(); + let target_field = arrow_type + .map(|arrow_type| DataType::from_pyarrow(&arrow_type.as_borrowed())) + .transpose()? + .map(|data_type| Field::new("", data_type, array.dtype().is_nullable())); if let Some(chunked_array) = array.as_opt::() { // We figure out a single Arrow Data Type to convert all chunks into, otherwise // the preferred type of each chunk may be different. - let arrow_field = session() - .arrow() - .to_arrow_field("", chunked_array.dtype())?; + let inferred_field; + let arrow_field = if let Some(target_field) = target_field.as_ref() { + target_field + } else { + inferred_field = session() + .arrow() + .to_arrow_field("", chunked_array.dtype())?; + &inferred_field + }; let chunks = chunked_array .iter_chunks() .map(|chunk| -> PyVortexResult<_> { Ok(session().arrow().execute_arrow( chunk.clone(), - Some(&arrow_field), + Some(arrow_field), &mut session().create_execution_ctx(), )?) }) @@ -491,7 +526,11 @@ impl PyArray { } else { Ok(session() .arrow() - .execute_arrow(array, None, &mut session().create_execution_ctx())? + .execute_arrow( + array, + target_field.as_ref(), + &mut session().create_execution_ctx(), + )? .into_data() .to_pyarrow(py)? .into_bound(py)) diff --git a/vortex-python/src/file.rs b/vortex-python/src/file.rs index 9550067b956..5f8a39862b7 100644 --- a/vortex-python/src/file.rs +++ b/vortex-python/src/file.rs @@ -4,6 +4,7 @@ use std::sync::Arc; use arrow_array::RecordBatchReader; +use arrow_schema::Schema; use pyo3::exceptions::PyTypeError; use pyo3::prelude::*; use pyo3::types::PyList; @@ -31,6 +32,7 @@ use vortex_arrow::ToArrowType; use crate::RUNTIME; use crate::arrays::PyArrayRef; +use crate::arrow::FromPyArrow; use crate::arrow::IntoPyArrow; use crate::dataset::PyVortexDataset; use crate::dtype::PyDType; @@ -153,15 +155,20 @@ impl PyVortexFile { }) } - #[pyo3(signature = (projection = None, *, expr = None, limit = None, batch_size = None))] + #[pyo3(signature = (projection = None, *, expr = None, limit = None, batch_size = None, schema = None))] fn to_arrow( slf: Bound, projection: Option, expr: Option, limit: Option, batch_size: Option, + schema: Option<&Bound>, ) -> PyVortexResult> { let vxf = slf.get().vxf.clone(); + let schema = schema + .map(|schema| Schema::from_pyarrow(&schema.as_borrowed())) + .transpose()? + .map(Arc::new); let reader = slf.py().detach(|| { let mut builder = vxf @@ -177,7 +184,10 @@ impl PyVortexFile { builder = builder.with_split_by(SplitBy::RowCount(batch_size)); } - let schema = Arc::new(builder.dtype()?.to_arrow_schema()?); + let schema = match schema { + Some(schema) => schema, + None => Arc::new(builder.dtype()?.to_arrow_schema()?), + }; builder.into_record_batch_reader(schema, &*RUNTIME) })?; diff --git a/vortex-python/test/test_array.py b/vortex-python/test/test_array.py index 7458f738e05..32409f87f33 100644 --- a/vortex-python/test/test_array.py +++ b/vortex-python/test/test_array.py @@ -25,6 +25,24 @@ def test_varbin_array_round_trip(): assert arr.to_arrow_array() == a +@pytest.mark.parametrize( + ("source_type", "target_type", "values"), + [ + (pa.string_view(), pa.string(), ["one", None, "three"]), + (pa.binary_view(), pa.binary(), [b"one", None, b"three"]), + ], +) +def test_varbin_offset_arrow_type( + source_type: pa.DataType, + target_type: pa.DataType, + values: list[str | bytes | None], +): + source = pa.array(values, type=source_type) + result = vortex.array(source).to_arrow_array(arrow_type=target_type) + assert result.type == target_type + assert result == pa.array(values, type=target_type) + + def test_varbin_array_take(): a = vortex.array(pa.array(["a", "b", "c", "d"], type=pa.string_view())) assert a.take(vortex.array(pa.array([0, 2]))).to_arrow_array() == pa.array( diff --git a/vortex-python/test/test_file.py b/vortex-python/test/test_file.py index 36196919f12..82fef155367 100644 --- a/vortex-python/test/test_file.py +++ b/vortex-python/test/test_file.py @@ -62,6 +62,16 @@ def test_to_arrow_columns(vxf: VortexFile): assert rbr.schema == pa.schema([("string", pa.string_view()), ("bool", pa.bool_())]) +def test_to_arrow_offset_string_schema(vxf: VortexFile): + schema = pa.schema([("string", pa.string())]) + table = vxf.to_arrow(projection=["string"], schema=schema).read_all() + assert table.schema == schema + assert table.column("string").combine_chunks() == pa.array( + (str(i) for i in range(1_000_000)), + type=pa.string(), + ) + + def test_empty_file(tmpdir_factory): # pyright: ignore[reportUnknownParameterType, reportMissingParameterType] # test for writing empty files with null columns # create an empty table with schema `empty: null` From 6c66449ff2e425dcd1f0e8f157406e6bfba81839 Mon Sep 17 00:00:00 2001 From: cl Date: Thu, 23 Jul 2026 14:26:03 +0800 Subject: [PATCH 02/10] perf(arrow): refine offset export paths Signed-off-by: cl --- Cargo.lock | 3 + encodings/runend/src/array.rs | 55 --------- encodings/sparse/src/lib.rs | 91 --------------- vortex-arrow/Cargo.toml | 3 + vortex-arrow/benches/to_arrow.rs | 190 ++++++++++++++++++++++++++++++- 5 files changed, 195 insertions(+), 147 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ab869e730b9..abd0e3dcba9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9746,9 +9746,12 @@ dependencies = [ "vortex-array", "vortex-buffer", "vortex-error", + "vortex-fsst", "vortex-mask", "vortex-runend", "vortex-session", + "vortex-sparse", + "vortex-zstd", ] [[package]] diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index b28a81b5375..3a073bb6556 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -15,7 +15,6 @@ use vortex_array::ArrayId; use vortex_array::ArrayParts; use vortex_array::ArrayRef; use vortex_array::ArrayView; -use vortex_array::Canonical; use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; @@ -23,17 +22,12 @@ use vortex_array::IntoArray; use vortex_array::TypedArrayRef; use vortex_array::VortexSessionExecute; use vortex_array::arrays::Primitive; -use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; -use vortex_array::arrays::varbin::VarBinArrayExt; use vortex_array::buffer::BufferHandle; -use vortex_array::builders::ArrayBuilder; -use vortex_array::builders::VarBinBufferBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::legacy_session; -use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::serde::ArrayChildren; use vortex_array::smallvec::smallvec; use vortex_array::validity::Validity; @@ -51,7 +45,6 @@ use crate::compress::runend_decode_primitive; use crate::compress::runend_decode_varbinview; use crate::compress::runend_encode; use crate::decompress_bool::runend_decode_bools; -use crate::iter::trimmed_ends_iter; use crate::ops::find_physical_index; use crate::ops::find_slice_end_index; use crate::rules::RULES; @@ -189,54 +182,6 @@ impl VTable for RunEnd { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { run_end_canonicalize(&array, ctx).map(ExecutionResult::done) } - - fn append_to_builder( - array: ArrayView<'_, Self>, - builder: &mut dyn ArrayBuilder, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - if matches!(array.dtype(), DType::Utf8(_) | DType::Binary(_)) - && let Some(builder) = builder.as_any_mut().downcast_mut::() - { - let ends = array - .ends() - .clone() - .execute_as::("ends", ctx)?; - let mut values_builder = VarBinBufferBuilder::with_capacity( - array.dtype().clone(), - builder.has_large_offsets(), - array.values().len(), - ); - array.values().append_to_builder(&mut values_builder, ctx)?; - let values = values_builder.finish_into_varbin(); - let validity = values - .as_ref() - .validity()? - .execute_mask(values.len(), ctx)?; - - match_each_unsigned_integer_ptype!(ends.ptype(), |E| { - let run_ends = trimmed_ends_iter(ends.as_slice::(), array.offset(), array.len()); - let mut previous_end = 0; - for (index, end) in run_ends.enumerate() { - let run_length = end - previous_end; - if validity.value(index) { - builder.append_n_values(values.bytes_at(index), run_length); - } else { - builder.append_nulls(run_length); - } - previous_end = end; - } - }); - return Ok(()); - } - - array - .array() - .clone() - .execute::(ctx)? - .into_array() - .append_to_builder(builder, ctx) - } } /// The run-end positions marking where each run terminates. diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index 282ac3e402d..9acf1b0be8f 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -26,17 +26,11 @@ use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; -use vortex_array::arrays::VarBinArray; use vortex_array::arrays::bool::BoolArrayExt; -use vortex_array::arrays::varbin::VarBinArrayExt; use vortex_array::buffer::BufferHandle; -use vortex_array::builders::ArrayBuilder; -use vortex_array::builders::VarBinBufferBuilder; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; -use vortex_array::dtype::IntegerPType; use vortex_array::dtype::Nullability; -use vortex_array::match_each_integer_ptype; use vortex_array::patches::PatchSlotIndices; use vortex_array::patches::Patches; use vortex_array::patches::PatchesData; @@ -356,91 +350,6 @@ impl VTable for Sparse { // TODO(joe): remove ctx from execute_sparse since all slots should be canonical. execute_sparse(parts, ctx).map(ExecutionResult::done) } - - fn append_to_builder( - array: ArrayView<'_, Self>, - builder: &mut dyn ArrayBuilder, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - if matches!(array.dtype(), DType::Utf8(_) | DType::Binary(_)) - && let Some(builder) = builder.as_any_mut().downcast_mut::() - { - return append_sparse_varbin(array, builder, ctx); - } - - array - .array() - .clone() - .execute::(ctx)? - .into_array() - .append_to_builder(builder, ctx) - } -} - -fn append_sparse_varbin( - array: ArrayView<'_, Sparse>, - builder: &mut VarBinBufferBuilder, - ctx: &mut ExecutionCtx, -) -> VortexResult<()> { - let patches = array.resolved_patches()?; - let indices = patches.indices().clone().execute::(ctx)?; - let mut values_builder = VarBinBufferBuilder::with_capacity( - array.dtype().clone(), - builder.has_large_offsets(), - patches.values().len(), - ); - patches - .values() - .append_to_builder(&mut values_builder, ctx)?; - let values = values_builder.finish_into_varbin(); - let validity = values - .as_ref() - .validity()? - .execute_mask(values.len(), ctx)?; - - match_each_integer_ptype!(indices.ptype(), |I| { - append_varbin_patches( - indices.as_slice::(), - &values, - &validity, - array.fill_scalar(), - array.len(), - builder, - ) - }) -} - -fn append_varbin_patches( - indices: &[I], - values: &VarBinArray, - validity: &Mask, - fill: &Scalar, - len: usize, - builder: &mut VarBinBufferBuilder, -) -> VortexResult<()> { - let mut row = 0; - let mut patch_index = 0; - while patch_index < indices.len() { - let index: usize = indices[patch_index].as_(); - let mut last_patch = patch_index; - while last_patch + 1 < indices.len() { - let next_index: usize = indices[last_patch + 1].as_(); - if next_index != index { - break; - } - last_patch += 1; - } - - builder.append_scalar_repeated(fill, index - row)?; - if validity.value(last_patch) { - builder.append_n_values(values.bytes_at(last_patch), 1); - } else { - builder.append_nulls(1); - } - row = index + 1; - patch_index = last_patch + 1; - } - builder.append_scalar_repeated(fill, len - row) } const PATCH_SLOTS: PatchSlotIndices = PatchSlotIndices { diff --git a/vortex-arrow/Cargo.toml b/vortex-arrow/Cargo.toml index 9d69a06fd7d..102a2bd7439 100644 --- a/vortex-arrow/Cargo.toml +++ b/vortex-arrow/Cargo.toml @@ -40,6 +40,9 @@ vortex-session = { workspace = true } divan = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-fsst = { workspace = true } +vortex-sparse = { workspace = true } +vortex-zstd = { workspace = true } [[bench]] name = "to_arrow" diff --git a/vortex-arrow/benches/to_arrow.rs b/vortex-arrow/benches/to_arrow.rs index 1cec274c28e..bfebde42fa0 100644 --- a/vortex-arrow/benches/to_arrow.rs +++ b/vortex-arrow/benches/to_arrow.rs @@ -6,20 +6,28 @@ use std::sync::Arc; use std::sync::LazyLock; +use arrow_schema::DataType; +use arrow_schema::Field; use divan::Bencher; +use divan::counter::ItemsCount; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; +use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::FilterArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; +use vortex_array::session::ArraySessionExt; #[expect( deprecated, reason = "benchmark comparing deprecated method with new one" @@ -28,14 +36,25 @@ use vortex_arrow::ArrowArrayExecutor; use vortex_arrow::ArrowSessionExt; #[allow(deprecated)] use vortex_arrow::dtype::ToArrowType as _; +use vortex_fsst::fsst_compress; +use vortex_fsst::fsst_train_compressor; +use vortex_mask::Mask; use vortex_session::VortexSession; +use vortex_sparse::Sparse; +use vortex_zstd::Zstd; fn main() { LazyLock::force(&SESSION); divan::main(); } -static SESSION: LazyLock = LazyLock::new(array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = array_session(); + vortex_fsst::initialize(&session); + vortex_sparse::initialize(&session); + session.arrays().register(Zstd); + session +}); fn schema() -> DType { let fields = StructFields::from_iter([ @@ -98,6 +117,175 @@ fn ArrowExportVTable_to_arrow_field(bencher: Bencher) { .bench_values(|dtype| SESSION.arrow().to_arrow_field("", &dtype).unwrap()) } +#[derive(Clone, Copy, Debug)] +enum OffsetStringEncoding { + View, + Fsst, + Sparse, + Zstd, + Dict, + DictFsst, + DictZstd, + FilterFsst, + FilterZstd, + FilterDictFsst, + ChunkedFsst, +} + +const OFFSET_STRING_ENCODINGS: &[OffsetStringEncoding] = &[ + OffsetStringEncoding::View, + OffsetStringEncoding::Fsst, + OffsetStringEncoding::Sparse, + OffsetStringEncoding::Zstd, + OffsetStringEncoding::Dict, + OffsetStringEncoding::DictFsst, + OffsetStringEncoding::DictZstd, + OffsetStringEncoding::FilterFsst, + OffsetStringEncoding::FilterZstd, + OffsetStringEncoding::FilterDictFsst, + OffsetStringEncoding::ChunkedFsst, +]; +const OFFSET_STRING_ROWS: usize = 100_000; +const OFFSET_STRING_CHUNKS: usize = 4; +const DICTIONARY_SIZE: usize = 2_048; + +fn structured_strings(len: usize) -> VarBinViewArray { + let values = (0..len) + .map(|index| format!("https://example.com/common/path/{index:06}/shared-suffix")) + .collect::>(); + VarBinViewArray::from_iter_str(values.iter().map(String::as_str)) +} + +fn dictionary_values() -> VarBinViewArray { + structured_strings(DICTIONARY_SIZE) +} + +fn dictionary_codes() -> ArrayRef { + PrimitiveArray::from_iter( + (0..OFFSET_STRING_ROWS).map(|index| u16::try_from(index % DICTIONARY_SIZE).unwrap()), + ) + .into_array() +} + +fn half_rows_mask() -> Mask { + Mask::from_iter((0..OFFSET_STRING_ROWS).map(|index| index.is_multiple_of(2))) +} + +fn filtered(array: ArrayRef) -> ArrayRef { + // Keep Filter as a lazy intermediate so benchmark setup cannot optimize it away. + FilterArray::new(array, half_rows_mask()).into_array() +} + +fn chunked_fsst(ctx: &mut vortex_array::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; + let chunks = (0..OFFSET_STRING_CHUNKS) + .map(|chunk_index| { + let start = chunk_index * chunk_size; + let end = if chunk_index + 1 == OFFSET_STRING_CHUNKS { + OFFSET_STRING_ROWS + } else { + start + chunk_size + }; + let chunk = source.slice(start..end).unwrap(); + fsst_compress(&chunk, &compressor, ctx) + .unwrap() + .into_array() + }) + .collect(); + ChunkedArray::try_new(chunks, source.dtype().clone()) + .unwrap() + .into_array() +} + +fn offset_string_array(encoding: OffsetStringEncoding) -> 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::Sparse => { + let values = (0..OFFSET_STRING_ROWS) + .map(|index| { + (index % 20 == 0) + .then(|| format!("https://example.com/sparse/{index:06}/shared-suffix")) + }) + .collect::>(); + let source = VarBinViewArray::from_iter( + values.iter().map(|value| value.as_deref()), + DType::Utf8(Nullability::Nullable), + ) + .into_array(); + let encoded = Sparse::encode(&source, None, &mut ctx).unwrap(); + assert!(encoded.is::()); + encoded + } + OffsetStringEncoding::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 => { + 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) + .unwrap() + .into_array() + } + OffsetStringEncoding::DictZstd => { + let values = dictionary_values(); + let compressed_values = + Zstd::from_var_bin_view_without_dict(&values, 3, 8_192, &mut ctx) + .unwrap() + .into_array(); + DictArray::try_new(dictionary_codes(), compressed_values) + .unwrap() + .into_array() + } + OffsetStringEncoding::FilterFsst => { + filtered(offset_string_array(OffsetStringEncoding::Fsst)) + } + OffsetStringEncoding::FilterZstd => { + filtered(offset_string_array(OffsetStringEncoding::Zstd)) + } + OffsetStringEncoding::FilterDictFsst => { + filtered(offset_string_array(OffsetStringEncoding::DictFsst)) + } + 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); + let field = Field::new("value", DataType::Utf8, array.dtype().is_nullable()); + + bencher + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .input_counter(|(array, _)| ItemsCount::new(array.len())) + .bench_values(|(array, mut ctx)| { + SESSION + .arrow() + .execute_arrow(array, Some(&field), &mut ctx) + .unwrap() + }); +} + #[divan::bench] fn to_arrow_array(bencher: Bencher) { bencher From 22940d16208e1b003d0ac90a5be8d52457d0925a Mon Sep 17 00:00:00 2001 From: cl Date: Thu, 23 Jul 2026 14:32:14 +0800 Subject: [PATCH 03/10] docs(onpair): retain decode boundary context Signed-off-by: cl --- encodings/onpair/src/canonical.rs | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/encodings/onpair/src/canonical.rs b/encodings/onpair/src/canonical.rs index e955bd3219b..5179def3bf1 100644 --- a/encodings/onpair/src/canonical.rs +++ b/encodings/onpair/src/canonical.rs @@ -61,6 +61,15 @@ pub(crate) fn onpair_decode_bytes( .sum() }); + // `codes_offsets` holds the per-row code boundaries and may itself be a + // sliced or filtered view of the original. Its first and last entries + // bound the contiguous run of `codes` belonging to the rows present in + // this array: `slice` keeps the full `codes` child and only narrows + // `codes_offsets` (so `code_start > 0` and/or `code_end < codes.len()`), + // while `filter` rebuilds both children so the window is the whole stream. + // OnPair has no `TakeExecute`, so a reordering take is served from the + // canonical `VarBinView` and never reaches this path. We only need those + // two boundaries, so point-look them up rather than decoding every offset. let codes_offsets = array.codes_offsets(); let code_start = code_boundary_at(codes_offsets, 0, ctx)?; let code_end = code_boundary_at(codes_offsets, array.len(), ctx)?; @@ -75,6 +84,10 @@ pub(crate) fn onpair_decode_bytes( array.codes().len() ); + // Slice the `codes` child to that window *before* unpacking it, so a sliced + // array materialises only its own codes rather than the whole column's. The + // contiguous decoder walks `codes` in order and never reads the per-row + // boundaries, so an empty boundary slice is sound. let codes = collect_widened::(&array.codes().slice(code_start..code_end)?, ctx)?; let dict = dict_view(array, ctx)?; let mut out_bytes = ByteBufferMut::with_capacity(total_size); @@ -90,6 +103,7 @@ pub(crate) fn onpair_decode_bytes( "OnPair codes decoded to {written} bytes but uncompressed_lengths records {total_size}" ); } + // SAFETY: `try_decode_into` initialised exactly `written` bytes. unsafe { out_bytes.set_len(written) }; Ok((out_bytes, lengths)) } From 6f03c1db6b1b3df236f1e942beb111ca6158bebb Mon Sep 17 00:00:00 2001 From: cl Date: Thu, 23 Jul 2026 14:40:04 +0800 Subject: [PATCH 04/10] docs: use Arrow array type names Signed-off-by: cl --- encodings/fsst/src/canonical.rs | 1 + vortex-array/src/arrays/varbin/builder.rs | 12 ++++++------ vortex-python/python/vortex/file.py | 4 ++-- vortex-python/src/arrays/mod.rs | 6 +++--- 4 files changed, 12 insertions(+), 11 deletions(-) diff --git a/encodings/fsst/src/canonical.rs b/encodings/fsst/src/canonical.rs index dd14090541b..9f68e130249 100644 --- a/encodings/fsst/src/canonical.rs +++ b/encodings/fsst/src/canonical.rs @@ -68,6 +68,7 @@ pub(crate) fn fsst_decode_bytes( len == total_size, "FSST decoded {len} bytes, expected {total_size}" ); + // SAFETY: `decompress_into` initialized the first `len` bytes. unsafe { uncompressed_bytes.set_len(len) }; Ok((uncompressed_bytes, uncompressed_lens_array)) } diff --git a/vortex-array/src/arrays/varbin/builder.rs b/vortex-array/src/arrays/varbin/builder.rs index 947d619769c..5a52d3f1835 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -191,11 +191,11 @@ impl VarBinBuilder { } } -/// Builder for offset-based UTF-8 and binary arrays with either 32-bit or 64-bit offsets. +/// Builder for UTF-8 and binary [`VarBinArray`] values. /// -/// Unlike [`crate::builders::VarBinViewBuilder`], this builder produces a [`VarBinArray`]. -/// Encodings can downcast an [`ArrayBuilder`] to this type and decode directly into Arrow's -/// offset-based physical representation without first materializing a view array. +/// 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 VarBinBufferBuilder { dtype: DType, storage: TypedBuilder, @@ -270,7 +270,7 @@ impl VarBinBufferBuilder { Ok(()) } - /// Appends an existing offset-based Vortex array without constructing views. + /// Appends an existing [`VarBinArray`] without constructing views. pub fn append_varbin( &mut self, array: crate::ArrayView<'_, VarBin>, @@ -296,7 +296,7 @@ impl VarBinBufferBuilder { Ok(()) } - /// Appends a Vortex view array into the offset-based output. + /// Appends a [`VarBinViewArray`](crate::arrays::VarBinViewArray). pub fn append_varbinview( &mut self, array: crate::ArrayView<'_, VarBinView>, diff --git a/vortex-python/python/vortex/file.py b/vortex-python/python/vortex/file.py index e077fdf852e..78f038f1985 100644 --- a/vortex-python/python/vortex/file.py +++ b/vortex-python/python/vortex/file.py @@ -216,8 +216,8 @@ def to_arrow( batch_size : :class:`int` | None The number of rows to read per chunk. schema : :class:`pyarrow.Schema` | None - The Arrow physical schema to return. Use ``pyarrow.string()`` and - ``pyarrow.binary()`` fields to request offset-based arrays instead of view arrays. + The Arrow schema to return. Use ``pyarrow.string()`` for ``StringArray`` fields. + Use ``pyarrow.binary()`` for ``BinaryArray`` fields. """ return self._file.to_arrow( diff --git a/vortex-python/src/arrays/mod.rs b/vortex-python/src/arrays/mod.rs index 51c8dde7033..5daf57310d3 100644 --- a/vortex-python/src/arrays/mod.rs +++ b/vortex-python/src/arrays/mod.rs @@ -435,8 +435,8 @@ impl PyArray { /// Parameters /// ---------- /// arrow_type : :class:`pyarrow.DataType`, optional - /// The Arrow physical type to return. When omitted, Vortex chooses its preferred type, - /// including ``string_view`` and ``binary_view`` for UTF-8 and binary arrays. + /// The Arrow type to return. By default, UTF-8 data returns a ``StringViewArray`` and + /// binary data returns a ``BinaryViewArray``. /// /// Returns /// ------- @@ -457,7 +457,7 @@ impl PyArray { /// 3 /// ] /// ``` - /// Export an offset-based Arrow string array instead of the default string-view array: + /// Export a ``StringArray`` instead of a ``StringViewArray``: /// /// ```python /// >>> import pyarrow From aa9c2caa8dba008dc3a09e83545680dcd56bded0 Mon Sep 17 00:00:00 2001 From: cl Date: Thu, 23 Jul 2026 14:45:29 +0800 Subject: [PATCH 05/10] test: follow Vortex naming conventions Signed-off-by: cl --- encodings/fsst/src/canonical.rs | 2 +- encodings/onpair/src/tests.rs | 2 +- encodings/sparse/src/lib.rs | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/encodings/fsst/src/canonical.rs b/encodings/fsst/src/canonical.rs index 9f68e130249..b8ff74eff84 100644 --- a/encodings/fsst/src/canonical.rs +++ b/encodings/fsst/src/canonical.rs @@ -250,7 +250,7 @@ mod tests { } #[test] - fn rejects_incorrect_uncompressed_lengths() -> VortexResult<()> { + fn test_rejects_incorrect_uncompressed_lengths() -> VortexResult<()> { let input = VarBinViewArray::from_iter_str(["hello"]).into_array(); let mut ctx = SESSION.create_execution_ctx(); let encoded = fsst_compress(&input, &fsst_train_compressor(&input, &mut ctx)?, &mut ctx)?; diff --git a/encodings/onpair/src/tests.rs b/encodings/onpair/src/tests.rs index 41c3ee73bce..ddc75fcd183 100644 --- a/encodings/onpair/src/tests.rs +++ b/encodings/onpair/src/tests.rs @@ -62,7 +62,7 @@ fn sample_input() -> VarBinArray { } #[test] -fn direct_offset_builder() -> vortex_error::VortexResult<()> { +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)?; diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index 9acf1b0be8f..e2b0192148d 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -908,7 +908,7 @@ mod test { } #[test] - fn append_utf8_to_varbin_buffer_builder() { + fn test_append_utf8_to_varbin_buffer_builder() { let mut ctx = SESSION.create_execution_ctx(); let values = VarBinViewArray::from_iter_nullable_str([Some("patched"), None, Some("last")]) .into_array(); @@ -930,7 +930,7 @@ mod test { } #[test] - fn append_utf8_with_duplicate_patch_indices() { + fn test_append_utf8_with_duplicate_patch_indices() { let mut ctx = SESSION.create_execution_ctx(); let values = VarBinViewArray::from_iter_str(["first", "second"]).into_array(); let array = Sparse::try_new( From 413501e9ef9a991bc59f57e94f30dc3689f08dca Mon Sep 17 00:00:00 2001 From: cl Date: Thu, 23 Jul 2026 18:01:23 +0800 Subject: [PATCH 06/10] test(array): cover direct varbin builder paths Signed-off-by: cl --- vortex-array/src/arrays/dict/array.rs | 18 +++++ vortex-array/src/arrays/varbin/builder.rs | 82 +++++++++++++++++++++-- 2 files changed, 95 insertions(+), 5 deletions(-) diff --git a/vortex-array/src/arrays/dict/array.rs b/vortex-array/src/arrays/dict/array.rs index cefc5a381fd..24922ea5898 100644 --- a/vortex-array/src/arrays/dict/array.rs +++ b/vortex-array/src/arrays/dict/array.rs @@ -294,7 +294,9 @@ mod test { use crate::arrays::ChunkedArray; use crate::arrays::DictArray; use crate::arrays::PrimitiveArray; + use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; + use crate::builders::VarBinBufferBuilder; use crate::builders::builder_with_capacity; use crate::dtype::DType; use crate::dtype::NativePType; @@ -440,6 +442,22 @@ mod test { .into_array() } + #[test] + fn test_dict_utf8_append_to_varbin_buffer_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 = + VarBinBufferBuilder::with_capacity(dict.dtype().clone(), false, dict.len()); + let mut ctx = array_session().create_execution_ctx(); + + dict.into_array() + .append_to_builder(&mut builder, &mut ctx)?; + + assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx); + Ok(()) + } + #[test] fn test_dict_array_from_primitive_chunks() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); diff --git a/vortex-array/src/arrays/varbin/builder.rs b/vortex-array/src/arrays/varbin/builder.rs index 5a52d3f1835..fb5e8518612 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -218,11 +218,6 @@ impl VarBinBufferBuilder { Self { dtype, storage } } - /// Returns whether this builder uses 64-bit offsets. - pub fn has_large_offsets(&self) -> bool { - matches!(self.storage, TypedBuilder::I64(_)) - } - /// Appends decompressed values represented by one contiguous byte buffer and per-row lengths. pub fn append_values( &mut self, @@ -416,12 +411,20 @@ impl ArrayBuilder for VarBinBufferBuilder { #[cfg(test)] mod tests { + use rstest::rstest; use vortex_error::VortexResult; + use vortex_mask::Mask; + use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; + use crate::arrays::VarBinArray; + use crate::arrays::VarBinViewArray; use crate::arrays::varbin::VarBinArrayExt; + use crate::arrays::varbin::builder::VarBinBufferBuilder; use crate::arrays::varbin::builder::VarBinBuilder; + use crate::assert_arrays_eq; + use crate::builders::ArrayBuilder; use crate::dtype::DType; use crate::dtype::Nullability::Nullable; use crate::expr::stats::Precision; @@ -453,6 +456,75 @@ mod tests { ); } + #[rstest] + #[case(false)] + #[case(true)] + fn test_append_varbin_to_buffer_builder(#[case] large_offsets: bool) -> VortexResult<()> { + let source = + VarBinArray::from_iter([Some("hello"), None, Some("world")], DType::Utf8(Nullable)); + let mut builder = + VarBinBufferBuilder::with_capacity(source.dtype().clone(), large_offsets, source.len()); + let mut ctx = array_session().create_execution_ctx(); + + source + .clone() + .into_array() + .append_to_builder(&mut builder, &mut ctx)?; + + assert_arrays_eq!(builder.finish_into_varbin(), source, &mut ctx); + Ok(()) + } + + #[rstest] + #[case(false)] + #[case(true)] + fn test_array_builder_methods(#[case] large_offsets: bool) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + for validity in [ + Mask::new_true(3), + Mask::new_false(3), + Mask::from_iter([true, false, true]), + ] { + let mut builder = + VarBinBufferBuilder::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(); + assert_eq!(result.validity()?.execute_mask(3, &mut ctx)?, validity); + } + Ok(()) + } + + #[test] + fn test_append_varbinview_validity_to_buffer_builder() -> VortexResult<()> { + 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)); + let expected = VarBinViewArray::from_iter( + [None, None, Some("hello"), None, Some("world")], + DType::Utf8(Nullable), + ); + let mut builder = + VarBinBufferBuilder::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)?; + + assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx); + Ok(()) + } + #[test] fn offsets_have_is_sorted_stat() -> VortexResult<()> { let mut builder = VarBinBuilder::::with_capacity(0); From 63698fa47d7b556252c67b3910b25b6b138eada3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=94=A1=E7=95=A5?= Date: Thu, 23 Jul 2026 23:48:41 +0800 Subject: [PATCH 07/10] fix(arrow): address direct byte export checks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 蔡略 --- Cargo.lock | 1 - encodings/fsst/src/array.rs | 20 +++-- encodings/onpair/src/array.rs | 20 +++-- encodings/zstd/src/array.rs | 20 +++-- vortex-array/src/arrays/varbin/builder.rs | 88 +++++++++++++++++---- vortex-arrow/Cargo.toml | 1 - vortex-arrow/benches/to_arrow.rs | 20 ----- vortex-python/python/vortex/_lib/arrays.pyi | 4 +- vortex-python/python/vortex/file.py | 4 +- 9 files changed, 112 insertions(+), 66 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index abd0e3dcba9..1a0886f7632 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9750,7 +9750,6 @@ dependencies = [ "vortex-mask", "vortex-runend", "vortex-session", - "vortex-sparse", "vortex-zstd", ] diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index 363180a4fec..0bb4ba73f7b 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -315,14 +315,18 @@ impl VTable for FSST { .validity()? .execute_mask(array.array().len(), ctx)?; match_each_integer_ptype!(lengths.ptype(), |P| { - builder.append_values( - bytes.as_slice(), - lengths - .as_slice::

() - .iter() - .map(|length| AsPrimitive::::as_(*length)), - &validity, - ); + // SAFETY: FSST decoding preserves the source values, and fsst_decode_bytes + // verifies that the decoded byte length matches the per-row lengths. + unsafe { + builder.append_values_unchecked( + bytes.as_slice(), + lengths + .as_slice::

() + .iter() + .map(|length| AsPrimitive::::as_(*length)), + &validity, + ); + } }); return Ok(()); } diff --git a/encodings/onpair/src/array.rs b/encodings/onpair/src/array.rs index b447a268781..2e33ba23e91 100644 --- a/encodings/onpair/src/array.rs +++ b/encodings/onpair/src/array.rs @@ -555,14 +555,18 @@ impl VTable for OnPair { .validity()? .execute_mask(array.array().len(), ctx)?; match_each_integer_ptype!(lengths.ptype(), |P| { - builder.append_values( - bytes.as_slice(), - lengths - .as_slice::

() - .iter() - .map(|length| AsPrimitive::::as_(*length)), - &validity, - ); + // SAFETY: OnPair decoding preserves the source values and verifies that the + // decoded byte length matches the per-row lengths. + unsafe { + builder.append_values_unchecked( + bytes.as_slice(), + lengths + .as_slice::

() + .iter() + .map(|length| AsPrimitive::::as_(*length)), + &validity, + ); + } }); return Ok(()); } diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index 53545b10955..e2843b8f8ba 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -1155,7 +1155,10 @@ impl ZstdData { match mask.indices() { AllOr::All => { for value in values { - builder.append_n_values(value, 1); + // SAFETY: Zstd stores complete source values, preserving UTF-8 validity. + unsafe { + builder.append_n_values_unchecked(value, 1); + } } } AllOr::None => builder.append_nulls(slice.n_rows), @@ -1163,12 +1166,15 @@ impl ZstdData { 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, - ); + // SAFETY: Zstd stores complete source values, preserving UTF-8 validity. + unsafe { + builder.append_n_values_unchecked( + values + .next() + .vortex_expect("Zstd value count must match validity"), + 1, + ); + } row = valid_index + 1; } builder.append_nulls(slice.n_rows - row); diff --git a/vortex-array/src/arrays/varbin/builder.rs b/vortex-array/src/arrays/varbin/builder.rs index fb5e8518612..7e9d5871e50 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -113,9 +113,24 @@ impl VarBinBuilder { fn append_values_with_lengths( &mut self, values: &[u8], - lengths: impl Iterator, + lengths: impl Iterator + Clone, validity: &Mask, ) { + let total_len = lengths + .clone() + .try_fold(0usize, |total, len| total.checked_add(len)) + .unwrap_or_else(|| vortex_panic!("VarBin byte length overflow")); + let final_end = self + .data + .len() + .checked_add(total_len) + .unwrap_or_else(|| vortex_panic!("VarBin byte offset overflow")); + O::from(final_end).unwrap_or_else(|| { + vortex_panic!( + "Failed to convert byte offset {final_end} to {}", + std::any::type_name::() + ) + }); let mut end = self.data.len(); let mut len = 0; for value_len in lengths { @@ -219,10 +234,16 @@ impl VarBinBufferBuilder { } /// Appends decompressed values represented by one contiguous byte buffer and per-row lengths. - pub fn append_values( + /// + /// # Safety + /// + /// `lengths` must contain exactly `validity.len()` items whose sum is `values.len()`. + /// Non-nullable builders require an all-valid mask. For UTF-8 builders, every valid value + /// identified by `lengths` must be valid UTF-8. + pub unsafe fn append_values_unchecked( &mut self, values: &[u8], - lengths: impl Iterator, + lengths: impl Iterator + Clone, validity: &Mask, ) { match &mut self.storage { @@ -236,7 +257,11 @@ impl VarBinBufferBuilder { } /// Appends the same non-null value `n` times. - pub fn append_n_values(&mut self, value: impl AsRef<[u8]>, n: usize) { + /// + /// # Safety + /// + /// For a UTF-8 builder, `value` must be valid UTF-8. + pub unsafe fn append_n_values_unchecked(&mut self, value: impl AsRef<[u8]>, n: usize) { let value = value.as_ref(); for _ in 0..n { self.append_value(value); @@ -253,11 +278,13 @@ impl VarBinBufferBuilder { ); match self.dtype() { DType::Utf8(_) => match scalar.as_utf8().value() { - Some(value) => self.append_n_values(value, n), + // SAFETY: Utf8Scalar values are valid UTF-8. + Some(value) => unsafe { self.append_n_values_unchecked(value, n) }, None => self.push_nulls(n), }, DType::Binary(_) => match scalar.as_binary().value() { - Some(value) => self.append_n_values(value, n), + // SAFETY: Binary builders accept arbitrary bytes. + Some(value) => unsafe { self.append_n_values_unchecked(value, n) }, None => self.push_nulls(n), }, dtype => vortex_bail!("VarBinBufferBuilder cannot append scalar of dtype {dtype}"), @@ -271,6 +298,12 @@ impl VarBinBufferBuilder { array: crate::ArrayView<'_, VarBin>, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { + vortex_ensure!( + array.dtype() == self.dtype(), + "VarBinBufferBuilder expected array with dtype {}, got {}", + self.dtype(), + array.dtype() + ); let offsets = array.offsets().clone().execute::(ctx)?; let bytes: ByteBuffer = array.sliced_bytes(); let validity = array @@ -278,15 +311,18 @@ impl VarBinBufferBuilder { .execute_mask(array.as_ref().len(), ctx)?; crate::match_each_integer_ptype!(offsets.ptype(), |P| { let offsets = offsets.as_slice::

(); - self.append_values( - bytes.as_slice(), - offsets.windows(2).map(|window| { - let start: usize = window[0].as_(); - let end: usize = window[1].as_(); - end - start - }), - &validity, - ); + // SAFETY: VarBinArray guarantees matching offsets, bytes, validity, and UTF-8. + unsafe { + self.append_values_unchecked( + bytes.as_slice(), + offsets.windows(2).map(|window| { + let start: usize = window[0].as_(); + let end: usize = window[1].as_(); + end - start + }), + &validity, + ); + } }); Ok(()) } @@ -297,6 +333,12 @@ impl VarBinBufferBuilder { array: crate::ArrayView<'_, VarBinView>, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { + vortex_ensure!( + array.dtype() == self.dtype(), + "VarBinBufferBuilder expected array with dtype {}, got {}", + self.dtype(), + array.dtype() + ); let validity = array .varbinview_validity() .execute_mask(array.as_ref().len(), ctx)?; @@ -475,6 +517,22 @@ mod tests { Ok(()) } + #[test] + fn test_safe_append_rejects_dtype_mismatch() { + let mut builder = VarBinBufferBuilder::with_capacity(DType::Utf8(Nullable), false, 1); + let mut ctx = array_session().create_execution_ctx(); + let varbin = VarBinArray::from_iter([Some(b"hello".as_slice())], DType::Binary(Nullable)); + assert!(builder.append_varbin(varbin.as_view(), &mut ctx).is_err()); + + let varbinview = + VarBinViewArray::from_iter([Some(b"hello".as_slice())], DType::Binary(Nullable)); + assert!( + builder + .append_varbinview(varbinview.as_view(), &mut ctx) + .is_err() + ); + } + #[rstest] #[case(false)] #[case(true)] diff --git a/vortex-arrow/Cargo.toml b/vortex-arrow/Cargo.toml index 102a2bd7439..540deb58a2f 100644 --- a/vortex-arrow/Cargo.toml +++ b/vortex-arrow/Cargo.toml @@ -41,7 +41,6 @@ divan = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } vortex-fsst = { workspace = true } -vortex-sparse = { workspace = true } vortex-zstd = { workspace = true } [[bench]] diff --git a/vortex-arrow/benches/to_arrow.rs b/vortex-arrow/benches/to_arrow.rs index bfebde42fa0..4a03cf3c6a9 100644 --- a/vortex-arrow/benches/to_arrow.rs +++ b/vortex-arrow/benches/to_arrow.rs @@ -40,7 +40,6 @@ use vortex_fsst::fsst_compress; use vortex_fsst::fsst_train_compressor; use vortex_mask::Mask; use vortex_session::VortexSession; -use vortex_sparse::Sparse; use vortex_zstd::Zstd; fn main() { @@ -51,7 +50,6 @@ fn main() { static SESSION: LazyLock = LazyLock::new(|| { let session = array_session(); vortex_fsst::initialize(&session); - vortex_sparse::initialize(&session); session.arrays().register(Zstd); session }); @@ -121,7 +119,6 @@ fn ArrowExportVTable_to_arrow_field(bencher: Bencher) { enum OffsetStringEncoding { View, Fsst, - Sparse, Zstd, Dict, DictFsst, @@ -135,7 +132,6 @@ enum OffsetStringEncoding { const OFFSET_STRING_ENCODINGS: &[OffsetStringEncoding] = &[ OffsetStringEncoding::View, OffsetStringEncoding::Fsst, - OffsetStringEncoding::Sparse, OffsetStringEncoding::Zstd, OffsetStringEncoding::Dict, OffsetStringEncoding::DictFsst, @@ -210,22 +206,6 @@ fn offset_string_array(encoding: OffsetStringEncoding) -> ArrayRef { .unwrap() .into_array() } - OffsetStringEncoding::Sparse => { - let values = (0..OFFSET_STRING_ROWS) - .map(|index| { - (index % 20 == 0) - .then(|| format!("https://example.com/sparse/{index:06}/shared-suffix")) - }) - .collect::>(); - let source = VarBinViewArray::from_iter( - values.iter().map(|value| value.as_deref()), - DType::Utf8(Nullability::Nullable), - ) - .into_array(); - let encoded = Sparse::encode(&source, None, &mut ctx).unwrap(); - assert!(encoded.is::()); - encoded - } OffsetStringEncoding::Zstd => { let source = structured_strings(OFFSET_STRING_ROWS); Zstd::from_var_bin_view_without_dict(&source, 3, 8_192, &mut ctx) diff --git a/vortex-python/python/vortex/_lib/arrays.pyi b/vortex-python/python/vortex/_lib/arrays.pyi index 44c7c824ac9..a0c8317e9db 100644 --- a/vortex-python/python/vortex/_lib/arrays.pyi +++ b/vortex-python/python/vortex/_lib/arrays.pyi @@ -25,9 +25,7 @@ class Array: ) -> Array: ... @staticmethod def from_range(obj: range, *, dtype: DType | None = None) -> Array: ... - def to_arrow_array( - self, *, arrow_type: pa.DataType | None = None - ) -> pa.Array[pa.Scalar[pa.DataType]]: ... + def to_arrow_array(self, *, arrow_type: pa.DataType | None = None) -> pa.Array[pa.Scalar[pa.DataType]]: ... def __vortex_array_metadata__(self) -> tuple[str, bytes, int, bytes, list[object], list[object]]: ... @property def id(self) -> str: ... diff --git a/vortex-python/python/vortex/file.py b/vortex-python/python/vortex/file.py index 78f038f1985..90336338e84 100644 --- a/vortex-python/python/vortex/file.py +++ b/vortex-python/python/vortex/file.py @@ -220,9 +220,7 @@ def to_arrow( Use ``pyarrow.binary()`` for ``BinaryArray`` fields. """ - return self._file.to_arrow( - projection, expr=expr, limit=limit, batch_size=batch_size, schema=schema - ) + return self._file.to_arrow(projection, expr=expr, limit=limit, batch_size=batch_size, schema=schema) def to_dataset(self) -> VortexDataset: """Scan the Vortex file using the :class:`pyarrow.dataset.Dataset` API.""" From d8ce25a2f3152edc9570a156e1ed4bef20ae3d4e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=94=A1=E7=95=A5?= Date: Thu, 23 Jul 2026 23:53:06 +0800 Subject: [PATCH 08/10] revert: drop varbin builder hardening MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 蔡略 --- encodings/fsst/src/array.rs | 20 +++--- encodings/onpair/src/array.rs | 20 +++--- encodings/zstd/src/array.rs | 20 ++---- vortex-array/src/arrays/varbin/builder.rs | 88 ++++------------------- 4 files changed, 38 insertions(+), 110 deletions(-) diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index 0bb4ba73f7b..363180a4fec 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -315,18 +315,14 @@ impl VTable for FSST { .validity()? .execute_mask(array.array().len(), ctx)?; match_each_integer_ptype!(lengths.ptype(), |P| { - // SAFETY: FSST decoding preserves the source values, and fsst_decode_bytes - // verifies that the decoded byte length matches the per-row lengths. - unsafe { - builder.append_values_unchecked( - bytes.as_slice(), - lengths - .as_slice::

() - .iter() - .map(|length| AsPrimitive::::as_(*length)), - &validity, - ); - } + builder.append_values( + bytes.as_slice(), + lengths + .as_slice::

() + .iter() + .map(|length| AsPrimitive::::as_(*length)), + &validity, + ); }); return Ok(()); } diff --git a/encodings/onpair/src/array.rs b/encodings/onpair/src/array.rs index 2e33ba23e91..b447a268781 100644 --- a/encodings/onpair/src/array.rs +++ b/encodings/onpair/src/array.rs @@ -555,18 +555,14 @@ impl VTable for OnPair { .validity()? .execute_mask(array.array().len(), ctx)?; match_each_integer_ptype!(lengths.ptype(), |P| { - // SAFETY: OnPair decoding preserves the source values and verifies that the - // decoded byte length matches the per-row lengths. - unsafe { - builder.append_values_unchecked( - bytes.as_slice(), - lengths - .as_slice::

() - .iter() - .map(|length| AsPrimitive::::as_(*length)), - &validity, - ); - } + builder.append_values( + bytes.as_slice(), + lengths + .as_slice::

() + .iter() + .map(|length| AsPrimitive::::as_(*length)), + &validity, + ); }); return Ok(()); } diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index e2843b8f8ba..53545b10955 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -1155,10 +1155,7 @@ impl ZstdData { match mask.indices() { AllOr::All => { for value in values { - // SAFETY: Zstd stores complete source values, preserving UTF-8 validity. - unsafe { - builder.append_n_values_unchecked(value, 1); - } + builder.append_n_values(value, 1); } } AllOr::None => builder.append_nulls(slice.n_rows), @@ -1166,15 +1163,12 @@ impl ZstdData { let mut row = 0; for &valid_index in valid_indices { builder.append_nulls(valid_index - row); - // SAFETY: Zstd stores complete source values, preserving UTF-8 validity. - unsafe { - builder.append_n_values_unchecked( - values - .next() - .vortex_expect("Zstd value count must match validity"), - 1, - ); - } + 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); diff --git a/vortex-array/src/arrays/varbin/builder.rs b/vortex-array/src/arrays/varbin/builder.rs index 7e9d5871e50..fb5e8518612 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -113,24 +113,9 @@ impl VarBinBuilder { fn append_values_with_lengths( &mut self, values: &[u8], - lengths: impl Iterator + Clone, + lengths: impl Iterator, validity: &Mask, ) { - let total_len = lengths - .clone() - .try_fold(0usize, |total, len| total.checked_add(len)) - .unwrap_or_else(|| vortex_panic!("VarBin byte length overflow")); - let final_end = self - .data - .len() - .checked_add(total_len) - .unwrap_or_else(|| vortex_panic!("VarBin byte offset overflow")); - O::from(final_end).unwrap_or_else(|| { - vortex_panic!( - "Failed to convert byte offset {final_end} to {}", - std::any::type_name::() - ) - }); let mut end = self.data.len(); let mut len = 0; for value_len in lengths { @@ -234,16 +219,10 @@ impl VarBinBufferBuilder { } /// Appends decompressed values represented by one contiguous byte buffer and per-row lengths. - /// - /// # Safety - /// - /// `lengths` must contain exactly `validity.len()` items whose sum is `values.len()`. - /// Non-nullable builders require an all-valid mask. For UTF-8 builders, every valid value - /// identified by `lengths` must be valid UTF-8. - pub unsafe fn append_values_unchecked( + pub fn append_values( &mut self, values: &[u8], - lengths: impl Iterator + Clone, + lengths: impl Iterator, validity: &Mask, ) { match &mut self.storage { @@ -257,11 +236,7 @@ impl VarBinBufferBuilder { } /// Appends the same non-null value `n` times. - /// - /// # Safety - /// - /// For a UTF-8 builder, `value` must be valid UTF-8. - pub unsafe fn append_n_values_unchecked(&mut self, value: impl AsRef<[u8]>, n: usize) { + 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); @@ -278,13 +253,11 @@ impl VarBinBufferBuilder { ); match self.dtype() { DType::Utf8(_) => match scalar.as_utf8().value() { - // SAFETY: Utf8Scalar values are valid UTF-8. - Some(value) => unsafe { self.append_n_values_unchecked(value, n) }, + Some(value) => self.append_n_values(value, n), None => self.push_nulls(n), }, DType::Binary(_) => match scalar.as_binary().value() { - // SAFETY: Binary builders accept arbitrary bytes. - Some(value) => unsafe { self.append_n_values_unchecked(value, n) }, + Some(value) => self.append_n_values(value, n), None => self.push_nulls(n), }, dtype => vortex_bail!("VarBinBufferBuilder cannot append scalar of dtype {dtype}"), @@ -298,12 +271,6 @@ impl VarBinBufferBuilder { array: crate::ArrayView<'_, VarBin>, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - vortex_ensure!( - array.dtype() == self.dtype(), - "VarBinBufferBuilder expected array with dtype {}, got {}", - self.dtype(), - array.dtype() - ); let offsets = array.offsets().clone().execute::(ctx)?; let bytes: ByteBuffer = array.sliced_bytes(); let validity = array @@ -311,18 +278,15 @@ impl VarBinBufferBuilder { .execute_mask(array.as_ref().len(), ctx)?; crate::match_each_integer_ptype!(offsets.ptype(), |P| { let offsets = offsets.as_slice::

(); - // SAFETY: VarBinArray guarantees matching offsets, bytes, validity, and UTF-8. - unsafe { - self.append_values_unchecked( - bytes.as_slice(), - offsets.windows(2).map(|window| { - let start: usize = window[0].as_(); - let end: usize = window[1].as_(); - end - start - }), - &validity, - ); - } + self.append_values( + bytes.as_slice(), + offsets.windows(2).map(|window| { + let start: usize = window[0].as_(); + let end: usize = window[1].as_(); + end - start + }), + &validity, + ); }); Ok(()) } @@ -333,12 +297,6 @@ impl VarBinBufferBuilder { array: crate::ArrayView<'_, VarBinView>, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - vortex_ensure!( - array.dtype() == self.dtype(), - "VarBinBufferBuilder expected array with dtype {}, got {}", - self.dtype(), - array.dtype() - ); let validity = array .varbinview_validity() .execute_mask(array.as_ref().len(), ctx)?; @@ -517,22 +475,6 @@ mod tests { Ok(()) } - #[test] - fn test_safe_append_rejects_dtype_mismatch() { - let mut builder = VarBinBufferBuilder::with_capacity(DType::Utf8(Nullable), false, 1); - let mut ctx = array_session().create_execution_ctx(); - let varbin = VarBinArray::from_iter([Some(b"hello".as_slice())], DType::Binary(Nullable)); - assert!(builder.append_varbin(varbin.as_view(), &mut ctx).is_err()); - - let varbinview = - VarBinViewArray::from_iter([Some(b"hello".as_slice())], DType::Binary(Nullable)); - assert!( - builder - .append_varbinview(varbinview.as_view(), &mut ctx) - .is_err() - ); - } - #[rstest] #[case(false)] #[case(true)] From 4157bd960bce59eaff6d9db8fe57ccf035fb8df4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=94=A1=E7=95=A5?= Date: Sat, 25 Jul 2026 18:25:13 +0800 Subject: [PATCH 09/10] fix(arrow): append directly into offset builders MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 蔡略 --- encodings/fsst/src/array.rs | 12 +- encodings/fsst/src/canonical.rs | 4 +- encodings/onpair/src/array.rs | 12 +- encodings/onpair/src/tests.rs | 4 +- encodings/runend/src/array.rs | 4 +- encodings/sparse/src/lib.rs | 8 +- encodings/zstd/src/array.rs | 76 ++++----- encodings/zstd/src/test.rs | 4 +- .../src/arrays/constant/vtable/mod.rs | 6 +- vortex-array/src/arrays/dict/array.rs | 7 +- vortex-array/src/arrays/dict/vtable/mod.rs | 6 +- vortex-array/src/arrays/varbin/builder.rs | 152 +++++++++--------- .../src/arrays/varbin/compute/filter.rs | 2 +- vortex-array/src/arrays/varbin/vtable/mod.rs | 4 +- .../src/arrays/varbinview/vtable/mod.rs | 4 +- vortex-array/src/builders/mod.rs | 2 +- vortex-arrow/src/executor/byte.rs | 4 +- 17 files changed, 153 insertions(+), 158 deletions(-) diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index 5e6b397d077..cf0f9848cf8 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -34,7 +34,7 @@ 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::VarBinBufferBuilder; +use vortex_array::builders::DynVarBinBuilder; use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -307,7 +307,7 @@ impl VTable for FSST { builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - if let Some(builder) = builder.as_any_mut().downcast_mut::() { + if let Some(builder) = builder.as_any_mut().downcast_mut::() { let (bytes, lengths) = fsst_decode_bytes(array, ctx)?; let validity = array .array() @@ -316,10 +316,10 @@ impl VTable for FSST { match_each_integer_ptype!(lengths.ptype(), |P| { builder.append_values( bytes.as_slice(), - lengths - .as_slice::

() - .iter() - .map(|length| AsPrimitive::::as_(*length)), + lengths.as_slice::

().iter().scan(0usize, |end, length| { + *end += AsPrimitive::::as_(*length); + Some(*end) + }), &validity, ); }); diff --git a/encodings/fsst/src/canonical.rs b/encodings/fsst/src/canonical.rs index 0b9c3713ca2..44288f48cf9 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::VarBinBufferBuilder; + use vortex_array::builders::DynVarBinBuilder; use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -212,7 +212,7 @@ mod tests { { let mut builder = - VarBinBufferBuilder::with_capacity(chunked_arr.dtype().clone(), false, data.len()); + DynVarBinBuilder::with_capacity(chunked_arr.dtype().clone(), false, data.len()); chunked_arr .into_array() .append_to_builder(&mut builder, &mut ctx)?; diff --git a/encodings/onpair/src/array.rs b/encodings/onpair/src/array.rs index b447a268781..b37ddc74cc7 100644 --- a/encodings/onpair/src/array.rs +++ b/encodings/onpair/src/array.rs @@ -26,7 +26,7 @@ use vortex_array::IntoArray; use vortex_array::array_slots; use vortex_array::buffer::BufferHandle; use vortex_array::builders::ArrayBuilder; -use vortex_array::builders::VarBinBufferBuilder; +use vortex_array::builders::DynVarBinBuilder; use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -548,7 +548,7 @@ impl VTable for OnPair { builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - if let Some(builder) = builder.as_any_mut().downcast_mut::() { + if let Some(builder) = builder.as_any_mut().downcast_mut::() { let (bytes, lengths) = onpair_decode_bytes(array, ctx)?; let validity = array .array() @@ -557,10 +557,10 @@ impl VTable for OnPair { match_each_integer_ptype!(lengths.ptype(), |P| { builder.append_values( bytes.as_slice(), - lengths - .as_slice::

() - .iter() - .map(|length| AsPrimitive::::as_(*length)), + lengths.as_slice::

().iter().scan(0usize, |end, length| { + *end += AsPrimitive::::as_(*length); + Some(*end) + }), &validity, ); }); diff --git a/encodings/onpair/src/tests.rs b/encodings/onpair/src/tests.rs index ddc75fcd183..61b5591f54d 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::VarBinBufferBuilder; +use vortex_array::builders::DynVarBinBuilder; 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 = VarBinBufferBuilder::with_capacity(input.dtype().clone(), false, input.len()); + let mut builder = DynVarBinBuilder::with_capacity(input.dtype().clone(), false, 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 bd7a668b285..2557f617aa7 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -504,7 +504,7 @@ mod tests { use vortex_array::arrays::DictArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; - use vortex_array::builders::VarBinBufferBuilder; + use vortex_array::builders::DynVarBinBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -562,7 +562,7 @@ mod tests { Some("c"), ]) .into_array(); - let mut builder = VarBinBufferBuilder::with_capacity(arr.dtype().clone(), false, arr.len()); + let mut builder = DynVarBinBuilder::with_capacity(arr.dtype().clone(), false, 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 e2b0192148d..fa791e3c6a9 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -699,7 +699,7 @@ mod test { use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; - use vortex_array::builders::VarBinBufferBuilder; + use vortex_array::builders::DynVarBinBuilder; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -908,7 +908,7 @@ mod test { } #[test] - fn test_append_utf8_to_varbin_buffer_builder() { + fn test_append_utf8_to_dyn_varbin_builder() { let mut ctx = SESSION.create_execution_ctx(); let values = VarBinViewArray::from_iter_nullable_str([Some("patched"), None, Some("last")]) .into_array(); @@ -924,7 +924,7 @@ mod test { ]) .into_array(); let mut builder = - VarBinBufferBuilder::with_capacity(array.dtype().clone(), false, array.len()); + DynVarBinBuilder::with_capacity(array.dtype().clone(), false, array.len()); array.append_to_builder(&mut builder, &mut ctx).unwrap(); assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx); } @@ -942,7 +942,7 @@ mod test { .unwrap(); let expected = VarBinViewArray::from_iter_str(["fill", "second"]).into_array(); let mut builder = - VarBinBufferBuilder::with_capacity(array.dtype().clone(), false, array.len()); + DynVarBinBuilder::with_capacity(array.dtype().clone(), false, 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 3f689e346c8..90efd99448f 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -30,7 +30,7 @@ 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::VarBinBufferBuilder; +use vortex_array::builders::DynVarBinBuilder; use vortex_array::dtype::DType; use vortex_array::scalar::Scalar; use vortex_array::serde::ArrayChildren; @@ -277,13 +277,42 @@ impl VTable for Zstd { ctx: &mut ExecutionCtx, ) -> VortexResult<()> { if matches!(array.dtype(), DType::Binary(_) | DType::Utf8(_)) - && let Some(builder) = builder.as_any_mut().downcast_mut::() + && let Some(builder) = builder.as_any_mut().downcast_mut::() { let unsliced_validity = child_to_validity(array.slots()[0].as_ref(), array.dtype().nullability()); - return array + let slice = array .data() - .append_varbin(array.dtype(), &unsliced_validity, builder, ctx); + .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); + } + } + return Ok(()); } array @@ -1141,45 +1170,6 @@ impl ZstdData { } } - fn append_varbin( - &self, - dtype: &DType, - unsliced_validity: &Validity, - builder: &mut VarBinBufferBuilder, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - let slice = self.decompress_slice(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); - } - } - Ok(()) - } - /// Returns the length of the array. #[inline] pub fn len(&self) -> usize { diff --git a/encodings/zstd/src/test.rs b/encodings/zstd/src/test.rs index 9552475fc17..fac5fcdc631 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::VarBinBufferBuilder; +use vortex_array::builders::DynVarBinBuilder; 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 = - VarBinBufferBuilder::with_capacity(compressed.dtype().clone(), false, compressed.len()); + DynVarBinBuilder::with_capacity(compressed.dtype().clone(), false, 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 e3255693d42..d9131e3a933 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -33,9 +33,9 @@ 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::VarBinBufferBuilder; use crate::builders::VarBinViewBuilder; use crate::canonical::Canonical; use crate::dtype::DType; @@ -223,7 +223,7 @@ impl VTable for Constant { }); } DType::Utf8(_) => { - if let Some(builder) = builder.as_any_mut().downcast_mut::() { + if let Some(builder) = builder.as_any_mut().downcast_mut::() { builder.append_scalar_repeated(scalar, n)?; } else { append_value_or_nulls::(builder, scalar.is_null(), n, |b| { @@ -236,7 +236,7 @@ impl VTable for Constant { } } DType::Binary(_) => { - if let Some(builder) = builder.as_any_mut().downcast_mut::() { + if let Some(builder) = builder.as_any_mut().downcast_mut::() { builder.append_scalar_repeated(scalar, n)?; } else { append_value_or_nulls::(builder, scalar.is_null(), n, |b| { diff --git a/vortex-array/src/arrays/dict/array.rs b/vortex-array/src/arrays/dict/array.rs index 24922ea5898..ba1201693b6 100644 --- a/vortex-array/src/arrays/dict/array.rs +++ b/vortex-array/src/arrays/dict/array.rs @@ -296,7 +296,7 @@ mod test { use crate::arrays::PrimitiveArray; use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; - use crate::builders::VarBinBufferBuilder; + use crate::builders::DynVarBinBuilder; use crate::builders::builder_with_capacity; use crate::dtype::DType; use crate::dtype::NativePType; @@ -443,12 +443,11 @@ mod test { } #[test] - fn test_dict_utf8_append_to_varbin_buffer_builder() -> VortexResult<()> { + fn test_dict_utf8_append_to_dyn_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 = - VarBinBufferBuilder::with_capacity(dict.dtype().clone(), false, dict.len()); + let mut builder = DynVarBinBuilder::with_capacity(dict.dtype().clone(), false, dict.len()); let mut ctx = array_session().create_execution_ctx(); dict.into_array() diff --git a/vortex-array/src/arrays/dict/vtable/mod.rs b/vortex-array/src/arrays/dict/vtable/mod.rs index 36adc11a37a..60ac5c49f68 100644 --- a/vortex-array/src/arrays/dict/vtable/mod.rs +++ b/vortex-array/src/arrays/dict/vtable/mod.rs @@ -43,7 +43,7 @@ use crate::arrays::dict::compute::rules::PARENT_RULES; use crate::arrays::dict::execute::take_canonical; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; -use crate::builders::VarBinBufferBuilder; +use crate::builders::DynVarBinBuilder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -219,10 +219,10 @@ impl VTable for Dict { builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - if let Some(builder) = builder.as_any_mut().downcast_mut::() { + if let Some(builder) = builder.as_any_mut().downcast_mut::() { // Dictionary values may exceed the target's offset width even when the selected // output does not. Keep the intermediate wide and narrow only after applying codes. - let mut values_builder = VarBinBufferBuilder::with_capacity( + let mut values_builder = DynVarBinBuilder::with_capacity( array.values().dtype().clone(), true, array.values().len(), diff --git a/vortex-array/src/arrays/varbin/builder.rs b/vortex-array/src/arrays/varbin/builder.rs index a2236edf86a..6025bfcd1f1 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -100,38 +100,38 @@ impl VarBinBuilder { } #[inline] - pub fn append_values(&mut self, values: &[u8], end_offsets: impl Iterator, num: usize) - where - O: 'static, - usize: AsPrimitive, - { - self.offsets - .extend(end_offsets.map(|offset| offset + self.data.len().as_())); - self.data.extend_from_slice(values); - self.validity.append_n(true, num); - } - - fn append_values_with_lengths( + pub fn append_values

( &mut self, values: &[u8], - lengths: impl Iterator, + end_offsets: impl Iterator, validity: &Mask, - ) { - let mut end = self.data.len(); + ) where + P: AsPrimitive, + { + let data_start = self.data.len(); + let mut last_end = data_start; let mut len = 0; - for value_len in lengths { - end += value_len; - self.offsets.push(O::from(end).unwrap_or_else(|| { + self.offsets.extend(end_offsets.map(|offset| { + let relative_end = offset.as_(); + let end = data_start.checked_add(relative_end).unwrap_or_else(|| { + vortex_panic!("Byte offset overflow: {data_start} + {relative_end}") + }); + last_end = end; + len += 1; + O::from(end).unwrap_or_else(|| { vortex_panic!( "Failed to convert byte offset {end} to {}", std::any::type_name::() ) - })); - len += 1; - } + }) + })); debug_assert_eq!(len, validity.len()); - debug_assert_eq!(end - self.data.len(), values.len()); + debug_assert_eq!(last_end - data_start, values.len()); self.data.extend_from_slice(values); + self.append_validity(validity); + } + + 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), @@ -197,42 +197,41 @@ impl VarBinBuilder { /// 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 VarBinBufferBuilder { +pub struct DynVarBinBuilder { dtype: DType, - storage: TypedBuilder, + storage: DynOffsets, } -enum TypedBuilder { +enum DynOffsets { I32(VarBinBuilder), I64(VarBinBuilder), } -impl VarBinBufferBuilder { +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(_))); let storage = if large_offsets { - TypedBuilder::I64(VarBinBuilder::with_capacity(capacity)) + DynOffsets::I64(VarBinBuilder::with_capacity(capacity)) } else { - TypedBuilder::I32(VarBinBuilder::with_capacity(capacity)) + DynOffsets::I32(VarBinBuilder::with_capacity(capacity)) }; Self { dtype, storage } } - /// Appends decompressed values represented by one contiguous byte buffer and per-row lengths. - pub fn append_values( + /// Appends decompressed values represented by one contiguous byte buffer and relative end + /// offsets. + pub fn append_values

( &mut self, values: &[u8], - lengths: impl Iterator, + end_offsets: impl Iterator, validity: &Mask, - ) { + ) where + P: AsPrimitive, + { match &mut self.storage { - TypedBuilder::I32(builder) => { - builder.append_values_with_lengths(values, lengths, validity) - } - TypedBuilder::I64(builder) => { - builder.append_values_with_lengths(values, lengths, validity) - } + DynOffsets::I32(builder) => builder.append_values(values, end_offsets, validity), + DynOffsets::I64(builder) => builder.append_values(values, end_offsets, validity), } } @@ -248,7 +247,7 @@ impl VarBinBufferBuilder { pub fn append_scalar_repeated(&mut self, scalar: &Scalar, n: usize) -> VortexResult<()> { vortex_ensure!( scalar.dtype() == self.dtype(), - "VarBinBufferBuilder expected scalar with dtype {}, got {}", + "DynVarBinBuilder expected scalar with dtype {}, got {}", self.dtype(), scalar.dtype() ); @@ -261,7 +260,7 @@ impl VarBinBufferBuilder { Some(value) => self.append_n_values(value, n), None => self.push_nulls(n), }, - dtype => vortex_bail!("VarBinBufferBuilder cannot append scalar of dtype {dtype}"), + dtype => vortex_bail!("DynVarBinBuilder cannot append scalar of dtype {dtype}"), } Ok(()) } @@ -279,13 +278,13 @@ impl VarBinBufferBuilder { .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.windows(2).map(|window| { - let start: usize = window[0].as_(); - let end: usize = window[1].as_(); - end - start - }), + offsets + .iter() + .skip(1) + .map(|offset| AsPrimitive::::as_(*offset) - first_offset), &validity, ); }); @@ -324,34 +323,34 @@ impl VarBinBufferBuilder { /// Finishes the current values as a [`VarBinArray`] and resets the builder. pub fn finish_into_varbin(&mut self) -> VarBinArray { match &mut self.storage { - TypedBuilder::I32(builder) => std::mem::take(builder).finish(self.dtype.clone()), - TypedBuilder::I64(builder) => std::mem::take(builder).finish(self.dtype.clone()), + 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 { - TypedBuilder::I32(builder) => builder.append_value(value), - TypedBuilder::I64(builder) => builder.append_value(value), + DynOffsets::I32(builder) => builder.append_value(value), + DynOffsets::I64(builder) => builder.append_value(value), } } fn push_null(&mut self) { match &mut self.storage { - TypedBuilder::I32(builder) => builder.append_null(), - TypedBuilder::I64(builder) => builder.append_null(), + DynOffsets::I32(builder) => builder.append_null(), + DynOffsets::I64(builder) => builder.append_null(), } } fn push_nulls(&mut self, n: usize) { match &mut self.storage { - TypedBuilder::I32(builder) => builder.append_n_nulls(n), - TypedBuilder::I64(builder) => builder.append_n_nulls(n), + DynOffsets::I32(builder) => builder.append_n_nulls(n), + DynOffsets::I64(builder) => builder.append_n_nulls(n), } } } -impl ArrayBuilder for VarBinBufferBuilder { +impl ArrayBuilder for DynVarBinBuilder { fn as_any(&self) -> &dyn Any { self } @@ -366,8 +365,8 @@ impl ArrayBuilder for VarBinBufferBuilder { fn len(&self) -> usize { match &self.storage { - TypedBuilder::I32(builder) => builder.len(), - TypedBuilder::I64(builder) => builder.len(), + DynOffsets::I32(builder) => builder.len(), + DynOffsets::I64(builder) => builder.len(), } } @@ -387,15 +386,15 @@ impl ArrayBuilder for VarBinBufferBuilder { fn reserve_exact(&mut self, additional: usize) { match &mut self.storage { - TypedBuilder::I32(builder) => builder.reserve_exact(additional), - TypedBuilder::I64(builder) => builder.reserve_exact(additional), + DynOffsets::I32(builder) => builder.reserve_exact(additional), + DynOffsets::I64(builder) => builder.reserve_exact(additional), } } unsafe fn set_validity_unchecked(&mut self, validity: Mask) { match &mut self.storage { - TypedBuilder::I32(builder) => builder.set_validity(validity), - TypedBuilder::I64(builder) => builder.set_validity(validity), + DynOffsets::I32(builder) => builder.set_validity(validity), + DynOffsets::I64(builder) => builder.set_validity(validity), } } @@ -422,7 +421,7 @@ mod tests { use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::arrays::varbin::VarBinArraySlotsExt; - use crate::arrays::varbin::builder::VarBinBufferBuilder; + use crate::arrays::varbin::builder::DynVarBinBuilder; use crate::arrays::varbin::builder::VarBinBuilder; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; @@ -460,17 +459,24 @@ mod tests { #[rstest] #[case(false)] #[case(true)] - fn test_append_varbin_to_buffer_builder(#[case] large_offsets: bool) -> VortexResult<()> { - let source = - VarBinArray::from_iter([Some("hello"), None, Some("world")], DType::Utf8(Nullable)); + fn test_append_varbin_to_builder(#[case] large_offsets: bool) -> VortexResult<()> { + let source = VarBinArray::from_iter( + [ + Some("prefix"), + Some("hello"), + None, + Some("world"), + Some("suffix"), + ], + DType::Utf8(Nullable), + ) + .into_array() + .slice(1..4)?; let mut builder = - VarBinBufferBuilder::with_capacity(source.dtype().clone(), large_offsets, source.len()); + DynVarBinBuilder::with_capacity(source.dtype().clone(), large_offsets, source.len()); let mut ctx = array_session().create_execution_ctx(); - source - .clone() - .into_array() - .append_to_builder(&mut builder, &mut ctx)?; + source.append_to_builder(&mut builder, &mut ctx)?; assert_arrays_eq!(builder.finish_into_varbin(), source, &mut ctx); Ok(()) @@ -487,8 +493,8 @@ mod tests { Mask::from_iter([true, false, true]), ] { let mut builder = - VarBinBufferBuilder::with_capacity(DType::Utf8(Nullable), large_offsets, 0); - assert!(builder.as_any().is::()); + 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))?; @@ -503,7 +509,7 @@ mod tests { } #[test] - fn test_append_varbinview_validity_to_buffer_builder() -> VortexResult<()> { + fn test_append_varbinview_validity_to_dyn_builder() -> VortexResult<()> { 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)); @@ -512,7 +518,7 @@ mod tests { DType::Utf8(Nullable), ); let mut builder = - VarBinBufferBuilder::with_capacity(expected.dtype().clone(), false, expected.len()); + DynVarBinBuilder::with_capacity(expected.dtype().clone(), false, expected.len()); let mut ctx = array_session().create_execution_ctx(); all_null diff --git a/vortex-array/src/arrays/varbin/compute/filter.rs b/vortex-array/src/arrays/varbin/compute/filter.rs index 97f6a615a07..dea708485a6 100644 --- a/vortex-array/src/arrays/varbin/compute/filter.rs +++ b/vortex-array/src/arrays/varbin/compute/filter.rs @@ -154,7 +154,7 @@ fn update_non_nullable_slice( .iter() .map(|o| *o - offsets[start]) .dropping(1); - builder.append_values(new_data, new_offsets, end - start) + builder.append_values(new_data, new_offsets, &Mask::new_true(end - start)) } fn filter_select_var_bin_by_index( diff --git a/vortex-array/src/arrays/varbin/vtable/mod.rs b/vortex-array/src/arrays/varbin/vtable/mod.rs index 438aa7a22d8..6a51e61e3d3 100644 --- a/vortex-array/src/arrays/varbin/vtable/mod.rs +++ b/vortex-array/src/arrays/varbin/vtable/mod.rs @@ -25,7 +25,7 @@ use crate::arrays::varbin::VarBinData; use crate::arrays::varbin::VarBinSlots; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; -use crate::builders::VarBinBufferBuilder; +use crate::builders::DynVarBinBuilder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -209,7 +209,7 @@ impl VTable for VarBin { builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - if let Some(builder) = builder.as_any_mut().downcast_mut::() { + if let Some(builder) = builder.as_any_mut().downcast_mut::() { return builder.append_varbin(array, ctx); } varbin_to_canonical(array, ctx)? diff --git a/vortex-array/src/arrays/varbinview/vtable/mod.rs b/vortex-array/src/arrays/varbinview/vtable/mod.rs index 6f731fc5613..b3ad6767f0e 100644 --- a/vortex-array/src/arrays/varbinview/vtable/mod.rs +++ b/vortex-array/src/arrays/varbinview/vtable/mod.rs @@ -29,7 +29,7 @@ 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::VarBinBufferBuilder; +use crate::builders::DynVarBinBuilder; use crate::builders::VarBinViewBuilder; use crate::dtype::DType; use crate::hash::ArrayEq; @@ -249,7 +249,7 @@ 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::() { + if let Some(builder) = builder.as_any_mut().downcast_mut::() { return builder.append_varbinview(array, ctx); } 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 343edd2aba0..ac1183177bc 100644 --- a/vortex-array/src/builders/mod.rs +++ b/vortex-array/src/builders/mod.rs @@ -75,7 +75,7 @@ pub use primitive::*; pub use struct_::*; pub use varbinview::*; -pub use crate::arrays::varbin::builder::VarBinBufferBuilder; +pub use crate::arrays::varbin::builder::DynVarBinBuilder; #[cfg(test)] mod tests; diff --git a/vortex-arrow/src/executor/byte.rs b/vortex-arrow/src/executor/byte.rs index 9ff8e24ec73..813f5d92ab5 100644 --- a/vortex-arrow/src/executor/byte.rs +++ b/vortex-arrow/src/executor/byte.rs @@ -16,7 +16,7 @@ use vortex_array::ExecutionCtx; use vortex_array::arrays::VarBin; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::varbin::VarBinArraySlotsExt; -use vortex_array::builders::VarBinBufferBuilder; +use vortex_array::builders::DynVarBinBuilder; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; @@ -53,7 +53,7 @@ where return varbin_to_byte_array::(array, ctx); } - let mut builder = VarBinBufferBuilder::with_capacity( + let mut builder = DynVarBinBuilder::with_capacity( array.dtype().clone(), T::Offset::PTYPE == PType::I64, array.len(), From 95074ff62cd55a005a43210d395f80d70eafa780 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=94=A1=E7=95=A5?= Date: Sun, 26 Jul 2026 19:14:45 +0800 Subject: [PATCH 10/10] fix(arrow): handle byte export builder errors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 蔡略 --- encodings/fsst/src/array.rs | 4 +- encodings/onpair/src/array.rs | 4 +- encodings/zstd/src/array.rs | 80 +++++++------ vortex-array/src/arrays/dict/vtable/mod.rs | 22 ---- vortex-array/src/arrays/varbin/builder.rs | 106 +++++++++++++----- .../src/arrays/varbin/compute/filter.rs | 28 +++-- vortex-arrow/src/executor/byte.rs | 28 +++++ 7 files changed, 165 insertions(+), 107 deletions(-) diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index cf0f9848cf8..6ad30b4576b 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -321,8 +321,8 @@ impl VTable for FSST { Some(*end) }), &validity, - ); - }); + ) + })?; return Ok(()); } diff --git a/encodings/onpair/src/array.rs b/encodings/onpair/src/array.rs index b37ddc74cc7..95f93d6eeae 100644 --- a/encodings/onpair/src/array.rs +++ b/encodings/onpair/src/array.rs @@ -562,8 +562,8 @@ impl VTable for OnPair { Some(*end) }), &validity, - ); - }); + ) + })?; return Ok(()); } diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index 90efd99448f..77c6866d563 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -276,51 +276,49 @@ impl VTable for Zstd { builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - if matches!(array.dtype(), DType::Binary(_) | DType::Utf8(_)) - && let Some(builder) = builder.as_any_mut().downcast_mut::() - { - 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); - } + 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); + } + 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); } - return Ok(()); } - - array - .array() - .clone() - .execute::(ctx)? - .into_array() - .append_to_builder(builder, ctx) + Ok(()) } fn reduce_parent( diff --git a/vortex-array/src/arrays/dict/vtable/mod.rs b/vortex-array/src/arrays/dict/vtable/mod.rs index 60ac5c49f68..d9ced9eeffa 100644 --- a/vortex-array/src/arrays/dict/vtable/mod.rs +++ b/vortex-array/src/arrays/dict/vtable/mod.rs @@ -5,7 +5,6 @@ use std::hash::Hasher; use prost::Message; use smallvec::smallvec; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -35,15 +34,12 @@ use crate::array::VTable; use crate::array::with_empty_buffers; use crate::arrays::ConstantArray; use crate::arrays::Primitive; -use crate::arrays::VarBin; use crate::arrays::dict::DictArrayExt; use crate::arrays::dict::DictArraySlotsExt; -use crate::arrays::dict::TakeExecute; use crate::arrays::dict::compute::rules::PARENT_RULES; use crate::arrays::dict::execute::take_canonical; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; -use crate::builders::DynVarBinBuilder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -219,24 +215,6 @@ impl VTable for Dict { builder: &mut dyn ArrayBuilder, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - if let Some(builder) = builder.as_any_mut().downcast_mut::() { - // Dictionary values may exceed the target's offset width even when the selected - // output does not. Keep the intermediate wide and narrow only after applying codes. - let mut values_builder = DynVarBinBuilder::with_capacity( - array.values().dtype().clone(), - true, - array.values().len(), - ); - array - .values() - .clone() - .append_to_builder(&mut values_builder, ctx)?; - let values = values_builder.finish_into_varbin(); - let taken = ::take(values.as_view(), array.codes(), ctx)? - .vortex_expect("taking dictionary values should produce an array"); - return builder.append_varbin(taken.as_::(), ctx); - } - if !array.is_empty() && let (Some(codes), Some(values)) = ( array.codes().as_opt::(), diff --git a/vortex-array/src/arrays/varbin/builder.rs b/vortex-array/src/arrays/varbin/builder.rs index 6025bfcd1f1..9a0a2a8922f 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -105,30 +105,56 @@ impl VarBinBuilder { values: &[u8], end_offsets: impl Iterator, validity: &Mask, - ) where + ) -> VortexResult<()> + where P: AsPrimitive, { + let offsets_start = self.offsets.len(); let data_start = self.data.len(); - let mut last_end = data_start; + let mut previous_end = data_start; let mut len = 0; - self.offsets.extend(end_offsets.map(|offset| { + + for offset in end_offsets { let relative_end = offset.as_(); - let end = data_start.checked_add(relative_end).unwrap_or_else(|| { - vortex_panic!("Byte offset overflow: {data_start} + {relative_end}") - }); - last_end = end; - len += 1; - O::from(end).unwrap_or_else(|| { - vortex_panic!( - "Failed to convert byte offset {end} to {}", + 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::() - ) - }) - })); - debug_assert_eq!(len, validity.len()); - debug_assert_eq!(last_end - data_start, values.len()); + ); + }; + 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() + ); + } + self.data.extend_from_slice(values); self.append_validity(validity); + Ok(()) } fn append_validity(&mut self, validity: &Mask) { @@ -161,14 +187,18 @@ impl VarBinBuilder { #[allow(clippy::disallowed_methods)] pub fn finish(self, dtype: DType) -> 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()); - // The builder guarantees offsets are monotonically increasing, so we can set - // this stat eagerly. This avoids an O(n) recomputation when the array is - // deserialized and VarBinArray::validate checks sortedness. + // The builder adds offsets in monotonically increasing order. Store this statistic to + // prevent VarBinArray::validate from recomputing it after deserialization. #[cfg(debug_assertions)] { let offsets_are_sorted = offsets @@ -210,7 +240,10 @@ enum DynOffsets { 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(_))); + 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 { @@ -219,14 +252,16 @@ impl DynVarBinBuilder { Self { dtype, storage } } - /// Appends decompressed values represented by one contiguous byte buffer and relative end - /// offsets. + /// Appends decompressed values from one contiguous byte buffer. + /// + /// Each offset in `end_offsets` marks the end of one value relative to the start of `values`. pub fn append_values

( &mut self, values: &[u8], end_offsets: impl Iterator, validity: &Mask, - ) where + ) -> VortexResult<()> + where P: AsPrimitive, { match &mut self.storage { @@ -286,8 +321,8 @@ impl DynVarBinBuilder { .skip(1) .map(|offset| AsPrimitive::::as_(*offset) - first_offset), &validity, - ); - }); + ) + })?; Ok(()) } @@ -482,6 +517,27 @@ mod tests { Ok(()) } + #[test] + fn append_values_offset_overflow_returns_error() { + let mut builder = VarBinBuilder::::new(); + let values = [0u8; 128]; + + let result = builder.append_values(&values, [values.len()].into_iter(), &Mask::new_true(1)); + + assert!(result.is_err()); + assert_eq!(builder.offsets.len(), 1); + assert!(builder.data.is_empty()); + assert_eq!(builder.validity.len(), 0); + } + + #[test] + #[should_panic(expected = "The offset count must be one more than the validity length")] + fn finish_rejects_mismatched_validity() { + let mut builder = VarBinBuilder::::new(); + builder.validity.append_true(); + drop(builder.finish(DType::Utf8(Nullable))); + } + #[rstest] #[case(false)] #[case(true)] diff --git a/vortex-array/src/arrays/varbin/compute/filter.rs b/vortex-array/src/arrays/varbin/compute/filter.rs index dea708485a6..a90bc87ef07 100644 --- a/vortex-array/src/arrays/varbin/compute/filter.rs +++ b/vortex-array/src/arrays/varbin/compute/filter.rs @@ -6,7 +6,6 @@ use num_traits::AsPrimitive; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; -use vortex_error::vortex_panic; use vortex_mask::AllOr; use vortex_mask::Mask; use vortex_mask::MaskIter; @@ -92,9 +91,9 @@ where let mut builder = VarBinBuilder::::with_capacity(selection_count); match logical_validity.bit_buffer() { AllOr::All => { - mask_slices.iter().for_each(|(start, end)| { - update_non_nullable_slice(data, offsets, &mut builder, *start, *end) - }); + for &(start, end) in mask_slices { + update_non_nullable_slice(data, offsets, &mut builder, start, end)?; + } } AllOr::None => { builder.append_n_nulls(selection_count); @@ -103,7 +102,7 @@ where for (start, end) in mask_slices.iter().copied() { let null_sl = validity.slice(start..end); if null_sl.true_count() == null_sl.len() { - update_non_nullable_slice(data, offsets, &mut builder, start, end) + update_non_nullable_slice(data, offsets, &mut builder, start, end)?; } else { for (idx, valid) in null_sl.iter().enumerate() { if valid { @@ -137,19 +136,18 @@ fn update_non_nullable_slice( builder: &mut VarBinBuilder, start: usize, end: usize, -) where +) -> VortexResult<()> +where O: IntegerPType, usize: AsPrimitive, { - let new_data = { - let offset_start = offsets[start].to_usize().unwrap_or_else(|| { - vortex_panic!("Failed to convert offset to usize: {}", offsets[start]) - }); - let offset_end = offsets[end].to_usize().unwrap_or_else(|| { - vortex_panic!("Failed to convert offset to usize: {}", offsets[end]) - }); - &data[offset_start..offset_end] - }; + let offset_start = offsets[start] + .to_usize() + .ok_or_else(|| vortex_err!("Failed to convert offset to usize: {}", offsets[start]))?; + let offset_end = offsets[end] + .to_usize() + .ok_or_else(|| vortex_err!("Failed to convert offset to usize: {}", offsets[end]))?; + let new_data = &data[offset_start..offset_end]; let new_offsets = offsets[start..end + 1] .iter() .map(|o| *o - offsets[start]) diff --git a/vortex-arrow/src/executor/byte.rs b/vortex-arrow/src/executor/byte.rs index 813f5d92ab5..d7a435cca91 100644 --- a/vortex-arrow/src/executor/byte.rs +++ b/vortex-arrow/src/executor/byte.rs @@ -24,6 +24,7 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_error::VortexError; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use crate::byte_view::execute_varbinview_to_arrow; use crate::executor::validity::to_arrow_null_buffer; @@ -36,6 +37,14 @@ pub(super) fn to_arrow_byte_array( where T::Offset: NativePType, { + if !matches!(array.dtype(), DType::Utf8(_) | DType::Binary(_)) { + vortex_bail!( + "Cannot convert Vortex array with dtype {} to Arrow byte array type {}", + array.dtype(), + T::DATA_TYPE + ); + } + 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 { @@ -97,6 +106,7 @@ mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; + use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -200,6 +210,24 @@ mod tests { assert!(!arrow.is_null(2)); } + #[rstest] + #[case(DataType::Utf8)] + #[case(DataType::LargeUtf8)] + #[case(DataType::Binary)] + #[case(DataType::LargeBinary)] + fn incompatible_vortex_dtype_returns_error(#[case] target_dtype: DataType) { + 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( + PrimitiveArray::from_iter([1i32, 2, 3]).into_array(), + Some(&field), + &mut ctx, + ); + + assert!(result.is_err()); + } + #[test] fn filtered_utf8_view_export_does_not_retain_unselected_buffers() -> VortexResult<()> { let unselected = "x".repeat(1 << 20);