diff --git a/cpp/src/arrow/ipc/read_write_test.cc b/cpp/src/arrow/ipc/read_write_test.cc index aa3f9fd77f23..cfa6a541d5cc 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 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 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 263689a648d0..855d1a2f2116 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::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 (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 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 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();