Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Added pop to utf8/binary/fixedSize MutableArray #966

Merged
merged 5 commits into from
Apr 29, 2022
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions src/array/binary/mutable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,22 @@ impl<O: Offset> MutableBinaryArray<O> {
self.try_push(value).unwrap()
}

/// Pop the last entry from [`MutableBinaryArray`].
/// Note If the values is empty, this method will return None.
ygf11 marked this conversation as resolved.
Show resolved Hide resolved
pub fn pop(&mut self) -> Option<Vec<u8>> {
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<P: AsRef<[u8]>, I: IntoIterator<Item = Option<P>>>(iter: I) -> Result<Self> {
let iterator = iter.into_iter();
let (lower, _) = iterator.size_hint();
Expand Down
15 changes: 15 additions & 0 deletions src/array/fixed_size_binary/mutable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,21 @@ impl MutableFixedSizeBinaryArray {
self.try_push(value).unwrap()
}

/// Pop the last entry from [`MutableFixedSizeBinaryArray`].
/// Note If the values is empty, this method will return None.
ygf11 marked this conversation as resolved.
Show resolved Hide resolved
pub fn pop(&mut self) -> Option<Vec<u8>> {
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.
Expand Down
18 changes: 18 additions & 0 deletions src/array/utf8/mutable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -164,6 +164,24 @@ impl<O: Offset> MutableUtf8Array<O> {
self.try_push(value).unwrap()
}

/// Pop the last entry from [`MutableUtf8Array`].
/// Note If the values is empty, this method will return None.
ygf11 marked this conversation as resolved.
Show resolved Hide resolved
pub fn pop(&mut self) -> Option<String> {
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);
Expand Down
48 changes: 48 additions & 0 deletions tests/it/array/binary/mutable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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"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::<i32>::new();
Expand Down
56 changes: 56 additions & 0 deletions tests/it/array/fixed_size_binary/mutable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
48 changes: 47 additions & 1 deletion tests/it/array/utf8/mutable.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use arrow2::array::{MutableUtf8Array, Utf8Array};
use arrow2::array::{MutableArray, MutableUtf8Array, Utf8Array};
use arrow2::bitmap::Bitmap;
use arrow2::datatypes::DataType;

Expand All @@ -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("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]
Expand Down