Improve coalesce

Author: Dandandan

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<ArrayRef> = batch
-        .columns()
-        .iter()
+fn gc_string_view_batch(batch: RecordBatch) -> RecordBatch {
+    let (schema, columns, num_rows) = batch.into_parts();
+    let new_columns: Vec<ArrayRef> = 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<u8> = Vec::with_capacity(ideal_buffer_size);
+
+                let views: Vec<u128> = 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()