ARROW-17813: [Python] Nested ExtensionArray conversion to/from pandas/numpy

python/pyarrow/src/arrow_to_pandas.cc

@@ -18,6 +18,7 @@
 // Functions for pandas conversion via NumPy
 
 #include "arrow_to_pandas.h"
+#include <arrow/compute/cast.h>
 #include "numpy_interop.h"  // IWYU pragma: expand
 
 #include <cmath>
@@ -191,6 +192,13 @@ static inline bool ListTypeSupported(const DataType& type) {
       const auto& list_type = checked_cast<const BaseListType&>(type);
       return ListTypeSupported(*list_type.value_type());
     }
+    case Type::EXTENSION: {
+      auto ext = std::dynamic_pointer_cast<ExtensionType>(type.GetSharedPtr());
+      if (ext == nullptr) {
+        return false;
+      }
+      return ListTypeSupported(*(ext->storage_type()));
+    }
     default:
       break;
   }
@@ -730,31 +738,38 @@ Status DecodeDictionaries(MemoryPool* pool, const std::shared_ptr<DataType>& den
 template <typename ListArrayT>
 Status ConvertListsLike(PandasOptions options, const ChunkedArray& data,
                         PyObject** out_values) {
+  using ListArrayType = typename ListArrayT::TypeClass;
+  const auto& list_type = checked_cast<const ListArrayType&>(*data.type());
+  const auto& value_type = list_type.value_type();
+
+  const auto& val_type = checked_cast<const ExtensionType&>(*value_type);
+  const auto& storage_ty = val_type.storage_type();
+
   // Get column of underlying value arrays
   ArrayVector value_arrays;
   for (int c = 0; c < data.num_chunks(); c++) {
     const auto& arr = checked_cast<const ListArrayT&>(*data.chunk(c));
-    value_arrays.emplace_back(arr.values());
+    if (arr.values()->type()->id() == Type::EXTENSION) {
+      const auto& arr_ext = checked_cast<const ExtensionArray&>(*arr.values());
+      value_arrays.emplace_back(arr_ext.storage());
+    } else {
+      value_arrays.emplace_back(arr.values());
+    }
   }
-  using ListArrayType = typename ListArrayT::TypeClass;
-  const auto& list_type = checked_cast<const ListArrayType&>(*data.type());
-  auto value_type = list_type.value_type();
 
-  auto flat_column = std::make_shared<ChunkedArray>(value_arrays, value_type);
+  auto flat_column = std::make_shared<ChunkedArray>(value_arrays, storage_ty);
 
   options = MakeInnerOptions(std::move(options));
 
   OwnedRefNoGIL owned_numpy_array;
   RETURN_NOT_OK(ConvertChunkedArrayToPandas(options, flat_column, nullptr,
                                             owned_numpy_array.ref()));
-
   PyObject* numpy_array = owned_numpy_array.obj();
   DCHECK(PyArray_Check(numpy_array));
 
   int64_t chunk_offset = 0;
   for (int c = 0; c < data.num_chunks(); c++) {
     const auto& arr = checked_cast<const ListArrayT&>(*data.chunk(c));
-
     const bool has_nulls = data.null_count() > 0;
     for (int64_t i = 0; i < arr.length(); ++i) {
       if (has_nulls && arr.IsNull(i)) {
@@ -1969,6 +1984,7 @@ static Status GetPandasWriterType(const ChunkedArray& data, const PandasOptions&
     case Type::LARGE_LIST:
     case Type::MAP: {
       auto list_type = std::static_pointer_cast<BaseListType>(data.type());
+
       if (!ListTypeSupported(*list_type->value_type())) {
         return Status::NotImplemented("Not implemented type for Arrow list to pandas: ",
                                       list_type->value_type()->ToString());

python/pyarrow/tests/test_extension_type.py

@@ -925,3 +925,28 @@ def test_empty_take():
     result = empty_arr.take(pa.array([], pa.int32()))
     assert len(result) == 0
     assert result.equals(empty_arr)
+
+
+class LabelType(pa.ExtensionType):
+
+         def __init__(self):
+             super(LabelType, self).__init__(pa.string(), "label")
+
+         def __arrow_ext_serialize__(self):
+             return b""
+
+         @classmethod
+         def __arrow_ext_deserialize__(cls, storage_type, serialized):
+             return LabelType()
+
+@pytest.mark.parametrize("data,ty", (
+    ([1, 2, 3], IntegerType),
+    (["cat", "dog", "horse"], LabelType)
+))
+def test_extension_array_to_numpy(data, ty):     
+     storage = pa.array(data)
+     ext_arr = pa.ExtensionArray.from_storage(ty(), storage)
+     offsets = pa.array([0, 1])
+     list_arr = pa.ListArray.from_arrays(offsets, ext_arr)
+     list_arr.to_numpy(zero_copy_only=False)
+