From c3648f4add54d9cdf5a3a9a52084eee2d176e08e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Thu, 5 Jun 2025 19:48:38 +0200 Subject: [PATCH 01/19] Improve coalesce --- arrow-select/src/coalesce.rs | 90 +++++++++++++++++++++++------------- 1 file changed, 59 insertions(+), 31 deletions(-) diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs index 5a5e7e141c84..20012fbbe983 100644 --- a/arrow-select/src/coalesce.rs +++ b/arrow-select/src/coalesce.rs @@ -21,9 +21,9 @@ //! [`filter`]: crate::filter::filter //! [`take`]: crate::take::take use crate::concat::concat_batches; -use arrow_array::{ - builder::StringViewBuilder, cast::AsArray, Array, ArrayRef, RecordBatch, RecordBatchOptions, -}; +use arrow_array::StringViewArray; +use arrow_array::{cast::AsArray, Array, ArrayRef, RecordBatch, RecordBatchOptions}; +use arrow_data::ByteView; use arrow_schema::{ArrowError, SchemaRef}; use std::collections::VecDeque; use std::sync::Arc; @@ -164,7 +164,7 @@ impl BatchCoalescer { return Ok(()); } - let mut batch = gc_string_view_batch(&batch); + let mut batch = gc_string_view_batch(batch); // If pushing this batch would exceed the target batch size, // finish the current batch and start a new one @@ -242,14 +242,14 @@ impl BatchCoalescer { /// However, after a while (e.g., after `FilterExec` or `HashJoinExec`) the /// `StringViewArray` may only refer to a small portion of the buffer, /// significantly increasing memory usage. -fn gc_string_view_batch(batch: &RecordBatch) -> RecordBatch { - let new_columns: Vec = batch - .columns() - .iter() +fn gc_string_view_batch(batch: RecordBatch) -> RecordBatch { + let (schema, columns, num_rows) = batch.into_parts(); + let new_columns: Vec = columns + .into_iter() .map(|c| { // Try to re-create the `StringViewArray` to prevent holding the underlying buffer too long. let Some(s) = c.as_string_view_opt() else { - return Arc::clone(c); + return c; }; let ideal_buffer_size: usize = s .views() @@ -264,41 +264,65 @@ fn gc_string_view_batch(batch: &RecordBatch) -> RecordBatch { }) .sum(); let actual_buffer_size = s.get_buffer_memory_size(); + let buffers = s.data_buffers(); // Re-creating the array copies data and can be time consuming. // We only do it if the array is sparse if actual_buffer_size > (ideal_buffer_size * 2) { // We set the block size to `ideal_buffer_size` so that the new StringViewArray only has one buffer, which accelerate later concat_batches. // See https://github.com/apache/arrow-rs/issues/6094 for more details. - let mut builder = StringViewBuilder::with_capacity(s.len()); - if ideal_buffer_size > 0 { - builder = builder.with_fixed_block_size(ideal_buffer_size as u32); - } - - for v in s.iter() { - builder.append_option(v); - } - - let gc_string = builder.finish(); - - debug_assert!(gc_string.data_buffers().len() <= 1); // buffer count can be 0 if the `ideal_buffer_size` is 0 + let mut buffer: Vec = Vec::with_capacity(ideal_buffer_size); + + let views: Vec = s + .views() + .as_ref() + .iter() + .cloned() + .map(|v| { + // SAFETY: ByteView has same memory layout as u128 + let mut b: ByteView = ByteView::from(v); + + if b.length > 12 { + let offset = buffer.len() as u32; + buffer.extend_from_slice( + buffers[b.buffer_index as usize] + .get(b.offset as usize..b.offset as usize + b.length as usize) + .expect("Invalid buffer slice"), + ); + b.offset = offset; + b.buffer_index = 0; // Set buffer index to 0, as we only have one buffer + } + + b.into() + }) + .collect(); + + let buffers = if buffer.is_empty() { + vec![] + } else { + vec![buffer.into()] + }; + + let gc_string = unsafe { + StringViewArray::new_unchecked(views.into(), buffers, s.nulls().cloned()) + }; Arc::new(gc_string) } else { - Arc::clone(c) + c } }) .collect(); let mut options = RecordBatchOptions::new(); - options = options.with_row_count(Some(batch.num_rows())); - RecordBatch::try_new_with_options(batch.schema(), new_columns, &options) + options = options.with_row_count(Some(num_rows)); + RecordBatch::try_new_with_options(schema, new_columns, &options) .expect("Failed to re-create the gc'ed record batch") } #[cfg(test)] mod tests { use super::*; - use arrow_array::builder::ArrayBuilder; + use arrow_array::builder::{ArrayBuilder, StringViewBuilder}; use arrow_array::{StringViewArray, UInt32Array}; use arrow_schema::{DataType, Field, Schema}; use std::ops::Range; @@ -518,9 +542,11 @@ mod tests { fn test_gc_string_view_test_batch_empty() { let schema = Schema::empty(); let batch = RecordBatch::new_empty(schema.into()); - let output_batch = gc_string_view_batch(&batch); - assert_eq!(batch.num_columns(), output_batch.num_columns()); - assert_eq!(batch.num_rows(), output_batch.num_rows()); + let cols = batch.num_columns(); + let num_rows = batch.num_rows(); + let output_batch = gc_string_view_batch(batch); + assert_eq!(cols, output_batch.num_columns()); + assert_eq!(num_rows, output_batch.num_rows()); } #[test] @@ -568,9 +594,11 @@ mod tests { /// and ensures the number of rows are the same fn do_gc(array: StringViewArray) -> StringViewArray { let batch = RecordBatch::try_from_iter(vec![("a", Arc::new(array) as ArrayRef)]).unwrap(); - let gc_batch = gc_string_view_batch(&batch); - assert_eq!(batch.num_rows(), gc_batch.num_rows()); - assert_eq!(batch.schema(), gc_batch.schema()); + let rows = batch.num_rows(); + let schema = batch.schema(); + let gc_batch = gc_string_view_batch(batch); + assert_eq!(rows, gc_batch.num_rows()); + assert_eq!(schema, gc_batch.schema()); gc_batch .column(0) .as_any() From d9d309f4638c8f0c92e6aea26b36e94822f36932 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Fri, 6 Jun 2025 09:13:17 +0200 Subject: [PATCH 02/19] Specialize concat --- .../src/builder/generic_bytes_view_builder.rs | 14 +++++++++++- arrow-select/src/concat.rs | 22 ++++++++++++++++++- 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/arrow-array/src/builder/generic_bytes_view_builder.rs b/arrow-array/src/builder/generic_bytes_view_builder.rs index 7268e751b149..942902604c92 100644 --- a/arrow-array/src/builder/generic_bytes_view_builder.rs +++ b/arrow-array/src/builder/generic_bytes_view_builder.rs @@ -28,7 +28,7 @@ use hashbrown::HashTable; use crate::builder::ArrayBuilder; use crate::types::bytes::ByteArrayNativeType; use crate::types::{BinaryViewType, ByteViewType, StringViewType}; -use crate::{ArrayRef, GenericByteViewArray}; +use crate::{Array, ArrayRef, GenericByteViewArray}; const STARTING_BLOCK_SIZE: u32 = 8 * 1024; // 8KiB const MAX_BLOCK_SIZE: u32 = 2 * 1024 * 1024; // 2MiB @@ -205,6 +205,18 @@ impl GenericByteViewBuilder { self.null_buffer_builder.append_non_null(); } + /// Appends an array + pub fn append_array(&mut self, array: &GenericByteViewArray) { + self.flush_in_progress(); + self.completed.extend(array.data_buffers().iter().cloned()); + self.views_builder.append_slice(array.views()); + if let Some(null_buffer) = array.nulls() { + self.null_buffer_builder.append_buffer(null_buffer); + } else { + self.null_buffer_builder.append_n_non_nulls(array.len()); + } + } + /// Try to append a view of the given `block`, `offset` and `length` /// /// See [`Self::append_block`] diff --git a/arrow-select/src/concat.rs b/arrow-select/src/concat.rs index 486afbd14467..e6a2cf88d237 100644 --- a/arrow-select/src/concat.rs +++ b/arrow-select/src/concat.rs @@ -31,7 +31,9 @@ //! ``` use crate::dictionary::{merge_dictionary_values, should_merge_dictionary_values}; -use arrow_array::builder::{BooleanBuilder, GenericByteBuilder, PrimitiveBuilder}; +use arrow_array::builder::{ + BinaryViewBuilder, BooleanBuilder, GenericByteBuilder, PrimitiveBuilder, StringViewBuilder, +}; use arrow_array::cast::AsArray; use arrow_array::types::*; use arrow_array::*; @@ -84,6 +86,22 @@ fn fixed_size_list_capacity(arrays: &[&dyn Array], data_type: &DataType) -> Capa } } +fn concat_byte_view(arrays: &[&dyn Array]) -> Result { + let mut builder = BinaryViewBuilder::with_capacity(arrays.iter().map(|a| a.len()).sum()); + for &array in arrays.iter() { + builder.append_array(array.as_binary_view()); + } + Ok(Arc::new(builder.finish())) +} + +fn concat_string_view(arrays: &[&dyn Array]) -> Result { + let mut builder = StringViewBuilder::with_capacity(arrays.iter().map(|a| a.len()).sum()); + for &array in arrays.iter() { + builder.append_array(array.as_string_view()); + } + Ok(Arc::new(builder.finish())) +} + fn concat_dictionaries( arrays: &[&dyn Array], ) -> Result { @@ -425,6 +443,8 @@ pub fn concat(arrays: &[&dyn Array]) -> Result { _ => unreachable!("Unsupported run end index type: {r:?}"), } } + DataType::Utf8View => concat_string_view(arrays), + DataType::BinaryView => concat_byte_view(arrays), _ => { let capacity = get_capacity(arrays, d); concat_fallback(arrays, capacity) From cda04cf2acc93dd774b233a740f30a150d2eb94e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Fri, 6 Jun 2025 09:24:50 +0200 Subject: [PATCH 03/19] Specialize concat --- .../src/builder/generic_bytes_view_builder.rs | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/arrow-array/src/builder/generic_bytes_view_builder.rs b/arrow-array/src/builder/generic_bytes_view_builder.rs index 942902604c92..fc8db10fe508 100644 --- a/arrow-array/src/builder/generic_bytes_view_builder.rs +++ b/arrow-array/src/builder/generic_bytes_view_builder.rs @@ -209,7 +209,23 @@ impl GenericByteViewBuilder { pub fn append_array(&mut self, array: &GenericByteViewArray) { self.flush_in_progress(); self.completed.extend(array.data_buffers().iter().cloned()); - self.views_builder.append_slice(array.views()); + + if self.completed.len() == 0 { + self.views_builder.append_slice(array.views()); + } else { + let starting_buffer = self.completed.len() as u32; + + self.views_builder.extend(array.views().iter().map(|v| { + let mut byte_view = ByteView::from(*v); + if byte_view.length > 12 { + // If the view is small enough, we can inline it + byte_view.buffer_index += starting_buffer; + }; + + byte_view.into() + })); + } + if let Some(null_buffer) = array.nulls() { self.null_buffer_builder.append_buffer(null_buffer); } else { From 675d9d6566eca92b2f38ef1936e5a535afc7d63b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Fri, 6 Jun 2025 09:30:31 +0200 Subject: [PATCH 04/19] Specialize concat --- .../src/builder/generic_bytes_view_builder.rs | 25 +++++++++---------- 1 file changed, 12 insertions(+), 13 deletions(-) diff --git a/arrow-array/src/builder/generic_bytes_view_builder.rs b/arrow-array/src/builder/generic_bytes_view_builder.rs index fc8db10fe508..330d4138605d 100644 --- a/arrow-array/src/builder/generic_bytes_view_builder.rs +++ b/arrow-array/src/builder/generic_bytes_view_builder.rs @@ -19,7 +19,7 @@ use std::any::Any; use std::marker::PhantomData; use std::sync::Arc; -use arrow_buffer::{Buffer, BufferBuilder, NullBufferBuilder, ScalarBuffer}; +use arrow_buffer::{Buffer, NullBufferBuilder, ScalarBuffer}; use arrow_data::ByteView; use arrow_schema::ArrowError; use hashbrown::hash_table::Entry; @@ -79,7 +79,7 @@ impl BlockSizeGrowthStrategy { /// using [`GenericByteViewBuilder::append_block`] and then views into this block appended /// using [`GenericByteViewBuilder::try_append_view`] pub struct GenericByteViewBuilder { - views_builder: BufferBuilder, + views_builder: Vec, null_buffer_builder: NullBufferBuilder, completed: Vec, in_progress: Vec, @@ -99,7 +99,7 @@ impl GenericByteViewBuilder { /// Creates a new [`GenericByteViewBuilder`] with space for `capacity` string values. pub fn with_capacity(capacity: usize) -> Self { Self { - views_builder: BufferBuilder::new(capacity), + views_builder: Vec::with_capacity(capacity), null_buffer_builder: NullBufferBuilder::new(capacity), completed: vec![], in_progress: vec![], @@ -201,7 +201,7 @@ impl GenericByteViewBuilder { let b = b.get_unchecked(start..end); let view = make_view(b, block, offset); - self.views_builder.append(view); + self.views_builder.push(view); self.null_buffer_builder.append_non_null(); } @@ -211,7 +211,7 @@ impl GenericByteViewBuilder { self.completed.extend(array.data_buffers().iter().cloned()); if self.completed.len() == 0 { - self.views_builder.append_slice(array.views()); + self.views_builder.extend_from_slice(array.views()); } else { let starting_buffer = self.completed.len() as u32; @@ -222,7 +222,7 @@ impl GenericByteViewBuilder { byte_view.buffer_index += starting_buffer; }; - byte_view.into() + byte_view.as_u128() })); } @@ -315,7 +315,7 @@ impl GenericByteViewBuilder { let mut view_buffer = [0; 16]; view_buffer[0..4].copy_from_slice(&length.to_le_bytes()); view_buffer[4..4 + v.len()].copy_from_slice(v); - self.views_builder.append(u128::from_le_bytes(view_buffer)); + self.views_builder.push(u128::from_le_bytes(view_buffer)); self.null_buffer_builder.append_non_null(); return; } @@ -340,7 +340,7 @@ impl GenericByteViewBuilder { // If the string already exists, we will directly use the view let idx = occupied.get(); self.views_builder - .append(self.views_builder.as_slice()[*idx]); + .push(self.views_builder[*idx]); self.null_buffer_builder.append_non_null(); self.string_tracker = Some((ht, hasher)); return; @@ -369,7 +369,7 @@ impl GenericByteViewBuilder { buffer_index: self.completed.len() as u32, offset, }; - self.views_builder.append(view.into()); + self.views_builder.push(view.into()); self.null_buffer_builder.append_non_null(); } @@ -386,21 +386,20 @@ impl GenericByteViewBuilder { #[inline] pub fn append_null(&mut self) { self.null_buffer_builder.append_null(); - self.views_builder.append(0); + self.views_builder.push(0); } /// Builds the [`GenericByteViewArray`] and reset this builder pub fn finish(&mut self) -> GenericByteViewArray { self.flush_in_progress(); let completed = std::mem::take(&mut self.completed); - let len = self.views_builder.len(); - let views = ScalarBuffer::new(self.views_builder.finish(), 0, len); let nulls = self.null_buffer_builder.finish(); if let Some((ref mut ht, _)) = self.string_tracker.as_mut() { ht.clear(); } + let views = std::mem::take(&mut self.views_builder); // SAFETY: valid by construction - unsafe { GenericByteViewArray::new_unchecked(views, completed, nulls) } + unsafe { GenericByteViewArray::new_unchecked(views.into(), completed, nulls) } } /// Builds the [`GenericByteViewArray`] without resetting the builder From b44d0ab30ff999d98cac93a4146ef639c4a59a8f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Fri, 6 Jun 2025 09:55:32 +0200 Subject: [PATCH 05/19] Clippy --- arrow-array/src/builder/generic_bytes_view_builder.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/arrow-array/src/builder/generic_bytes_view_builder.rs b/arrow-array/src/builder/generic_bytes_view_builder.rs index 330d4138605d..056c20067a96 100644 --- a/arrow-array/src/builder/generic_bytes_view_builder.rs +++ b/arrow-array/src/builder/generic_bytes_view_builder.rs @@ -210,7 +210,7 @@ impl GenericByteViewBuilder { self.flush_in_progress(); self.completed.extend(array.data_buffers().iter().cloned()); - if self.completed.len() == 0 { + if self.completed.is_empty() == 0 { self.views_builder.extend_from_slice(array.views()); } else { let starting_buffer = self.completed.len() as u32; From 294e1887adc446667329ad04ba2f25850145daeb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Fri, 6 Jun 2025 09:57:01 +0200 Subject: [PATCH 06/19] Remove comment --- arrow-select/src/coalesce.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs index 20012fbbe983..642f872f48ae 100644 --- a/arrow-select/src/coalesce.rs +++ b/arrow-select/src/coalesce.rs @@ -279,7 +279,6 @@ fn gc_string_view_batch(batch: RecordBatch) -> RecordBatch { .iter() .cloned() .map(|v| { - // SAFETY: ByteView has same memory layout as u128 let mut b: ByteView = ByteView::from(v); if b.length > 12 { From 9936a8a7659f516f3729af83c86f18b25b183518 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Fri, 6 Jun 2025 10:04:04 +0200 Subject: [PATCH 07/19] Fix --- arrow-array/src/builder/generic_bytes_view_builder.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/arrow-array/src/builder/generic_bytes_view_builder.rs b/arrow-array/src/builder/generic_bytes_view_builder.rs index 056c20067a96..b4d0222474f7 100644 --- a/arrow-array/src/builder/generic_bytes_view_builder.rs +++ b/arrow-array/src/builder/generic_bytes_view_builder.rs @@ -210,7 +210,7 @@ impl GenericByteViewBuilder { self.flush_in_progress(); self.completed.extend(array.data_buffers().iter().cloned()); - if self.completed.is_empty() == 0 { + if self.completed.is_empty() { self.views_builder.extend_from_slice(array.views()); } else { let starting_buffer = self.completed.len() as u32; From abe62b97f9e53b173721a8200fae1c5bb035f617 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Fri, 6 Jun 2025 10:29:38 +0200 Subject: [PATCH 08/19] Fmt --- arrow-array/src/builder/generic_bytes_view_builder.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/arrow-array/src/builder/generic_bytes_view_builder.rs b/arrow-array/src/builder/generic_bytes_view_builder.rs index b4d0222474f7..3f5c823ab74e 100644 --- a/arrow-array/src/builder/generic_bytes_view_builder.rs +++ b/arrow-array/src/builder/generic_bytes_view_builder.rs @@ -339,8 +339,7 @@ impl GenericByteViewBuilder { Entry::Occupied(occupied) => { // If the string already exists, we will directly use the view let idx = occupied.get(); - self.views_builder - .push(self.views_builder[*idx]); + self.views_builder.push(self.views_builder[*idx]); self.null_buffer_builder.append_non_null(); self.string_tracker = Some((ht, hasher)); return; From ab2878e993a5d933f78afad526868e1e8b562039 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Fri, 6 Jun 2025 12:52:20 +0200 Subject: [PATCH 09/19] Rename --- .../src/builder/generic_bytes_view_builder.rs | 34 +++++++++---------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/arrow-array/src/builder/generic_bytes_view_builder.rs b/arrow-array/src/builder/generic_bytes_view_builder.rs index 3f5c823ab74e..49c0d01792c6 100644 --- a/arrow-array/src/builder/generic_bytes_view_builder.rs +++ b/arrow-array/src/builder/generic_bytes_view_builder.rs @@ -79,7 +79,7 @@ impl BlockSizeGrowthStrategy { /// using [`GenericByteViewBuilder::append_block`] and then views into this block appended /// using [`GenericByteViewBuilder::try_append_view`] pub struct GenericByteViewBuilder { - views_builder: Vec, + views_buffer: Vec, null_buffer_builder: NullBufferBuilder, completed: Vec, in_progress: Vec, @@ -99,7 +99,7 @@ impl GenericByteViewBuilder { /// Creates a new [`GenericByteViewBuilder`] with space for `capacity` string values. pub fn with_capacity(capacity: usize) -> Self { Self { - views_builder: Vec::with_capacity(capacity), + views_buffer: Vec::with_capacity(capacity), null_buffer_builder: NullBufferBuilder::new(capacity), completed: vec![], in_progress: vec![], @@ -148,7 +148,7 @@ impl GenericByteViewBuilder { pub fn with_deduplicate_strings(self) -> Self { Self { string_tracker: Some(( - HashTable::with_capacity(self.views_builder.capacity()), + HashTable::with_capacity(self.views_buffer.capacity()), Default::default(), )), ..self @@ -201,7 +201,7 @@ impl GenericByteViewBuilder { let b = b.get_unchecked(start..end); let view = make_view(b, block, offset); - self.views_builder.push(view); + self.views_buffer.push(view); self.null_buffer_builder.append_non_null(); } @@ -211,11 +211,11 @@ impl GenericByteViewBuilder { self.completed.extend(array.data_buffers().iter().cloned()); if self.completed.is_empty() { - self.views_builder.extend_from_slice(array.views()); + self.views_buffer.extend_from_slice(array.views()); } else { let starting_buffer = self.completed.len() as u32; - self.views_builder.extend(array.views().iter().map(|v| { + self.views_buffer.extend(array.views().iter().map(|v| { let mut byte_view = ByteView::from(*v); if byte_view.length > 12 { // If the view is small enough, we can inline it @@ -283,7 +283,7 @@ impl GenericByteViewBuilder { /// Useful if we want to know what value has been inserted to the builder /// The index has to be smaller than `self.len()`, otherwise it will panic pub fn get_value(&self, index: usize) -> &[u8] { - let view = self.views_builder.as_slice().get(index).unwrap(); + let view = self.views_buffer.as_slice().get(index).unwrap(); let len = *view as u32; if len <= 12 { // # Safety @@ -315,7 +315,7 @@ impl GenericByteViewBuilder { let mut view_buffer = [0; 16]; view_buffer[0..4].copy_from_slice(&length.to_le_bytes()); view_buffer[4..4 + v.len()].copy_from_slice(v); - self.views_builder.push(u128::from_le_bytes(view_buffer)); + self.views_buffer.push(u128::from_le_bytes(view_buffer)); self.null_buffer_builder.append_non_null(); return; } @@ -339,7 +339,7 @@ impl GenericByteViewBuilder { Entry::Occupied(occupied) => { // If the string already exists, we will directly use the view let idx = occupied.get(); - self.views_builder.push(self.views_builder[*idx]); + self.views_buffer.push(self.views_buffer[*idx]); self.null_buffer_builder.append_non_null(); self.string_tracker = Some((ht, hasher)); return; @@ -347,7 +347,7 @@ impl GenericByteViewBuilder { Entry::Vacant(vacant) => { // o.w. we insert the (string hash -> view index) // the idx is current length of views_builder, as we are inserting a new view - vacant.insert(self.views_builder.len()); + vacant.insert(self.views_buffer.len()); } } self.string_tracker = Some((ht, hasher)); @@ -368,7 +368,7 @@ impl GenericByteViewBuilder { buffer_index: self.completed.len() as u32, offset, }; - self.views_builder.push(view.into()); + self.views_buffer.push(view.into()); self.null_buffer_builder.append_non_null(); } @@ -385,7 +385,7 @@ impl GenericByteViewBuilder { #[inline] pub fn append_null(&mut self) { self.null_buffer_builder.append_null(); - self.views_builder.push(0); + self.views_buffer.push(0); } /// Builds the [`GenericByteViewArray`] and reset this builder @@ -396,7 +396,7 @@ impl GenericByteViewBuilder { if let Some((ref mut ht, _)) = self.string_tracker.as_mut() { ht.clear(); } - let views = std::mem::take(&mut self.views_builder); + let views = std::mem::take(&mut self.views_buffer); // SAFETY: valid by construction unsafe { GenericByteViewArray::new_unchecked(views.into(), completed, nulls) } } @@ -407,8 +407,8 @@ impl GenericByteViewBuilder { if !self.in_progress.is_empty() { completed.push(Buffer::from_slice_ref(&self.in_progress)); } - let len = self.views_builder.len(); - let views = Buffer::from_slice_ref(self.views_builder.as_slice()); + let len = self.views_buffer.len(); + let views = Buffer::from_slice_ref(self.views_buffer.as_slice()); let views = ScalarBuffer::new(views, 0, len); let nulls = self.null_buffer_builder.finish_cloned(); // SAFETY: valid by construction @@ -422,7 +422,7 @@ impl GenericByteViewBuilder { /// Return the allocated size of this builder in bytes, useful for memory accounting. pub fn allocated_size(&self) -> usize { - let views = self.views_builder.capacity() * std::mem::size_of::(); + let views = self.views_buffer.capacity() * std::mem::size_of::(); let null = self.null_buffer_builder.allocated_size(); let buffer_size = self.completed.iter().map(|b| b.capacity()).sum::(); let in_progress = self.in_progress.capacity(); @@ -444,7 +444,7 @@ impl std::fmt::Debug for GenericByteViewBuilder { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}ViewBuilder", T::PREFIX)?; f.debug_struct("") - .field("views_builder", &self.views_builder) + .field("views_builder", &self.views_buffer) .field("in_progress", &self.in_progress) .field("completed", &self.completed) .field("null_buffer_builder", &self.null_buffer_builder) From 06101a1af31692dec9648dc1aa35f1be0cfd5389 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Fri, 6 Jun 2025 13:45:49 +0200 Subject: [PATCH 10/19] comment --- arrow-array/src/builder/generic_bytes_view_builder.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/arrow-array/src/builder/generic_bytes_view_builder.rs b/arrow-array/src/builder/generic_bytes_view_builder.rs index 49c0d01792c6..7d27bf0f5f53 100644 --- a/arrow-array/src/builder/generic_bytes_view_builder.rs +++ b/arrow-array/src/builder/generic_bytes_view_builder.rs @@ -205,7 +205,9 @@ impl GenericByteViewBuilder { self.null_buffer_builder.append_non_null(); } - /// Appends an array + /// Appends an array to the builder. + /// This will flush any in-progress block and append the data buffers + /// and add the (adapted) views. pub fn append_array(&mut self, array: &GenericByteViewArray) { self.flush_in_progress(); self.completed.extend(array.data_buffers().iter().cloned()); From 7b408ad051c9130410f40ddce7cc9fde922274a0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Fri, 6 Jun 2025 14:28:41 +0200 Subject: [PATCH 11/19] Style suggestion --- arrow-array/src/builder/generic_bytes_view_builder.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/arrow-array/src/builder/generic_bytes_view_builder.rs b/arrow-array/src/builder/generic_bytes_view_builder.rs index 7d27bf0f5f53..1e06cd769872 100644 --- a/arrow-array/src/builder/generic_bytes_view_builder.rs +++ b/arrow-array/src/builder/generic_bytes_view_builder.rs @@ -446,7 +446,7 @@ impl std::fmt::Debug for GenericByteViewBuilder { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}ViewBuilder", T::PREFIX)?; f.debug_struct("") - .field("views_builder", &self.views_buffer) + .field("views_buffer", &self.views_buffer) .field("in_progress", &self.in_progress) .field("completed", &self.completed) .field("null_buffer_builder", &self.null_buffer_builder) From 16ff18cafe9824fc6b61d16e67dc91e91ee3c665 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Fri, 6 Jun 2025 20:06:05 +0200 Subject: [PATCH 12/19] Update arrow-array/src/builder/generic_bytes_view_builder.rs Co-authored-by: Andrew Lamb --- arrow-array/src/builder/generic_bytes_view_builder.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/arrow-array/src/builder/generic_bytes_view_builder.rs b/arrow-array/src/builder/generic_bytes_view_builder.rs index 1e06cd769872..9ebf959b8dc4 100644 --- a/arrow-array/src/builder/generic_bytes_view_builder.rs +++ b/arrow-array/src/builder/generic_bytes_view_builder.rs @@ -220,7 +220,7 @@ impl GenericByteViewBuilder { self.views_buffer.extend(array.views().iter().map(|v| { let mut byte_view = ByteView::from(*v); if byte_view.length > 12 { - // If the view is small enough, we can inline it + // Small views (<=12 bytes) are inlined, so only need to update large views byte_view.buffer_index += starting_buffer; }; From 97d93aa64e1025285841b9a5abc80933b024d96a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Sat, 7 Jun 2025 00:16:43 +0200 Subject: [PATCH 13/19] Fix and extend fast path --- arrow-array/src/builder/generic_bytes_view_builder.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/arrow-array/src/builder/generic_bytes_view_builder.rs b/arrow-array/src/builder/generic_bytes_view_builder.rs index 9ebf959b8dc4..f50986e2bf9c 100644 --- a/arrow-array/src/builder/generic_bytes_view_builder.rs +++ b/arrow-array/src/builder/generic_bytes_view_builder.rs @@ -210,9 +210,12 @@ impl GenericByteViewBuilder { /// and add the (adapted) views. pub fn append_array(&mut self, array: &GenericByteViewArray) { self.flush_in_progress(); + // keep original views if the array is empty or if there are no data buffers (all inline views) + let keep_views = self.completed.is_empty() || array.data_buffers().is_empty(); + self.completed.extend(array.data_buffers().iter().cloned()); - if self.completed.is_empty() { + if keep_views { self.views_buffer.extend_from_slice(array.views()); } else { let starting_buffer = self.completed.len() as u32; From c4222161fbca96ce5e543d2e94a74b1120e01ea2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Sat, 7 Jun 2025 00:30:55 +0200 Subject: [PATCH 14/19] Adjust comment --- arrow-array/src/builder/generic_bytes_view_builder.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/arrow-array/src/builder/generic_bytes_view_builder.rs b/arrow-array/src/builder/generic_bytes_view_builder.rs index f50986e2bf9c..3237c5f9639e 100644 --- a/arrow-array/src/builder/generic_bytes_view_builder.rs +++ b/arrow-array/src/builder/generic_bytes_view_builder.rs @@ -210,7 +210,7 @@ impl GenericByteViewBuilder { /// and add the (adapted) views. pub fn append_array(&mut self, array: &GenericByteViewArray) { self.flush_in_progress(); - // keep original views if the array is empty or if there are no data buffers (all inline views) + // keep original views if this array is the first to be added or if there are no data buffers (all inline views) let keep_views = self.completed.is_empty() || array.data_buffers().is_empty(); self.completed.extend(array.data_buffers().iter().cloned()); From 7a506141f821c812c7bd6087f910e59091c80e1b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Sat, 7 Jun 2025 00:37:14 +0200 Subject: [PATCH 15/19] Address feedback --- arrow-select/src/concat.rs | 21 +++++++-------------- 1 file changed, 7 insertions(+), 14 deletions(-) diff --git a/arrow-select/src/concat.rs b/arrow-select/src/concat.rs index f1db901dc8be..33fed0f95bdd 100644 --- a/arrow-select/src/concat.rs +++ b/arrow-select/src/concat.rs @@ -32,7 +32,7 @@ use crate::dictionary::{merge_dictionary_values, should_merge_dictionary_values}; use arrow_array::builder::{ - BinaryViewBuilder, BooleanBuilder, GenericByteBuilder, PrimitiveBuilder, StringViewBuilder, + BooleanBuilder, GenericByteBuilder, GenericByteViewBuilder, PrimitiveBuilder, }; use arrow_array::cast::AsArray; use arrow_array::types::*; @@ -86,18 +86,11 @@ fn fixed_size_list_capacity(arrays: &[&dyn Array], data_type: &DataType) -> Capa } } -fn concat_byte_view(arrays: &[&dyn Array]) -> Result { - let mut builder = BinaryViewBuilder::with_capacity(arrays.iter().map(|a| a.len()).sum()); +fn concat_byte_view(arrays: &[&dyn Array]) -> Result { + let mut builder = + GenericByteViewBuilder::::with_capacity(arrays.iter().map(|a| a.len()).sum()); for &array in arrays.iter() { - builder.append_array(array.as_binary_view()); - } - Ok(Arc::new(builder.finish())) -} - -fn concat_string_view(arrays: &[&dyn Array]) -> Result { - let mut builder = StringViewBuilder::with_capacity(arrays.iter().map(|a| a.len()).sum()); - for &array in arrays.iter() { - builder.append_array(array.as_string_view()); + builder.append_array(array.as_byte_view()); } Ok(Arc::new(builder.finish())) } @@ -443,8 +436,8 @@ pub fn concat(arrays: &[&dyn Array]) -> Result { _ => unreachable!("Unsupported run end index type: {r:?}"), } } - DataType::Utf8View => concat_string_view(arrays), - DataType::BinaryView => concat_byte_view(arrays), + DataType::Utf8View => concat_byte_view::(arrays), + DataType::BinaryView => concat_byte_view::(arrays), _ => { let capacity = get_capacity(arrays, d); concat_fallback(arrays, capacity) From d2305b97362cc6d7d9d7f3033cd6713a6b66808c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Sat, 7 Jun 2025 00:56:40 +0200 Subject: [PATCH 16/19] new_unchecked --- arrow-select/src/coalesce.rs | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs index 642f872f48ae..0624dfa8e804 100644 --- a/arrow-select/src/coalesce.rs +++ b/arrow-select/src/coalesce.rs @@ -22,7 +22,7 @@ //! [`take`]: crate::take::take use crate::concat::concat_batches; use arrow_array::StringViewArray; -use arrow_array::{cast::AsArray, Array, ArrayRef, RecordBatch, RecordBatchOptions}; +use arrow_array::{cast::AsArray, Array, ArrayRef, RecordBatch}; use arrow_data::ByteView; use arrow_schema::{ArrowError, SchemaRef}; use std::collections::VecDeque; @@ -312,10 +312,7 @@ fn gc_string_view_batch(batch: RecordBatch) -> RecordBatch { } }) .collect(); - let mut options = RecordBatchOptions::new(); - options = options.with_row_count(Some(num_rows)); - RecordBatch::try_new_with_options(schema, new_columns, &options) - .expect("Failed to re-create the gc'ed record batch") + unsafe { RecordBatch::new_unchecked(schema, new_columns, num_rows) } } #[cfg(test)] From a94a09628e3ce0afcdb4ae42fa0866b86b42682d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Sat, 7 Jun 2025 01:37:43 +0200 Subject: [PATCH 17/19] Update coalesce.rs --- arrow-select/src/coalesce.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs index 0624dfa8e804..d123802033ba 100644 --- a/arrow-select/src/coalesce.rs +++ b/arrow-select/src/coalesce.rs @@ -319,7 +319,7 @@ fn gc_string_view_batch(batch: RecordBatch) -> RecordBatch { mod tests { use super::*; use arrow_array::builder::{ArrayBuilder, StringViewBuilder}; - use arrow_array::{StringViewArray, UInt32Array}; + use arrow_array::{RecordBatchOptions, StringViewArray, UInt32Array}; use arrow_schema::{DataType, Field, Schema}; use std::ops::Range; From cf26d8606ea4fe317ba2d781c786a6da670086d1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Sat, 7 Jun 2025 09:51:09 +0200 Subject: [PATCH 18/19] add another fast path --- arrow-select/src/coalesce.rs | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs index d123802033ba..86d7944d8c34 100644 --- a/arrow-select/src/coalesce.rs +++ b/arrow-select/src/coalesce.rs @@ -269,6 +269,17 @@ fn gc_string_view_batch(batch: RecordBatch) -> RecordBatch { // Re-creating the array copies data and can be time consuming. // We only do it if the array is sparse if actual_buffer_size > (ideal_buffer_size * 2) { + if ideal_buffer_size == 0 { + // If the ideal buffer size is 0, all views are inlined + // so just reuse the views + return Arc::new(unsafe { + StringViewArray::new_unchecked( + s.views().clone(), + vec![], + s.nulls().cloned(), + ) + }); + } // We set the block size to `ideal_buffer_size` so that the new StringViewArray only has one buffer, which accelerate later concat_batches. // See https://github.com/apache/arrow-rs/issues/6094 for more details. let mut buffer: Vec = Vec::with_capacity(ideal_buffer_size); From 289aeda9a532a2c3c8c1c8cd2b2e9317cffec37d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Sat, 7 Jun 2025 09:58:47 +0200 Subject: [PATCH 19/19] add another fast path --- arrow-select/src/coalesce.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs index 86d7944d8c34..41fd0ff16859 100644 --- a/arrow-select/src/coalesce.rs +++ b/arrow-select/src/coalesce.rs @@ -251,6 +251,10 @@ fn gc_string_view_batch(batch: RecordBatch) -> RecordBatch { let Some(s) = c.as_string_view_opt() else { return c; }; + if s.data_buffers().is_empty() { + // If there are no data buffers, we can just return the array as is + return c; + } let ideal_buffer_size: usize = s .views() .iter()