use arrow2::{
array::{
growable::{Growable, GrowableList},
Array, ListArray, MutableListArray, MutablePrimitiveArray, TryExtend,
},
datatypes::DataType,
};
fn create_list_array(data: Vec<Option<Vec<Option<i32>>>>) -> ListArray<i32> {
let mut array = MutableListArray::<i32, MutablePrimitiveArray<i32>>::new();
array.try_extend(data).unwrap();
array.into()
}
#[test]
fn extension() {
let data = vec![
Some(vec![Some(1i32), Some(2), Some(3)]),
Some(vec![Some(4), Some(5)]),
Some(vec![Some(6i32), Some(7), Some(8)]),
];
let array = create_list_array(data);
let data_type =
DataType::Extension("ext".to_owned(), Box::new(array.data_type().clone()), None);
let array_ext = ListArray::new(
data_type,
array.offsets().clone(),
array.values().clone(),
array.validity().cloned(),
);
let mut a = GrowableList::new(vec![&array_ext], false, 0);
a.extend(0, 0, 1);
assert_eq!(a.len(), 1);
let result: ListArray<i32> = a.into();
assert_eq!(array_ext.data_type(), result.data_type());
dbg!(result);
}
#[test]
fn basic() {
let data = vec![
Some(vec![Some(1i32), Some(2), Some(3)]),
Some(vec![Some(4), Some(5)]),
Some(vec![Some(6i32), Some(7), Some(8)]),
];
let array = create_list_array(data);
let mut a = GrowableList::new(vec![&array], false, 0);
a.extend(0, 0, 1);
assert_eq!(a.len(), 1);
let result: ListArray<i32> = a.into();
let expected = vec![Some(vec![Some(1i32), Some(2), Some(3)])];
let expected = create_list_array(expected);
assert_eq!(result, expected)
}
#[test]
fn null_offset() {
let data = vec![
Some(vec![Some(1i32), Some(2), Some(3)]),
None,
Some(vec![Some(6i32), Some(7), Some(8)]),
];
let array = create_list_array(data);
let array = array.sliced(1, 2);
let mut a = GrowableList::new(vec![&array], false, 0);
a.extend(0, 1, 1);
assert_eq!(a.len(), 1);
let result: ListArray<i32> = a.into();
let expected = vec![Some(vec![Some(6i32), Some(7), Some(8)])];
let expected = create_list_array(expected);
assert_eq!(result, expected)
}
#[test]
fn null_offsets() {
let data = vec![
Some(vec![Some(1i32), Some(2), Some(3)]),
None,
Some(vec![Some(6i32), None, Some(8)]),
];
let array = create_list_array(data);
let array = array.sliced(1, 2);
let mut a = GrowableList::new(vec![&array], false, 0);
a.extend(0, 1, 1);
assert_eq!(a.len(), 1);
let result: ListArray<i32> = a.into();
let expected = vec![Some(vec![Some(6i32), None, Some(8)])];
let expected = create_list_array(expected);
assert_eq!(result, expected)
}
#[test]
fn test_from_two_lists() {
let data_1 = vec![
Some(vec![Some(1i32), Some(2), Some(3)]),
None,
Some(vec![Some(6i32), None, Some(8)]),
];
let array_1 = create_list_array(data_1);
let data_2 = vec![
Some(vec![Some(8i32), Some(7), Some(6)]),
Some(vec![Some(5i32), None, Some(4)]),
Some(vec![Some(2i32), Some(1), Some(0)]),
];
let array_2 = create_list_array(data_2);
let mut a = GrowableList::new(vec![&array_1, &array_2], false, 6);
a.extend(0, 0, 2);
a.extend(1, 1, 1);
assert_eq!(a.len(), 3);
let result: ListArray<i32> = a.into();
let expected = vec![
Some(vec![Some(1i32), Some(2), Some(3)]),
None,
Some(vec![Some(5i32), None, Some(4)]),
];
let expected = create_list_array(expected);
assert_eq!(result, expected);
}