Skip to content

Commit

Permalink
test: Add unit test for extending slice of list array (#5948)
Browse files Browse the repository at this point in the history
* test: Add unit test for extending slice of list array

* For review
  • Loading branch information
viirya authored Jun 25, 2024
1 parent 66bada5 commit 460fd55
Showing 1 changed file with 38 additions and 1 deletion.
39 changes: 38 additions & 1 deletion arrow-array/src/ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1229,13 +1229,14 @@ mod tests_from_ffi {
use arrow_data::ArrayData;
use arrow_schema::{DataType, Field};

use crate::types::Int32Type;
use crate::{
array::{
Array, BooleanArray, DictionaryArray, FixedSizeBinaryArray, FixedSizeListArray,
Int32Array, Int64Array, StringArray, StructArray, UInt32Array, UInt64Array,
},
ffi::{from_ffi, FFI_ArrowArray, FFI_ArrowSchema},
make_array, ArrayRef,
make_array, ArrayRef, ListArray,
};

use super::{ImportedArrowArray, Result};
Expand Down Expand Up @@ -1514,4 +1515,40 @@ mod tests_from_ffi {
&imported
);
}

fn roundtrip_list_array(array: ListArray) -> ListArray {
let data = array.into_data();

let array = FFI_ArrowArray::new(&data);
let schema = FFI_ArrowSchema::try_from(data.data_type()).unwrap();

let array = unsafe { from_ffi(array, &schema) }.unwrap();
ListArray::from(array)
}

#[test]
fn test_extend_imported_list_slice() {
let mut data = vec![];

for i in 0..1000 {
let mut list = vec![];
for j in 0..100 {
list.push(Some(i * 1000 + j));
}
data.push(Some(list));
}

let list_array = ListArray::from_iter_primitive::<Int32Type, _, _>(data);

let slice = list_array.slice(500, 500);
let imported = roundtrip_list_array(slice.clone());
assert_eq!(imported.len(), 500);
assert_eq!(&slice, &imported);

let copied = extend_array(&imported);
assert_eq!(
copied.as_any().downcast_ref::<ListArray>().unwrap(),
&imported
);
}
}

0 comments on commit 460fd55

Please sign in to comment.