Add pop test about some MutableArray

ygf11 · Apr 27, 2022 · 898c7c8 · 898c7c8
1 parent b8d2002
commit 898c7c8
Show file tree

Hide file tree

Showing 4 changed files with 151 additions and 2 deletions.
diff --git a/src/array/fixed_size_binary/mutable.rs b/src/array/fixed_size_binary/mutable.rs
@@ -118,7 +118,6 @@ impl MutableFixedSizeBinaryArray {
         if self.values.len() < self.size {
             return None;
         }
-
         let value_start = self.values.len() - self.size;
         let value = self.values.split_off(value_start);
         self.validity

diff --git a/tests/it/array/binary/mutable.rs b/tests/it/array/binary/mutable.rs
@@ -73,6 +73,54 @@ fn push_null() {
     assert_eq!(array.validity(), Some(&Bitmap::from([false])));
 }
 
+#[test]
+fn pop() {
+    let mut a = MutableBinaryArray::<i32>::new();
+    a.push(Some(b"first"));
+    a.push(Some(b"second"));
+    a.push::<Vec<u8>>(None);
+    a.push_null();
+
+    assert_eq!(a.pop(), None);
+    assert_eq!(a.len(), 3);
+    assert_eq!(a.pop(), None);
+    assert_eq!(a.len(), 2);
+    assert_eq!(a.pop(), Some(b"second".to_vec()));
+    assert_eq!(a.len(), 1);
+    assert_eq!(a.pop(), Some(b"first".to_vec()));
+    assert_eq!(a.len(), 0);
+    assert_eq!(a.pop(), None);
+    assert_eq!(a.len(), 0);
+}
+
+#[test]
+fn pop_all_some() {
+    let mut a = MutableBinaryArray::<i32>::new();
+    a.push(Some(b"first"));
+    a.push(Some(b"second"));
+    a.push(Some(b"third"));
+    a.push(Some(b"firth"));
+
+    for _ in 0..4 {
+        a.push(Some(b"aaaa"));
+    }
+
+    a.push(Some(b"bbbb"));
+
+    assert_eq!(a.pop(), Some(b"bbbb".to_vec()));
+    assert_eq!(a.pop(), Some(b"aaaa".to_vec()));
+    assert_eq!(a.pop(), Some(b"aaaa".to_vec()));
+    assert_eq!(a.pop(), Some(b"aaaa".to_vec()));
+    assert_eq!(a.len(), 5);
+    assert_eq!(a.pop(), Some(b"aaaa".to_vec()));
+    assert_eq!(a.pop(), Some(b"firth".to_vec()));
+    assert_eq!(a.pop(), Some(b"third".to_vec()));
+    assert_eq!(a.pop(), Some(b"second".to_vec()));
+    assert_eq!(a.pop(), Some(b"first".to_vec()));
+    assert!(a.is_empty());
+    assert_eq!(a.pop(), None);
+}
+
 #[test]
 fn extend_trusted_len_values() {
     let mut array = MutableBinaryArray::<i32>::new();

diff --git a/tests/it/array/fixed_size_binary/mutable.rs b/tests/it/array/fixed_size_binary/mutable.rs
@@ -67,6 +67,62 @@ fn push_null() {
     assert_eq!(array.validity(), Some(&Bitmap::from([false])));
 }
 
+#[test]
+fn pop() {
+    let mut a = MutableFixedSizeBinaryArray::new(2);
+    a.push(Some(b"aa"));
+    a.push::<&[u8]>(None);
+    a.push(Some(b"bb"));
+    a.push::<&[u8]>(None);
+
+    assert_eq!(a.pop(), None);
+    assert_eq!(a.len(), 3);
+    assert_eq!(a.pop(), Some(b"bb".to_vec()));
+    assert_eq!(a.len(), 2);
+    assert_eq!(a.pop(), None);
+    assert_eq!(a.len(), 1);
+    assert_eq!(a.pop(), Some(b"aa".to_vec()));
+    assert!(a.is_empty());
+    assert_eq!(a.pop(), None);
+    assert!(a.is_empty());
+}
+
+#[test]
+fn pop_all_some() {
+    let mut a = MutableFixedSizeBinaryArray::new(2);
+    a.push(Some(b"aa"));
+    a.push(Some(b"bb"));
+    a.push(Some(b"cc"));
+    a.push(Some(b"dd"));
+
+    for _ in 0..4 {
+        a.push(Some(b"11"));
+    }
+
+    a.push(Some(b"22"));
+
+    assert_eq!(a.pop(), Some(b"22".to_vec()));
+    assert_eq!(a.pop(), Some(b"11".to_vec()));
+    assert_eq!(a.pop(), Some(b"11".to_vec()));
+    assert_eq!(a.pop(), Some(b"11".to_vec()));
+    assert_eq!(a.len(), 5);
+
+    assert_eq!(
+        a,
+        MutableFixedSizeBinaryArray::try_from_iter(
+            vec![
+                Some(b"aa"),
+                Some(b"bb"),
+                Some(b"cc"),
+                Some(b"dd"),
+                Some(b"11"),
+            ],
+            2,
+        )
+        .unwrap()
+    );
+}
+
 #[test]
 fn as_arc() {
     let mut array = MutableFixedSizeBinaryArray::try_from_iter(

diff --git a/tests/it/array/utf8/mutable.rs b/tests/it/array/utf8/mutable.rs
@@ -1,4 +1,4 @@
-use arrow2::array::{MutableUtf8Array, Utf8Array};
+use arrow2::array::{MutableArray, MutableUtf8Array, Utf8Array};
 use arrow2::bitmap::Bitmap;
 use arrow2::datatypes::DataType;
 
@@ -19,6 +19,52 @@ fn push_null() {
     assert_eq!(array.validity(), Some(&Bitmap::from([false])));
 }
 
+#[test]
+fn pop() {
+    let mut a = MutableUtf8Array::<i32>::new();
+    a.push(Some("first"));
+    a.push(Some("second"));
+    a.push(Some("third"));
+    a.push::<&str>(None);
+
+    assert_eq!(a.pop(), None);
+    assert_eq!(a.len(), 3);
+    assert_eq!(a.pop(), Some("third".to_owned()));
+    assert_eq!(a.len(), 2);
+    assert_eq!(a.pop(), Some("second".to_string()));
+    assert_eq!(a.len(), 1);
+    assert_eq!(a.pop(), Some("first".to_string()));
+    assert!(a.is_empty());
+    assert_eq!(a.pop(), None);
+    assert!(a.is_empty());
+}
+
+#[test]
+fn pop_all_some() {
+    let mut a = MutableUtf8Array::<i32>::new();
+    a.push(Some("first"));
+    a.push(Some("second"));
+    a.push(Some("third"));
+    a.push(Some("firth"));
+    for _ in 0..4 {
+        a.push(Some("aaaa"));
+    }
+    a.push(Some("bbbb"));
+
+    assert_eq!(a.pop(), Some("bbbb".to_string()));
+    assert_eq!(a.pop(), Some("aaaa".to_string()));
+    assert_eq!(a.pop(), Some("aaaa".to_string()));
+    assert_eq!(a.pop(), Some("aaaa".to_string()));
+    assert_eq!(a.len(), 5);
+    assert_eq!(a.pop(), Some("aaaa".to_string()));
+    assert_eq!(a.pop(), Some("firth".to_string()));
+    assert_eq!(a.pop(), Some("third".to_string()));
+    assert_eq!(a.pop(), Some("second".to_string()));
+    assert_eq!(a.pop(), Some("first".to_string()));
+    assert!(a.is_empty());
+    assert_eq!(a.pop(), None);
+}
+
 /// Safety guarantee
 #[test]
 #[should_panic]