From a2bcd42f27460b16399490609be42fa1e244c2c9 Mon Sep 17 00:00:00 2001
From: Ritchie Vink <ritchie46@gmail.com>
Date: Wed, 29 Sep 2021 13:09:26 +0200
Subject: [PATCH 1/3] expose shrink_to_fit to mutable arrays

---
 src/array/binary/mutable.rs            |  8 ++++++++
 src/array/boolean/mutable.rs           |  4 ++++
 src/array/dictionary/mutable.rs        |  8 ++++++++
 src/array/fixed_size_binary/mutable.rs |  9 +++++++++
 src/array/fixed_size_list/mutable.rs   |  8 ++++++++
 src/array/list/mutable.rs              |  9 +++++++++
 src/array/mod.rs                       |  3 +++
 src/array/primitive/mutable.rs         |  9 +++++++++
 src/array/utf8/mutable.rs              | 10 ++++++++++
 9 files changed, 68 insertions(+)

diff --git a/src/array/binary/mutable.rs b/src/array/binary/mutable.rs
index 1eee5e10406..3c905d27f6c 100644
--- a/src/array/binary/mutable.rs
+++ b/src/array/binary/mutable.rs
@@ -135,6 +135,10 @@ impl<O: Offset> MutableBinaryArray<O> {
         let a: BinaryArray<O> = self.into();
         Arc::new(a)
     }
+    /// Shrinks the capacity of the [`MutableBinaryArray`] to fit its current length.
+    pub fn shrink_to_fit(&mut self) {
+        self.values.shrink_to_fit();
+    }
 }
 
 impl<O: Offset> MutableArray for MutableBinaryArray<O> {
@@ -180,6 +184,10 @@ impl<O: Offset> MutableArray for MutableBinaryArray<O> {
     fn push_null(&mut self) {
         self.push::<&[u8]>(None)
     }
+
+    fn shrink_to_fit(&mut self) {
+        self.shrink_to_fit()
+    }
 }
 
 impl<O: Offset, P: AsRef<[u8]>> FromIterator<Option<P>> for MutableBinaryArray<O> {
diff --git a/src/array/boolean/mutable.rs b/src/array/boolean/mutable.rs
index b27e7b2bb2d..2473cf8d6ec 100644
--- a/src/array/boolean/mutable.rs
+++ b/src/array/boolean/mutable.rs
@@ -461,6 +461,10 @@ impl MutableArray for MutableBooleanArray {
     fn push_null(&mut self) {
         self.push(None)
     }
+
+    fn shrink_to_fit(&mut self) {
+        // no-op still have to implement this for bitmaps
+    }
 }
 
 impl Extend<Option<bool>> for MutableBooleanArray {
diff --git a/src/array/dictionary/mutable.rs b/src/array/dictionary/mutable.rs
index 17f6f4caa4d..8d54b60ede0 100644
--- a/src/array/dictionary/mutable.rs
+++ b/src/array/dictionary/mutable.rs
@@ -104,6 +104,11 @@ impl<K: DictionaryKey, M: MutableArray> MutableDictionaryArray<K, M> {
         let a: DictionaryArray<K> = self.into();
         Arc::new(a)
     }
+
+    /// Shrinks the capacity of the [`MutableDictionaryArray`] to fit its current length.
+    pub fn shrink_to_fit(&mut self) {
+        self.values.shrink_to_fit();
+    }
 }
 
 impl<K: DictionaryKey, M: 'static + MutableArray> MutableArray for MutableDictionaryArray<K, M> {
@@ -144,6 +149,9 @@ impl<K: DictionaryKey, M: 'static + MutableArray> MutableArray for MutableDictio
     fn push_null(&mut self) {
         self.keys.push(None)
     }
+    fn shrink_to_fit(&mut self) {
+        self.shrink_to_fit()
+    }
 }
 
 impl<K, M, T: Hash> TryExtend<Option<T>> for MutableDictionaryArray<K, M>
diff --git a/src/array/fixed_size_binary/mutable.rs b/src/array/fixed_size_binary/mutable.rs
index 9f3872c704b..4d6b364d733 100644
--- a/src/array/fixed_size_binary/mutable.rs
+++ b/src/array/fixed_size_binary/mutable.rs
@@ -149,6 +149,11 @@ impl MutableFixedSizeBinaryArray {
     pub unsafe fn value_unchecked(&self, i: usize) -> &[u8] {
         std::slice::from_raw_parts(self.values.as_ptr().add(i * self.size), self.size)
     }
+
+    /// Shrinks the capacity of the [`MutablePrimitive`] to fit its current length.
+    pub fn shrink_to_fit(&mut self) {
+        self.values.shrink_to_fit();
+    }
 }
 
 /// Accessors
@@ -204,6 +209,10 @@ impl MutableArray for MutableFixedSizeBinaryArray {
     fn push_null(&mut self) {
         self.values.extend_constant(self.size, 0);
     }
+
+    fn shrink_to_fit(&mut self) {
+        self.shrink_to_fit()
+    }
 }
 
 impl FixedSizeBinaryValues for MutableFixedSizeBinaryArray {
diff --git a/src/array/fixed_size_list/mutable.rs b/src/array/fixed_size_list/mutable.rs
index dcbc902b9dc..5780c29d0b3 100644
--- a/src/array/fixed_size_list/mutable.rs
+++ b/src/array/fixed_size_list/mutable.rs
@@ -74,6 +74,10 @@ impl<M: MutableArray> MutableFixedSizeListArray<M> {
             None => self.init_validity(),
         }
     }
+    /// Shrinks the capacity of the [`MutableFixedSizeList`] to fit its current length.
+    pub fn shrink_to_fit(&mut self) {
+        self.values.shrink_to_fit()
+    }
 }
 
 impl<M: MutableArray + 'static> MutableArray for MutableFixedSizeListArray<M> {
