GH-50623: [C++][IPC] Fix extension-wrapped union IPC roundtrip

[Alb3e3]

Rationale for this change

IPC schema serialization already treats extension arrays as their storage type. Record batch serialization still handled extension arrays according to the logical extension type when deciding whether to emit a validity buffer. That mismatches storage types whose IPC layout differs from the default, including unions.

For extension-wrapped dense and sparse unions, the writer emitted a spurious top-level validity-buffer slot. The resulting IPC payload dropped the union buffers from the positions the reader expects, producing invalid data on read.

What changes are included in this PR?

• teach RecordBatchSerializer::VisitArray() to use the physical storage array when deciding common IPC buffer handling for extension arrays
• keep field-node accounting on the logical extension array while delegating layout-sensitive buffer decisions to the storage array
• add dense and sparse union extension roundtrip regression tests for both file and stream IPC writers

This follows the same storage-type rule already used by IPC schema serialization.

AI-assisted contribution note: this patch was prepared with AI assistance under Arrow's AI-generated code guidance. I reproduced the bug locally, validated the fix, and reviewed the final diff myself.

Are these changes tested?

Yes.

Locally verified with:

• arrow-ipc-read-write-test --gtest_filter='TestFileFormat.DenseUnionExtensionRoundTrip:TestFileFormat.SparseUnionExtensionRoundTrip:TestStreamFormat.DenseUnionExtensionRoundTrip:TestStreamFormat.SparseUnionExtensionRoundTrip'
• arrow-ipc-read-write-test --gtest_filter='-TestSchemaMetadata.MetadataVersionForwardCompatibility'
• arrow-extension-type-test

The excluded TestSchemaMetadata.MetadataVersionForwardCompatibility requires external test data and is unrelated to this patch.

Are there any user-facing changes?

No API changes.

This PR contains a "Critical Fix" in Arrow's template sense: it fixes IPC output that could otherwise contain invalid union payloads when an extension type wraps union storage, which can fail validation and crash on readback.

• GitHub Issue: [C++][Python] IPC serialization drops a union's buffers when the union is wrapped in an extension type (invalid data; segfault on read) #50623

[github-actions]

⚠️ GitHub issue #50623 has been automatically assigned in GitHub to PR creator.

[pitrou]

Thanks for the PR @Alb3e3 . This looks good on the principle, here are some suggestions to make this cleaner and more reusable.

[pitrou]

Can we move these two classes to arrow/testing/extension_type.h? They're probably going to be useful for other tests at some point.

[Alb3e3]

Moved to arrow/testing/extension_type.h in 4c4ecb9390, with the ExtensionEquals / MakeArray / Deserialize bodies in testing/gtest_util.cc next to the other example extension types.

I kept UnionExtensionType parameterized on (storage_type, extension_name) rather than splitting it into two classes, so the dense- and sparse-backed variants can be registered at the same time. Accessors are dense_union_extension_type() and sparse_union_extension_type(), matching the dict_extension_type() / complex128() naming. Deserialize() now also validates the storage type against storage_type_, which the throwaway version did not.

[pitrou]

Perhaps move these two functions to arrow/ipc/test_common.h?

[Alb3e3]

Moved to arrow/ipc/test_common.{h,cc} as MakeDenseUnionExtension / MakeSparseUnionExtension, sitting right after MakeDictExtension. Both delegate to a file-local MakeUnionExtension(type, out).

While moving them I also made them match the shape of the neighbouring makers — two fields (f0 nullable, f1 non-nullable) instead of one, and the nullable one now carries a null ([[0, 1.5], [1, null]]), which the originals did not exercise.

[pitrou]

Why not extend kBatchCases instead of adding dedicated test functions?

[Alb3e3]

Done — folded into kBatchCases and the four TEST_F blocks are gone.

Worth noting this is a strict improvement in coverage, not just tidier: the two cases now run through 22 test instances instead of 4, and that picks up TestIpcRoundTrip.SliceRoundTrip and TestIpcRoundTrip.ZeroLengthArrays, plus the four StreamDecoder* variants and TestFileFormatGenerator{,Coalesced} — none of which the dedicated tests touched.

It also made the regression sharper. Reverting writer.cc alone and rerunning, the new cases fail in TestFileFormat, TestIpcRoundTrip.RoundTrip, TestIpcRoundTrip.ZeroLengthArrays, and then TestFileFormatGenerator aborts outright on an out-of-bounds ArrayData child access.

[pitrou]

This is still using arr while switching to physical_arr in other places (e.g. arr.length() above vs. physical_arr->length() below), with no obvious rationale. Can we make the code more consistent?

[Alb3e3]

Fair — that inconsistency was not deliberate, it was left over from patching the call sites one at a time.

Rewritten to unwrap once into a reference and then use it everywhere:

const Array& physical_arr = arr.type_id() == Type::EXTENSION
                                ? *checked_cast<const ExtensionArray&>(arr).storage()
                                : arr;

arr is not referenced again after that line. This is safe because ExtensionArray::SetData() builds the storage from data->Copy() with only the type swapped, so length, offset, null_count and the buffers are shared — the only thing that differs is the type id, which is exactly what the layout decisions below need. Added a comment saying so.

[Alb3e3]

Thanks @pitrou — all four addressed in `4c4ecb9390`, replies inline. Summary:

Ask	Done
Move the two classes to `arrow/testing/extension_type.h`	Yes — declarations there, bodies in `testing/gtest_util.cc`
Move the two functions to `arrow/ipc/test_common.h`	Yes — as `MakeDenseUnionExtension` / `MakeSparseUnionExtension`
Extend `kBatchCases` instead of dedicated tests	Yes — the four `TEST_F` blocks are gone
`writer.cc` `arr` vs `physical_arr` consistency	Yes — unwrapped once into a reference, `arr` unused after

Net effect on `read_write_test.cc` is -87/+20: the file now only registers the two types in `ExtensionTypesMixin` and lists the two makers in `kBatchCases`.

Verification — `arrow-ipc-read-write-test`: 431/432, the one failure being `TestSchemaMetadata.MetadataVersionForwardCompatibility`, which is my checkout missing the `testing/` submodule (`Test resources not found, set ARROW_TEST_DATA`), not related to this change. `arrow-extension-type-test` 6/6 and `arrow-feather-test` 83/83 also pass, since both link `arrow_testing`.

Non-vacuousness: reverting `writer.cc` and keeping the tests reproduces the bug in `TestFileFormat.RoundTrip/{27,28}`, `TestIpcRoundTrip.RoundTrip/{27,28}`, `TestIpcRoundTrip.ZeroLengthArrays/{27,28}`, and aborts `TestFileFormatGenerator` on std::vector bounds-checking of the ArrayData children.

On the red CI: `C++ Format` was mine — `read_write_test.cc` was not clang-format-18.1.8 clean at the previous head. That code is deleted now, and all six touched files pass `clang-format --dry-run -Werror`. The other three failures look unrelated to me — `Python (Cython) Lint` dies on `AttributeError: 'DictComprehensionAppendNode' object has no attribute 'value_expr'`, and `AMD64 Conda C++ AVX2` fails only `arrow-filesystem-test` with `Attempt to initialize S3 after it has been finalized`. Happy to look if you think either is on me.

For transparency, and since it came up on another PR I have open: this work is AI-assisted. The evidence above is all locally reproducible and I would rather you have that context up front.

[github-actions]

⚠️ GitHub issue #50623 has been automatically assigned in GitHub to PR creator.

[pitrou]

+1, thank you @Alb3e3

[pitrou]

FTR @paleolimbot