From eb7a75445441c78dadaec9456bd240a144d4cd2b Mon Sep 17 00:00:00 2001
From: "Jorge C. Leitao" <jorgecarleitao@gmail.com>
Date: Thu, 16 Dec 2021 17:29:08 +0000
Subject: [PATCH] Added MutablePrimitiveArray::extend_constant

---
 src/array/primitive/mutable.rs      | 21 ++++++++++++++
 tests/it/array/primitive/mutable.rs | 45 +++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+)

diff --git a/src/array/primitive/mutable.rs b/src/array/primitive/mutable.rs
index 15c7b9cd4a3..ac5613b754d 100644
--- a/src/array/primitive/mutable.rs
+++ b/src/array/primitive/mutable.rs
@@ -142,6 +142,27 @@ impl<T: NativeType> MutablePrimitiveArray<T> {
         }
     }
 
+    /// Extends the [`MutablePrimitiveArray`] with a constant
+    #[inline]
+    pub fn extend_constant(&mut self, additional: usize, value: Option<T>) {
+        if let Some(value) = value {
+            self.values.extend_constant(additional, value);
+            if let Some(validity) = &mut self.validity {
+                validity.extend_constant(additional, true)
+            }
+        } else {
+            if let Some(validity) = &mut self.validity {
+                validity.extend_constant(additional, false)
+            } else {
+                let mut validity = MutableBitmap::with_capacity(self.values.capacity());
+                validity.extend_constant(self.len(), true);
+                validity.extend_constant(additional, false);
+                self.validity = Some(validity)
+            }
+            self.values.extend_constant(additional, T::default());
+        }
+    }
+
     /// Extends the [`MutablePrimitiveArray`] from an iterator of trusted len.
     #[inline]
     pub fn extend_trusted_len<P, I>(&mut self, iterator: I)
diff --git a/tests/it/array/primitive/mutable.rs b/tests/it/array/primitive/mutable.rs
index f54167eb3f4..9615aa022a5 100644
--- a/tests/it/array/primitive/mutable.rs
+++ b/tests/it/array/primitive/mutable.rs
@@ -126,6 +126,51 @@ fn extend_trusted_len() {
     assert_eq!(a.values(), &MutableBuffer::<i32>::from([1, 2, 0, 4]));
 }
 
+#[test]
+fn extend_constant_no_validity() {
+    let mut a = MutablePrimitiveArray::<i32>::new();
+    a.push(Some(1));
+    a.extend_constant(2, Some(3));
+    assert_eq!(a.validity(), None);
+    assert_eq!(a.values(), &MutableBuffer::<i32>::from([1, 3, 3]));
+}
+
+#[test]
+fn extend_constant_validity() {
+    let mut a = MutablePrimitiveArray::<i32>::new();
+    a.push(Some(1));
+    a.extend_constant(2, None);
+    assert_eq!(
+        a.validity(),
+        Some(&MutableBitmap::from([true, false, false]))
+    );
+    assert_eq!(a.values(), &MutableBuffer::<i32>::from([1, 0, 0]));
+}
+
+#[test]
+fn extend_constant_validity_inverse() {
+    let mut a = MutablePrimitiveArray::<i32>::new();
+    a.push(None);
+    a.extend_constant(2, Some(1));
+    assert_eq!(
+        a.validity(),
+        Some(&MutableBitmap::from([false, true, true]))
+    );
+    assert_eq!(a.values(), &MutableBuffer::<i32>::from([0, 1, 1]));
+}
+
+#[test]
+fn extend_constant_validity_none() {
+    let mut a = MutablePrimitiveArray::<i32>::new();
+    a.push(None);
+    a.extend_constant(2, None);
+    assert_eq!(
+        a.validity(),
+        Some(&MutableBitmap::from([false, false, false]))
+    );
+    assert_eq!(a.values(), &MutableBuffer::<i32>::from([0, 0, 0]));
+}
+
 #[test]
 fn extend_trusted_len_values() {
     let mut a = MutablePrimitiveArray::<i32>::new();