@@ -124,6 +128,10 @@ impl<M: MutableArray + 'static> MutableArray for MutableFixedSizeListArray<M> {
             self.init_validity()
         }
     }
+
+    fn shrink_to_fit(&mut self) {
+        self.shrink_to_fit()
+    }
 }
 
 impl<M, I, T> TryExtend<Option<I>> for MutableFixedSizeListArray<M>
diff --git a/src/array/list/mutable.rs b/src/array/list/mutable.rs
index 1de8430ef8a..93d75c71d89 100644
--- a/src/array/list/mutable.rs
+++ b/src/array/list/mutable.rs
@@ -42,6 +42,12 @@ impl<O: Offset, M: MutableArray + Default> MutableListArray<O, M> {
             validity: None,
         }
     }
+
+    /// Shrinks the capacity of the [`MutableList`] to fit its current length.
+    pub fn shrink_to_fit(&mut self) {
+        self.values.shrink_to_fit();
+        self.offsets.shrink_to_fit();
+    }
 }
 
 impl<O: Offset, M: MutableArray + Default> Default for MutableListArray<O, M> {
@@ -222,4 +228,7 @@ impl<O: Offset, M: MutableArray + 'static> MutableArray for MutableListArray<O,
     fn push_null(&mut self) {
         self.push_null()
     }
+    fn shrink_to_fit(&mut self) {
+        self.shrink_to_fit();
+    }
 }
diff --git a/src/array/mod.rs b/src/array/mod.rs
index 3d1afd382aa..1068db6c293 100644
--- a/src/array/mod.rs
+++ b/src/array/mod.rs
@@ -155,6 +155,9 @@ pub trait MutableArray: std::fmt::Debug {
             .map(|x| x.get(index))
             .unwrap_or(true)
     }
+
+    /// Shrink the array to fit its length.
+    fn shrink_to_fit(&mut self);
 }
 
 macro_rules! general_dyn {
diff --git a/src/array/primitive/mutable.rs b/src/array/primitive/mutable.rs
index bb5b611d638..9b271d52bd1 100644
--- a/src/array/primitive/mutable.rs
+++ b/src/array/primitive/mutable.rs
@@ -237,6 +237,11 @@ impl<T: NativeType> MutablePrimitiveArray<T> {
         let a: PrimitiveArray<T> = self.into();
         Arc::new(a)
     }
+
+    /// Shrinks the capacity of the [`MutablePrimitive`] to fit its current length.
+    pub fn shrink_to_fit(&mut self) {
+        self.values.shrink_to_fit();
+    }
 }
 
 /// Accessors
