From 72de99c5f6cd2973494af48f49d93b0411ee1186 Mon Sep 17 00:00:00 2001
From: Jorge Leitao <jorgecarleitao@gmail.com>
Date: Wed, 3 Aug 2022 05:30:03 +0200
Subject: [PATCH] Added MutableArray::reserve (#1202)

---
 src/array/binary/mutable.rs                  |  4 ++++
 src/array/boolean/mutable.rs                 |  4 ++++
 src/array/dictionary/mutable.rs              | 10 ++++++++++
 src/array/fixed_size_binary/mutable.rs       | 12 ++++++++++++
 src/array/fixed_size_list/mutable.rs         | 13 +++++++++++++
 src/array/list/mutable.rs                    | 13 +++++++++++++
 src/array/mod.rs                             |  3 +++
 src/array/primitive/mutable.rs               |  4 ++++
 src/array/utf8/mutable.rs                    |  4 ++++
 src/io/avro/read/nested.rs                   | 12 ++++++++++++
 src/io/parquet/read/statistics/dictionary.rs |  4 ++++
 src/io/parquet/read/statistics/list.rs       |  4 ++++
 src/io/parquet/read/statistics/map.rs        |  4 ++++
 src/io/parquet/read/statistics/struct_.rs    |  4 ++++
 14 files changed, 95 insertions(+)

diff --git a/src/array/binary/mutable.rs b/src/array/binary/mutable.rs
index 7c44a5ff446..9990a542c2d 100644
--- a/src/array/binary/mutable.rs
+++ b/src/array/binary/mutable.rs
@@ -230,6 +230,10 @@ impl<O: Offset> MutableArray for MutableBinaryArray<O> {
         self.push::<&[u8]>(None)
     }
 
+    fn reserve(&mut self, additional: usize) {
+        self.reserve(additional)
+    }
+
     fn shrink_to_fit(&mut self) {
         self.shrink_to_fit()
     }
diff --git a/src/array/boolean/mutable.rs b/src/array/boolean/mutable.rs
index 39acab03e07..83fac4b7ef5 100644
--- a/src/array/boolean/mutable.rs
+++ b/src/array/boolean/mutable.rs
@@ -495,6 +495,10 @@ impl MutableArray for MutableBooleanArray {
         self.push(None)
     }
 
+    fn reserve(&mut self, additional: usize) {
+        self.reserve(additional)
+    }
+
     fn shrink_to_fit(&mut self) {
         self.shrink_to_fit()
     }
diff --git a/src/array/dictionary/mutable.rs b/src/array/dictionary/mutable.rs
index c5ca7efb310..665d6fc9234 100644
--- a/src/array/dictionary/mutable.rs
+++ b/src/array/dictionary/mutable.rs
@@ -135,6 +135,11 @@ impl<K: DictionaryKey, M: MutableArray> MutableDictionaryArray<K, M> {
         Box::new(a)
     }
 
+    /// Reserves `additional` slots.
+    pub fn reserve(&mut self, additional: usize) {
+        self.keys.reserve(additional);
+    }
+
     /// Shrinks the capacity of the [`MutableDictionaryArray`] to fit its current length.
     pub fn shrink_to_fit(&mut self) {
         self.values.shrink_to_fit();
@@ -196,6 +201,11 @@ impl<K: DictionaryKey, M: 'static + MutableArray> MutableArray for MutableDictio
     fn push_null(&mut self) {
         self.keys.push(None)
     }
+
+    fn reserve(&mut self, additional: usize) {
+        self.reserve(additional)
+    }
+
     fn shrink_to_fit(&mut self) {
         self.shrink_to_fit()
     }
diff --git a/src/array/fixed_size_binary/mutable.rs b/src/array/fixed_size_binary/mutable.rs
index 528689ae4bc..dc70e49e4ac 100644
--- a/src/array/fixed_size_binary/mutable.rs
+++ b/src/array/fixed_size_binary/mutable.rs
@@ -175,6 +175,14 @@ impl MutableFixedSizeBinaryArray {
         std::slice::from_raw_parts(self.values.as_ptr().add(i * self.size), self.size)
     }
 
+    /// Reserves `additional` slots.
+    pub fn reserve(&mut self, additional: usize) {
+        self.values.reserve(additional * self.size);
+        if let Some(x) = self.validity.as_mut() {
+            x.reserve(additional)
+        }
+    }
+
     /// Shrinks the capacity of the [`MutableFixedSizeBinaryArray`] to fit its current length.
     pub fn shrink_to_fit(&mut self) {
         self.values.shrink_to_fit();
@@ -238,6 +246,10 @@ impl MutableArray for MutableFixedSizeBinaryArray {
         self.push::<&[u8]>(None);
     }
 
+    fn reserve(&mut self, additional: usize) {
+        self.reserve(additional)
+    }
+
     fn shrink_to_fit(&mut self) {
         self.shrink_to_fit()
     }
diff --git a/src/array/fixed_size_list/mutable.rs b/src/array/fixed_size_list/mutable.rs
index 5798e4472e2..c638e84033b 100644
--- a/src/array/fixed_size_list/mutable.rs
+++ b/src/array/fixed_size_list/mutable.rs
@@ -99,6 +99,15 @@ impl<M: MutableArray> MutableFixedSizeListArray<M> {
             None => self.init_validity(),
         }
     }
+
+    /// Reserves `additional` slots.
+    pub fn reserve(&mut self, additional: usize) {
+        self.values.reserve(additional);
+        if let Some(x) = self.validity.as_mut() {
+            x.reserve(additional)
+        }
+    }
+
     /// Shrinks the capacity of the [`MutableFixedSizeListArray`] to fit its current length.
     pub fn shrink_to_fit(&mut self) {
         self.values.shrink_to_fit();
@@ -157,6 +166,10 @@ impl<M: MutableArray + 'static> MutableArray for MutableFixedSizeListArray<M> {
         }
     }
 
+    fn reserve(&mut self, additional: usize) {
+        self.reserve(additional)
+    }
+
     fn shrink_to_fit(&mut self) {
         self.shrink_to_fit()
     }
diff --git a/src/array/list/mutable.rs b/src/array/list/mutable.rs
index 2f6d8d1ebf4..8d0dc3b05e8 100644
--- a/src/array/list/mutable.rs
+++ b/src/array/list/mutable.rs
@@ -193,6 +193,14 @@ impl<O: Offset, M: MutableArray> MutableListArray<O, M> {
         Box::new(a)
     }
 
+    /// Reserves `additional` slots.
+    pub fn reserve(&mut self, additional: usize) {
+        self.offsets.reserve(additional);
+        if let Some(x) = self.validity.as_mut() {
+            x.reserve(additional)
+        }
+    }
+
     /// Shrinks the capacity of the [`MutableListArray`] to fit its current length.
     pub fn shrink_to_fit(&mut self) {
         self.values.shrink_to_fit();
@@ -254,6 +262,11 @@ impl<O: Offset, M: MutableArray + 'static> MutableArray for MutableListArray<O,
     fn push_null(&mut self) {
         self.push_null()
     }
+
+    fn reserve(&mut self, additional: usize) {
+        self.reserve(additional)
+    }
+
     fn shrink_to_fit(&mut self) {
         self.shrink_to_fit();
     }
diff --git a/src/array/mod.rs b/src/array/mod.rs
index cc5038fa0f4..8010b5cf920 100644
--- a/src/array/mod.rs
+++ b/src/array/mod.rs
@@ -163,6 +163,9 @@ pub trait MutableArray: std::fmt::Debug + Send + Sync {
             .unwrap_or(true)
     }
 
+    /// Reserves additional slots to its capacity.
+    fn reserve(&mut self, additional: usize);
+
     /// Shrink the array to fit its length.
     fn shrink_to_fit(&mut self);
 }
diff --git a/src/array/primitive/mutable.rs b/src/array/primitive/mutable.rs
index c390690d3ba..5f2014db319 100644
--- a/src/array/primitive/mutable.rs
+++ b/src/array/primitive/mutable.rs
@@ -419,6 +419,10 @@ impl<T: NativeType> MutableArray for MutablePrimitiveArray<T> {
         self.push(None)
     }
 
+    fn reserve(&mut self, additional: usize) {
+        self.reserve(additional)
+    }
+
     fn shrink_to_fit(&mut self) {
         self.shrink_to_fit()
     }
diff --git a/src/array/utf8/mutable.rs b/src/array/utf8/mutable.rs
index 52150942370..0b87b97c25c 100644
--- a/src/array/utf8/mutable.rs
+++ b/src/array/utf8/mutable.rs
@@ -318,6 +318,10 @@ impl<O: Offset> MutableArray for MutableUtf8Array<O> {
         self.push::<&str>(None)
     }
 
+    fn reserve(&mut self, additional: usize) {
+        self.reserve(additional, 0)
+    }
+
     fn shrink_to_fit(&mut self) {
         self.shrink_to_fit()
     }
diff --git a/src/io/avro/read/nested.rs b/src/io/avro/read/nested.rs
index 450930cad73..981b6894f9b 100644
--- a/src/io/avro/read/nested.rs
+++ b/src/io/avro/read/nested.rs
@@ -112,6 +112,10 @@ impl<O: Offset> MutableArray for DynMutableListArray<O> {
         self.push_null()
     }
 
+    fn reserve(&mut self, _: usize) {
+        todo!();
+    }
+
     fn shrink_to_fit(&mut self) {
         todo!();
     }
@@ -195,6 +199,10 @@ impl MutableArray for FixedItemsUtf8Dictionary {
         self.push_null()
     }
 
+    fn reserve(&mut self, _: usize) {
+        todo!();
+    }
+
     fn shrink_to_fit(&mut self) {
         todo!();
     }
@@ -287,6 +295,10 @@ impl MutableArray for DynMutableStructArray {
         self.push_null()
     }
 
+    fn reserve(&mut self, _: usize) {
+        todo!();
+    }
+
     fn shrink_to_fit(&mut self) {
         todo!();
     }
diff --git a/src/io/parquet/read/statistics/dictionary.rs b/src/io/parquet/read/statistics/dictionary.rs
index b2553c62f1b..a4bc349741f 100644
--- a/src/io/parquet/read/statistics/dictionary.rs
+++ b/src/io/parquet/read/statistics/dictionary.rs
@@ -59,6 +59,10 @@ impl MutableArray for DynMutableDictionary {
         todo!()
     }
 
+    fn reserve(&mut self, _: usize) {
+        todo!();
+    }
+
     fn shrink_to_fit(&mut self) {
         todo!()
     }
diff --git a/src/io/parquet/read/statistics/list.rs b/src/io/parquet/read/statistics/list.rs
index 6b4bb2d7d3d..496c7deb9d3 100644
--- a/src/io/parquet/read/statistics/list.rs
+++ b/src/io/parquet/read/statistics/list.rs
@@ -73,6 +73,10 @@ impl MutableArray for DynMutableListArray {
         todo!()
     }
 
+    fn reserve(&mut self, _: usize) {
+        todo!();
+    }
+
     fn shrink_to_fit(&mut self) {
         todo!()
     }
diff --git a/src/io/parquet/read/statistics/map.rs b/src/io/parquet/read/statistics/map.rs
index 71ef8154dd0..c31dfb9d9e7 100644
--- a/src/io/parquet/read/statistics/map.rs
+++ b/src/io/parquet/read/statistics/map.rs
@@ -58,6 +58,10 @@ impl MutableArray for DynMutableMapArray {
         todo!()
     }
 
+    fn reserve(&mut self, _: usize) {
+        todo!();
+    }
+
     fn shrink_to_fit(&mut self) {
         todo!()
     }
diff --git a/src/io/parquet/read/statistics/struct_.rs b/src/io/parquet/read/statistics/struct_.rs
index 7ae80af3e79..c19a22b1dbd 100644
--- a/src/io/parquet/read/statistics/struct_.rs
+++ b/src/io/parquet/read/statistics/struct_.rs
@@ -55,6 +55,10 @@ impl MutableArray for DynMutableStructArray {
         todo!()
     }
 
+    fn reserve(&mut self, _: usize) {
+        todo!();
+    }
+
     fn shrink_to_fit(&mut self) {
         todo!()
     }