diff --git a/Cargo.lock b/Cargo.lock index 651aa144bda..c328cf9345d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9746,9 +9746,11 @@ dependencies = [ "vortex-array", "vortex-buffer", "vortex-error", + "vortex-fsst", "vortex-mask", "vortex-runend", "vortex-session", + "vortex-zstd", ] [[package]] diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index 120d007054d..6ad30b4576b 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; @@ -33,11 +34,13 @@ use vortex_array::arrays::VarBinArray; use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::buffer::BufferHandle; use vortex_array::builders::ArrayBuilder; +use vortex_array::builders::DynVarBinBuilder; use vortex_array::builders::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::validity::Validity; use vortex_array::vtable::VTable; @@ -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; @@ -303,6 +307,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().scan(0usize, |end, length| { + *end += AsPrimitive::::as_(*length); + Some(*end) + }), + &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 b40bbf5d0bb..44288f48cf9 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; @@ -41,20 +42,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() @@ -69,14 +61,25 @@ 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}" + ); + // SAFETY: `decompress_into` initialized the first `len` bytes. 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, @@ -98,15 +101,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::DynVarBinBuilder; 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; @@ -200,6 +209,20 @@ mod tests { .collect::>(); assert_eq!(data, res2) }; + + { + let mut builder = + DynVarBinBuilder::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(()) } @@ -226,4 +249,22 @@ mod tests { let _result = builder.finish_into_varbinview(); Ok(()) } + + #[test] + 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)?; + 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..95f93d6eeae 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::DynVarBinBuilder; 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().scan(0usize, |end, length| { + *end += AsPrimitive::::as_(*length); + Some(*end) + }), + &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..5179def3bf1 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() @@ -91,7 +90,6 @@ pub(crate) fn onpair_decode_views( // 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()) { @@ -107,7 +105,15 @@ pub(crate) fn onpair_decode_views( } // 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..61b5591f54d 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::DynVarBinBuilder; 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 test_direct_offset_builder() -> vortex_error::VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = sample_input(); + let encoded = compress_onpair(input.as_ref(), &mut ctx)?; + let mut builder = DynVarBinBuilder::with_capacity(input.dtype().clone(), false, input.len()); + 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 9b7556322b6..2557f617aa7 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -504,6 +504,7 @@ mod tests { use vortex_array::arrays::DictArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; + use vortex_array::builders::DynVarBinBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -542,14 +543,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 = 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 4d4ec2623de..fa791e3c6a9 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -697,7 +697,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::DynVarBinBuilder; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -905,6 +907,46 @@ mod test { ); } + #[test] + 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(); + 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 = + 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); + } + + #[test] + 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( + 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 = + 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); + } + #[test] #[should_panic] fn test_invalid_length() { diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index 781e0df348f..77c6866d563 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -29,6 +29,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::DynVarBinBuilder; use vortex_array::dtype::DType; use vortex_array::scalar::Scalar; use vortex_array::serde::ArrayChildren; @@ -269,6 +271,56 @@ impl VTable for Zstd { .map(ExecutionResult::done) } + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let Some(builder) = builder.as_any_mut().downcast_mut::() else { + return array + .array() + .clone() + .execute::(ctx)? + .into_array() + .append_to_builder(builder, ctx); + }; + + let unsliced_validity = + child_to_validity(array.slots()[0].as_ref(), array.dtype().nullability()); + let slice = array + .data() + .decompress_slice(array.dtype(), &unsliced_validity, ctx)?; + let value_start = slice.value_idx_start - slice.n_skipped_values; + let value_count = slice.value_idx_stop - slice.value_idx_start; + let mut values = zstd_values(slice.bytes.as_slice()) + .skip(value_start) + .take(value_count); + let mask = slice.validity.execute_mask(slice.n_rows, ctx)?; + match mask.indices() { + AllOr::All => { + for value in values { + builder.append_n_values(value, 1); + } + } + AllOr::None => builder.append_nulls(slice.n_rows), + AllOr::Some(valid_indices) => { + let mut row = 0; + for &valid_index in valid_indices { + builder.append_nulls(valid_index - row); + builder.append_n_values( + values + .next() + .vortex_expect("Zstd value count must match validity"), + 1, + ); + row = valid_index + 1; + } + builder.append_nulls(slice.n_rows - row); + } + } + Ok(()) + } + fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, @@ -517,6 +569,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( @@ -900,12 +980,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); @@ -1000,29 +1080,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` @@ -1031,24 +1129,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 } @@ -1059,7 +1157,7 @@ impl ZstdData { views.freeze(), Arc::from(buffers), dtype.clone(), - slice_validity, + slice.validity, ) } .into_array()) diff --git a/encodings/zstd/src/test.rs b/encodings/zstd/src/test.rs index cf45eb97c04..fac5fcdc631 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::DynVarBinBuilder; 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 = + DynVarBinBuilder::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..d9131e3a933 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -33,6 +33,7 @@ use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::BoolBuilder; use crate::builders::DecimalBuilder; +use crate::builders::DynVarBinBuilder; use crate::builders::NullBuilder; use crate::builders::PrimitiveBuilder; use crate::builders::VarBinViewBuilder; @@ -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/array.rs b/vortex-array/src/arrays/dict/array.rs index cefc5a381fd..ba1201693b6 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::DynVarBinBuilder; use crate::builders::builder_with_capacity; use crate::dtype::DType; use crate::dtype::NativePType; @@ -440,6 +442,21 @@ mod test { .into_array() } + #[test] + 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 = DynVarBinBuilder::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 eaac8e6f600..9a0a2a8922f 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -1,24 +1,41 @@ // 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::varbin::VarBinArraySlotsExt; +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, @@ -83,27 +100,105 @@ impl VarBinBuilder { } #[inline] - pub fn append_values(&mut self, values: &[u8], end_offsets: impl Iterator, num: usize) + pub fn append_values

