Various fixes for nested Extension types

src/array/growable/structure.rs

@@ -25,6 +25,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) {
@@ -69,7 +71,7 @@ impl<'a> GrowableStruct<'a> {
         let values = values.into_iter().map(|mut x| x.as_box()).collect();
 
         StructArray::new(
-            DataType::Struct(self.arrays[0].fields().to_vec()),
+            self.arrays[0].data_type().clone(),
             values,
             validity.into(),
         )
@@ -121,7 +123,7 @@ impl<'a> From<GrowableStruct<'a>> 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(),
         )

src/array/struct_/mod.rs

@@ -131,7 +131,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()))

src/array/union/mod.rs

@@ -152,7 +152,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()))

src/ffi/schema.rs

@@ -46,7 +46,7 @@ impl ArrowSchema {
         let mut flags = field.is_nullable as i64 * 2;
 
         // allocate (and hold) the children
-        let children_vec = match field.data_type() {
+        let children_vec = match field.data_type().to_logical_type() {
             DataType::List(field) => {
                 vec![Box::new(ArrowSchema::new(field.as_ref()))]
             }
@@ -473,7 +473,7 @@ fn to_format(data_type: &DataType) -> String {
 }
 
 pub(super) fn get_child(data_type: &DataType, index: usize) -> Result<DataType> {
-    match (index, data_type) {
+    match (index, data_type.to_logical_type()) {
         (0, DataType::List(field)) => Ok(field.data_type().clone()),
         (0, DataType::FixedSizeList(field, _)) => Ok(field.data_type().clone()),
         (0, DataType::LargeList(field)) => Ok(field.data_type().clone()),

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::{self, DataType},
 };
 
 fn create_list_array(data: Vec<Option<Vec<Option<i32>>>>) -> ListArray<i32> {
@@ -9,6 +12,33 @@ fn create_list_array(data: Vec<Option<Vec<Option<i32>>>>) -> ListArray<i32> {
     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::from_data(
+        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<i32> = a.into();
+    assert_eq!(array_ext.data_type(), result.data_type());
+    dbg!(result);
+}
+
 #[test]
 fn basic() {
     let data = vec![

tests/it/array/growable/mod.rs

@@ -10,9 +10,9 @@ mod struct_;
 mod union;
 mod utf8;
 
-use arrow2::array::growable::make_growable;
+use arrow2::array::growable::{make_growable, GrowableStruct};
 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);
 }

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![

tests/it/ffi/data.rs

@@ -346,3 +346,22 @@ fn extension() -> Result<()> {
     );
     test_round_trip_schema(field)
 }
+
+#[test]
+fn extension_struct() -> Result<()> {
+    let data_type = DataType::Extension(
+        "ext".to_owned(),
+        Box::new(DataType::Struct(vec![Field::new(
+            "a",
+            DataType::Int32,
+            true,
+        )])),
+        None,
+    );
+
+    let values = vec![Int32Array::from([Some(1), None, Some(3)]).boxed()];
+    let validity = Bitmap::from([true, false, true]);
+
+    let array = StructArray::from_data(data_type, values, validity.into());
+    test_round_trip(array)
+}