pass down batch_size from ArrowReaderBuilder through ArrayReaderBuilder.Ensure internal buffers to be pre-allocated.

lyang24 · lyang24 · commit 224798a92b39 · 2026-01-10T22:31:04.000-08:00
diff --git a/parquet/src/arrow/array_reader/builder.rs b/parquet/src/arrow/array_reader/builder.rs
@@ -95,6 +95,8 @@ pub struct ArrayReaderBuilder<'a> {
     parquet_metadata: Option<&'a ParquetMetaData>,
     /// metrics
     metrics: &'a ArrowReaderMetrics,
+    /// Batch size for pre-allocating internal buffers
+    batch_size: Option<usize>,
 }
 
 impl<'a> ArrayReaderBuilder<'a> {
@@ -104,6 +106,7 @@ impl<'a> ArrayReaderBuilder<'a> {
             cache_options: None,
             parquet_metadata: None,
             metrics,
+            batch_size: None,
         }
     }
 
@@ -119,6 +122,15 @@ impl<'a> ArrayReaderBuilder<'a> {
         self
     }
 
+    /// Set the batch size for pre-allocating internal buffers
+    ///
+    /// This allows the reader to pre-allocate buffers with the expected capacity,
+    /// avoiding reallocations when reading the first batch of data.
+    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
+        self.batch_size = Some(batch_size);
+        self
+    }
+
     /// Create [`ArrayReader`] from parquet schema, projection mask, and parquet file reader.
     pub fn build_array_reader(
         &self,
@@ -410,18 +422,18 @@ impl<'a> ArrayReaderBuilder<'a> {
             )?) as _,
             PhysicalType::BYTE_ARRAY => match arrow_type {
                 Some(DataType::Dictionary(_, _)) => {
-                    make_byte_array_dictionary_reader(page_iterator, column_desc, arrow_type)?
+                    make_byte_array_dictionary_reader(page_iterator, column_desc, arrow_type, self.batch_size)?
                 }
                 Some(DataType::Utf8View | DataType::BinaryView) => {
-                    make_byte_view_array_reader(page_iterator, column_desc, arrow_type)?
+                    make_byte_view_array_reader(page_iterator, column_desc, arrow_type, self.batch_size)?
                 }
-                _ => make_byte_array_reader(page_iterator, column_desc, arrow_type)?,
+                _ => make_byte_array_reader(page_iterator, column_desc, arrow_type, self.batch_size)?,
             },
             PhysicalType::FIXED_LEN_BYTE_ARRAY => match arrow_type {
                 Some(DataType::Dictionary(_, _)) => {
-                    make_byte_array_dictionary_reader(page_iterator, column_desc, arrow_type)?
+                    make_byte_array_dictionary_reader(page_iterator, column_desc, arrow_type, self.batch_size)?
                 }
-                _ => make_fixed_len_byte_array_reader(page_iterator, column_desc, arrow_type)?,
+                _ => make_fixed_len_byte_array_reader(page_iterator, column_desc, arrow_type, self.batch_size)?,
             },
         };
         Ok(Some(reader))
diff --git a/parquet/src/arrow/array_reader/byte_array.rs b/parquet/src/arrow/array_reader/byte_array.rs
@@ -38,10 +38,14 @@ use std::any::Any;
 use std::sync::Arc;
 
 /// Returns an [`ArrayReader`] that decodes the provided byte array column