( + &mut self, + values: &[u8], + end_offsets: impl Iterator, + validity: &Mask, + ) -> VortexResult<()> where - O: 'static, - usize: AsPrimitive, + P: AsPrimitive, { - self.offsets - .extend(end_offsets.map(|offset| offset + self.data.len().as_())); + let offsets_start = self.offsets.len(); + let data_start = self.data.len(); + let mut previous_end = data_start; + let mut len = 0; + + for offset in end_offsets { + let relative_end = offset.as_(); + let Some(end) = data_start.checked_add(relative_end) else { + self.offsets.truncate(offsets_start); + vortex_bail!("Byte offset overflow: {data_start} + {relative_end}"); + }; + if end < previous_end { + self.offsets.truncate(offsets_start); + vortex_bail!("End offsets must be monotonically increasing"); + } + let Some(end_offset) = O::from(end) else { + self.offsets.truncate(offsets_start); + vortex_bail!( + "Byte offset {end} does not fit in {}", + std::any::type_name::() + ); + }; + self.offsets.push(end_offset); + previous_end = end; + len += 1; + } + + if len != validity.len() { + self.offsets.truncate(offsets_start); + vortex_bail!( + "End offset count {len} does not match validity length {}", + validity.len() + ); + } + if previous_end - data_start != values.len() { + self.offsets.truncate(offsets_start); + vortex_bail!( + "Final relative end offset {} does not match values length {}", + previous_end - data_start, + values.len() + ); + } + self.data.extend_from_slice(values); - self.validity.append_n(true, num); + self.append_validity(validity); + Ok(()) + } + + fn append_validity(&mut self, validity: &Mask) { + match validity { + Mask::AllTrue(len) => self.validity.append_n(true, *len), + Mask::AllFalse(len) => self.validity.append_n(false, *len), + Mask::Values(values) => self.validity.append_buffer(values.bit_buffer()), + } + } + + 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 { + 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 @@ -127,14 +222,244 @@ impl VarBinBuilder { } } +/// Builder for UTF-8 and binary [`VarBinArray`] values. +/// +/// Unlike [`crate::builders::VarBinViewBuilder`], this builder stores 32-bit or 64-bit offsets. +/// Encodings can decode into this builder without first creating a +/// [`VarBinViewArray`](crate::arrays::VarBinViewArray). +pub struct DynVarBinBuilder { + dtype: DType, + storage: DynOffsets, +} + +enum DynOffsets { + I32(VarBinBuilder), + I64(VarBinBuilder), +} + +impl DynVarBinBuilder { + /// Creates a builder with 32-bit offsets, or 64-bit offsets when `large_offsets` is true. + pub fn with_capacity(dtype: DType, large_offsets: bool, capacity: usize) -> Self { + assert!( + matches!(dtype, DType::Utf8(_) | DType::Binary(_)), + "DynVarBinBuilder dtype must be Utf8 or Binary" + ); + let storage = if large_offsets { + DynOffsets::I64(VarBinBuilder::with_capacity(capacity)) + } else { + DynOffsets::I32(VarBinBuilder::with_capacity(capacity)) + }; + Self { dtype, storage } + } + + /// 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, + ) -> VortexResult<()> + where + P: AsPrimitive, + { + match &mut self.storage { + DynOffsets::I32(builder) => builder.append_values(values, end_offsets, validity), + DynOffsets::I64(builder) => builder.append_values(values, end_offsets, validity), + } + } + + /// Appends the same non-null value `n` times. + pub fn append_n_values(&mut self, value: impl AsRef<[u8]>, n: usize) { + let value = value.as_ref(); + for _ in 0..n { + self.append_value(value); + } + } + + /// Appends the same UTF-8 or binary scalar `n` times. + pub fn append_scalar_repeated(&mut self, scalar: &Scalar, n: usize) -> VortexResult<()> { + vortex_ensure!( + scalar.dtype() == self.dtype(), + "DynVarBinBuilder expected scalar with dtype {}, got {}", + self.dtype(), + scalar.dtype() + ); + match self.dtype() { + DType::Utf8(_) => match scalar.as_utf8().value() { + Some(value) => self.append_n_values(value, n), + None => self.push_nulls(n), + }, + DType::Binary(_) => match scalar.as_binary().value() { + Some(value) => self.append_n_values(value, n), + None => self.push_nulls(n), + }, + dtype => vortex_bail!("DynVarBinBuilder cannot append scalar of dtype {dtype}"), + } + Ok(()) + } + + /// Appends an existing [`VarBinArray`] without constructing views. + pub fn append_varbin( + &mut self, + array: crate::ArrayView<'_, VarBin>, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let offsets = array.offsets().clone().execute::(ctx)?; + let bytes: ByteBuffer = array.sliced_bytes(); + let validity = array + .varbin_validity() + .execute_mask(array.as_ref().len(), ctx)?; + crate::match_each_integer_ptype!(offsets.ptype(), |P| { + let offsets = offsets.as_slice::

(); + let first_offset: usize = offsets[0].as_(); + self.append_values( + bytes.as_slice(), + offsets + .iter() + .skip(1) + .map(|offset| AsPrimitive::::as_(*offset) - first_offset), + &validity, + ) + })?; + Ok(()) + } + + /// Appends a [`VarBinViewArray`](crate::arrays::VarBinViewArray). + 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 { + DynOffsets::I32(builder) => std::mem::take(builder).finish(self.dtype.clone()), + DynOffsets::I64(builder) => std::mem::take(builder).finish(self.dtype.clone()), + } + } + + fn append_value(&mut self, value: impl AsRef<[u8]>) { + match &mut self.storage { + DynOffsets::I32(builder) => builder.append_value(value), + DynOffsets::I64(builder) => builder.append_value(value), + } + } + + fn push_null(&mut self) { + match &mut self.storage { + DynOffsets::I32(builder) => builder.append_null(), + DynOffsets::I64(builder) => builder.append_null(), + } + } + + fn push_nulls(&mut self, n: usize) { + match &mut self.storage { + DynOffsets::I32(builder) => builder.append_n_nulls(n), + DynOffsets::I64(builder) => builder.append_n_nulls(n), + } + } +} + +impl ArrayBuilder for DynVarBinBuilder { + 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 { + DynOffsets::I32(builder) => builder.len(), + DynOffsets::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 { + 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 { + DynOffsets::I32(builder) => builder.set_validity(validity), + DynOffsets::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 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::VarBinArraySlotsExt; + use crate::arrays::varbin::builder::DynVarBinBuilder; 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; @@ -166,6 +491,103 @@ mod tests { ); } + #[rstest] + #[case(false)] + #[case(true)] + 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 = + DynVarBinBuilder::with_capacity(source.dtype().clone(), large_offsets, source.len()); + let mut ctx = array_session().create_execution_ctx(); + + source.append_to_builder(&mut builder, &mut ctx)?; + + assert_arrays_eq!(builder.finish_into_varbin(), source, &mut ctx); + 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)] + 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 = + DynVarBinBuilder::with_capacity(DType::Utf8(Nullable), large_offsets, 0); + assert!(builder.as_any().is::()); + builder.reserve_exact(3); + builder.append_zero(); + builder.append_scalar(&Scalar::utf8("hello", Nullable))?; + builder.append_null(); + assert_eq!(builder.len(), 3); + builder.set_validity(validity.clone()); + + let result = builder.finish_into_canonical(&mut ctx).into_array(); + assert_eq!(result.validity()?.execute_mask(3, &mut ctx)?, validity); + } + Ok(()) + } + + #[test] + 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)); + let expected = VarBinViewArray::from_iter( + [None, None, Some("hello"), None, Some("world")], + DType::Utf8(Nullable), + ); + let mut builder = + DynVarBinBuilder::with_capacity(expected.dtype().clone(), false, expected.len()); + let mut ctx = array_session().create_execution_ctx(); + + all_null + .into_array() + .append_to_builder(&mut builder, &mut ctx)?; + mixed + .into_array() + .append_to_builder(&mut builder, &mut ctx)?; + + 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); diff --git a/vortex-array/src/arrays/varbin/compute/filter.rs b/vortex-array/src/arrays/varbin/compute/filter.rs index 97f6a615a07..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,24 +136,23 @@ 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]) .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 caf88b55358..6a51e61e3d3 100644 --- a/vortex-array/src/arrays/varbin/vtable/mod.rs +++ b/vortex-array/src/arrays/varbin/vtable/mod.rs @@ -24,6 +24,8 @@ use crate::arrays::varbin::VarBinArraySlotsExt; use crate::arrays::varbin::VarBinData; use crate::arrays::varbin::VarBinSlots; use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::builders::DynVarBinBuilder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -202,6 +204,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 f6659e03ee4..b3ad6767f0e 100644 --- a/vortex-array/src/arrays/varbinview/vtable/mod.rs +++ b/vortex-array/src/arrays/varbinview/vtable/mod.rs @@ -29,6 +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::DynVarBinBuilder; 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..ac1183177bc 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::DynVarBinBuilder; + #[cfg(test)] mod tests; diff --git a/vortex-arrow/Cargo.toml b/vortex-arrow/Cargo.toml index 9d69a06fd7d..540deb58a2f 100644 --- a/vortex-arrow/Cargo.toml +++ b/vortex-arrow/Cargo.toml @@ -40,6 +40,8 @@ vortex-session = { workspace = true } divan = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-fsst = { 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..4a03cf3c6a9 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,23 @@ 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_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); + session.arrays().register(Zstd); + session +}); fn schema() -> DType { let fields = StructFields::from_iter([ @@ -98,6 +115,157 @@ 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, + Zstd, + Dict, + DictFsst, + DictZstd, + FilterFsst, + FilterZstd, + FilterDictFsst, + ChunkedFsst, +} + +const OFFSET_STRING_ENCODINGS: &[OffsetStringEncoding] = &[ + OffsetStringEncoding::View, + OffsetStringEncoding::Fsst, + OffsetStringEncoding::Zstd, + OffsetStringEncoding::Dict, + OffsetStringEncoding::DictFsst, + OffsetStringEncoding::DictZstd, + OffsetStringEncoding::FilterFsst, + OffsetStringEncoding::FilterZstd, + OffsetStringEncoding::FilterDictFsst, + OffsetStringEncoding::ChunkedFsst, +]; +const 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::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 diff --git a/vortex-arrow/src/executor/byte.rs b/vortex-arrow/src/executor/byte.rs index e4b37f84da3..d7a435cca91 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,12 +16,15 @@ use vortex_array::ExecutionCtx; use vortex_array::arrays::VarBin; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::varbin::VarBinArraySlotsExt; +use vortex_array::builders::DynVarBinBuilder; 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; +use vortex_error::vortex_bail; use crate::byte_view::execute_varbinview_to_arrow; use crate::executor::validity::to_arrow_null_buffer; @@ -33,20 +37,38 @@ 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 { + 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 = DynVarBinBuilder::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 +106,14 @@ 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; 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"]) @@ -187,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); diff --git a/vortex-python/python/vortex/_lib/arrays.pyi b/vortex-python/python/vortex/_lib/arrays.pyi index 8da62190291..a0c8317e9db 100644 --- a/vortex-python/python/vortex/_lib/arrays.pyi +++ b/vortex-python/python/vortex/_lib/arrays.pyi @@ -25,7 +25,7 @@ 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..90336338e84 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,12 @@ 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 schema to return. Use ``pyarrow.string()`` for ``StringArray`` fields. + Use ``pyarrow.binary()`` for ``BinaryArray`` fields. """ - 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..5daf57310d3 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 type to return. By default, UTF-8 data returns a ``StringViewArray`` and + /// binary data returns a ``BinaryViewArray``. + /// /// Returns /// ------- /// :class:`pyarrow.Array` @@ -448,24 +457,50 @@ impl PyArray { /// 3 /// ] /// ``` - fn to_arrow_array<'py>(self_: &'py Bound<'py, Self>) -> PyVortexResult> { + /// Export a ``StringArray`` instead of a ``StringViewArray``: + /// + /// ```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`