From 1daaedbf1613f87d36f940d4b5007c5cb25ccfa1 Mon Sep 17 00:00:00 2001 From: Alberto Maschietto Date: Thu, 20 Aug 2026 17:12:03 +0400 Subject: [PATCH 1/2] GH-50623: [C++][IPC] fix extension-wrapped union IPC roundtrip --- cpp/src/arrow/ipc/read_write_test.cc | 87 +++++++++++++++++++++++++++- cpp/src/arrow/ipc/writer.cc | 15 +++-- 2 files changed, 97 insertions(+), 5 deletions(-) diff --git a/cpp/src/arrow/ipc/read_write_test.cc b/cpp/src/arrow/ipc/read_write_test.cc index aa3f9fd77f23..3cbc52db6244 100644 --- a/cpp/src/arrow/ipc/read_write_test.cc +++ b/cpp/src/arrow/ipc/read_write_test.cc @@ -71,6 +71,73 @@ using MetadataVector = std::vector>; namespace test { +class UnionExtensionArray : public ExtensionArray { + public: + using ExtensionArray::ExtensionArray; +}; + +class UnionExtensionType : public ExtensionType { + public: + UnionExtensionType(std::shared_ptr 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 other.extension_name() == extension_name_; + } + + std::shared_ptr MakeArray(std::shared_ptr data) const override { + return std::make_shared(std::move(data)); + } + + Result> Deserialize( + std::shared_ptr storage_type, + const std::string& serialized) const override { + if (serialized != extension_name_) { + return Status::Invalid("Type identifier did not match"); + } + return std::make_shared(std::move(storage_type), + extension_name_); + } + + std::string Serialize() const override { return extension_name_; } + + private: + std::string extension_name_; +}; + +std::shared_ptr dense_union_extension_type() { + return std::make_shared( + dense_union({field("floats", float64()), field("strings", large_utf8())}, {0, 1}), + "dense-union-extension"); +} + +std::shared_ptr sparse_union_extension_type() { + return std::make_shared( + sparse_union({field("floats", float64()), field("strings", large_utf8())}, {0, 1}), + "sparse-union-extension"); +} + +Status MakeDenseUnionExtension(std::shared_ptr* out) { + auto type = dense_union_extension_type(); + auto storage_type = checked_cast(*type).storage_type(); + auto storage = ArrayFromJSON(storage_type, R"([[0, 1.5], [1, "abc"]])"); + auto array = ExtensionType::WrapArray(type, storage); + *out = RecordBatch::Make(schema({field("f0", type)}), array->length(), {array}); + return Status::OK(); +} + +Status MakeSparseUnionExtension(std::shared_ptr* out) { + auto type = sparse_union_extension_type(); + auto storage_type = checked_cast(*type).storage_type(); + auto storage = ArrayFromJSON(storage_type, R"([[0, 1.5], [1, "abc"]])"); + auto array = ExtensionType::WrapArray(type, storage); + *out = RecordBatch::Make(schema({field("f0", type)}), array->length(), {array}); + return Status::OK(); +} + const std::vector kMetadataVersions = {MetadataVersion::V4, MetadataVersion::V5}; @@ -404,7 +471,9 @@ 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_; @@ -1954,6 +2023,22 @@ TEST_P(TestFileFormatGeneratorCoalesced, RoundTrip) { TEST_P(TestStreamFormat, RoundTrip) { TestRoundTripWithOptions(*GetParam()); } +TEST_F(TestFileFormat, DenseUnionExtensionRoundTrip) { + TestRoundTrip(MakeDenseUnionExtension, IpcWriteOptions::Defaults()); +} + +TEST_F(TestFileFormat, SparseUnionExtensionRoundTrip) { + TestRoundTrip(MakeSparseUnionExtension, IpcWriteOptions::Defaults()); +} + +TEST_F(TestStreamFormat, DenseUnionExtensionRoundTrip) { + TestRoundTrip(MakeDenseUnionExtension, IpcWriteOptions::Defaults()); +} + +TEST_F(TestStreamFormat, SparseUnionExtensionRoundTrip) { + TestRoundTrip(MakeSparseUnionExtension, IpcWriteOptions::Defaults()); +} + TEST_P(TestStreamDecoderData, RoundTrip) { TestRoundTripWithOptions(*GetParam()); } TEST_P(TestStreamDecoderBuffer, RoundTrip) { TestRoundTripWithOptions(*GetParam()); } diff --git a/cpp/src/arrow/ipc/writer.cc b/cpp/src/arrow/ipc/writer.cc index 263689a648d0..33a21053b100 100644 --- a/cpp/src/arrow/ipc/writer.cc +++ b/cpp/src/arrow/ipc/writer.cc @@ -146,6 +146,11 @@ class RecordBatchSerializer { Status VisitArray(const Array& arr) { static std::shared_ptr kNullBuffer = std::make_shared(nullptr, 0); + const Array* physical_arr = &arr; + if (arr.type_id() == Type::EXTENSION) { + physical_arr = checked_cast(arr).storage().get(); + } + if (max_recursion_depth_ <= 0) { return Status::Invalid("Max recursion depth reached"); } @@ -154,7 +159,8 @@ class RecordBatchSerializer { 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"); } @@ -164,10 +170,11 @@ 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)) { + if (internal::HasValidityBitmap(physical_arr->type_id(), options_.metadata_version)) { if (arr.null_count() > 0) { std::shared_ptr bitmap; - RETURN_NOT_OK(GetTruncatedBitmap(arr.offset(), arr.length(), arr.null_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 { @@ -175,7 +182,7 @@ class RecordBatchSerializer { out_->body_buffers.emplace_back(kNullBuffer); } } - return VisitType(arr); + return VisitType(*physical_arr); } // Override this for writing dictionary metadata From 4c4ecb93903a52c95aef436f93f1694daeb050f2 Mon Sep 17 00:00:00 2001 From: Alberto Maschietto Date: Tue, 25 Aug 2026 18:17:47 +0400 Subject: [PATCH 2/2] =?UTF-8?q?GH-50623:=20address=20review=20=E2=80=94=20?= =?UTF-8?q?move=20test=20helpers=20to=20arrow/testing=20and=20arrow/ipc/te?= =?UTF-8?q?st=5Fcommon?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - UnionExtensionArray / UnionExtensionType now live in arrow/testing/extension_type.h (implementations in testing/gtest_util.cc, alongside the other example extension types), with dense_union_extension_type() / sparse_union_extension_type() accessors. - MakeDenseUnionExtension / MakeSparseUnionExtension moved to arrow/ipc/test_common, next to MakeUuid / MakeComplex128 / MakeDictExtension. - Added both to kBatchCases instead of keeping four dedicated TEST_F cases. - writer.cc: unwrap once into a physical_arr reference and use it consistently. --- cpp/src/arrow/ipc/read_write_test.cc | 87 +------------------------- cpp/src/arrow/ipc/test_common.cc | 29 +++++++++ cpp/src/arrow/ipc/test_common.h | 6 ++ cpp/src/arrow/ipc/writer.cc | 35 ++++++----- cpp/src/arrow/testing/extension_type.h | 38 +++++++++++ cpp/src/arrow/testing/gtest_util.cc | 35 +++++++++++ 6 files changed, 131 insertions(+), 99 deletions(-) diff --git a/cpp/src/arrow/ipc/read_write_test.cc b/cpp/src/arrow/ipc/read_write_test.cc index 3cbc52db6244..cfa6a541d5cc 100644 --- a/cpp/src/arrow/ipc/read_write_test.cc +++ b/cpp/src/arrow/ipc/read_write_test.cc @@ -71,73 +71,6 @@ using MetadataVector = std::vector>; namespace test { -class UnionExtensionArray : public ExtensionArray { - public: - using ExtensionArray::ExtensionArray; -}; - -class UnionExtensionType : public ExtensionType { - public: - UnionExtensionType(std::shared_ptr 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 other.extension_name() == extension_name_; - } - - std::shared_ptr MakeArray(std::shared_ptr data) const override { - return std::make_shared(std::move(data)); - } - - Result> Deserialize( - std::shared_ptr storage_type, - const std::string& serialized) const override { - if (serialized != extension_name_) { - return Status::Invalid("Type identifier did not match"); - } - return std::make_shared(std::move(storage_type), - extension_name_); - } - - std::string Serialize() const override { return extension_name_; } - - private: - std::string extension_name_; -}; - -std::shared_ptr dense_union_extension_type() { - return std::make_shared( - dense_union({field("floats", float64()), field("strings", large_utf8())}, {0, 1}), - "dense-union-extension"); -} - -std::shared_ptr sparse_union_extension_type() { - return std::make_shared( - sparse_union({field("floats", float64()), field("strings", large_utf8())}, {0, 1}), - "sparse-union-extension"); -} - -Status MakeDenseUnionExtension(std::shared_ptr* out) { - auto type = dense_union_extension_type(); - auto storage_type = checked_cast(*type).storage_type(); - auto storage = ArrayFromJSON(storage_type, R"([[0, 1.5], [1, "abc"]])"); - auto array = ExtensionType::WrapArray(type, storage); - *out = RecordBatch::Make(schema({field("f0", type)}), array->length(), {array}); - return Status::OK(); -} - -Status MakeSparseUnionExtension(std::shared_ptr* out) { - auto type = sparse_union_extension_type(); - auto storage_type = checked_cast(*type).storage_type(); - auto storage = ArrayFromJSON(storage_type, R"([[0, 1.5], [1, "abc"]])"); - auto array = ExtensionType::WrapArray(type, storage); - *out = RecordBatch::Make(schema({field("f0", type)}), array->length(), {array}); - return Status::OK(); -} - const std::vector kMetadataVersions = {MetadataVersion::V4, MetadataVersion::V5}; @@ -464,7 +397,9 @@ const std::vector kBatchCases = { &MakeIntervals, &MakeUuid, &MakeComplex128, - &MakeDictExtension}; + &MakeDictExtension, + &MakeDenseUnionExtension, + &MakeSparseUnionExtension}; static int g_file_number = 0; @@ -2023,22 +1958,6 @@ TEST_P(TestFileFormatGeneratorCoalesced, RoundTrip) { TEST_P(TestStreamFormat, RoundTrip) { TestRoundTripWithOptions(*GetParam()); } -TEST_F(TestFileFormat, DenseUnionExtensionRoundTrip) { - TestRoundTrip(MakeDenseUnionExtension, IpcWriteOptions::Defaults()); -} - -TEST_F(TestFileFormat, SparseUnionExtensionRoundTrip) { - TestRoundTrip(MakeSparseUnionExtension, IpcWriteOptions::Defaults()); -} - -TEST_F(TestStreamFormat, DenseUnionExtensionRoundTrip) { - TestRoundTrip(MakeDenseUnionExtension, IpcWriteOptions::Defaults()); -} - -TEST_F(TestStreamFormat, SparseUnionExtensionRoundTrip) { - TestRoundTrip(MakeSparseUnionExtension, IpcWriteOptions::Defaults()); -} - TEST_P(TestStreamDecoderData, RoundTrip) { TestRoundTripWithOptions(*GetParam()); } TEST_P(TestStreamDecoderBuffer, RoundTrip) { TestRoundTripWithOptions(*GetParam()); } diff --git a/cpp/src/arrow/ipc/test_common.cc b/cpp/src/arrow/ipc/test_common.cc index ceca6d9e4340..2b9cf1154219 100644 --- a/cpp/src/arrow/ipc/test_common.cc +++ b/cpp/src/arrow/ipc/test_common.cc @@ -1164,6 +1164,35 @@ Status MakeDictExtension(std::shared_ptr* out) { namespace { +Status MakeUnionExtension(const std::shared_ptr& type, + std::shared_ptr* out) { + auto storage_type = checked_cast(*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* out) { + return MakeUnionExtension(dense_union_extension_type(), out); +} + +Status MakeSparseUnionExtension(std::shared_ptr* out) { + return MakeUnionExtension(sparse_union_extension_type(), out); +} + +namespace { + template void FillRandomData(CValueType* data, size_t n, CValueType min, CValueType max, SeedType seed) { diff --git a/cpp/src/arrow/ipc/test_common.h b/cpp/src/arrow/ipc/test_common.h index 6044ef207bc7..dd0abc1ff495 100644 --- a/cpp/src/arrow/ipc/test_common.h +++ b/cpp/src/arrow/ipc/test_common.h @@ -185,6 +185,12 @@ Status MakeComplex128(std::shared_ptr* out); ARROW_TESTING_EXPORT Status MakeDictExtension(std::shared_ptr* out); +ARROW_TESTING_EXPORT +Status MakeDenseUnionExtension(std::shared_ptr* out); + +ARROW_TESTING_EXPORT +Status MakeSparseUnionExtension(std::shared_ptr* out); + ARROW_TESTING_EXPORT Status MakeRandomTensor(const std::shared_ptr& type, const std::vector& shape, bool row_major_p, diff --git a/cpp/src/arrow/ipc/writer.cc b/cpp/src/arrow/ipc/writer.cc index 33a21053b100..855d1a2f2116 100644 --- a/cpp/src/arrow/ipc/writer.cc +++ b/cpp/src/arrow/ipc/writer.cc @@ -146,43 +146,48 @@ class RecordBatchSerializer { Status VisitArray(const Array& arr) { static std::shared_ptr kNullBuffer = std::make_shared(nullptr, 0); - const Array* physical_arr = &arr; - if (arr.type_id() == Type::EXTENSION) { - physical_arr = checked_cast(arr).storage().get(); - } - if (max_recursion_depth_ <= 0) { return Status::Invalid("Max recursion depth reached"); } - if (!options_.allow_64bit && arr.length() > std::numeric_limits::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(arr).storage() + : arr; + + if (!options_.allow_64bit && + physical_arr.length() > std::numeric_limits::max()) { return Status::CapacityError("Cannot write arrays larger than 2^31 - 1 in length"); } - if (physical_arr->offset() != 0 && - physical_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(physical_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 bitmap; - RETURN_NOT_OK(GetTruncatedBitmap(physical_arr->offset(), physical_arr->length(), - physical_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(*physical_arr); + return VisitType(physical_arr); } // Override this for writing dictionary metadata diff --git a/cpp/src/arrow/testing/extension_type.h b/cpp/src/arrow/testing/extension_type.h index 9b4492a543a3..e5d6b597c5be 100644 --- a/cpp/src/arrow/testing/extension_type.h +++ b/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 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 MakeArray(std::shared_ptr data) const override; + + Result> Deserialize( + std::shared_ptr 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 binary_view_extension_type(); ARROW_TESTING_EXPORT std::shared_ptr complex128(); +ARROW_TESTING_EXPORT +std::shared_ptr dense_union_extension_type(); + +ARROW_TESTING_EXPORT +std::shared_ptr sparse_union_extension_type(); + ARROW_TESTING_EXPORT std::shared_ptr ExampleUuid(); diff --git a/cpp/src/arrow/testing/gtest_util.cc b/cpp/src/arrow/testing/gtest_util.cc index 3ea7d9bb22ec..daadfe9c2cdc 100644 --- a/cpp/src/arrow/testing/gtest_util.cc +++ b/cpp/src/arrow/testing/gtest_util.cc @@ -977,6 +977,29 @@ Result> BinaryViewExtensionType::Deserialize( return std::make_shared(); } +bool UnionExtensionType::ExtensionEquals(const ExtensionType& other) const { + return (other.extension_name() == this->extension_name()); +} + +std::shared_ptr UnionExtensionType::MakeArray( + std::shared_ptr data) const { + DCHECK_EQ(data->type->id(), Type::EXTENSION); + DCHECK(ExtensionEquals(checked_cast(*data->type))); + return std::make_shared(data); +} + +Result> UnionExtensionType::Deserialize( + std::shared_ptr 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(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 dict_extension_type() { std::shared_ptr complex128() { return std::make_shared(); } +std::shared_ptr dense_union_extension_type() { + return std::make_shared( + dense_union({field("floats", float64()), field("strings", large_utf8())}, {0, 1}), + "dense-union-extension"); +} + +std::shared_ptr sparse_union_extension_type() { + return std::make_shared( + sparse_union({field("floats", float64()), field("strings", large_utf8())}, {0, 1}), + "sparse-union-extension"); +} + std::shared_ptr MakeComplex128(const std::shared_ptr& real, const std::shared_ptr& imag) { auto type = complex128();