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

Added capacity to some mutable arrays and tests #913

Merged
merged 4 commits into from
Mar 14, 2022
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
7 changes: 2 additions & 5 deletions src/array/fixed_size_binary/iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -60,10 +60,7 @@ impl<'a> FixedSizeBinaryArray {
pub fn iter(
&'a self,
) -> ZipValidity<'a, &'a [u8], FixedSizeBinaryValuesIter<'a, FixedSizeBinaryArray>> {
zip_validity(
FixedSizeBinaryValuesIter::new(self),
self.validity.as_ref().map(|x| x.iter()),
)
zip_validity(self.iter_values(), self.validity.as_ref().map(|x| x.iter()))
}

/// Returns iterator over the values of [`FixedSizeBinaryArray`]
Expand All @@ -88,7 +85,7 @@ impl<'a> MutableFixedSizeBinaryArray {
&'a self,
) -> ZipValidity<'a, &'a [u8], FixedSizeBinaryValuesIter<'a, MutableFixedSizeBinaryArray>> {
zip_validity(
FixedSizeBinaryValuesIter::new(self),
self.iter_values(),
self.validity().as_ref().map(|x| x.iter()),
)
}
Expand Down
3 changes: 2 additions & 1 deletion src/array/fixed_size_binary/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -81,9 +81,10 @@ impl FixedSizeBinaryArray {

/// Returns a new null [`FixedSizeBinaryArray`].
pub fn new_null(data_type: DataType, length: usize) -> Self {
let size = Self::maybe_get_size(&data_type).unwrap();
Self::new(
data_type,
Buffer::new_zeroed(length),
Buffer::new_zeroed(length * size),
Some(Bitmap::new_zeroed(length)),
)
}
Expand Down
5 changes: 5 additions & 0 deletions src/array/fixed_size_binary/mutable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,11 @@ impl MutableFixedSizeBinaryArray {
self.size
}

/// Returns the capacity of this array
pub fn capacity(&self) -> usize {
self.values.capacity() / self.size
}

fn init_validity(&mut self) {
let mut validity = MutableBitmap::new();
validity.extend_constant(self.len(), true);
Expand Down
11 changes: 11 additions & 0 deletions src/array/primitive/mutable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -262,6 +262,11 @@ impl<T: NativeType> MutablePrimitiveArray<T> {
validity.shrink_to_fit()
}
}

/// Returns the capacity of this [`MutablePrimitiveArray`].
pub fn capacity(&self) -> usize {
self.values.capacity()
}
}

/// Accessors
Expand Down Expand Up @@ -477,6 +482,12 @@ impl<T: NativeType> MutablePrimitiveArray<T> {
}
}

/// Creates a (non-null) [`MutablePrimitiveArray`] from a vector of values.
/// This does not have memcopy and is the fastest way to create a [`PrimitiveArray`].
pub fn from_vec(values: Vec<T>) -> Self {
Self::from_data(T::PRIMITIVE.into(), values, None)
}

