add bytes_estimate for binary push in parquet deserialize (jorgecarle…

…itao#1308)
ritchie46 · Dec 12, 2022 · 1fcfd7c · 1fcfd7c
1 parent 1417f88
commit 1fcfd7c
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/src/io/parquet/read/deserialize/binary/basic.rs b/src/io/parquet/read/deserialize/binary/basic.rs
@@ -477,9 +477,13 @@ impl<'a, O: Offset> utils::Decoder<'a> for BinaryDecoder<O> {
 
 pub(super) fn finish<O: Offset, A: TraitBinaryArray<O>>(
     data_type: &DataType,
-    values: Binary<O>,
-    validity: MutableBitmap,
+    mut values: Binary<O>,
+    mut validity: MutableBitmap,
 ) -> Result<A> {
+    values.offsets.shrink_to_fit();
+    values.values.shrink_to_fit();
+    validity.shrink_to_fit();
+
     A::try_new(
         data_type.clone(),
         values.offsets.into(),

diff --git a/src/io/parquet/read/deserialize/binary/utils.rs b/src/io/parquet/read/deserialize/binary/utils.rs
@@ -39,12 +39,20 @@ impl<O: Offset> Binary<O> {
     pub fn with_capacity(capacity: usize) -> Self {
         Self {
             offsets: Offsets::with_capacity(capacity),
-            values: Vec::with_capacity(capacity * 24),
+            values: Vec::with_capacity(capacity.min(100) * 24),
         }
     }
 
     #[inline]
     pub fn push(&mut self, v: &[u8]) {
+        if self.offsets.len() == 100 && self.offsets.capacity() > 100 {
+            let bytes_per_row = self.values.len() / 100 + 1;
+            let bytes_estimate = bytes_per_row * self.offsets.capacity();
+            if bytes_estimate > self.values.capacity() {
+                self.values.reserve(bytes_estimate - self.values.capacity());
+            }
+        }
+
         self.values.extend(v);
         self.offsets.try_push_usize(v.len()).unwrap()
     }