diff --git a/cpp/src/arrow/array/array_view_test.cc b/cpp/src/arrow/array/array_view_test.cc index dfa8021dde1..ca1c82bfc90 100644 --- a/cpp/src/arrow/array/array_view_test.cc +++ b/cpp/src/arrow/array/array_view_test.cc @@ -27,6 +27,7 @@ #include "arrow/extension_type.h" #include "arrow/result.h" #include "arrow/status.h" +#include "arrow/testing/extension_type.h" #include "arrow/testing/gtest_util.h" #include "arrow/type.h" #include "arrow/util/endian.h" @@ -480,6 +481,32 @@ TEST(TestArrayView, ExtensionType) { CheckView(expected, arr); } +TEST(TestArrayView, ExtensionTypeNestedStorage) { + auto ty1 = list_extension_type(); + const auto& ext_type = static_cast(*ty1); + auto data = ArrayFromJSON(ext_type.storage_type(), "[[0, -1, 42], [5, 6]]")->data(); + data->type = ty1; + auto arr = ext_type.MakeArray(data); + ASSERT_OK(arr->ValidateFull()); + + CheckView(arr, arr); +} + +TEST(TestArrayView, StructAsStructWithExtensionField) { + auto ty1 = list_extension_type(); + const auto& ext_type = static_cast(*ty1); + auto ext_data = ArrayFromJSON(ext_type.storage_type(), "[[0, -1, 42], [5, 6]]")->data(); + ext_data->type = ty1; + auto ext_arr = ext_type.MakeArray(ext_data); + + auto sibling = ArrayFromJSON(int8(), "[1, 2]"); + ASSERT_OK_AND_ASSIGN(auto arr, StructArray::Make({ext_arr, sibling}, + std::vector{"a", "b"})); + ASSERT_OK(arr->ValidateFull()); + + CheckView(arr, arr); +} + TEST(TestArrayView, NonZeroOffset) { auto arr = ArrayFromJSON(int16(), "[10, 11, 12, 13]"); diff --git a/cpp/src/arrow/array/data.cc b/cpp/src/arrow/array/data.cc index dcc659e7094..edeb1c31190 100644 --- a/cpp/src/arrow/array/data.cc +++ b/cpp/src/arrow/array/data.cc @@ -744,7 +744,12 @@ namespace { void AccumulateLayouts(const std::shared_ptr& type, std::vector* layouts) { layouts->push_back(type->layout()); - for (const auto& child : type->fields()) { + const DataType* type_for_children = type.get(); + if (type->id() == Type::EXTENSION) { + const auto& ext_type = checked_cast(*type); + type_for_children = ext_type.storage_type().get(); + } + for (const auto& child : type_for_children->fields()) { AccumulateLayouts(child->type(), layouts); } } @@ -916,8 +921,14 @@ struct ViewDataImpl { out_type, out_length, std::move(out_buffers), out_null_count, out_offset); out_data->dictionary = dictionary; + const DataType* type_for_children = out_type.get(); + if (out_type->id() == Type::EXTENSION) { + const auto& ext_type = checked_cast(*out_type); + type_for_children = ext_type.storage_type().get(); + } + // Process children recursively, depth-first - for (const auto& child_field : out_type->fields()) { + for (const auto& child_field : type_for_children->fields()) { std::shared_ptr child_data; RETURN_NOT_OK(MakeDataView(child_field, &child_data)); out_data->child_data.push_back(std::move(child_data)); diff --git a/python/pyarrow/tests/test_extension_type.py b/python/pyarrow/tests/test_extension_type.py index 35b801eca87..ce2870f795d 100644 --- a/python/pyarrow/tests/test_extension_type.py +++ b/python/pyarrow/tests/test_extension_type.py @@ -784,6 +784,38 @@ def test_cast_to_extension_with_nested_storage(): assert result.equals(expected) +def test_cast_table_with_extension_type_in_struct(): + # https://github.com/apache/arrow/issues/37004 + + # fixed-size list + array = pa.array([[1, 2], [3, 4], [5, 6]], pa.list_(pa.int32(), 2)) + ext_type = MyFixedListType(pa.list_(pa.int32(), 2)) + ext_array = pa.ExtensionArray.from_storage(ext_type, array) + struct_array = pa.StructArray.from_arrays([ext_array], "x") + + table = pa.Table.from_arrays([struct_array], ["field1"]) + result = table.cast(table.schema) + assert result.schema == table.schema + assert result.equals(table) + + result = struct_array.cast(struct_array.type) + assert result.equals(struct_array) + + # variable-size list + array = pa.array([[1, 2], [3, 4, 5], [6]], pa.list_(pa.int32())) + ext_type = MyListType(pa.list_(pa.int32())) + ext_array = pa.ExtensionArray.from_storage(ext_type, array) + struct_array = pa.StructArray.from_arrays([ext_array], "x") + + table = pa.Table.from_arrays([struct_array], ["field1"]) + result = table.cast(table.schema) + assert result.schema == table.schema + assert result.equals(table) + + result = struct_array.cast(struct_array.type) + assert result.equals(struct_array) + + def test_concat(): arr1 = pa.array([1, 2, 3], IntegerType()) arr2 = pa.array([4, 5, 6], IntegerType())