From 8aab4065eedfdee458ceb315b65e23cc8e9f06f5 Mon Sep 17 00:00:00 2001
From: Ritchie Vink <ritchie46@gmail.com>
Date: Mon, 6 Sep 2021 12:59:21 +0200
Subject: [PATCH] improve MutablePrimitiveArray

adds

* extend_trusted_len_values
* extend_trusted_len_values_unchecked
* extend_from_slice
---
 src/array/primitive/mutable.rs      | 46 +++++++++++++++++++++++++++++
 src/compute/like.rs                 |  5 +++-
 tests/it/array/primitive/mutable.rs | 32 ++++++++++++++++++++
 3 files changed, 82 insertions(+), 1 deletion(-)

diff --git a/src/array/primitive/mutable.rs b/src/array/primitive/mutable.rs
index ee330188126..0735347a860 100644
--- a/src/array/primitive/mutable.rs
+++ b/src/array/primitive/mutable.rs
@@ -172,6 +172,52 @@ impl<T: NativeType> MutablePrimitiveArray<T> {
             }
         }
     }
+    /// Extends the [`MutablePrimitiveArray`] from an iterator of values of trusted len.
+    /// This differs from `extend_trusted_len` which accepts in iterator of optional values.
+    #[inline]
+    pub fn extend_trusted_len_values<I>(&mut self, iterator: I)
+    where
+        I: TrustedLen<Item = T>,
+    {
+        unsafe { self.extend_trusted_len_values_unchecked(iterator) }
+    }
+
+    /// Extends the [`MutablePrimitiveArray`] from an iterator of values of trusted len.
+    /// This differs from `extend_trusted_len_unchecked` which accepts in iterator of optional values.
+    /// # Safety
+    /// The iterator must be trusted len.
+    #[inline]
+    pub unsafe fn extend_trusted_len_values_unchecked<I>(&mut self, iterator: I)
+    where
+        I: Iterator<Item = T>,
+    {
+        let size = iterator
+            .size_hint()
+            .1
+            .expect("upper bound should be set by trusted len iter");
+        self.values.extend_from_trusted_len_iter_unchecked(iterator);
+        self.update_all_valid(size);
+    }
+
+    #[inline]
+    /// Extends the [`MutablePrimitiveArray`] from a slice
+    pub fn extend_from_slice(&mut self, items: &[T]) {
+        self.values.extend_from_slice(items);
+        self.update_all_valid(items.len());
+    }
+
+    fn update_all_valid(&mut self, additional: usize) {
+        if let Some(validity) = self.validity.as_mut() {
+            validity.extend_constant(additional, true);
+        } else {
+            let validity = MutableBitmap::from_trusted_len_iter(
+                std::iter::repeat(true).take(self.len() + additional),
+            );
+            if validity.null_count() > 0 {
+                self.validity = Some(validity);
+            }
+        }
+    }
 
     fn init_validity(&mut self) {
         let mut validity = MutableBitmap::new();
diff --git a/src/compute/like.rs b/src/compute/like.rs
index bcd2ab71a96..1160721cef7 100644
--- a/src/compute/like.rs
+++ b/src/compute/like.rs
@@ -232,7 +232,10 @@ fn a_like_binary_scalar<O: Offset, F: Fn(bool) -> bool>(
 ) -> Result<BooleanArray> {
     let validity = lhs.validity();
     let pattern = std::str::from_utf8(rhs).map_err(|e| {
-        ArrowError::InvalidArgumentError(format!("Unable to convert the LIKE pattern to string: {}", e))
+        ArrowError::InvalidArgumentError(format!(
+            "Unable to convert the LIKE pattern to string: {}",
+            e
+        ))
     })?;
 
     let values = if !pattern.contains(is_like_pattern) {
diff --git a/tests/it/array/primitive/mutable.rs b/tests/it/array/primitive/mutable.rs
index f5f4cd381d7..ecc1fd8f994 100644
--- a/tests/it/array/primitive/mutable.rs
+++ b/tests/it/array/primitive/mutable.rs
@@ -126,6 +126,38 @@ fn extend_trusted_len() {
     assert_eq!(a.values(), &MutableBuffer::<i32>::from([1, 2, 0, 4]));
 }
 
+#[test]
+fn extend_trusted_len_values() {
+    let mut a = MutablePrimitiveArray::<i32>::new();
+    a.extend_trusted_len_values(vec![1, 2, 3].into_iter());
+    assert_eq!(a.validity(), &None);
+    assert_eq!(a.values(), &MutableBuffer::<i32>::from([1, 2, 3]));
+
+    let mut a = MutablePrimitiveArray::<i32>::new();
+    a.push(None);
+    a.extend_trusted_len_values(vec![1, 2].into_iter());
+    assert_eq!(
+        a.validity(),
+        &Some(MutableBitmap::from([false, true, true]))
+    );
+}
+
+#[test]
+fn extend_from_slice() {
+    let mut a = MutablePrimitiveArray::<i32>::new();
+    a.extend_from_slice(&[1, 2, 3]);
+    assert_eq!(a.validity(), &None);
+    assert_eq!(a.values(), &MutableBuffer::<i32>::from([1, 2, 3]));
+
+    let mut a = MutablePrimitiveArray::<i32>::new();
+    a.push(None);
+    a.extend_from_slice(&[1, 2]);
+    assert_eq!(
+        a.validity(),
+        &Some(MutableBitmap::from([false, true, true]))
+    );
+}
+
 #[test]
 fn set_validity() {
     let mut a = MutablePrimitiveArray::<i32>::new();