From b7e991366104d1647b955a828e0551256ef2e7c9 Mon Sep 17 00:00:00 2001
From: Ritchie Vink <ritchie46@gmail.com>
Date: Sat, 2 Oct 2021 09:28:10 +0200
Subject: [PATCH] expose `shrink_to_fit` to mutable arrays (#467)

Co-authored-by: Jorge Leitao <jorgecarleitao@gmail.com>
---
 src/array/binary/mutable.rs            | 12 ++++++++++++
 src/array/boolean/mutable.rs           | 12 ++++++++++++
 src/array/dictionary/mutable.rs        |  9 +++++++++
 src/array/fixed_size_binary/mutable.rs | 12 ++++++++++++
 src/array/fixed_size_list/mutable.rs   | 11 +++++++++++
 src/array/list/mutable.rs              | 14 +++++++++++++-
 src/array/mod.rs                       |  3 +++
 src/array/primitive/mutable.rs         | 12 ++++++++++++
 src/array/utf8/mutable.rs              | 13 +++++++++++++
 9 files changed, 97 insertions(+), 1 deletion(-)

diff --git a/src/array/binary/mutable.rs b/src/array/binary/mutable.rs
index 1eee5e10406..81f7294324f 100644
--- a/src/array/binary/mutable.rs
+++ b/src/array/binary/mutable.rs
@@ -135,6 +135,14 @@ 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();
+        self.offsets.shrink_to_fit();
+        if let Some(validity) = &mut self.validity {
+            validity.shrink_to_fit()
+        }
+    }
 }
 
 impl<O: Offset> MutableArray for MutableBinaryArray<O> {
@@ -180,6 +188,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..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>`.
@@ -461,6 +469,10 @@ impl MutableArray for MutableBooleanArray {
     fn push_null(&mut self) {
         self.push(None)
     }
+
+    fn shrink_to_fit(&mut self) {
+        self.shrink_to_fit()
+    }
 }
 
 impl Extend<Option<bool>> for MutableBooleanArray {
diff --git a/src/array/dictionary/mutable.rs b/src/array/dictionary/mutable.rs
index 17f6f4caa4d..895bec0f1bb 100644
--- a/src/array/dictionary/mutable.rs
+++ b/src/array/dictionary/mutable.rs
@@ -104,6 +104,12 @@ 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();
+        self.keys.shrink_to_fit();
+    }
 }
 
 impl<K: DictionaryKey, M: 'static + MutableArray> MutableArray for MutableDictionaryArray<K, M> {
@@ -144,6 +150,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..7717e594dda 100644
--- a/src/array/fixed_size_binary/mutable.rs
+++ b/src/array/fixed_size_binary/mutable.rs
@@ -149,6 +149,14 @@ 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();
+        if let Some(validity) = &mut self.validity {
+            validity.shrink_to_fit()
+        }
+    }
 }
 
 /// Accessors
@@ -204,6 +212,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..13924106f8e 100644
--- a/src/array/fixed_size_list/mutable.rs
+++ b/src/array/fixed_size_list/mutable.rs
@@ -74,6 +74,13 @@ 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();
+        if let Some(validity) = &mut self.validity {
+            validity.shrink_to_fit()
+        }
+    }
 }
 
 impl<M: MutableArray + 'static> MutableArray for MutableFixedSizeListArray<M> {
@@ -124,6 +131,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..1892dc50955 100644
--- a/src/array/list/mutable.rs
+++ b/src/array/list/mutable.rs
@@ -42,6 +42,15 @@ 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();
+        if let Some(validity) = &mut self.validity {
+            validity.shrink_to_fit()
+        }
+    }
 }
 
 impl<O: Offset, M: MutableArray + Default> Default for MutableListArray<O, M> {
@@ -179,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
     }
@@ -222,4 +231,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..e8d764fddc3 100644
--- a/src/array/primitive/mutable.rs
+++ b/src/array/primitive/mutable.rs
@@ -237,6 +237,14 @@ 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();
+        if let Some(validity) = &mut self.validity {
+            validity.shrink_to_fit()
+        }
+    }
 }
 
 /// Accessors
@@ -357,6 +365,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..25d4521d883 100644
--- a/src/array/utf8/mutable.rs
+++ b/src/array/utf8/mutable.rs
@@ -164,6 +164,15 @@ 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();
+        if let Some(validity) = &mut self.validity {
+            validity.shrink_to_fit()
+        }
+    }
 }
 
 impl<O: Offset> MutableUtf8Array<O> {
@@ -224,6 +233,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> {