diff --git a/arrow-array/src/builder/generic_bytes_view_builder.rs b/arrow-array/src/builder/generic_bytes_view_builder.rs index 6ff209ef7445..3c21b8624b99 100644 --- a/arrow-array/src/builder/generic_bytes_view_builder.rs +++ b/arrow-array/src/builder/generic_bytes_view_builder.rs @@ -224,6 +224,10 @@ impl GenericByteViewBuilder { /// Appends an array to the builder. /// This will flush any in-progress block and append the data buffers /// and add the (adapted) views. + /// + /// # Panics + /// + /// Panics if the builder would hold more than `u32::MAX` data buffers pub fn append_array(&mut self, array: &GenericByteViewArray) { self.flush_in_progress(); // keep original views if this array is the first to be added or if there are no data buffers (all inline views) @@ -231,18 +235,28 @@ impl GenericByteViewBuilder { let starting_buffer = self.completed.len() as u32; self.completed.extend(array.data_buffers().iter().cloned()); + assert!( + u32::try_from(self.completed.len()).is_ok(), + "Too many blocks" + ); if keep_views { self.views_buffer.extend_from_slice(array.views()); } else { self.views_buffer.extend(array.views().iter().map(|v| { - let mut byte_view = ByteView::from(*v); - if byte_view.length > MAX_INLINE_VIEW_LEN { - // Small views (<=12 bytes) are inlined, so only need to update large views - byte_view.buffer_index += starting_buffer; - } - - byte_view.as_u128() + // A view of a string of 12 bytes or fewer stores the string + // inline and has no buffer index, so adding 0 leaves it + // unchanged. + let offset = if *v as u32 > MAX_INLINE_VIEW_LEN { + starting_buffer + } else { + 0 + }; + // The buffer index is bits 64 to 95 of a view. Adding + // `starting_buffer` gives the string's buffer index in the + // builder, which fits in 32 bits, so the sum never carries + // into bit 96. + v + ((offset as u128) << 64) })); } diff --git a/arrow-select/src/coalesce/byte_view.rs b/arrow-select/src/coalesce/byte_view.rs index a6697bfdb599..ca3d2280a141 100644 --- a/arrow-select/src/coalesce/byte_view.rs +++ b/arrow-select/src/coalesce/byte_view.rs @@ -194,8 +194,12 @@ impl InProgressByteViewArray { self.size_of_completed_buffers_from_current_source += buffers_size; } - let starting_buffer: u32 = self.completed.len().try_into().expect("too many buffers"); + let starting_buffer = self.completed.len() as u32; self.completed.extend_from_slice(buffers); + assert!( + u32::try_from(self.completed.len()).is_ok(), + "too many buffers" + ); if starting_buffer == 0 { // If there are no buffers, we can just use the views as is @@ -203,12 +207,19 @@ impl InProgressByteViewArray { } else { // If there are buffers, we need to update the buffer index let updated_views = views.iter().map(|v| { - let mut byte_view = ByteView::from(*v); - if byte_view.length > MAX_INLINE_VIEW_LEN { - // Small views (<=12 bytes) are inlined, so only need to update large views - byte_view.buffer_index += starting_buffer; - } - byte_view.as_u128() + // A view of a string of 12 bytes or fewer stores the string + // inline and has no buffer index, so adding 0 leaves it + // unchanged. + let offset = if *v as u32 > MAX_INLINE_VIEW_LEN { + starting_buffer + } else { + 0 + }; + // The buffer index is bits 64 to 95 of a view. Adding + // `starting_buffer` gives the string's buffer index in the + // output, which fits in 32 bits, so the sum never carries into + // bit 96. + v + ((offset as u128) << 64) }); self.views.extend(updated_views); diff --git a/arrow/benches/coalesce_kernels.rs b/arrow/benches/coalesce_kernels.rs index 9cdd71cdc3bc..1f0c524cb6c5 100644 --- a/arrow/benches/coalesce_kernels.rs +++ b/arrow/benches/coalesce_kernels.rs @@ -1447,10 +1447,17 @@ impl DataStreamBuilder { ) .with_timezone(Arc::clone(tz)), ), - DataType::Utf8 => Arc::new(create_string_array::( - self.batch_size, - self.null_density, - )), // TODO seed + // Strings of 0 to 399 bytes; `max_string_len` doesn't apply + DataType::Utf8 => Arc::new( + create_string_array_with_len_range_and_prefix_and_seed::( + self.batch_size, + self.null_density, + 0, + 399, + "", + seed, + ), + ), DataType::Binary => Arc::new(create_binary_array_with_len_range_and_prefix_and_seed::< i32, >( @@ -1482,11 +1489,18 @@ impl DataStreamBuilder { if key_type.as_ref() == &DataType::Int32 && value_type.as_ref() == &DataType::Utf8 => { - Arc::new(create_string_dict_array::( - self.batch_size, - self.null_density, - self.max_string_len, - )) // TODO seed + Arc::new( + create_string_array_with_len_range_and_prefix_and_seed::( + self.batch_size, + self.null_density, + self.max_string_len, + self.max_string_len, + "", + seed, + ) + .iter() + .collect::>(), + ) } DataType::FixedSizeBinary(size) => Arc::new(create_fixed_size_binary_array( self.batch_size,