Added MutablePrimitiveArray::extend_constant

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)

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();