diff --git a/cpp/src/arrow/c/bridge.cc b/cpp/src/arrow/c/bridge.cc index 4391d4cbc248..287226c6339d 100644 --- a/cpp/src/arrow/c/bridge.cc +++ b/cpp/src/arrow/c/bridge.cc @@ -579,16 +579,19 @@ struct ArrayExporter { // This is because ARROW-9037 is in version 0.17 and 0.17.1, and they are // not able to import arrays without a null bitmap and null_count == -1. data->GetNullCount(); + + const auto physical_type_id = data->type->storage_id(); + // Store buffer pointers size_t n_buffers = data->buffers.size(); auto buffers_begin = data->buffers.begin(); - if (n_buffers > 0 && !internal::may_have_validity_bitmap(data->type->id())) { + if (n_buffers > 0 && !internal::may_have_validity_bitmap(physical_type_id)) { --n_buffers; ++buffers_begin; } - bool need_variadic_buffer_sizes = data->type->storage_id() == Type::BINARY_VIEW || - data->type->storage_id() == Type::STRING_VIEW; + bool need_variadic_buffer_sizes = physical_type_id == Type::BINARY_VIEW || + physical_type_id == Type::STRING_VIEW; if (need_variadic_buffer_sizes) { ++n_buffers; } diff --git a/cpp/src/arrow/c/bridge_test.cc b/cpp/src/arrow/c/bridge_test.cc index 437273267548..9a4d104d0004 100644 --- a/cpp/src/arrow/c/bridge_test.cc +++ b/cpp/src/arrow/c/bridge_test.cc @@ -575,9 +575,15 @@ struct ArrayExportChecker { ASSERT_EQ(c_export->null_count, expected_data.null_count); ASSERT_EQ(c_export->offset, expected_data.offset); + const DataType* physical_type = expected_data.type.get(); + if (physical_type->id() == Type::EXTENSION) { + physical_type = + checked_cast(*physical_type).storage_type().get(); + } + auto expected_n_buffers = static_cast(expected_data.buffers.size()); auto expected_buffers = expected_data.buffers.data(); - if (!internal::may_have_validity_bitmap(expected_data.type->id())) { + if (!internal::may_have_validity_bitmap(physical_type->id())) { --expected_n_buffers; ++expected_buffers; } @@ -1173,6 +1179,8 @@ TEST_F(TestArrayExport, Extension) { TestPrimitive(ExampleUuid); TestPrimitive(ExampleSmallint); TestPrimitive(ExampleComplex128); + TestPrimitive(ExampleDenseUnionExtension); + TestPrimitive(ExampleSparseUnionExtension); } TEST_F(TestArrayExport, MovePrimitive) { diff --git a/cpp/src/arrow/flight/integration_tests/test_integration_client.cc b/cpp/src/arrow/flight/integration_tests/test_integration_client.cc index c5c3f10576d3..f39a4f04eb79 100644 --- a/cpp/src/arrow/flight/integration_tests/test_integration_client.cc +++ b/cpp/src/arrow/flight/integration_tests/test_integration_client.cc @@ -145,8 +145,9 @@ class IntegrationTestScenario : public Scenario { Status RunClient(std::unique_ptr client) override { // Make sure the required extension types are registered. - ExtensionTypeGuard uuid_ext_guard(uuid()); - ExtensionTypeGuard dict_ext_guard(dict_extension_type()); + ExtensionTypeGuard ext_guard({uuid(), dict_extension_type(), + dense_union_extension_type(), + sparse_union_extension_type()}); FlightDescriptor descr{FlightDescriptor::PATH, "", {FLAGS_path}}; diff --git a/cpp/src/arrow/integration/c_data_integration_internal.cc b/cpp/src/arrow/integration/c_data_integration_internal.cc index b21a0cc13b20..80b17580866e 100644 --- a/cpp/src/arrow/integration/c_data_integration_internal.cc +++ b/cpp/src/arrow/integration/c_data_integration_internal.cc @@ -38,7 +38,8 @@ namespace { // Make sure the extension types referenced in test data are registered. [[nodiscard]] auto RequireExtensionTypes() { - return ExtensionTypeGuard({uuid(), dict_extension_type()}); + return ExtensionTypeGuard({uuid(), dict_extension_type(), dense_union_extension_type(), + sparse_union_extension_type()}); } template diff --git a/cpp/src/arrow/integration/json_integration_test.cc b/cpp/src/arrow/integration/json_integration_test.cc index 700551c23b79..d10d72c1c7e2 100644 --- a/cpp/src/arrow/integration/json_integration_test.cc +++ b/cpp/src/arrow/integration/json_integration_test.cc @@ -226,7 +226,9 @@ Status RunCommand(const std::string& json_path, const std::string& arrow_path, const std::string& command) { // Make sure the required extension types are registered, as they will be // referenced in test data. - ExtensionTypeGuard ext_guard({uuid(), dict_extension_type()}); + ExtensionTypeGuard ext_guard({uuid(), dict_extension_type(), + dense_union_extension_type(), + sparse_union_extension_type()}); if (json_path == "") { return Status::Invalid("Must specify json file name"); diff --git a/cpp/src/arrow/integration/json_internal.cc b/cpp/src/arrow/integration/json_internal.cc index abf48b9df767..788ccd9c1134 100644 --- a/cpp/src/arrow/integration/json_internal.cc +++ b/cpp/src/arrow/integration/json_internal.cc @@ -1953,7 +1953,11 @@ class ArrayReader { Result> Parse() { ARROW_ASSIGN_OR_RAISE(length_, GetMemberInt(obj_, "count")); - if (::arrow::internal::may_have_validity_bitmap(type_->id())) { + const auto physical_type_id = + type_->id() == Type::EXTENSION + ? checked_cast(*type_).storage_type()->id() + : type_->id(); + if (::arrow::internal::may_have_validity_bitmap(physical_type_id)) { // Null and union types don't have a validity bitmap RETURN_NOT_OK(ParseValidityBitmap()); } diff --git a/cpp/src/arrow/testing/extension_type.h b/cpp/src/arrow/testing/extension_type.h index e5d6b597c5be..f931206a0d68 100644 --- a/cpp/src/arrow/testing/extension_type.h +++ b/cpp/src/arrow/testing/extension_type.h @@ -266,6 +266,12 @@ std::shared_ptr ExampleDictExtension(); ARROW_TESTING_EXPORT std::shared_ptr ExampleComplex128(); +ARROW_TESTING_EXPORT +std::shared_ptr ExampleDenseUnionExtension(); + +ARROW_TESTING_EXPORT +std::shared_ptr ExampleSparseUnionExtension(); + ARROW_TESTING_EXPORT std::shared_ptr MakeComplex128(const std::shared_ptr& real, const std::shared_ptr& imag); diff --git a/cpp/src/arrow/testing/gtest_util.cc b/cpp/src/arrow/testing/gtest_util.cc index daadfe9c2cdc..b7d2a963d0de 100644 --- a/cpp/src/arrow/testing/gtest_util.cc +++ b/cpp/src/arrow/testing/gtest_util.cc @@ -1092,6 +1092,20 @@ std::shared_ptr ExampleComplex128() { return ExtensionType::WrapArray(complex128(), arr); } +std::shared_ptr ExampleDenseUnionExtension() { + auto type = dense_union_extension_type(); + auto storage_type = checked_cast(*type).storage_type(); + return ExtensionType::WrapArray( + type, ArrayFromJSON(storage_type, R"([[0, 1.5], [1, "abc"]])")); +} + +std::shared_ptr ExampleSparseUnionExtension() { + auto type = sparse_union_extension_type(); + auto storage_type = checked_cast(*type).storage_type(); + return ExtensionType::WrapArray( + type, ArrayFromJSON(storage_type, R"([[0, 1.5], [1, "abc"]])")); +} + ExtensionTypeGuard::ExtensionTypeGuard(const std::shared_ptr& type) : ExtensionTypeGuard(DataTypeVector{type}) {} diff --git a/dev/archery/archery/integration/datagen.py b/dev/archery/archery/integration/datagen.py index ed50919dfcd6..adffb2a1b21d 100644 --- a/dev/archery/archery/integration/datagen.py +++ b/dev/archery/archery/integration/datagen.py @@ -1938,6 +1938,29 @@ def generate_extension_case(): dictionaries=[dict0]) +def generate_extension_wrapped_union_case(): + # Unions wrapped in an extension type, exercising the extension/union + # interaction across implementations (see the C++ fix in GH-50623). + sparse_union_type = ExtensionType( + 'sparse-union-extension', 'sparse-union-extension', + SparseUnionField('', [get_field('floats', 'float64'), + get_field('strings', 'largeutf8')], + type_ids=[0, 1])) + dense_union_type = ExtensionType( + 'dense-union-extension', 'dense-union-extension', + DenseUnionField('', [get_field('floats', 'float64'), + get_field('strings', 'largeutf8')], + type_ids=[0, 1])) + + fields = [ + ExtensionField('sparse_union_ext', sparse_union_type), + ExtensionField('dense_union_ext', dense_union_type), + ] + + batch_sizes = [0, 7] + return _generate_file("extension_union", fields, batch_sizes) + + def get_generated_json_files(tempdir=None): tempdir = tempdir or tempfile.mkdtemp(prefix='arrow-integration-') @@ -2047,6 +2070,8 @@ def get_generated_json_files(tempdir=None): # TODO(https://github.com/apache/arrow/issues/38045) .skip_format(SKIP_FLIGHT, '.NET') .skip_tester('Ruby'), + + generate_extension_wrapped_union_case(), ] generated_paths = []