This is an automated email from the ASF dual-hosted git repository.
pitrou pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new 3e6f730bd0d GH-50623: [C++][IPC] Fix extension-wrapped union IPC
roundtrip (#50927)
3e6f730bd0d is described below
commit 3e6f730bd0d04d972d60d4efee12b6be9c53cf25
Author: alb3e3 <[email protected]>
AuthorDate: Tue Aug 25 19:30:25 2026 +0400
GH-50623: [C++][IPC] Fix extension-wrapped union IPC roundtrip (#50927)
### 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: #50623
Authored-by: Alberto Maschietto <[email protected]>
Signed-off-by: Antoine Pitrou <[email protected]>
---
cpp/src/arrow/ipc/read_write_test.cc | 8 +++++--
cpp/src/arrow/ipc/test_common.cc | 29 ++++++++++++++++++++++++++
cpp/src/arrow/ipc/test_common.h | 6 ++++++
cpp/src/arrow/ipc/writer.cc | 28 ++++++++++++++++++-------
cpp/src/arrow/testing/extension_type.h | 38 ++++++++++++++++++++++++++++++++++
cpp/src/arrow/testing/gtest_util.cc | 35 +++++++++++++++++++++++++++++++
6 files changed, 134 insertions(+), 10 deletions(-)
diff --git a/cpp/src/arrow/ipc/read_write_test.cc
b/cpp/src/arrow/ipc/read_write_test.cc
index aa3f9fd77f2..cfa6a541d5c 100644
--- a/cpp/src/arrow/ipc/read_write_test.cc
+++ b/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_;
diff --git a/cpp/src/arrow/ipc/test_common.cc b/cpp/src/arrow/ipc/test_common.cc
index ceca6d9e434..2b9cf115421 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<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) {
diff --git a/cpp/src/arrow/ipc/test_common.h b/cpp/src/arrow/ipc/test_common.h
index 6044ef207bc..dd0abc1ff49 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<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,
diff --git a/cpp/src/arrow/ipc/writer.cc b/cpp/src/arrow/ipc/writer.cc
index 263689a648d..855d1a2f211 100644
--- a/cpp/src/arrow/ipc/writer.cc
+++ b/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
diff --git a/cpp/src/arrow/testing/extension_type.h
b/cpp/src/arrow/testing/extension_type.h
index 9b4492a543a..e5d6b597c5b 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<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();
diff --git a/cpp/src/arrow/testing/gtest_util.cc
b/cpp/src/arrow/testing/gtest_util.cc
index 3ea7d9bb22e..daadfe9c2cd 100644
--- a/cpp/src/arrow/testing/gtest_util.cc
+++ b/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();
