Skip to content

GH-50993: [CI][Integration] Add extension-wrapped union to integration data #51027

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Alb3e3 wants to merge 7 commits into
base: main
Choose a base branch
from
Open

GH-50993: [CI][Integration] Add extension-wrapped union to integration data #51027

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
3c71b93
GH-50993: [CI][Integration] Add extension-wrapped union to integratio…
Alb3e3 Aug 28, 2026
5028fb5
GH-50993: flag extension_union as allow-unregistered
Alb3e3 Aug 28, 2026
de746b3
GH-50993: exercise registered union extensions
Alb3e3 Aug 31, 2026
f215787
GH-50993: Register union extension types in C++ integration consumers
Alb3e3 Aug 31, 2026
b2b42d4
GH-50993: Export union extension arrays with storage layout
Alb3e3 Aug 31, 2026
b6634be
GH-50993: Format bridge.cc
Alb3e3 Aug 31, 2026
3abb59a
GH-50993: Address review: use storage_id() and add union extension ex…
Alb3e3 Sep 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions cpp/src/arrow/c/bridge.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
10 changes: 9 additions & 1 deletion cpp/src/arrow/c/bridge_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<const ExtensionType&>(*physical_type).storage_type().get();
}

auto expected_n_buffers = static_cast<int64_t>(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;
}
Expand Down Expand Up @@ -1173,6 +1179,8 @@ TEST_F(TestArrayExport, Extension) {
TestPrimitive(ExampleUuid);
TestPrimitive(ExampleSmallint);
TestPrimitive(ExampleComplex128);
TestPrimitive(ExampleDenseUnionExtension);
TestPrimitive(ExampleSparseUnionExtension);
}

TEST_F(TestArrayExport, MovePrimitive) {
Expand Down
5 changes: 3 additions & 2 deletions cpp/src/arrow/flight/integration_tests/test_integration_client.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,8 +145,9 @@ class IntegrationTestScenario : public Scenario {

Status RunClient(std::unique_ptr<FlightClient> 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}};

Expand Down
3 changes: 2 additions & 1 deletion cpp/src/arrow/integration/c_data_integration_internal.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 <typename Func>
Expand Down
4 changes: 3 additions & 1 deletion cpp/src/arrow/integration/json_integration_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()});
Comment on lines +229 to +231

@pitrou pitrou Sep 1, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@paleolimbot @zeroshade I think similar guards will have to be added to arrow-go and nanoarrow so that the extension type wrapping happens in integration testing.


if (json_path == "") {
return Status::Invalid("Must specify json file name");
Expand Down
6 changes: 5 additions & 1 deletion cpp/src/arrow/integration/json_internal.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1953,7 +1953,11 @@ class ArrayReader {
Result<std::shared_ptr<ArrayData>> Parse() {
ARROW_ASSIGN_OR_RAISE(length_, GetMemberInt<int32_t>(obj_, "count"));

if (::arrow::internal::may_have_validity_bitmap(type_->id())) {
const auto physical_type_id =
type_->id() == Type::EXTENSION
? checked_cast<const ExtensionType&>(*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());
}
Expand Down
6 changes: 6 additions & 0 deletions cpp/src/arrow/testing/extension_type.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -266,6 +266,12 @@ std::shared_ptr<Array> ExampleDictExtension();
ARROW_TESTING_EXPORT
std::shared_ptr<Array> ExampleComplex128();

ARROW_TESTING_EXPORT
std::shared_ptr<Array> ExampleDenseUnionExtension();

ARROW_TESTING_EXPORT
std::shared_ptr<Array> ExampleSparseUnionExtension();

ARROW_TESTING_EXPORT
std::shared_ptr<Array> MakeComplex128(const std::shared_ptr<Array>& real,
const std::shared_ptr<Array>& imag);
Expand Down
14 changes: 14 additions & 0 deletions cpp/src/arrow/testing/gtest_util.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1092,6 +1092,20 @@ std::shared_ptr<Array> ExampleComplex128() {
return ExtensionType::WrapArray(complex128(), arr);
}

std::shared_ptr<Array> ExampleDenseUnionExtension() {
auto type = dense_union_extension_type();
auto storage_type = checked_cast<const ExtensionType&>(*type).storage_type();
return ExtensionType::WrapArray(
type, ArrayFromJSON(storage_type, R"([[0, 1.5], [1, "abc"]])"));
}

std::shared_ptr<Array> ExampleSparseUnionExtension() {
auto type = sparse_union_extension_type();
auto storage_type = checked_cast<const ExtensionType&>(*type).storage_type();
return ExtensionType::WrapArray(
type, ArrayFromJSON(storage_type, R"([[0, 1.5], [1, "abc"]])"));
}

ExtensionTypeGuard::ExtensionTypeGuard(const std::shared_ptr<DataType>& type)
: ExtensionTypeGuard(DataTypeVector{type}) {}

Expand Down
25 changes: 25 additions & 0 deletions dev/archery/archery/integration/datagen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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-')

Expand Down Expand Up @@ -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(),

@alamb alamb Aug 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

so does this mean that all the other implementations correctly handle this extension wrapped union? If so, maybe that implies that whatever the problem was in #50927 is C++ specific and maybe not worth additional cross implementation coverage

@pitrou pitrou Aug 31, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@alamb It depends how other implementations handle extension types. If they need to be registered as in Arrow C++ for wrapping to occur, then a bug could be hidden by lack of registration.

@Alb3e3 Alb3e3 Aug 31, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed. After removing the bypass, there were still two C++ integration consumers missing extension-type registration: flight-test-integration-client and c_data_integration_internal.

With just those registrations reverted locally, archery integration --run-flight --run-c-data -x -k extension_union fails in Flight with Extension type 'sparse-union-extension' not found while opening generated_extension_union.json.

I pushed f215787 to register the dense/sparse union extension types in both places, and rerunning the same focused extension_union Flight + C Data command passes locally.

]

generated_paths = []
Expand Down
Loading