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

Added more tests #591

Merged
merged 1 commit into from
Nov 10, 2021
Merged
Show file tree
Hide file tree
Changes from all 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
46 changes: 43 additions & 3 deletions src/array/binary/mutable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,19 @@ impl<O: Offset> MutableBinaryArray<O> {
}
}

/// Initializes a new [`MutableBinaryArray`] with a pre-allocated capacity of slots and values.
pub fn with_capacities(capacity: usize, values: usize) -> Self {
let mut offsets = MutableBuffer::<O>::with_capacity(capacity + 1);
offsets.push(O::default());

Self {
data_type: Self::default_data_type(),
offsets,
values: MutableBuffer::<u8>::with_capacity(values),
validity: None,
}
}

/// Reserves `additional` slots.
pub fn reserve(&mut self, additional: usize) {
self.offsets.reserve(additional);
Expand Down Expand Up @@ -145,6 +158,18 @@ impl<O: Offset> MutableBinaryArray<O> {
}
}

impl<O: Offset> MutableBinaryArray<O> {
/// returns its values.
pub fn values(&self) -> &MutableBuffer<u8> {
&self.values
}

/// returns its offsets.
pub fn offsets(&self) -> &MutableBuffer<O> {
&self.offsets
}
}

impl<O: Offset> MutableArray for MutableBinaryArray<O> {
fn len(&self) -> usize {
self.offsets.len() - 1
Expand Down Expand Up @@ -228,15 +253,26 @@ impl<O: Offset> MutableBinaryArray<O> {
}

/// Creates a new [`BinaryArray`] from a [`TrustedLen`] of `&[u8]`.
/// # Safety
/// The iterator must be [`TrustedLen`](https://doc.rust-lang.org/std/iter/trait.TrustedLen.html).
/// I.e. that `size_hint().1` correctly reports its length.
#[inline]
pub fn from_trusted_len_values_iter<T: AsRef<[u8]>, I: TrustedLen<Item = T>>(
pub unsafe fn from_trusted_len_values_iter_unchecked<T: AsRef<[u8]>, I: Iterator<Item = T>>(
iterator: I,
) -> Self {
// soundness: I is `TrustedLen`
let (offsets, values) = unsafe { trusted_len_values_iter(iterator) };
Self::from_data(Self::default_data_type(), offsets, values, None)
}

/// Creates a new [`BinaryArray`] from a [`TrustedLen`] of `&[u8]`.
#[inline]
pub fn from_trusted_len_values_iter<T: AsRef<[u8]>, I: TrustedLen<Item = T>>(
iterator: I,
) -> Self {
// soundness: I is `TrustedLen`
unsafe { Self::from_trusted_len_values_iter_unchecked(iterator) }
}

/// Creates a [`MutableBinaryArray`] from an falible iterator of trusted length.
/// # Safety
/// The iterator must be [`TrustedLen`](https://doc.rust-lang.org/std/iter/trait.TrustedLen.html).
Expand All @@ -252,7 +288,11 @@ impl<O: Offset> MutableBinaryArray<O> {
let iterator = iterator.into_iter();

// soundness: assumed trusted len
let (validity, offsets, values) = try_trusted_len_unzip(iterator)?;
let (mut validity, offsets, values) = try_trusted_len_unzip(iterator)?;

if validity.as_mut().unwrap().null_count() == 0 {
validity = None;
}

Ok(Self::from_data(
Self::default_data_type(),
Expand Down
65 changes: 64 additions & 1 deletion tests/it/array/binary/mutable.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,68 @@
use arrow2::array::{BinaryArray, MutableBinaryArray};
use std::ops::Deref;

use arrow2::array::{BinaryArray, MutableArray, MutableBinaryArray};
use arrow2::bitmap::Bitmap;
use arrow2::error::ArrowError;

#[test]
fn new() {
assert_eq!(MutableBinaryArray::<i32>::new().len(), 0);

let a = MutableBinaryArray::<i32>::with_capacity(2);
assert_eq!(a.len(), 0);
assert!(a.offsets().capacity() >= 3);
assert_eq!(a.values().capacity(), 0);

let a = MutableBinaryArray::<i32>::with_capacities(2, 60);
assert_eq!(a.len(), 0);
assert!(a.offsets().capacity() >= 3);
assert!(a.values().capacity() >= 60);
}

#[test]
fn from_iter() {
let iter = (0..3u8).map(|x| Some(vec![x; x as usize]));
let a: MutableBinaryArray<i32> = iter.clone().collect();
assert_eq!(a.values().deref(), &[1u8, 2, 2]);
assert_eq!(a.offsets().deref(), &[0, 0, 1, 3]);
assert_eq!(a.validity(), None);

let a = unsafe { MutableBinaryArray::<i32>::from_trusted_len_iter_unchecked(iter) };
assert_eq!(a.values().deref(), &[1u8, 2, 2]);
assert_eq!(a.offsets().deref(), &[0, 0, 1, 3]);
assert_eq!(a.validity(), None);
}

#[test]
fn from_trusted_len_iter() {
let iter = (0..3u8).map(|x| vec![x; x as usize]);
let a: MutableBinaryArray<i32> = iter.clone().map(Some).collect();
assert_eq!(a.values().deref(), &[1u8, 2, 2]);
assert_eq!(a.offsets().deref(), &[0, 0, 1, 3]);
assert_eq!(a.validity(), None);

let a = unsafe {
MutableBinaryArray::<i32>::from_trusted_len_iter_unchecked(iter.clone().map(Some))
};
assert_eq!(a.values().deref(), &[1u8, 2, 2]);
assert_eq!(a.offsets().deref(), &[0, 0, 1, 3]);
assert_eq!(a.validity(), None);

let a = unsafe {
MutableBinaryArray::<i32>::try_from_trusted_len_iter_unchecked::<ArrowError, _, _>(
iter.clone().map(Some).map(Ok),
)
}
.unwrap();
assert_eq!(a.values().deref(), &[1u8, 2, 2]);
assert_eq!(a.offsets().deref(), &[0, 0, 1, 3]);
assert_eq!(a.validity(), None);

let a = unsafe { MutableBinaryArray::<i32>::from_trusted_len_values_iter_unchecked(iter) };
assert_eq!(a.values().deref(), &[1u8, 2, 2]);
assert_eq!(a.offsets().deref(), &[0, 0, 1, 3]);
assert_eq!(a.validity(), None);
}

#[test]
fn push_null() {
Expand Down