[arrow-rs] branch master updated: Preallocate buffers for FixedSizeBinary array creation (#3793)

[tustvold]

This is an automated email from the ASF dual-hosted git repository.

tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new b9fcd7fa2 Preallocate buffers for FixedSizeBinary array creation 
(#3793)
b9fcd7fa2 is described below

commit b9fcd7fa2154f848823521a11aeeba4e687025b8
Author: Max Burke <m...@urbanlogiq.com>
AuthorDate: Fri Mar 3 03:32:57 2023 -0800

    Preallocate buffers for FixedSizeBinary array creation (#3793)
    
    * Preallocate buffers for FixedSizeBinary array creation
    
    * fix build
    
    * code review fixup
---
 arrow-array/src/array/fixed_size_binary_array.rs | 29 ++++++++++++++++++------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/arrow-array/src/array/fixed_size_binary_array.rs 
b/arrow-array/src/array/fixed_size_binary_array.rs
index e927c8d8a..87f1b9557 100644
--- a/arrow-array/src/array/fixed_size_binary_array.rs
+++ b/arrow-array/src/array/fixed_size_binary_array.rs
@@ -141,8 +141,11 @@ impl FixedSizeBinaryArray {
         let mut len = 0;
         let mut size = None;
         let mut byte = 0;
-        let mut null_buf = MutableBuffer::from_len_zeroed(0);
-        let mut buffer = MutableBuffer::from_len_zeroed(0);
+
+        let iter_size_hint = iter.size_hint().0;
+        let mut null_buf = MutableBuffer::new(bit_util::ceil(iter_size_hint, 
8));
+        let mut buffer = MutableBuffer::new(0);
+
         let mut prepend = 0;
         iter.try_for_each(|item| -> Result<(), ArrowError> {
             // extend null bitmask by one byte per each 8 items
@@ -163,7 +166,12 @@ impl FixedSizeBinaryArray {
                         )));
                     }
                 } else {
-                    size = Some(slice.len());
+                    let len = slice.len();
+                    size = Some(len);
+                    // Now that we know how large each element is we can 
reserve
+                    // sufficient capacity in the underlying mutable buffer for
+                    // the data.
+                    buffer.reserve(iter_size_hint * len);
                     buffer.extend_zeros(slice.len() * prepend);
                 }
                 bit_util::set_bit(null_buf.as_slice_mut(), len);
@@ -234,8 +242,10 @@ impl FixedSizeBinaryArray {
     {
         let mut len = 0;
         let mut byte = 0;
-        let mut null_buf = MutableBuffer::from_len_zeroed(0);
-        let mut buffer = MutableBuffer::from_len_zeroed(0);
+
+        let iter_size_hint = iter.size_hint().0;
+        let mut null_buf = MutableBuffer::new(bit_util::ceil(iter_size_hint, 
8));
+        let mut buffer = MutableBuffer::new(iter_size_hint * (size as usize));
 
         iter.try_for_each(|item| -> Result<(), ArrowError> {
             // extend null bitmask by one byte per each 8 items
@@ -304,7 +314,9 @@ impl FixedSizeBinaryArray {
     {
         let mut len = 0;
         let mut size = None;
-        let mut buffer = MutableBuffer::from_len_zeroed(0);
+        let iter_size_hint = iter.size_hint().0;
+        let mut buffer = MutableBuffer::new(0);
+
         iter.try_for_each(|item| -> Result<(), ArrowError> {
             let slice = item.as_ref();
             if let Some(size) = size {
@@ -316,8 +328,11 @@ impl FixedSizeBinaryArray {
                     )));
                 }
             } else {
-                size = Some(slice.len());
+                let len = slice.len();
+                size = Some(len);
+                buffer.reserve(iter_size_hint * len);
             }
+
             buffer.extend_from_slice(slice);
 
             len += 1;