use re_arrow2::array::*; use re_arrow2::bitmap::Bitmap; use re_arrow2::datatypes::{DataType, Field, TimeUnit}; use re_arrow2::{error::Result, ffi}; use std::collections::BTreeMap; use std::sync::Arc; fn _test_round_trip(array: Box, expected: Box) -> Result<()> { let field = Field::new("a", array.data_type().clone(), true); // export array and corresponding data_type let array_ffi = ffi::export_array_to_c(array); let schema_ffi = ffi::export_field_to_c(&field); // import references let result_field = unsafe { ffi::import_field_from_c(&schema_ffi)? }; let result_array = unsafe { ffi::import_array_from_c(array_ffi, result_field.data_type.clone())? }; assert_eq!(&result_array, &expected); assert_eq!(result_field, field); Ok(()) } fn test_round_trip(expected: impl Array + Clone + 'static) -> Result<()> { let array: Box = Box::new(expected.clone()); let expected = Box::new(expected) as Box; _test_round_trip(array.clone(), clone(expected.as_ref()))?; // sliced _test_round_trip(array.sliced(1, 2), expected.sliced(1, 2)) } fn test_round_trip_schema(field: Field) -> Result<()> { let schema_ffi = ffi::export_field_to_c(&field); let result = unsafe { ffi::import_field_from_c(&schema_ffi)? }; assert_eq!(result, field); Ok(()) } #[test] fn bool_nullable() -> Result<()> { let data = BooleanArray::from(&[Some(true), None, Some(false), None]); test_round_trip(data) } #[test] fn bool() -> Result<()> { let data = BooleanArray::from_slice([true, true, false]); test_round_trip(data) } #[test] fn bool_nullable_sliced() -> Result<()> { let bitmap = Bitmap::from([true, false, false, true]).sliced(1, 3); let data = BooleanArray::try_new(DataType::Boolean, [true, true, false].into(), Some(bitmap))?; test_round_trip(data) } #[test] fn u32_nullable() -> Result<()> { let data = Int32Array::from(&[Some(2), None, Some(1), None]); test_round_trip(data) } #[test] fn u32() -> Result<()> { let data = Int32Array::from_slice([2, 0, 1, 0]); test_round_trip(data) } #[test] fn u32_sliced() -> Result<()> { let bitmap = Bitmap::from([true, false, false, true]).sliced(1, 3); let data = Int32Array::try_new(DataType::Int32, vec![1, 2, 3].into(), Some(bitmap))?; test_round_trip(data) } #[test] fn decimal() -> Result<()> { let data = Int128Array::from_slice([1, 0, 2, 0]); test_round_trip(data) } #[test] fn decimal_nullable() -> Result<()> { let data = Int128Array::from(&[Some(1), None, Some(2), None]); test_round_trip(data) } #[test] fn timestamp_tz() -> Result<()> { let data = Int64Array::from(&vec![Some(2), None, None]).to(DataType::Timestamp( TimeUnit::Second, Some(Arc::new("UTC".to_string())), )); test_round_trip(data) } #[test] fn utf8_nullable() -> Result<()> { let data = Utf8Array::::from([Some("a"), None, Some("bb"), None]); test_round_trip(data) } #[test] fn utf8() -> Result<()> { let data = Utf8Array::::from_slice(["a", "", "bb", ""]); test_round_trip(data) } #[test] fn utf8_sliced() -> Result<()> { let bitmap = Bitmap::from([true, false, false, true]).sliced(1, 3); let data = Utf8Array::::try_new( DataType::Utf8, vec![0, 1, 1, 2].try_into().unwrap(), b"ab".to_vec().into(), Some(bitmap), )?; test_round_trip(data) } #[test] fn large_utf8() -> Result<()> { let data = Utf8Array::::from([Some("a"), None, Some("bb"), None]); test_round_trip(data) } #[test] fn binary_nullable() -> Result<()> { let data = BinaryArray::::from([Some(b"a".as_ref()), None, Some(b"bb".as_ref()), None]); test_round_trip(data) } #[test] fn binary() -> Result<()> { let data = BinaryArray::::from_slice([b"a".as_ref(), b"", b"bb", b""]); test_round_trip(data) } #[test] fn binary_sliced() -> Result<()> { let bitmap = Bitmap::from([true, false, false, true]).sliced(1, 3); let data = BinaryArray::::try_new( DataType::Binary, vec![0, 1, 1, 2].try_into().unwrap(), b"ab".to_vec().into(), Some(bitmap), )?; test_round_trip(data) } #[test] fn large_binary() -> Result<()> { let data = BinaryArray::::from([Some(b"a".as_ref()), None, Some(b"bb".as_ref()), None]); test_round_trip(data) } #[test] fn fixed_size_binary() -> Result<()> { let data = FixedSizeBinaryArray::new( DataType::FixedSizeBinary(2), vec![1, 2, 3, 4, 5, 6].into(), None, ); test_round_trip(data) } #[test] fn fixed_size_binary_nullable() -> Result<()> { let data = FixedSizeBinaryArray::new( DataType::FixedSizeBinary(2), vec![1, 2, 3, 4, 5, 6].into(), Some([true, true, false].into()), ); test_round_trip(data) } #[test] fn fixed_size_binary_sliced() -> Result<()> { let bitmap = Bitmap::from([true, false, false, true]).sliced(1, 3); let data = FixedSizeBinaryArray::try_new( DataType::FixedSizeBinary(2), b"ababab".to_vec().into(), Some(bitmap), )?; test_round_trip(data) } #[test] fn list() -> Result<()> { let data = vec![ Some(vec![Some(1i32), Some(2), Some(3)]), None, Some(vec![Some(4), None, Some(6)]), ]; let mut array = MutableListArray::>::new(); array.try_extend(data)?; let array: ListArray = array.into(); test_round_trip(array) } #[test] fn list_sliced() -> Result<()> { let bitmap = Bitmap::from([true, false, false, true]).sliced(1, 3); let array = ListArray::::try_new( DataType::List(std::sync::Arc::new(Field::new("a", DataType::Int32, true))), vec![0, 1, 1, 2].try_into().unwrap(), Box::new(PrimitiveArray::::from_slice([1, 2])), Some(bitmap), )?; test_round_trip(array) } #[test] fn large_list() -> Result<()> { let data = vec![ Some(vec![Some(1i32), Some(2), Some(3)]), None, Some(vec![Some(4), None, Some(6)]), ]; let mut array = MutableListArray::>::new(); array.try_extend(data)?; let array: ListArray = array.into(); test_round_trip(array) } #[test] fn fixed_size_list() -> Result<()> { let data = vec![ Some(vec![Some(1i32), Some(2), Some(3)]), None, Some(vec![Some(4), None, Some(6)]), ]; let mut array = MutableFixedSizeListArray::new(MutablePrimitiveArray::::new(), 3); array.try_extend(data)?; let array: FixedSizeListArray = array.into(); test_round_trip(array) } #[test] fn fixed_size_list_sliced() -> Result<()> { let bitmap = Bitmap::from([true, false, false, true]).sliced(1, 3); let array = FixedSizeListArray::try_new( DataType::FixedSizeList( std::sync::Arc::new(Field::new("a", DataType::Int32, true)), 2, ), Box::new(PrimitiveArray::::from_vec(vec![1, 2, 3, 4, 5, 6])), Some(bitmap), )?; test_round_trip(array) } #[test] fn list_list() -> Result<()> { let data = vec![ Some(vec![ Some(vec![None]), Some(vec![Some(2)]), Some(vec![Some(3)]), ]), None, Some(vec![Some(vec![Some(4), None, Some(6)])]), ]; let mut array = MutableListArray::>>::new(); array.try_extend(data)?; let array: ListArray = array.into(); test_round_trip(array) } #[test] fn struct_() -> Result<()> { let data_type = DataType::Struct(Arc::new(vec![Field::new("a", DataType::Int32, true)])); let values = vec![Int32Array::from([Some(1), None, Some(3)]).boxed()]; let validity = Bitmap::from([true, false, true]); let array = StructArray::new(data_type, values, validity.into()); test_round_trip(array) } #[test] fn dict() -> Result<()> { let data = vec![Some("a"), Some("a"), None, Some("b")]; let mut array = MutableDictionaryArray::>::new(); array.try_extend(data)?; let array: DictionaryArray = array.into(); test_round_trip(array) } #[test] fn schema() -> Result<()> { let field = Field::new( "a", DataType::List(std::sync::Arc::new(Field::new("a", DataType::UInt32, true))), true, ); test_round_trip_schema(field)?; let field = Field::new( "a", DataType::Dictionary(u32::KEY_TYPE, Arc::new(DataType::Utf8), false), true, ); test_round_trip_schema(field)?; let field = Field::new("a", DataType::Int32, true); let mut metadata = BTreeMap::new(); metadata.insert("some".to_string(), "stuff".to_string()); let field = field.with_metadata(metadata); test_round_trip_schema(field) } #[test] fn extension() -> Result<()> { let field = Field::new( "a", DataType::Extension( "a".to_string(), Arc::new(DataType::Int32), Some(Arc::new("bla".to_string())), ), true, ); test_round_trip_schema(field) } #[test] fn extension_children() -> Result<()> { let field = Field::new( "a", DataType::Extension( "b".to_string(), Arc::new(DataType::Struct(Arc::new(vec![Field::new( "c", DataType::Int32, true, )]))), None, ), true, ); test_round_trip_schema(field) }