Improve coalesce and concat performance for views
#7614
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -21,9 +21,9 @@ | |
| //! [`filter`]: crate::filter::filter | ||
| //! [`take`]: crate::take::take | ||
| use crate::concat::concat_batches; | ||
| use arrow_array::StringViewArray; | ||
| use arrow_array::{cast::AsArray, Array, ArrayRef, RecordBatch}; | ||
| 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); | ||
|
|
||
| // If pushing this batch would exceed the target batch size, | ||
| // finish the current batch and start a new one | ||
|
|
@@ -242,15 +242,19 @@ 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 (schema, columns, num_rows) = batch.into_parts(); | ||
| let new_columns: Vec<ArrayRef> = columns | ||
|
Contributor
There was a problem hiding this comment. I don't know if it matters but in #7620 I simply modified the existing Vec rather than collecting into a new one to save maybe an allocation for mut c in columns.iter_mut() {
...
}
Contributor
Author
There was a problem hiding this comment. AFAIK there is a Rust compiler optimization that will reuse a
Contributor
Author
There was a problem hiding this comment. This is one of the reasons I think using
Contributor
There was a problem hiding this comment. Yeah I think using |
||
| .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 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() | ||
|
|
@@ -264,42 +268,73 @@ 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) { | ||
| 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<u8> = Vec::with_capacity(ideal_buffer_size); | ||
|
|
||
| let views: Vec<u128> = s | ||
| .views() | ||
| .as_ref() | ||
| .iter() | ||
| .cloned() | ||
| .map(|v| { | ||
| 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(); | ||
|
Comment on lines
+291
to
+312
Contributor
There was a problem hiding this comment. I'm curious it is better if implement like this: let mut views = Vec::with_capacity(s.views().len());
for &v in s.views().as_ref() {
let mut b: ByteView = ByteView::from(v);
if b.length > 12 {
...
}
views.push(b.into());
}vector pre-allocates and no clone 🤔
Contributor
Author
There was a problem hiding this comment. No that is generally worse -
the |
||
|
|
||
| let buffers = if buffer.is_empty() { | ||
|
Contributor
There was a problem hiding this comment. I think the only way this could happen is if if ideal_buffer_size == 0 {
// just reuse views as is...
} |
||
| vec![] | ||
| } else { | ||
| vec![buffer.into()] | ||
| }; | ||
|
|
||
| let gc_string = unsafe { | ||
| StringViewArray::new_unchecked(views.into(), buffers, s.nulls().cloned()) | ||
| }; | ||
|
|
||
| Arc::new(gc_string) | ||
| } else { | ||
| c | ||
| } | ||
| }) | ||
| .collect(); | ||
| unsafe { RecordBatch::new_unchecked(schema, new_columns, num_rows) } | ||
| } | ||
|
|
||
| #[cfg(test)] | ||
| mod tests { | ||
| use super::*; | ||
| use arrow_array::builder::{ArrayBuilder, StringViewBuilder}; | ||
| use arrow_array::{RecordBatchOptions, StringViewArray, UInt32Array}; | ||
| use arrow_schema::{DataType, Field, Schema}; | ||
| use std::ops::Range; | ||
|
|
||
|
|
@@ -518,9 +553,11 @@ mod tests { | |
| fn test_gc_string_view_test_batch_empty() { | ||
| let schema = Schema::empty(); | ||
| let batch = RecordBatch::new_empty(schema.into()); | ||
| 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 +605,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 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() | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
this is a great idea 💯