diff --git a/src/array/growable/structure.rs b/src/array/growable/structure.rs index e7f997024de..8e91cf03ce0 100644 --- a/src/array/growable/structure.rs +++ b/src/array/growable/structure.rs @@ -3,7 +3,6 @@ use std::sync::Arc; use crate::{ array::{Array, StructArray}, bitmap::MutableBitmap, - datatypes::DataType, }; use super::{ @@ -25,6 +24,8 @@ impl<'a> GrowableStruct<'a> { /// # Panics /// If `arrays` is empty. pub fn new(arrays: Vec<&'a StructArray>, mut use_validity: bool, capacity: usize) -> Self { + assert!(!arrays.is_empty()); + // if any of the arrays has nulls, insertions from any array requires setting bits // as there is at least one array with nulls. if arrays.iter().any(|array| array.null_count() > 0) { @@ -68,11 +69,7 @@ impl<'a> GrowableStruct<'a> { let values = std::mem::take(&mut self.values); let values = values.into_iter().map(|mut x| x.as_box()).collect(); - StructArray::new( - DataType::Struct(self.arrays[0].fields().to_vec()), - values, - validity.into(), - ) + StructArray::new(self.arrays[0].data_type().clone(), values, validity.into()) } } @@ -121,7 +118,7 @@ impl<'a> From> for StructArray { let values = val.values.into_iter().map(|mut x| x.as_box()).collect(); StructArray::new( - DataType::Struct(val.arrays[0].fields().to_vec()), + val.arrays[0].data_type().clone(), values, val.validity.into(), ) diff --git a/src/array/struct_/mod.rs b/src/array/struct_/mod.rs index 19383386592..a8848bb5578 100644 --- a/src/array/struct_/mod.rs +++ b/src/array/struct_/mod.rs @@ -122,7 +122,7 @@ impl StructArray { /// Creates an empty [`StructArray`]. pub fn new_empty(data_type: DataType) -> Self { - if let DataType::Struct(fields) = &data_type { + if let DataType::Struct(fields) = &data_type.to_logical_type() { let values = fields .iter() .map(|field| new_empty_array(field.data_type().clone())) diff --git a/src/array/union/mod.rs b/src/array/union/mod.rs index 3aa963c00a8..e767aac2b21 100644 --- a/src/array/union/mod.rs +++ b/src/array/union/mod.rs @@ -195,7 +195,7 @@ impl UnionArray { /// Creates a new empty [`UnionArray`]. pub fn new_empty(data_type: DataType) -> Self { - if let DataType::Union(f, _, mode) = &data_type { + if let DataType::Union(f, _, mode) = data_type.to_logical_type() { let fields = f .iter() .map(|x| new_empty_array(x.data_type().clone())) diff --git a/tests/it/array/growable/list.rs b/tests/it/array/growable/list.rs index ef183654579..9aa56b109aa 100644 --- a/tests/it/array/growable/list.rs +++ b/tests/it/array/growable/list.rs @@ -1,6 +1,9 @@ -use arrow2::array::{ - growable::{Growable, GrowableList}, - ListArray, MutableListArray, MutablePrimitiveArray, TryExtend, +use arrow2::{ + array::{ + growable::{Growable, GrowableList}, + Array, ListArray, MutableListArray, MutablePrimitiveArray, TryExtend, + }, + datatypes::DataType, }; fn create_list_array(data: Vec>>>) -> ListArray { @@ -9,6 +12,33 @@ fn create_list_array(data: Vec>>>) -> ListArray { 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); + + let result: ListArray = a.into(); + assert_eq!(array_ext.data_type(), result.data_type()); + dbg!(result); +} + #[test] fn basic() { let data = vec![ diff --git a/tests/it/array/growable/mod.rs b/tests/it/array/growable/mod.rs index a6e25ecba37..480bcf06c31 100644 --- a/tests/it/array/growable/mod.rs +++ b/tests/it/array/growable/mod.rs @@ -12,7 +12,7 @@ mod utf8; use arrow2::array::growable::make_growable; use arrow2::array::*; -use arrow2::datatypes::DataType; +use arrow2::datatypes::{DataType, Field}; #[test] fn test_make_growable() { @@ -37,11 +37,36 @@ fn test_make_growable() { let array = FixedSizeBinaryArray::new(DataType::FixedSizeBinary(2), b"abcd".to_vec().into(), None); make_growable(&[&array], false, 2); +} +#[test] +fn test_make_growable_extension() { let array = DictionaryArray::try_from_keys( Int32Array::from_slice([1, 0]), Int32Array::from_slice([1, 2]).boxed(), ) .unwrap(); make_growable(&[&array], false, 2); + + let data_type = DataType::Extension("ext".to_owned(), Box::new(DataType::Int32), None); + let array = Int32Array::from_slice([1, 2]).to(data_type.clone()); + let array_grown = make_growable(&[&array], false, 2).as_box(); + assert_eq!(array_grown.data_type(), &data_type); + + let data_type = DataType::Extension( + "ext".to_owned(), + Box::new(DataType::Struct(vec![Field::new( + "a", + DataType::Int32, + false, + )])), + None, + ); + let array = StructArray::new( + data_type.clone(), + vec![Int32Array::from_slice([1, 2]).boxed()], + None, + ); + let array_grown = make_growable(&[&array], false, 2).as_box(); + assert_eq!(array_grown.data_type(), &data_type); } diff --git a/tests/it/array/mod.rs b/tests/it/array/mod.rs index 9df7c53d02e..85318ba628a 100644 --- a/tests/it/array/mod.rs +++ b/tests/it/array/mod.rs @@ -58,6 +58,11 @@ fn empty() { DataType::Utf8, DataType::Binary, DataType::List(Box::new(Field::new("a", DataType::Binary, true))), + DataType::List(Box::new(Field::new( + "a", + DataType::Extension("ext".to_owned(), Box::new(DataType::Int32), None), + true, + ))), DataType::Union( vec![Field::new("a", DataType::Binary, true)], None, @@ -68,11 +73,42 @@ fn empty() { None, UnionMode::Dense, ), + DataType::Struct(vec![Field::new("a", DataType::Int32, true)]), ]; let a = datatypes.into_iter().all(|x| new_empty_array(x).len() == 0); assert!(a); } +#[test] +fn empty_extension() { + let datatypes = vec![ + DataType::Int32, + DataType::Float64, + DataType::Utf8, + DataType::Binary, + DataType::List(Box::new(Field::new("a", DataType::Binary, true))), + DataType::Union( + vec![Field::new("a", DataType::Binary, true)], + None, + UnionMode::Sparse, + ), + DataType::Union( + vec![Field::new("a", DataType::Binary, true)], + None, + UnionMode::Dense, + ), + DataType::Struct(vec![Field::new("a", DataType::Int32, true)]), + ]; + let a = datatypes + .into_iter() + .map(|dt| DataType::Extension("ext".to_owned(), Box::new(dt), None)) + .all(|x| { + let a = new_empty_array(x); + a.len() == 0 && matches!(a.data_type(), DataType::Extension(_, _, _)) + }); + assert!(a); +} + #[test] fn test_clone() { let datatypes = vec![