Avoid overallocating arrays in coalesce primitives / views

arrow-select/src/coalesce.rs

@@ -615,6 +615,7 @@ mod tests {
     use crate::concat::concat_batches;
     use arrow_array::builder::StringViewBuilder;
     use arrow_array::cast::AsArray;
+    use arrow_array::types::Int32Type;
     use arrow_array::{
         BinaryViewArray, Int32Array, Int64Array, RecordBatchOptions, StringArray, StringViewArray,
         TimestampNanosecondArray, UInt32Array, UInt64Array,
@@ -1625,6 +1626,11 @@ mod tests {
         // Now should have a completed batch (100 rows total)
         assert!(coalescer.has_completed_batch());
         let output_batch = coalescer.next_completed_batch().unwrap();
+        let size = output_batch
+            .column(0)
+            .as_primitive::<Int32Type>()
+            .get_buffer_memory_size();
+        assert_eq!(size, 400); // 100 rows * 4 bytes each
         assert_eq!(output_batch.num_rows(), 100);
 
         assert_eq!(coalescer.get_buffered_rows(), 0);

arrow-select/src/coalesce/byte_view.rs

@@ -98,7 +98,10 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
     /// This is done on write (when we know it is necessary) rather than
     /// eagerly to avoid allocations that are not used.
     fn ensure_capacity(&mut self) {
-        self.views.reserve(self.batch_size);
+        if self.views.capacity() == 0 {
+            self.views.reserve(self.batch_size);
+        }
+        debug_assert_eq!(self.views.capacity(), self.batch_size);
     }
 
     /// Finishes in progress buffer, if any

arrow-select/src/coalesce/primitive.rs

@@ -55,7 +55,10 @@ impl<T: ArrowPrimitiveType> InProgressPrimitiveArray<T> {
     /// This is done on write (when we know it is necessary) rather than
     /// eagerly to avoid allocations that are not used.
     fn ensure_capacity(&mut self) {
-        self.current.reserve(self.batch_size);
+        if self.current.capacity() == 0 {
+            self.current.reserve(self.batch_size);
+        }
+        debug_assert_eq!(self.current.capacity(), self.batch_size);
     }
 }