@@ -357,6 +362,10 @@ impl<T: NativeType> MutableArray for MutablePrimitiveArray<T> {
     fn push_null(&mut self) {
         self.push(None)
     }
+
+    fn shrink_to_fit(&mut self) {
+        self.shrink_to_fit()
+    }
 }
 
 impl<T: NativeType + NaturalDataType> MutablePrimitiveArray<T> {
diff --git a/src/array/utf8/mutable.rs b/src/array/utf8/mutable.rs
index 78b3a7ba7d0..b7a3c7ea0ec 100644
--- a/src/array/utf8/mutable.rs
+++ b/src/array/utf8/mutable.rs
@@ -164,6 +164,12 @@ impl<O: Offset> MutableUtf8Array<O> {
         let a: Utf8Array<O> = self.into();
         Arc::new(a)
     }
+
+    /// Shrinks the capacity of the [`MutableUtf8`] to fit its current length.
+    pub fn shrink_to_fit(&mut self) {
+        self.values.shrink_to_fit();
+        self.offsets.shrink_to_fit();
+    }
 }
 
 impl<O: Offset> MutableUtf8Array<O> {
@@ -224,6 +230,10 @@ impl<O: Offset> MutableArray for MutableUtf8Array<O> {
     fn push_null(&mut self) {
         self.push::<&str>(None)
     }
+
+    fn shrink_to_fit(&mut self) {
+        self.shrink_to_fit()
+    }
 }
 
 impl<O: Offset, P: AsRef<str>> FromIterator<Option<P>> for MutableUtf8Array<O> {

From f70122bf2d608e395f16c328ec315fa4fdcc6e20 Mon Sep 17 00:00:00 2001
From: Ritchie Vink <ritchie46@gmail.com>
Date: Thu, 30 Sep 2021 20:12:39 +0200
Subject: [PATCH 2/3] also shrink the validity bitmaps

---
 src/array/binary/mutable.rs            |  3 +++
 src/array/boolean/mutable.rs           | 10 +++++++++-
 src/array/dictionary/mutable.rs        |  1 +
 src/array/fixed_size_binary/mutable.rs |  3 +++
 src/array/fixed_size_list/mutable.rs   |  5 ++++-
 src/array/list/mutable.rs              |  5 ++++-
 src/array/primitive/mutable.rs         |  3 +++
 src/array/utf8/mutable.rs              |  3 +++
 8 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/src/array/binary/mutable.rs b/src/array/binary/mutable.rs
index 3c905d27f6c..b56a932ed77 100644
--- a/src/array/binary/mutable.rs
+++ b/src/array/binary/mutable.rs
@@ -138,6 +138,9 @@ impl<O: Offset> MutableBinaryArray<O> {
     /// Shrinks the capacity of the [`MutableBinaryArray`] to fit its current length.
     pub fn shrink_to_fit(&mut self) {
         self.values.shrink_to_fit();
+        if let Some(validity) = &mut self.validity {
+            validity.shrink_to_fit()
+        }
     }
 }
 
diff --git a/src/array/boolean/mutable.rs b/src/array/boolean/mutable.rs
index 2473cf8d6ec..a8f0a2bc9df 100644
--- a/src/array/boolean/mutable.rs
+++ b/src/array/boolean/mutable.rs
@@ -293,6 +293,14 @@ impl MutableBooleanArray {
         // Safety: `I` is `TrustedLen`
         unsafe { Self::try_from_trusted_len_iter_unchecked(iterator) }
     }
+
+    /// Shrinks the capacity of the [`MutableBooleanArray`] to fit its current length.
+    pub fn shrink_to_fit(&mut self) {
+        self.values.shrink_to_fit();
+        if let Some(validity) = &mut self.validity {
+            validity.shrink_to_fit()
+        }
+    }
 }
 
 /// Creates a Bitmap and an optional [`MutableBitmap`] from an iterator of `Option<bool>`.
@@ -463,7 +471,7 @@ impl MutableArray for MutableBooleanArray {
     }
 
     fn shrink_to_fit(&mut self) {
-        // no-op still have to implement this for bitmaps
+        self.shrink_to_fit()
     }
 }
 
diff --git a/src/array/dictionary/mutable.rs b/src/array/dictionary/mutable.rs
index 8d54b60ede0..895bec0f1bb 100644
--- a/src/array/dictionary/mutable.rs
+++ b/src/array/dictionary/mutable.rs
@@ -108,6 +108,7 @@ impl<K: DictionaryKey, M: MutableArray> MutableDictionaryArray<K, M> {
     /// Shrinks the capacity of the [`MutableDictionaryArray`] to fit its current length.
     pub fn shrink_to_fit(&mut self) {
         self.values.shrink_to_fit();
+        self.keys.shrink_to_fit();
     }
 }
 