+///
+/// The optional `batch_size` parameter is used to pre-allocate internal buffers,
+/// avoiding reallocations when reading the first batch of data.
 pub fn make_byte_array_reader(
     pages: Box<dyn PageIterator>,
     column_desc: ColumnDescPtr,
     arrow_type: Option<ArrowType>,
+    batch_size: Option<usize>,
 ) -> Result<Box<dyn ArrayReader>> {
     // Check if Arrow type is specified, else create it from Parquet type
     let data_type = match arrow_type {
@@ -56,13 +60,23 @@ pub fn make_byte_array_reader(
         | ArrowType::Utf8
         | ArrowType::Decimal128(_, _)
         | ArrowType::Decimal256(_, _) => {
-            let reader = GenericRecordReader::new(column_desc);
+            let reader = match batch_size {
+                Some(capacity) => {
+                    GenericRecordReader::new_with_capacity(column_desc, capacity)
+                }
+                None => GenericRecordReader::new(column_desc),
+            };
             Ok(Box::new(ByteArrayReader::<i32>::new(
                 pages, data_type, reader,
             )))
         }
         ArrowType::LargeUtf8 | ArrowType::LargeBinary => {
-            let reader = GenericRecordReader::new(column_desc);
+            let reader = match batch_size {
+                Some(capacity) => {
+                    GenericRecordReader::new_with_capacity(column_desc, capacity)
+                }
+                None => GenericRecordReader::new(column_desc),
+            };
             Ok(Box::new(ByteArrayReader::<i64>::new(
                 pages, data_type, reader,
             )))
diff --git a/parquet/src/arrow/array_reader/byte_array_dictionary.rs b/parquet/src/arrow/array_reader/byte_array_dictionary.rs
@@ -40,14 +40,17 @@ use crate::util::bit_util::FromBytes;
 /// A macro to reduce verbosity of [`make_byte_array_dictionary_reader`]
 macro_rules! make_reader {
     (
-        ($pages:expr, $column_desc:expr, $data_type:expr) => match ($k:expr, $v:expr) {
+        ($pages:expr, $column_desc:expr, $data_type:expr, $batch_size:expr) => match ($k:expr, $v:expr) {
             $(($key_arrow:pat, $value_arrow:pat) => ($key_type:ty, $value_type:ty),)+
         }
     ) => {
         match (($k, $v)) {
             $(
                 ($key_arrow, $value_arrow) => {
-                    let reader = GenericRecordReader::new($column_desc);
+                    let reader = match $batch_size {
+                        Some(capacity) => GenericRecordReader::new_with_capacity($column_desc, capacity),
+                        None => GenericRecordReader::new($column_desc),
+                    };
                     Ok(Box::new(ByteArrayDictionaryReader::<$key_type, $value_type>::new(
                         $pages, $data_type, reader,
                     )))
@@ -73,10 +76,13 @@ macro_rules! make_reader {
 /// It is therefore recommended that if `pages` contains data from multiple column chunks,
 /// that the read batch size used is a divisor of the row group size
 ///
+/// The optional `batch_size` parameter is used to pre-allocate internal buffers,
+/// avoiding reallocations when reading the first batch of data.
 pub fn make_byte_array_dictionary_reader(
     pages: Box<dyn PageIterator>,
     column_desc: ColumnDescPtr,
     arrow_type: Option<ArrowType>,
+    batch_size: Option<usize>,
 ) -> Result<Box<dyn ArrayReader>> {
     // Check if Arrow type is specified, else create it from Parquet type
     let data_type = match arrow_type {
@@ -89,7 +95,7 @@ pub fn make_byte_array_dictionary_reader(
     match &data_type {
         ArrowType::Dictionary(key_type, value_type) => {
             make_reader! {
-                (pages, column_desc, data_type) => match (key_type.as_ref(), value_type.as_ref()) {
+                (pages, column_desc, data_type, batch_size) => match (key_type.as_ref(), value_type.as_ref()) {
                     (ArrowType::UInt8, ArrowType::Binary | ArrowType::Utf8 | ArrowType::FixedSizeBinary(_)) => (u8, i32),
                     (ArrowType::UInt8, ArrowType::LargeBinary | ArrowType::LargeUtf8) => (u8, i64),
                     (ArrowType::Int8, ArrowType::Binary | ArrowType::Utf8 | ArrowType::FixedSizeBinary(_)) => (i8, i32),
diff --git a/parquet/src/arrow/array_reader/byte_view_array.rs b/parquet/src/arrow/array_reader/byte_view_array.rs
@@ -36,10 +36,14 @@ use bytes::Bytes;
 use std::any::Any;
 
 /// Returns an [`ArrayReader`] that decodes the provided byte array column to view types.
+///
+/// The optional `batch_size` parameter is used to pre-allocate internal buffers,
+/// avoiding reallocations when reading the first batch of data.
 pub fn make_byte_view_array_reader(
     pages: Box<dyn PageIterator>,
     column_desc: ColumnDescPtr,
     arrow_type: Option<ArrowType>,
+    batch_size: Option<usize>,
 ) -> Result<Box<dyn ArrayReader>> {
     // Check if Arrow type is specified, else create it from Parquet type
     let data_type = match arrow_type {
@@ -52,7 +56,10 @@ pub fn make_byte_view_array_reader(
 
     match data_type {
         ArrowType::BinaryView | ArrowType::Utf8View => {
-            let reader = GenericRecordReader::new(column_desc);
+            let reader = match batch_size {
+                Some(capacity) => GenericRecordReader::new_with_capacity(column_desc, capacity),
+                None => GenericRecordReader::new(column_desc),
+            };
             Ok(Box::new(ByteViewArrayReader::new(pages, data_type, reader)))
         }
 
diff --git a/parquet/src/arrow/array_reader/fixed_len_byte_array.rs b/parquet/src/arrow/array_reader/fixed_len_byte_array.rs
@@ -40,10 +40,14 @@ use std::ops::Range;
 use std::sync::Arc;
 
 /// Returns an [`ArrayReader`] that decodes the provided fixed length byte array column
+///
+/// The optional `batch_size` parameter is used to pre-allocate internal buffers,
+/// avoiding reallocations when reading the first batch of data.
 pub fn make_fixed_len_byte_array_reader(
     pages: Box<dyn PageIterator>,
     column_desc: ColumnDescPtr,
     arrow_type: Option<ArrowType>,
+    batch_size: Option<usize>,
 ) -> Result<Box<dyn ArrayReader>> {
     // Check if Arrow type is specified, else create it from Parquet type
     let data_type = match arrow_type {
@@ -126,6 +130,7 @@ pub fn make_fixed_len_byte_array_reader(
         column_desc,
         data_type,
         byte_length,
+        batch_size,
     )))
 }
 
@@ -144,14 +149,19 @@ impl FixedLenByteArrayReader {
         column_desc: ColumnDescPtr,
         data_type: ArrowType,
         byte_length: usize,
+        batch_size: Option<usize>,
     ) -> Self {
+        let record_reader = match batch_size {
+            Some(capacity) => GenericRecordReader::new_with_capacity(column_desc, capacity),
+            None => GenericRecordReader::new(column_desc),
+        };
         Self {
             data_type,
             byte_length,
             pages,
             def_levels_buffer: None,
             rep_levels_buffer: None,
-            record_reader: GenericRecordReader::new(column_desc),
+            record_reader,
         }
     }
 }
@@ -284,6 +294,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,
diff --git a/parquet/src/arrow/arrow_reader/mod.rs b/parquet/src/arrow/arrow_reader/mod.rs
@@ -1066,6 +1066,7 @@ impl<T: ChunkReader + 'static> ParquetRecordBatchReaderBuilder<T> {
 
                 let array_reader = ArrayReaderBuilder::new(&reader, &metrics)
                     .with_parquet_metadata(&reader.metadata)
+                    .with_batch_size(batch_size)
                     .build_array_reader(fields.as_deref(), predicate.projection())?;
 
                 plan_builder = plan_builder.with_predicate(array_reader, predicate.as_mut())?;
@@ -1074,6 +1075,7 @@ impl<T: ChunkReader + 'static> ParquetRecordBatchReaderBuilder<T> {
 
         let array_reader = ArrayReaderBuilder::new(&reader, &metrics)
             .with_parquet_metadata(&reader.metadata)
+            .with_batch_size(batch_size)
             .build_array_reader(fields.as_deref(), &projection)?;
 
         let read_plan = plan_builder
@@ -1381,6 +1383,7 @@ impl ParquetRecordBatchReader {
         let metrics = ArrowReaderMetrics::disabled();
         let array_reader = ArrayReaderBuilder::new(row_groups, &metrics)
             .with_parquet_metadata(row_groups.metadata())
+            .with_batch_size(batch_size)
             .build_array_reader(levels.levels.as_ref(), &ProjectionMask::all())?;
 
         let read_plan = ReadPlanBuilder::new(batch_size)
diff --git a/parquet/src/arrow/buffer/dictionary_buffer.rs b/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,
diff --git a/parquet/src/arrow/buffer/offset_buffer.rs b/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,
diff --git a/parquet/src/arrow/buffer/view_buffer.rs b/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,
diff --git a/parquet/src/arrow/push_decoder/reader_builder/mod.rs b/parquet/src/arrow/push_decoder/reader_builder/mod.rs
@@ -435,6 +435,7 @@ impl RowGroupReaderBuilder {
                 let array_reader = ArrayReaderBuilder::new(&row_group, &self.metrics)
                     .with_cache_options(Some(&cache_options))
                     .with_parquet_metadata(&self.metadata)
+                    .with_batch_size(self.batch_size)
                     .build_array_reader(self.fields.as_deref(), predicate.projection())?;
 
                 plan_builder =
@@ -592,7 +593,8 @@ impl RowGroupReaderBuilder {
 
                 // if we have any cached results, connect them up
                 let array_reader_builder = ArrayReaderBuilder::new(&row_group, &self.metrics)
-                    .with_parquet_metadata(&self.metadata);
+                    .with_parquet_metadata(&self.metadata)
+                    .with_batch_size(self.batch_size);
                 let array_reader = if let Some(cache_info) = cache_info.as_ref() {
                     let cache_options = cache_info.builder().consumer();
                     array_reader_builder
diff --git a/parquet/src/arrow/record_reader/buffer.rs b/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,
diff --git a/parquet/src/arrow/record_reader/mod.rs b/parquet/src/arrow/record_reader/mod.rs
