Added trait TryPush.

jorgecarleitao · Aug 22, 2021 · a7f23ca · a7f23ca
1 parent 83256a8
commit a7f23ca
Show file tree

Hide file tree

Showing 7 changed files with 115 additions and 70 deletions.
diff --git a/src/array/binary/mutable.rs b/src/array/binary/mutable.rs
@@ -1,7 +1,7 @@
 use std::{iter::FromIterator, sync::Arc};
 
 use crate::{
-    array::{Array, MutableArray, Offset, TryExtend},
+    array::{Array, MutableArray, Offset, TryExtend, TryPush},
     bitmap::MutableBitmap,
     buffer::MutableBuffer,
     datatypes::DataType,
@@ -61,34 +61,9 @@ impl<O: Offset> MutableBinaryArray<O> {
         *self.offsets.last().unwrap()
     }
 
-    pub fn try_push<T: AsRef<[u8]>>(&mut self, value: Option<T>) -> Result<()> {
-        match value {
-            Some(value) => {
-                let bytes = value.as_ref();
-
-                let size = O::from_usize(self.values.len() + bytes.len())
-                    .ok_or(ArrowError::KeyOverflowError)?;
-
-                self.values.extend_from_slice(bytes);
-
-                self.offsets.push(size);
-
-                match &mut self.validity {
-                    Some(validity) => validity.push(true),
-                    None => {}
-                }
-            }
-            None => {
-                self.offsets.push(self.last_offset());
-                match &mut self.validity {
-                    Some(validity) => validity.push(false),
-                    None => self.init_validity(),
-                }
-            }
-        }
-        Ok(())
-    }
-
+    /// Pushes a new element to the array.
+    /// # Panic
+    /// This operation panics iff the length of all values (in bytes) exceeds `O` maximum value.
     pub fn push<T: AsRef<[u8]>>(&mut self, value: Option<T>) {
         self.try_push(value).unwrap()
     }
@@ -176,3 +151,33 @@ impl<O: Offset, T: AsRef<[u8]>> TryExtend<Option<T>> for MutableBinaryArray<O> {
         iter.try_for_each(|x| self.try_push(x))
     }
 }
+
+impl<O: Offset, T: AsRef<[u8]>> TryPush<Option<T>> for MutableBinaryArray<O> {
+    fn try_push(&mut self, value: Option<T>) -> Result<()> {
+        match value {
+            Some(value) => {
+                let bytes = value.as_ref();
+
+                let size = O::from_usize(self.values.len() + bytes.len())
+                    .ok_or(ArrowError::KeyOverflowError)?;
+
+                self.values.extend_from_slice(bytes);
+
+                self.offsets.push(size);
+
+                match &mut self.validity {
+                    Some(validity) => validity.push(true),
+                    None => {}
+                }
+            }
+            None => {
+                self.offsets.push(self.last_offset());
+                match &mut self.validity {
+                    Some(validity) => validity.push(false),
+                    None => self.init_validity(),
+                }
+            }
+        }
+        Ok(())
+    }
+}
diff --git a/src/array/boolean/mutable.rs b/src/array/boolean/mutable.rs
@@ -1,9 +1,8 @@
 use std::iter::FromIterator;
 use std::sync::Arc;
 
-use crate::array::TryExtend;
 use crate::{
-    array::{Array, MutableArray},
+    array::{Array, MutableArray, TryExtend, TryPush},
     bitmap::MutableBitmap,
     datatypes::DataType,
     error::Result,
@@ -352,6 +351,14 @@ impl TryExtend<Option<bool>> for MutableBooleanArray {
     }
 }
 
