From d023c09a4b506db77832c70d4f5784b0add3f91a Mon Sep 17 00:00:00 2001
From: "Jorge C. Leitao" <jorgecarleitao@gmail.com>
Date: Tue, 9 Nov 2021 08:03:16 +0000
Subject: [PATCH] Added more tests.

---
 src/array/binary/mutable.rs      | 46 ++++++++++++++++++++--
 tests/it/array/binary/mutable.rs | 65 +++++++++++++++++++++++++++++++-
 2 files changed, 107 insertions(+), 4 deletions(-)

diff --git a/src/array/binary/mutable.rs b/src/array/binary/mutable.rs
index 28b1ebb08df..ae631ab7fda 100644
--- a/src/array/binary/mutable.rs
+++ b/src/array/binary/mutable.rs
@@ -92,6 +92,19 @@ impl<O: Offset> MutableBinaryArray<O> {
         }
     }
 
+    /// Initializes a new [`MutableBinaryArray`] with a pre-allocated capacity of slots and values.
+    pub fn with_capacities(capacity: usize, values: usize) -> Self {
+        let mut offsets = MutableBuffer::<O>::with_capacity(capacity + 1);
+        offsets.push(O::default());
+
+        Self {
+            data_type: Self::default_data_type(),
+            offsets,
+            values: MutableBuffer::<u8>::with_capacity(values),
+            validity: None,
+        }
+    }
+
     /// Reserves `additional` slots.
     pub fn reserve(&mut self, additional: usize) {
         self.offsets.reserve(additional);
@@ -145,6 +158,18 @@ impl<O: Offset> MutableBinaryArray<O> {
     }
 }
 
