From 8a7ab318f99f8e317ebf07f11933085611a78066 Mon Sep 17 00:00:00 2001
From: Igor Stadnyk <fallenmi@users.noreply.github.com>
Date: Sun, 16 Aug 2026 07:03:51 +0100
Subject: [PATCH] GH-50623: [C++][IPC] Use extension storage layout for
 validity buffers

---
 cpp/src/arrow/ipc/read_write_test.cc | 67 ++++++++++++++++++++++++++++
 cpp/src/arrow/ipc/writer.cc          |  4 +-
 2 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/cpp/src/arrow/ipc/read_write_test.cc b/cpp/src/arrow/ipc/read_write_test.cc
index aa3f9fd77f23..74ecc4d06f34 100644
--- a/cpp/src/arrow/ipc/read_write_test.cc
+++ b/cpp/src/arrow/ipc/read_write_test.cc
@@ -410,6 +410,40 @@ class ExtensionTypesMixin {
   ExtensionTypeGuard ext_guard_;
 };
 
+class IpcStorageExtensionType : public ExtensionType {
+ public:
+  IpcStorageExtensionType(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 {
+    return extension_name_ == other.extension_name() &&
+           storage_type()->Equals(*other.storage_type());
+  }
+
+  std::shared_ptr<Array> MakeArray(std::shared_ptr<ArrayData> data) const override {
+    return std::make_shared<ExtensionArray>(std::move(data));
+  }
+
+  Result<std::shared_ptr<DataType>> Deserialize(
+      std::shared_ptr<DataType> storage_type,
+      const std::string& serialized_data) const override {
+    if (serialized_data != extension_name_) {
+      return Status::Invalid("Unexpected extension metadata: ", serialized_data);
+    }
+    return std::make_shared<IpcStorageExtensionType>(std::move(storage_type),
+                                                     extension_name_);
+  }
+
+  std::string Serialize() const override { return extension_name_; }
+
+ private:
+  std::string extension_name_;
+};
+
 class IpcTestFixture : public io::MemoryMapFixture, public ExtensionTypesMixin {
  public:
   void SetUp() {
@@ -695,6 +729,39 @@ TEST_P(TestIpcRoundTrip, ZeroLengthArrays) {
   CheckRoundtrip(bin_array2);
 }
 
+TEST_F(TestIpcRoundTrip, ExtensionWithUnionStorage) {
+  std::shared_ptr<RecordBatch> batch;
+  ASSERT_OK(MakeUnion(&batch));
+
+  for (int i = 0; i < batch->num_columns(); ++i) {
+    auto extension_type = std::make_shared<IpcStorageExtensionType>(
+        batch->column(i)->type(), "ipc.union." + std::to_string(i));
+    ExtensionTypeGuard extension_guard(extension_type);
+    for (const auto version : kMetadataVersions) {
+      options_.metadata_version = version;
+      CheckRoundtrip(ExtensionType::WrapArray(extension_type, batch->column(i)),
+                     options_);
+    }
+  }
+}
+
+TEST_F(TestIpcRoundTrip, ExtensionWithNullStorage) {
+  auto extension_type = std::make_shared<IpcStorageExtensionType>(null(), "ipc.null");
+  ExtensionTypeGuard extension_guard(extension_type);
+
+  auto extension =
+      ExtensionType::WrapArray(extension_type, std::make_shared<NullArray>(3));
+  auto values = ArrayFromJSON(int64(), "[1, 2, 3]");
+  auto batch = RecordBatch::Make(
+      schema({field("extension", extension_type), field("values", int64())}), 3,
+      {extension, values});
+
+  for (const auto version : kMetadataVersions) {
+    options_.metadata_version = version;
+    CheckRoundtrip(*batch, options_);
+  }
+}
+
 TEST_F(TestIpcRoundTrip, SparseUnionOfStructsWithReusedBuffers) {
   auto storage_type = struct_({
       field("i", int32()),
diff --git a/cpp/src/arrow/ipc/writer.cc b/cpp/src/arrow/ipc/writer.cc
index 263689a648d0..c24e1a5283c3 100644
--- a/cpp/src/arrow/ipc/writer.cc
+++ b/cpp/src/arrow/ipc/writer.cc
@@ -164,7 +164,9 @@ class RecordBatchSerializer {
 
     // 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)) {
+    // Extension arrays use the physical layout of their storage type.
+    if (internal::HasValidityBitmap(arr.type()->storage_id(),
+                                    options_.metadata_version)) {
       if (arr.null_count() > 0) {
         std::shared_ptr<Buffer> bitmap;
         RETURN_NOT_OK(GetTruncatedBitmap(arr.offset(), arr.length(), arr.null_bitmap(),