Repository: jorgecarleitao/arrow2
Pull request #1334 (2022-12-17 13:14:37)

Fix nested extension types
This PR fixes various issues with Extension types at various nesting levels. Rebased onto `main`.

codecov[bot]: # [Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1334?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) Report
Base: **83.08**% // Head: **83.11**% // Increases project coverage by **`+0.02%`** :tada:
> Coverage data is based on head [(`739dfeb`)](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1334?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) compared to base [(`df95803`)](https://codecov.io/gh/jorgecarleitao/arrow2/commit/df95803b9a9a7d8a4daa8380e822293faa6aa9a6?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao).
> Patch coverage: 100.00% of modified lines in pull request are covered.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #1334      +/-   ##
==========================================
+ Coverage   83.08%   83.11%   +0.02%     
==========================================
  Files         370      370              
  Lines       40170    40170              
==========================================
+ Hits        33376    33387      +11     
+ Misses       6794     6783      -11     
```


| [Impacted Files](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1334?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) | Coverage Δ | |
|---|---|---|
| [src/array/list/mod.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1334/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L2xpc3QvbW9kLnJz) | `83.66% <100.00%> (ø)` | |
| [src/array/struct\_/mod.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1334/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L3N0cnVjdF8vbW9kLnJz) | `64.73% <100.00%> (+4.62%)` | :arrow_up: |
| [src/array/union/mod.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1334/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L3VuaW9uL21vZC5ycw==) | `76.55% <100.00%> (ø)` | |
| [src/io/ipc/read/file.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1334/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL2lwYy9yZWFkL2ZpbGUucnM=) | `96.42% <0.00%> (-0.45%)` | :arrow_down: |
| [src/ffi/schema.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1334/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2ZmaS9zY2hlbWEucnM=) | `90.28% <0.00%> (+0.20%)` | :arrow_up: |
| [src/offset.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1334/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL29mZnNldC5ycw==) | `84.89% <0.00%> (+0.35%)` | :arrow_up: |
| [src/array/mod.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1334/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L21vZC5ycw==) | `65.58% <0.00%> (+0.64%)` | :arrow_up: |
| [src/array/growable/mod.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1334/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L2dyb3dhYmxlL21vZC5ycw==) | `74.64% <0.00%> (+1.40%)` | :arrow_up: |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1334?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao).

jondo2010: Seems like a codecov submission failure. @jorgecarleitao, can you re-trigger the build on this please?

Patch:
diff --git a/src/array/growable/structure.rs b/src/array/growable/structure.rs
index e7f997024d..8e91cf03ce 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<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(),
         )
diff --git a/src/array/struct_/mod.rs b/src/array/struct_/mod.rs
index 1938338659..a8848bb557 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 3aa963c00a..e767aac2b2 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/src/io/parquet/read/indexes/mod.rs b/src/io/parquet/read/indexes/mod.rs
index e890f27d69..045482a831 100644
--- a/src/io/parquet/read/indexes/mod.rs
+++ b/src/io/parquet/read/indexes/mod.rs
@@ -245,7 +245,7 @@ pub fn has_indexes(row_group: &RowGroupMetaData) -> bool {
 /// This function is IO-bounded and calls `reader.read_exact` exactly once.
 /// # Error
 /// Errors iff the indexes can't be read or their deserialization to arrow is incorrect (e.g. invalid utf-8)
-fn read_columns_indexes<R: Read + Seek>(
+pub fn read_columns_indexes<R: Read + Seek>(
     reader: &mut R,
     chunks: &[ColumnChunkMetaData],
     fields: &[Field],
@@ -264,7 +264,7 @@ fn read_columns_indexes<R: Read + Seek>(
 }
 
 /// Returns the set of (row) intervals of the pages.
-fn compute_page_row_intervals(
+pub fn compute_page_row_intervals(
     locations: &[PageLocation],
     num_rows: usize,
 ) -> Result<Vec<Interval>, ParquetError> {


Test patch:
diff --git a/tests/it/array/growable/list.rs b/tests/it/array/growable/list.rs
index ef18365457..9aa56b109a 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<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::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<i32> = 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 a6e25ecba3..480bcf06c3 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 9df7c53d02..85318ba628 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![