+impl<O: Offset> MutableBinaryArray<O> {
+    /// returns its values.
+    pub fn values(&self) -> &MutableBuffer<u8> {
+        &self.values
+    }
+
+    /// returns its offsets.
+    pub fn offsets(&self) -> &MutableBuffer<O> {
+        &self.offsets
+    }
+}
+
 impl<O: Offset> MutableArray for MutableBinaryArray<O> {
     fn len(&self) -> usize {
         self.offsets.len() - 1
@@ -228,15 +253,26 @@ impl<O: Offset> MutableBinaryArray<O> {
     }
 
     /// Creates a new [`BinaryArray`] from a [`TrustedLen`] of `&[u8]`.
+    /// # Safety
+    /// The iterator must be [`TrustedLen`](https://doc.rust-lang.org/std/iter/trait.TrustedLen.html).
+    /// I.e. that `size_hint().1` correctly reports its length.
     #[inline]
-    pub fn from_trusted_len_values_iter<T: AsRef<[u8]>, I: TrustedLen<Item = T>>(
+    pub unsafe fn from_trusted_len_values_iter_unchecked<T: AsRef<[u8]>, I: Iterator<Item = T>>(
         iterator: I,
     ) -> Self {
-        // soundness: I is `TrustedLen`
         let (offsets, values) = unsafe { trusted_len_values_iter(iterator) };
         Self::from_data(Self::default_data_type(), offsets, values, None)
     }
 
+    /// Creates a new [`BinaryArray`] from a [`TrustedLen`] of `&[u8]`.
+    #[inline]
+    pub fn from_trusted_len_values_iter<T: AsRef<[u8]>, I: TrustedLen<Item = T>>(
+        iterator: I,
+    ) -> Self {
+        // soundness: I is `TrustedLen`
+        unsafe { Self::from_trusted_len_values_iter_unchecked(iterator) }
+    }
+
     /// Creates a [`MutableBinaryArray`] from an falible iterator of trusted length.
     /// # Safety
     /// The iterator must be [`TrustedLen`](https://doc.rust-lang.org/std/iter/trait.TrustedLen.html).
@@ -252,7 +288,11 @@ impl<O: Offset> MutableBinaryArray<O> {
         let iterator = iterator.into_iter();
 
         // soundness: assumed trusted len
-        let (validity, offsets, values) = try_trusted_len_unzip(iterator)?;
+        let (mut validity, offsets, values) = try_trusted_len_unzip(iterator)?;
+
+        if validity.as_mut().unwrap().null_count() == 0 {
+            validity = None;
+        }
 
         Ok(Self::from_data(
             Self::default_data_type(),
diff --git a/tests/it/array/binary/mutable.rs b/tests/it/array/binary/mutable.rs
index bbf79431449..7258dea1332 100644
--- a/tests/it/array/binary/mutable.rs
+++ b/tests/it/array/binary/mutable.rs
@@ -1,5 +1,68 @@
-use arrow2::array::{BinaryArray, MutableBinaryArray};
+use std::ops::Deref;
+
+use arrow2::array::{BinaryArray, MutableArray, MutableBinaryArray};
 use arrow2::bitmap::Bitmap;
+use arrow2::error::ArrowError;
+
+#[test]
+fn new() {
+    assert_eq!(MutableBinaryArray::<i32>::new().len(), 0);
+
+    let a = MutableBinaryArray::<i32>::with_capacity(2);
+    assert_eq!(a.len(), 0);
+    assert!(a.offsets().capacity() >= 3);
+    assert_eq!(a.values().capacity(), 0);
+
+    let a = MutableBinaryArray::<i32>::with_capacities(2, 60);
+    assert_eq!(a.len(), 0);
+    assert!(a.offsets().capacity() >= 3);
+    assert!(a.values().capacity() >= 60);
+}
+
+#[test]
+fn from_iter() {
+    let iter = (0..3u8).map(|x| Some(vec![x; x as usize]));
+    let a: MutableBinaryArray<i32> = iter.clone().collect();
+    assert_eq!(a.values().deref(), &[1u8, 2, 2]);
+    assert_eq!(a.offsets().deref(), &[0, 0, 1, 3]);
+    assert_eq!(a.validity(), None);
+
+    let a = unsafe { MutableBinaryArray::<i32>::from_trusted_len_iter_unchecked(iter) };
+    assert_eq!(a.values().deref(), &[1u8, 2, 2]);
+    assert_eq!(a.offsets().deref(), &[0, 0, 1, 3]);
+    assert_eq!(a.validity(), None);
+}
+
+#[test]
+fn from_trusted_len_iter() {
+    let iter = (0..3u8).map(|x| vec![x; x as usize]);
+    let a: MutableBinaryArray<i32> = iter.clone().map(Some).collect();
+    assert_eq!(a.values().deref(), &[1u8, 2, 2]);
+    assert_eq!(a.offsets().deref(), &[0, 0, 1, 3]);
+    assert_eq!(a.validity(), None);
+
+    let a = unsafe {
+        MutableBinaryArray::<i32>::from_trusted_len_iter_unchecked(iter.clone().map(Some))
+    };
+    assert_eq!(a.values().deref(), &[1u8, 2, 2]);
+    assert_eq!(a.offsets().deref(), &[0, 0, 1, 3]);
+    assert_eq!(a.validity(), None);
+
+    let a = unsafe {
+        MutableBinaryArray::<i32>::try_from_trusted_len_iter_unchecked::<ArrowError, _, _>(
+            iter.clone().map(Some).map(Ok),
+        )
+    }
+    .unwrap();
+    assert_eq!(a.values().deref(), &[1u8, 2, 2]);
+    assert_eq!(a.offsets().deref(), &[0, 0, 1, 3]);
+    assert_eq!(a.validity(), None);
+
+    let a = unsafe { MutableBinaryArray::<i32>::from_trusted_len_values_iter_unchecked(iter) };
+    assert_eq!(a.values().deref(), &[1u8, 2, 2]);
+    assert_eq!(a.offsets().deref(), &[0, 0, 1, 3]);
+    assert_eq!(a.validity(), None);
+}
 
 #[test]
 fn push_null() {