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

Commit

Permalink
Made FixedSizeList::try_push_valid public and added new_with_field (
Browse files Browse the repository at this point in the history
  • Loading branch information
ncpenke committed Mar 6, 2022
1 parent 1d5661c commit 81bfadd
Show file tree
Hide file tree
Showing 2 changed files with 58 additions and 2 deletions.
21 changes: 19 additions & 2 deletions src/array/fixed_size_list/mutable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use std::sync::Arc;
use crate::{
array::{Array, MutableArray, TryExtend, TryPush},
bitmap::MutableBitmap,
datatypes::DataType,
datatypes::{DataType, Field},
error::{ArrowError, Result},
};

Expand Down Expand Up @@ -41,6 +41,21 @@ impl<M: MutableArray> MutableFixedSizeListArray<M> {
}
}

/// Creates a new [`MutableFixedSizeListArray`] from a [`MutableArray`] and size.
pub fn new_with_field(values: M, name: &str, nullable: bool, size: usize) -> Self {
let data_type = DataType::FixedSizeList(
Box::new(Field::new(name, values.data_type().clone(), nullable)),
size,
);
assert_eq!(values.len(), 0);
Self {
size,
data_type,
values,
validity: None,
}
}

/// The inner values
pub fn values(&self) -> &M {
&self.values
Expand All @@ -61,7 +76,9 @@ impl<M: MutableArray> MutableFixedSizeListArray<M> {
}

#[inline]
fn try_push_valid(&mut self) -> Result<()> {
/// Needs to be called when a valid value was extended to this array.
/// This is a relatively low level function, prefer `try_push` when you can.
pub fn try_push_valid(&mut self) -> Result<()> {
if self.values.len() % self.size != 0 {
return Err(ArrowError::Overflow);
};
Expand Down
39 changes: 39 additions & 0 deletions tests/it/array/fixed_size_list/mutable.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use arrow2::array::*;
use arrow2::datatypes::{DataType, Field};

#[test]
fn primitive() {
Expand All @@ -24,3 +25,41 @@ fn primitive() {
let expected = Int32Array::from(vec![None, None, None]);
assert_eq!(a, &expected)
}

#[test]
fn new_with_field() {
let data = vec![
Some(vec![Some(1i32), Some(2), Some(3)]),
Some(vec![None, None, None]),
Some(vec![Some(4), None, Some(6)]),
];

let mut list = MutableFixedSizeListArray::new_with_field(
MutablePrimitiveArray::<i32>::new(),
"custom_items",
false,
3,
);
list.try_extend(data).unwrap();
let list: FixedSizeListArray = list.into();

assert_eq!(
list.data_type(),
&DataType::FixedSizeList(
Box::new(Field::new("custom_items", DataType::Int32, false)),
3
)
);

let a = list.value(0);
let a = a.as_any().downcast_ref::<Int32Array>().unwrap();

let expected = Int32Array::from(vec![Some(1i32), Some(2), Some(3)]);
assert_eq!(a, &expected);

let a = list.value(1);
let a = a.as_any().downcast_ref::<Int32Array>().unwrap();

let expected = Int32Array::from(vec![None, None, None]);
assert_eq!(a, &expected)
}

0 comments on commit 81bfadd

Please sign in to comment.