/// Creates a new [`MutablePrimitiveArray`] from an iterator over values
/// # Safety
/// The iterator must be [`TrustedLen`](https://doc.rust-lang.org/std/iter/trait.TrustedLen.html).
Expand Down
42 changes: 24 additions & 18 deletions src/array/utf8/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -341,24 +341,30 @@ impl<O: Offset> Utf8Array<O> {
}
} else {
match (self.values.into_mut(), self.offsets.into_mut()) {
(Left(immutable_values), Left(immutable_offsets)) => Left(Utf8Array::from_data(
self.data_type,
immutable_offsets,
immutable_values,
None,
)),
(Left(immutable_values), Right(mutable_offsets)) => Left(Utf8Array::from_data(
self.data_type,
mutable_offsets.into(),
immutable_values,
None,
)),
(Right(mutable_values), Left(immutable_offsets)) => Left(Utf8Array::from_data(
self.data_type,
immutable_offsets,
mutable_values.into(),
None,
)),
(Left(immutable_values), Left(immutable_offsets)) => Left(unsafe {
Utf8Array::new_unchecked(
self.data_type,
immutable_offsets,
immutable_values,
None,
)
}),
(Left(immutable_values), Right(mutable_offsets)) => Left(unsafe {
Utf8Array::new_unchecked(
self.data_type,
mutable_offsets.into(),
immutable_values,
None,
)
}),
(Right(mutable_values), Left(immutable_offsets)) => Left(unsafe {
Utf8Array::from_data(
self.data_type,
immutable_offsets,
mutable_values.into(),
None,
)
}),
(Right(mutable_values), Right(mutable_offsets)) => {
Right(MutableUtf8Array::from_data(
self.data_type,
Expand Down
41 changes: 40 additions & 1 deletion tests/it/array/fixed_size_binary/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ fn basics() {
#[test]
fn with_validity() {
let values = Buffer::from_slice([1, 2, 3, 4, 5, 6]);
let a = FixedSizeBinaryArray::from_data(DataType::FixedSizeBinary(2), values, None);
let a = FixedSizeBinaryArray::new(DataType::FixedSizeBinary(2), values, None);
let a = a.with_validity(Some(Bitmap::from([true, false, true])));
assert!(a.validity().is_some());
}
Expand All @@ -50,9 +50,48 @@ fn empty() {
assert_eq!(array.validity(), None);
}

#[test]
fn null() {
let array = FixedSizeBinaryArray::new_null(DataType::FixedSizeBinary(2), 2);
assert_eq!(array.values().len(), 4);
assert_eq!(array.validity().cloned(), Some([false, false].into()));
}

#[test]
fn from_iter() {
let iter = std::iter::repeat(vec![1u8, 2]).take(2).map(Some);
let a = FixedSizeBinaryArray::from_iter(iter, 2);
assert_eq!(a.len(), 2);
}

#[test]
fn wrong_size() {
let values = Buffer::from_slice(b"abb");
assert!(FixedSizeBinaryArray::try_new(DataType::FixedSizeBinary(2), values, None).is_err());
}

#[test]
fn wrong_len() {
let values = Buffer::from_slice(b"abba");
let validity = Some([true, false, false].into()); // it should be 2
assert!(FixedSizeBinaryArray::try_new(DataType::FixedSizeBinary(2), values, validity).is_err());
}

#[test]
fn wrong_data_type() {
let values = Buffer::from_slice(b"abba");
assert!(FixedSizeBinaryArray::try_new(DataType::Binary, values, None).is_err());
}

#[test]
fn to() {
let values = Buffer::from_slice(b"abba");
let a = FixedSizeBinaryArray::new(DataType::FixedSizeBinary(2), values, None);

let extension = DataType::Extension(
"a".to_string(),
Box::new(DataType::FixedSizeBinary(2)),
None,
);
let _ = a.to(extension);
}
32 changes: 32 additions & 0 deletions tests/it/array/fixed_size_binary/mutable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -66,3 +66,35 @@ fn push_null() {
let array: FixedSizeBinaryArray = array.into();
assert_eq!(array.validity(), Some(&Bitmap::from([false])));
}

#[test]
fn as_arc() {
let mut array = MutableFixedSizeBinaryArray::try_from_iter(
vec![Some(b"ab"), Some(b"bc"), None, Some(b"fh")],
2,
)
.unwrap();

let array = array.as_arc();
assert_eq!(array.len(), 4);
}

#[test]
fn as_box() {
let mut array = MutableFixedSizeBinaryArray::try_from_iter(
vec![Some(b"ab"), Some(b"bc"), None, Some(b"fh")],
2,
)
.unwrap();

let array = array.as_box();
assert_eq!(array.len(), 4);
}

#[test]
fn shrink_to_fit_and_capacity() {
let mut array = MutableFixedSizeBinaryArray::with_capacity(2, 100);
array.push(Some([1, 2]));
array.shrink_to_fit();
assert_eq!(array.capacity(), 1);
}
36 changes: 36 additions & 0 deletions tests/it/array/growable/fixed_binary.rs
Original file line number Diff line number Diff line change
Expand Up @@ -87,3 +87,39 @@ fn sized_offsets() {
let expected = FixedSizeBinaryArray::from_iter(vec![Some(&[0, 2]), Some(&[0, 1])], 2);
assert_eq!(result, expected);
}

/// to, as_box, as_arc
#[test]
fn as_box() {
let array =
FixedSizeBinaryArray::from_iter(vec![Some(b"ab"), Some(b"bc"), None, Some(b"de")], 2);
let mut a = GrowableFixedSizeBinary::new(vec![&array], false, 0);
a.extend(0, 1, 2);

let result = a.as_box();
let result = result
.as_any()
.downcast_ref::<FixedSizeBinaryArray>()
.unwrap();

let expected = FixedSizeBinaryArray::from_iter(vec![Some("bc"), None], 2);
assert_eq!(&expected, result);
}

/// as_arc
#[test]
fn as_arc() {
let array =
FixedSizeBinaryArray::from_iter(vec![Some(b"ab"), Some(b"bc"), None, Some(b"de")], 2);
let mut a = GrowableFixedSizeBinary::new(vec![&array], false, 0);
a.extend(0, 1, 2);

let result = a.as_arc();
let result = result
.as_any()
.downcast_ref::<FixedSizeBinaryArray>()
.unwrap();

let expected = FixedSizeBinaryArray::from_iter(vec![Some("bc"), None], 2);
assert_eq!(&expected, result);
}
Loading