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

cpp/src/arrow/ipc/read_write_test.cc

@@ -397,14 +397,18 @@ const std::vector<test::MakeRecordBatch*> kBatchCases = {
     &MakeIntervals,
     &MakeUuid,
     &MakeComplex128,
-    &MakeDictExtension};
+    &MakeDictExtension,
+    &MakeDenseUnionExtension,
+    &MakeSparseUnionExtension};
 
 static int g_file_number = 0;
 
 class ExtensionTypesMixin {
  public:
   // Register the extension types required to ensure roundtripping
-  ExtensionTypesMixin() : ext_guard_({uuid(), dict_extension_type(), complex128()}) {}
+  ExtensionTypesMixin()
+      : ext_guard_({uuid(), dict_extension_type(), complex128(),
+                    dense_union_extension_type(), sparse_union_extension_type()}) {}
 
  protected:
   ExtensionTypeGuard ext_guard_;

cpp/src/arrow/ipc/test_common.cc

@@ -1164,6 +1164,35 @@ Status MakeDictExtension(std::shared_ptr<RecordBatch>* out) {
 
 namespace {
 
+Status MakeUnionExtension(const std::shared_ptr<DataType>& type,
+                          std::shared_ptr<RecordBatch>* out) {
+  auto storage_type = checked_cast<const ExtensionType&>(*type).storage_type();
+
+  auto f0 = field("f0", type);
+  auto f1 = field("f1", type, /*nullable=*/false);
+  auto schema = ::arrow::schema({f0, f1});
+
+  auto a0 = ExtensionType::WrapArray(
+      type, ArrayFromJSON(storage_type, R"([[0, 1.5], [1, null]])"));
+  auto a1 = ExtensionType::WrapArray(
+      type, ArrayFromJSON(storage_type, R"([[0, 1.5], [1, "abc"]])"));
+
+  *out = RecordBatch::Make(schema, a1->length(), {a0, a1});
+  return Status::OK();
+}
+
+}  // namespace
+
+Status MakeDenseUnionExtension(std::shared_ptr<RecordBatch>* out) {
+  return MakeUnionExtension(dense_union_extension_type(), out);
+}
+
+Status MakeSparseUnionExtension(std::shared_ptr<RecordBatch>* out) {
+  return MakeUnionExtension(sparse_union_extension_type(), out);
+}
+
+namespace {
+
 template <typename CValueType, typename SeedType, typename DistributionType>
 void FillRandomData(CValueType* data, size_t n, CValueType min, CValueType max,
                     SeedType seed) {

cpp/src/arrow/ipc/test_common.h

@@ -185,6 +185,12 @@ Status MakeComplex128(std::shared_ptr<RecordBatch>* out);
 ARROW_TESTING_EXPORT
 Status MakeDictExtension(std::shared_ptr<RecordBatch>* out);
 
+ARROW_TESTING_EXPORT
+Status MakeDenseUnionExtension(std::shared_ptr<RecordBatch>* out);
+
+ARROW_TESTING_EXPORT
+Status MakeSparseUnionExtension(std::shared_ptr<RecordBatch>* out);
+
 ARROW_TESTING_EXPORT
 Status MakeRandomTensor(const std::shared_ptr<DataType>& type,
                         const std::vector<int64_t>& shape, bool row_major_p,

cpp/src/arrow/ipc/writer.cc

@@ -150,32 +150,44 @@ class RecordBatchSerializer {
       return Status::Invalid("Max recursion depth reached");
     }
 
-    if (!options_.allow_64bit && arr.length() > std::numeric_limits<int32_t>::max()) {
+    // An extension array is serialized as its storage array: the extension type
+    // itself is carried in the schema, not in the record batch body.  The storage
+    // array shares the extension array's ArrayData, so length, offset and null
+    // count are unchanged; only the type id differs, and it is the storage type id
+    // that decides the buffer layout below.
+    const Array& physical_arr = arr.type_id() == Type::EXTENSION
+                                    ? *checked_cast<const ExtensionArray&>(arr).storage()
+                                    : arr;
+
+    if (!options_.allow_64bit &&
+        physical_arr.length() > std::numeric_limits<int32_t>::max()) {
       return Status::CapacityError("Cannot write arrays larger than 2^31 - 1 in length");
     }
 
-    if (arr.offset() != 0 && arr.device_type() != DeviceAllocationType::kCPU) {
+    if (physical_arr.offset() != 0 &&
+        physical_arr.device_type() != DeviceAllocationType::kCPU) {
       // https://github.com/apache/arrow/issues/43029
       return Status::NotImplemented("Cannot compute null count for non-cpu sliced array");
     }
 
     // push back all common elements
-    field_nodes_.push_back({arr.length(), arr.null_count(), 0});
+    field_nodes_.push_back({physical_arr.length(), physical_arr.null_count(), 0});
 
     // In V4, null types have no validity bitmap
     // In V5 and later, null and union types have no validity bitmap
-    if (internal::HasValidityBitmap(arr.type_id(), options_.metadata_version)) {
-      if (arr.null_count() > 0) {
+    if (internal::HasValidityBitmap(physical_arr.type_id(), options_.metadata_version)) {
+      if (physical_arr.null_count() > 0) {
         std::shared_ptr<Buffer> bitmap;
-        RETURN_NOT_OK(GetTruncatedBitmap(arr.offset(), arr.length(), arr.null_bitmap(),
-                                         options_.memory_pool, &bitmap));
+        RETURN_NOT_OK(GetTruncatedBitmap(physical_arr.offset(), physical_arr.length(),
+                                         physical_arr.null_bitmap(), options_.memory_pool,
+                                         &bitmap));
         out_->body_buffers.emplace_back(std::move(bitmap));
       } else {
         // Push a dummy zero-length buffer, not to be copied
         out_->body_buffers.emplace_back(kNullBuffer);
       }
     }
-    return VisitType(arr);
+    return VisitType(physical_arr);
   }
 
   // Override this for writing dictionary metadata

cpp/src/arrow/testing/extension_type.h

@@ -168,6 +168,38 @@ class ARROW_TESTING_EXPORT MetadataOptionalExtensionType : public ExtensionType
   }
 };
 
+class ARROW_TESTING_EXPORT UnionExtensionArray : public ExtensionArray {
+ public:
+  using ExtensionArray::ExtensionArray;
+};
+
+/// \brief An extension type over an arbitrary union storage type.
+///
+/// The storage type and extension name are given at construction, so that a
+/// dense-union-backed and a sparse-union-backed extension type can coexist in
+/// the same extension type registry.
+class ARROW_TESTING_EXPORT UnionExtensionType : public ExtensionType {
+ public:
+  UnionExtensionType(std::shared_ptr<DataType> storage_type, std::string extension_name)
+      : ExtensionType(std::move(storage_type)),
+        extension_name_(std::move(extension_name)) {}
+
+  std::string extension_name() const override { return extension_name_; }
+
+  bool ExtensionEquals(const ExtensionType& other) const override;
+
+  std::shared_ptr<Array> MakeArray(std::shared_ptr<ArrayData> data) const override;
+
+  Result<std::shared_ptr<DataType>> Deserialize(
+      std::shared_ptr<DataType> storage_type,
+      const std::string& serialized) const override;
+
+  std::string Serialize() const override { return extension_name_; }
+
+ private:
+  std::string extension_name_;
+};
+
 class ARROW_TESTING_EXPORT Complex128Array : public ExtensionArray {
  public:
   using ExtensionArray::ExtensionArray;
@@ -213,6 +245,12 @@ std::shared_ptr<DataType> binary_view_extension_type();
 ARROW_TESTING_EXPORT
 std::shared_ptr<DataType> complex128();
 
+ARROW_TESTING_EXPORT
+std::shared_ptr<DataType> dense_union_extension_type();
+
+ARROW_TESTING_EXPORT
+std::shared_ptr<DataType> sparse_union_extension_type();
+
 ARROW_TESTING_EXPORT
 std::shared_ptr<Array> ExampleUuid();
 

cpp/src/arrow/testing/gtest_util.cc

@@ -977,6 +977,29 @@ Result<std::shared_ptr<DataType>> BinaryViewExtensionType::Deserialize(
   return std::make_shared<BinaryViewExtensionType>();
 }
 
+bool UnionExtensionType::ExtensionEquals(const ExtensionType& other) const {
+  return (other.extension_name() == this->extension_name());
+}
+
+std::shared_ptr<Array> UnionExtensionType::MakeArray(
+    std::shared_ptr<ArrayData> data) const {
+  DCHECK_EQ(data->type->id(), Type::EXTENSION);
+  DCHECK(ExtensionEquals(checked_cast<const ExtensionType&>(*data->type)));
+  return std::make_shared<UnionExtensionArray>(data);
+}
+
+Result<std::shared_ptr<DataType>> UnionExtensionType::Deserialize(
+    std::shared_ptr<DataType> storage_type, const std::string& serialized) const {
+  if (serialized != extension_name_) {
+    return Status::Invalid("Type identifier did not match: '", serialized, "'");
+  }
+  if (!storage_type->Equals(*storage_type_)) {
+    return Status::Invalid("Invalid storage type for ", extension_name_, ": ",
+                           storage_type->ToString());
+  }
+  return std::make_shared<UnionExtensionType>(std::move(storage_type), extension_name_);
+}
+
 bool Complex128Type::ExtensionEquals(const ExtensionType& other) const {
   return (other.extension_name() == this->extension_name());
 }
@@ -1019,6 +1042,18 @@ std::shared_ptr<DataType> dict_extension_type() {
 
 std::shared_ptr<DataType> complex128() { return std::make_shared<Complex128Type>(); }
 
+std::shared_ptr<DataType> dense_union_extension_type() {
+  return std::make_shared<UnionExtensionType>(
+      dense_union({field("floats", float64()), field("strings", large_utf8())}, {0, 1}),
+      "dense-union-extension");
+}
+
+std::shared_ptr<DataType> sparse_union_extension_type() {
+  return std::make_shared<UnionExtensionType>(
+      sparse_union({field("floats", float64()), field("strings", large_utf8())}, {0, 1}),
+      "sparse-union-extension");
+}
+
 std::shared_ptr<Array> MakeComplex128(const std::shared_ptr<Array>& real,
                                       const std::shared_ptr<Array>& imag) {
   auto type = complex128();