diff --git a/src/array/fixed_size_binary/mutable.rs b/src/array/fixed_size_binary/mutable.rs
index 4d6b364d733..7717e594dda 100644
--- a/src/array/fixed_size_binary/mutable.rs
+++ b/src/array/fixed_size_binary/mutable.rs
@@ -153,6 +153,9 @@ impl MutableFixedSizeBinaryArray {
     /// Shrinks the capacity of the [`MutablePrimitive`] to fit its current length.
     pub fn shrink_to_fit(&mut self) {
         self.values.shrink_to_fit();
+        if let Some(validity) = &mut self.validity {
+            validity.shrink_to_fit()
+        }
     }
 }
 
diff --git a/src/array/fixed_size_list/mutable.rs b/src/array/fixed_size_list/mutable.rs
index 5780c29d0b3..13924106f8e 100644
--- a/src/array/fixed_size_list/mutable.rs
+++ b/src/array/fixed_size_list/mutable.rs
@@ -76,7 +76,10 @@ impl<M: MutableArray> MutableFixedSizeListArray<M> {
     }
     /// Shrinks the capacity of the [`MutableFixedSizeList`] to fit its current length.
     pub fn shrink_to_fit(&mut self) {
-        self.values.shrink_to_fit()
+        self.values.shrink_to_fit();
+        if let Some(validity) = &mut self.validity {
+            validity.shrink_to_fit()
+        }
     }
 }
 
diff --git a/src/array/list/mutable.rs b/src/array/list/mutable.rs
index 93d75c71d89..1892dc50955 100644
--- a/src/array/list/mutable.rs
+++ b/src/array/list/mutable.rs
@@ -47,6 +47,9 @@ impl<O: Offset, M: MutableArray + Default> MutableListArray<O, M> {
     pub fn shrink_to_fit(&mut self) {
         self.values.shrink_to_fit();
         self.offsets.shrink_to_fit();
+        if let Some(validity) = &mut self.validity {
+            validity.shrink_to_fit()
+        }
     }
 }
 
@@ -185,7 +188,7 @@ impl<O: Offset, M: MutableArray> MutableListArray<O, M> {
     }
 }
 
-impl<O: Offset, M: MutableArray + 'static> MutableArray for MutableListArray<O, M> {
+impl<O: Offset, M: MutableArray + Default + 'static> MutableArray for MutableListArray<O, M> {
     fn len(&self) -> usize {
         self.offsets.len() - 1
     }
diff --git a/src/array/primitive/mutable.rs b/src/array/primitive/mutable.rs
index 9b271d52bd1..e8d764fddc3 100644
--- a/src/array/primitive/mutable.rs
+++ b/src/array/primitive/mutable.rs
@@ -241,6 +241,9 @@ impl<T: NativeType> MutablePrimitiveArray<T> {
     /// Shrinks the capacity of the [`MutablePrimitive`] to fit its current length.
     pub fn shrink_to_fit(&mut self) {
         self.values.shrink_to_fit();
+        if let Some(validity) = &mut self.validity {
+            validity.shrink_to_fit()
+        }
     }
 }
 
diff --git a/src/array/utf8/mutable.rs b/src/array/utf8/mutable.rs
index b7a3c7ea0ec..25d4521d883 100644
--- a/src/array/utf8/mutable.rs
+++ b/src/array/utf8/mutable.rs
@@ -169,6 +169,9 @@ impl<O: Offset> MutableUtf8Array<O> {
     pub fn shrink_to_fit(&mut self) {
         self.values.shrink_to_fit();
         self.offsets.shrink_to_fit();
+        if let Some(validity) = &mut self.validity {
+            validity.shrink_to_fit()
+        }
     }
 }
 

From 15330236c34a3759b6035cb13fb8e9e1177cf99c Mon Sep 17 00:00:00 2001
From: Ritchie Vink <ritchie46@gmail.com>
Date: Fri, 1 Oct 2021 09:04:08 +0200
Subject: [PATCH 3/3] Update src/array/binary/mutable.rs

Co-authored-by: Jorge Leitao <jorgecarleitao@gmail.com>
---
 src/array/binary/mutable.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/array/binary/mutable.rs b/src/array/binary/mutable.rs
index b56a932ed77..81f7294324f 100644
--- a/src/array/binary/mutable.rs
+++ b/src/array/binary/mutable.rs
@@ -138,6 +138,7 @@ impl<O: Offset> MutableBinaryArray<O> {
     /// Shrinks the capacity of the [`MutableBinaryArray`] to fit its current length.
     pub fn shrink_to_fit(&mut self) {
         self.values.shrink_to_fit();
+        self.offsets.shrink_to_fit();
         if let Some(validity) = &mut self.validity {
             validity.shrink_to_fit()
         }