Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
67 changes: 67 additions & 0 deletions cpp/src/arrow/ipc/read_write_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down Expand Up @@ -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()),
Expand Down
4 changes: 3 additions & 1 deletion cpp/src/arrow/ipc/writer.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down