From b93333cfb603903dacaabdfd191d6f456c19a92b Mon Sep 17 00:00:00 2001 From: ygf11 Date: Fri, 29 Apr 2022 15:02:03 +0800 Subject: [PATCH] Added pop to utf8/binary/fixedSize MutableArray (#966) --- src/array/binary/mutable.rs | 16 ++++++ src/array/fixed_size_binary/mutable.rs | 15 ++++++ src/array/utf8/mutable.rs | 18 +++++++ tests/it/array/binary/mutable.rs | 48 ++++++++++++++++++ tests/it/array/fixed_size_binary/mutable.rs | 56 +++++++++++++++++++++ tests/it/array/utf8/mutable.rs | 48 +++++++++++++++++- 6 files changed, 200 insertions(+), 1 deletion(-) diff --git a/src/array/binary/mutable.rs b/src/array/binary/mutable.rs index 1a824193dcf..945ea337f17 100644 --- a/src/array/binary/mutable.rs +++ b/src/array/binary/mutable.rs @@ -125,6 +125,22 @@ impl MutableBinaryArray { self.try_push(value).unwrap() } + /// Pop the last entry from [`MutableBinaryArray`]. + /// This function returns `None` iff this array is empty + pub fn pop(&mut self) -> Option> { + if self.offsets.len() < 2 { + return None; + } + self.offsets.pop()?; + let value_start = self.offsets.iter().last().cloned()?.to_usize(); + let value = self.values.split_off(value_start); + self.validity + .as_mut() + .map(|x| x.pop()?.then(|| ())) + .unwrap_or_else(|| Some(())) + .map(|_| value) + } + fn try_from_iter, I: IntoIterator>>(iter: I) -> Result { let iterator = iter.into_iter(); let (lower, _) = iterator.size_hint(); diff --git a/src/array/fixed_size_binary/mutable.rs b/src/array/fixed_size_binary/mutable.rs index 787996267bc..579509bded4 100644 --- a/src/array/fixed_size_binary/mutable.rs +++ b/src/array/fixed_size_binary/mutable.rs @@ -112,6 +112,21 @@ impl MutableFixedSizeBinaryArray { self.try_push(value).unwrap() } + /// Pop the last entry from [`MutableFixedSizeBinaryArray`]. + /// This function returns `None` iff this array is empty + pub fn pop(&mut self) -> Option> { + 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 + .as_mut() + .map(|x| x.pop()?.then(|| ())) + .unwrap_or_else(|| Some(())) + .map(|_| value) + } + /// Creates a new [`MutableFixedSizeBinaryArray`] from an iterator of values. /// # Errors /// Errors iff the size of any of the `value` is not equal to its own size. diff --git a/src/array/utf8/mutable.rs b/src/array/utf8/mutable.rs index 27bd94a3976..f9f53a68462 100644 --- a/src/array/utf8/mutable.rs +++ b/src/array/utf8/mutable.rs @@ -164,6 +164,24 @@ impl MutableUtf8Array { self.try_push(value).unwrap() } + /// Pop the last entry from [`MutableUtf8Array`]. + /// This function returns `None` iff this array is empty. + pub fn pop(&mut self) -> Option { + if self.offsets.len() < 2 { + return None; + } + self.offsets.pop()?; + let value_start = self.offsets.iter().last().cloned()?.to_usize(); + let value = self.values.split_off(value_start); + self.validity + .as_mut() + .map(|x| x.pop()?.then(|| ())) + .unwrap_or_else(|| Some(())) + .map(|_| + // soundness: we always check for utf8 soundness on constructors. + unsafe { String::from_utf8_unchecked(value) }) + } + fn init_validity(&mut self) { let mut validity = MutableBitmap::with_capacity(self.offsets.capacity()); validity.extend_constant(self.len(), true); diff --git a/tests/it/array/binary/mutable.rs b/tests/it/array/binary/mutable.rs index 7258dea1332..67fd731d57f 100644 --- 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::::new(); + a.push(Some(b"first")); + a.push(Some(b"second")); + a.push::>(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::::new(); + a.push(Some(b"first")); + a.push(Some(b"second")); + a.push(Some(b"third")); + a.push(Some(b"fourth")); + + 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"fourth".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::::new(); diff --git a/tests/it/array/fixed_size_binary/mutable.rs b/tests/it/array/fixed_size_binary/mutable.rs index b5f5babae00..068ed0681c4 100644 --- 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 index 548409f68dc..8ee76476cbb 100644 --- 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::::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::::new(); + a.push(Some("first")); + a.push(Some("second")); + a.push(Some("third")); + a.push(Some("fourth")); + for _ in 0..4 { + a.push(Some("aaaa")); + } + a.push(Some("こんにちは")); + + assert_eq!(a.pop(), Some("こんにちは".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("fourth".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]