Add with_capacity to ValuesBuffer trait

Author: lyang24

parquet/src/arrow/array_reader/fixed_len_byte_array.rs

@@ -284,6 +284,15 @@ fn move_values<F>(
 }
 
 impl ValuesBuffer for FixedLenByteArrayBuffer {
+    fn with_capacity(_capacity: usize) -> Self {
+        // byte_length is not known at trait level, so we return a default buffer
+        // The decoder will pre-allocate when it knows both capacity and byte_length
+        Self {
+            buffer: Vec::new(),
+            byte_length: None,
+        }
+    }
+
     fn pad_nulls(
         &mut self,
         read_offset: usize,

parquet/src/arrow/buffer/dictionary_buffer.rs

@@ -194,6 +194,12 @@ impl<K: ArrowNativeType + Ord, V: OffsetSizeTrait> DictionaryBuffer<K, V> {
 }
 
 impl<K: ArrowNativeType, V: OffsetSizeTrait> ValuesBuffer for DictionaryBuffer<K, V> {
+    fn with_capacity(capacity: usize) -> Self {
+        Self::Values {
+            values: OffsetBuffer::with_capacity(capacity),
+        }
+    }
+
     fn pad_nulls(
         &mut self,
         read_offset: usize,

parquet/src/arrow/buffer/offset_buffer.rs

@@ -44,6 +44,19 @@ impl<I: OffsetSizeTrait> Default for OffsetBuffer<I> {
 }
 
 impl<I: OffsetSizeTrait> OffsetBuffer<I> {
+    /// Create a new `OffsetBuffer` with capacity for at least `capacity` elements
+    ///
+    /// Pre-allocates the offsets vector to avoid reallocations during reading.
+    /// The values vector is not pre-allocated as its size is unpredictable.
+    pub fn with_capacity(capacity: usize) -> Self {
+        let mut offsets = Vec::with_capacity(capacity + 1);
+        offsets.push(I::default());
+        Self {
+            offsets,
+            values: Vec::new(),
+        }
+    }
+
     /// Returns the number of byte arrays in this buffer
     pub fn len(&self) -> usize {
         self.offsets.len() - 1
@@ -139,6 +152,10 @@ impl<I: OffsetSizeTrait> OffsetBuffer<I> {
 }
 
 impl<I: OffsetSizeTrait> ValuesBuffer for OffsetBuffer<I> {
+    fn with_capacity(capacity: usize) -> Self {
+        Self::with_capacity(capacity)
+    }
+
     fn pad_nulls(
         &mut self,
         read_offset: usize,

parquet/src/arrow/buffer/view_buffer.rs

@@ -105,6 +105,10 @@ impl ViewBuffer {
 }
 
 impl ValuesBuffer for ViewBuffer {
+    fn with_capacity(capacity: usize) -> Self {
+        Self::with_capacity(capacity)
+    }
+
     fn pad_nulls(
         &mut self,
         read_offset: usize,

parquet/src/arrow/record_reader/buffer.rs

@@ -19,6 +19,12 @@ use crate::arrow::buffer::bit_util::iter_set_bits_rev;
 
 /// A buffer that supports padding with nulls
 pub trait ValuesBuffer: Default {
+    /// Create a new buffer with capacity for at least `capacity` elements
+    ///
+    /// This allows pre-allocating buffers to avoid reallocations during reading,
+    /// improving performance when the number of values is known in advance.
+    fn with_capacity(capacity: usize) -> Self;
+
     /// If a column contains nulls, more level data may be read than value data, as null
     /// values are not encoded. Therefore, first the levels data is read, the null count
     /// determined, and then the corresponding number of values read to a [`ValuesBuffer`].
@@ -43,6 +49,10 @@ pub trait ValuesBuffer: Default {
 }
 
 impl<T: Copy + Default> ValuesBuffer for Vec<T> {
+    fn with_capacity(capacity: usize) -> Self {
+        Vec::with_capacity(capacity)
+    }
+
     fn pad_nulls(
         &mut self,
         read_offset: usize,

parquet/src/arrow/record_reader/mod.rs

@@ -58,6 +58,8 @@ pub struct GenericRecordReader<V, CV> {
     num_values: usize,
     /// Number of buffered records
     num_records: usize,
+    /// Capacity hint for pre-allocating buffers based on batch size
+    capacity_hint: usize,
 }
 
 impl<V, CV> GenericRecordReader<V, CV>
@@ -67,19 +69,23 @@ where
 {
     /// Create a new [`GenericRecordReader`]
     pub fn new(desc: ColumnDescPtr) -> Self {
+        // Start with a reasonable default capacity to avoid zero reallocations on first batch
+        const DEFAULT_CAPACITY: usize = 1024;
+
         let def_levels = (desc.max_def_level() > 0)
             .then(|| DefinitionLevelBuffer::new(&desc, packed_null_mask(&desc)));
 
         let rep_levels = (desc.max_rep_level() > 0).then(Vec::new);
 
         Self {
-            values: V::default(),
+            values: V::with_capacity(DEFAULT_CAPACITY),
             def_levels,
             rep_levels,
             column_reader: None,
             column_desc: desc,
             num_values: 0,
             num_records: 0,
+            capacity_hint: DEFAULT_CAPACITY,
         }
     }
 
@@ -169,7 +175,9 @@ where
     /// Returns currently stored buffer data.
     /// The side effect is similar to `consume_def_levels`.
     pub fn consume_record_data(&mut self) -> V {
-        std::mem::take(&mut self.values)
+        // Replace the buffer with a new one that has the same capacity
+        // This avoids reallocations on subsequent batches
+        std::mem::replace(&mut self.values, V::with_capacity(self.capacity_hint))
     }
 
     /// Returns currently stored null bitmap data for nullable columns.
@@ -208,6 +216,11 @@ where
 
     /// Try to read one batch of data returning the number of records read
     fn read_one_batch(&mut self, batch_size: usize) -> Result<usize> {
+        // Update capacity hint to the largest batch size seen
+        if batch_size > self.capacity_hint {
+            self.capacity_hint = batch_size;
+        }
+
         let (records_read, values_read, levels_read) =
             self.column_reader.as_mut().unwrap().read_records(
                 batch_size,