+impl TryPush<Option<bool>> for MutableBooleanArray {
+    /// This is infalible and is implemented for consistency with all other types
+    fn try_push(&mut self, item: Option<bool>) -> Result<()> {
+        self.push(item);
+        Ok(())
+    }
+}
+
 impl PartialEq for MutableBooleanArray {
     fn eq(&self, other: &Self) -> bool {
         self.iter().eq(other.iter())

diff --git a/src/array/list/mutable.rs b/src/array/list/mutable.rs
@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use crate::{
-    array::{Array, MutableArray, Offset, TryExtend},
+    array::{Array, MutableArray, Offset, TryExtend, TryPush},
     bitmap::MutableBitmap,
     buffer::MutableBuffer,
     datatypes::{DataType, Field},
@@ -67,13 +67,25 @@ where
 {
     fn try_extend<II: IntoIterator<Item = Option<I>>>(&mut self, iter: II) -> Result<()> {
         for items in iter {
-            if let Some(items) = items {
-                let values = self.mut_values();
-                values.try_extend(items)?;
-                self.try_push_valid()?;
-            } else {
-                self.push_null();
-            }
+            self.try_push(items)?;
+        }
+        Ok(())
+    }
+}
+
+impl<O, M, I, T> TryPush<Option<I>> for MutableListArray<O, M>
+where
+    O: Offset,
+    M: MutableArray + TryExtend<Option<T>>,
+    I: IntoIterator<Item = Option<T>>,
+{
+    fn try_push(&mut self, item: Option<I>) -> Result<()> {
+        if let Some(items) = item {
+            let values = self.mut_values();
+            values.try_extend(items)?;
+            self.try_push_valid()?;
+        } else {
+            self.push_null();
         }
         Ok(())
     }

diff --git a/src/array/mod.rs b/src/array/mod.rs
@@ -414,6 +414,12 @@ pub trait TryExtend<A> {
     fn try_extend<I: IntoIterator<Item = A>>(&mut self, iter: I) -> Result<()>;
 }
 
+/// A trait describing the ability of a struct to receive new items.
+pub trait TryPush<A> {
+    /// Tries to push a new element.
+    fn try_push(&mut self, item: A) -> Result<()>;
+}
+
 fn display_helper<T: std::fmt::Display, I: IntoIterator<Item = Option<T>>>(iter: I) -> Vec<String> {
     let iterator = iter.into_iter();
     let len = iterator.size_hint().0;

diff --git a/src/array/primitive/mutable.rs b/src/array/primitive/mutable.rs
@@ -1,7 +1,7 @@
 use std::{iter::FromIterator, sync::Arc};
 
 use crate::{
-    array::{Array, MutableArray, TryExtend},
+    array::{Array, MutableArray, TryExtend, TryPush},
     bitmap::MutableBitmap,
     buffer::MutableBuffer,
     datatypes::DataType,
@@ -266,6 +266,14 @@ impl<T: NativeType> TryExtend<Option<T>> for MutablePrimitiveArray<T> {
     }
 }
 
+impl<T: NativeType> TryPush<Option<T>> for MutablePrimitiveArray<T> {
+    /// This is infalible and is implemented for consistency with all other types
+    fn try_push(&mut self, item: Option<T>) -> Result<()> {
+        self.push(item);
+        Ok(())
+    }
+}
+
 impl<T: NativeType> MutableArray for MutablePrimitiveArray<T> {
     fn len(&self) -> usize {
         self.values.len()

diff --git a/src/array/utf8/mutable.rs b/src/array/utf8/mutable.rs
@@ -3,7 +3,7 @@ use std::{iter::FromIterator, sync::Arc};
 use crate::{
     array::{
         specification::{check_offsets, check_offsets_and_utf8},
-        Array, MutableArray, Offset, TryExtend,
+        Array, MutableArray, Offset, TryExtend, TryPush,
     },
     bitmap::MutableBitmap,
     buffer::MutableBuffer,
@@ -126,35 +126,6 @@ impl<O: Offset> MutableUtf8Array<O> {
         *self.offsets.last().unwrap()
     }
 
-    /// Tries to push a new element to the array.
-    /// # Error
-    /// This operation errors iff the length of all values (in bytes) exceeds `O` maximum value.
-    pub fn try_push<T: AsRef<str>>(&mut self, value: Option<T>) -> Result<()> {
-        match value {
-            Some(value) => {
-                let bytes = value.as_ref().as_bytes();
-                self.values.extend_from_slice(bytes);
-
-                let size = O::from_usize(self.values.len()).ok_or(ArrowError::KeyOverflowError)?;
-
-                self.offsets.push(size);
-
-                match &mut self.validity {
-                    Some(validity) => validity.push(true),
-                    None => {}
-                }
-            }
-            None => {
-                self.offsets.push(self.last_offset());
-                match &mut self.validity {
-                    Some(validity) => validity.push(false),
-                    None => self.init_validity(),
-                }
-            }
-        }
-        Ok(())
-    }
-
     /// Pushes a new element to the array.
     /// # Panic
     /// This operation panics iff the length of all values (in bytes) exceeds `O` maximum value.
@@ -340,6 +311,34 @@ impl<O: Offset, T: AsRef<str>> TryExtend<Option<T>> for MutableUtf8Array<O> {
     }
 }
 
+impl<O: Offset, T: AsRef<str>> TryPush<Option<T>> for MutableUtf8Array<O> {
+    fn try_push(&mut self, value: Option<T>) -> Result<()> {
+        match value {
+            Some(value) => {
+                let bytes = value.as_ref().as_bytes();
+                self.values.extend_from_slice(bytes);
+
+                let size = O::from_usize(self.values.len()).ok_or(ArrowError::KeyOverflowError)?;
+
+                self.offsets.push(size);
+
+                match &mut self.validity {
+                    Some(validity) => validity.push(true),
+                    None => {}
+                }
+            }
+            None => {
+                self.offsets.push(self.last_offset());
+                match &mut self.validity {
+                    Some(validity) => validity.push(false),
+                    None => self.init_validity(),
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
 /// Creates [`MutableBitmap`] and two [`MutableBuffer`]s from an iterator of `Option`.
 /// The first buffer corresponds to a offset buffer, the second one
 /// corresponds to a values buffer.

diff --git a/tests/it/array/list/mutable.rs b/tests/it/array/list/mutable.rs
@@ -29,3 +29,11 @@ fn basics() {
     );
     assert_eq!(expected, array);
 }
+
+#[test]
+fn push() {
+    let mut array = MutableListArray::<i32, MutablePrimitiveArray<i32>>::new();
+    array
+        .try_push(Some(vec![Some(1i32), Some(2), Some(3)]))
+        .unwrap();
+}