public:

void SetUp() {

version_ = GetParam();

fb_version_ = internal::MetadataVersionToFlatbuffer(version_);

options_ = IpcWriteOptions::Defaults();

options_.metadata_version = version_;

}

protected:

MetadataVersion version_;

flatbuf::MetadataVersion fb_version_;

IpcWriteOptions options_;

std::string metadata = "foo";

std::string body = "bar";

auto b1 = std::make_shared<Buffer>(metadata);

auto b2 = std::make_shared<Buffer>(metadata);

auto b3 = std::make_shared<Buffer>(body);

auto b4 = std::make_shared<Buffer>(body);

Message msg1(b1, b3);

Message msg2(b2, b4);

Message msg3(b1, nullptr);

Message msg4(b2, nullptr);

ASSERT_TRUE(msg1.Equals(msg2));

ASSERT_TRUE(msg3.Equals(msg4));

ASSERT_FALSE(msg1.Equals(msg3));

ASSERT_FALSE(msg3.Equals(msg1));

// same metadata as msg1, different body

Message msg5(b2, b1);

ASSERT_FALSE(msg1.Equals(msg5));

ASSERT_FALSE(msg5.Equals(msg1));

const int64_t body_length = 64;

flatbuffers::FlatBufferBuilder fbb;

fbb.Finish(flatbuf::CreateMessage(fbb, fb_version_, flatbuf::MessageHeader::RecordBatch,

0 /* header */, body_length));

std::shared_ptr<Buffer> metadata;

ASSERT_OK_AND_ASSIGN(metadata, internal::WriteFlatbufferBuilder(fbb));

std::string body = "abcdef";

ASSERT_OK_AND_ASSIGN(std::unique_ptr<Message> message,

Message::Open(metadata, std::make_shared<Buffer>(body)));

auto CheckWithAlignment = [&](int32_t alignment) {

options_.alignment = alignment;

const int32_t prefix_size = 8;

int64_t output_length = 0;

ASSERT_OK_AND_ASSIGN(auto stream, io::BufferOutputStream::Create(1 << 10));

ASSERT_OK(message->SerializeTo(stream.get(), options_, &output_length));

ASSERT_EQ(bit_util::RoundUp(metadata->size() + prefix_size, alignment) + body_length,

output_length);

ASSERT_OK_AND_EQ(output_length, stream->Tell());

ASSERT_OK_AND_ASSIGN(auto buffer, stream->Finish());

// chech whether length is written in little endian

auto buffer_ptr = buffer.get()->data();

ASSERT_EQ(output_length - body_length - prefix_size,

bit_util::FromLittleEndian(*(uint32_t*)(buffer_ptr + 4)));

};

CheckWithAlignment(8);

CheckWithAlignment(64);

std::vector<std::shared_ptr<KeyValueMetadata>> cases = {

nullptr, key_value_metadata({}, {}),

key_value_metadata({"foo", "bar"}, {"fizz", "buzz"})};

for (auto metadata : cases) {

std::shared_ptr<Buffer> serialized;

ASSERT_OK(internal::WriteRecordBatchMessage(

/*length=*/0, /*body_length=*/0, metadata,

/*nodes=*/{},

/*buffers=*/{}, options_, &serialized));

ASSERT_OK_AND_ASSIGN(std::unique_ptr<Message> message,

Message::Open(serialized, /*body=*/nullptr));

if (metadata) {

ASSERT_TRUE(message->custom_metadata()->Equals(*metadata));

} else {

ASSERT_EQ(nullptr, message->custom_metadata());

}

}

ASSERT_GT(left.size(), 0);

ASSERT_GT(right.size(), 0);

ASSERT_TRUE(left.Equals(right, std::min(left.size(), right.size())));

std::shared_ptr<RecordBatch> batch;

ASSERT_OK(MakeIntBatchSized(36, &batch));

auto RoundtripWithOptions = [&](std::shared_ptr<Buffer>* out_serialized,

std::unique_ptr<Message>* out) {

IpcPayload payload;

ASSERT_OK(GetRecordBatchPayload(*batch, options_, &payload));

ASSERT_OK_AND_ASSIGN(auto stream, io::BufferOutputStream::Create(1 << 20));

int32_t metadata_length = -1;

ASSERT_OK(WriteIpcPayload(payload, options_, stream.get(), &metadata_length));

ASSERT_OK_AND_ASSIGN(*out_serialized, stream->Finish());

io::BufferReader io_reader(*out_serialized);

ASSERT_OK(ReadMessage(&io_reader).Value(out));

};

std::shared_ptr<Buffer> serialized, legacy_serialized;

std::unique_ptr<Message> message, legacy_message;

RoundtripWithOptions(&serialized, &message);

// First 4 bytes 0xFFFFFFFF Continuation marker

ASSERT_EQ(-1, util::SafeLoadAs<int32_t>(serialized->data()));

options_.write_legacy_ipc_format = true;

RoundtripWithOptions(&legacy_serialized, &legacy_message);

// Check that the continuation marker is not written

ASSERT_NE(-1, util::SafeLoadAs<int32_t>(legacy_serialized->data()));

// Have to use the smaller size to exclude padding

BuffersOverlapEquals(*legacy_message->metadata(), *message->metadata());

ASSERT_TRUE(legacy_message->body()->Equals(*message->body()));

std::string metadata = "invalid";

std::string body = "abcdef";

Message message(std::make_shared<Buffer>(metadata), std::make_shared<Buffer>(body));

ASSERT_FALSE(message.Verify());

public:

void SetUp() {}

void CheckSchemaRoundtrip(const Schema& schema) {

ASSERT_OK_AND_ASSIGN(std::shared_ptr<Buffer> buffer, SerializeSchema(schema));

io::BufferReader reader(buffer);

DictionaryMemo in_memo;

ASSERT_OK_AND_ASSIGN(auto actual_schema, ReadSchema(&reader, &in_memo));

AssertSchemaEqual(schema, *actual_schema);

}

auto f0 = field("f0", std::make_shared<Int8Type>());

auto f1 = field("f1", std::make_shared<Int16Type>(), false);

auto f2 = field("f2", std::make_shared<Int32Type>());

auto f3 = field("f3", std::make_shared<Int64Type>());

auto f4 = field("f4", std::make_shared<UInt8Type>());

auto f5 = field("f5", std::make_shared<UInt16Type>());

auto f6 = field("f6", std::make_shared<UInt32Type>());

auto f7 = field("f7", std::make_shared<UInt64Type>());

auto f8 = field("f8", std::make_shared<FloatType>());

auto f9 = field("f9", std::make_shared<DoubleType>(), false);

auto f10 = field("f10", std::make_shared<BooleanType>());

Schema schema({f0, f1, f2, f3, f4, f5, f6, f7, f8, f9, f10});

CheckSchemaRoundtrip(schema);

auto type = list(int32());

auto f0 = field("f0", type);

std::shared_ptr<StructType> type2(

new StructType({field("k1", INT32), field("k2", INT32), field("k3", INT32)}));

auto f1 = field("f1", type2);

Schema schema({f0, f1});

CheckSchemaRoundtrip(schema);

{

auto dict_type = dictionary(int8(), int32(), true /* ordered */);

auto f0 = field("f0", dict_type);

auto f1 = field("f1", list(dict_type));

Schema schema({f0, f1});

CheckSchemaRoundtrip(schema);

}

{

auto dict_type = dictionary(int8(), list(int32()));

auto f0 = field("f0", dict_type);

Schema schema({f0});

CheckSchemaRoundtrip(schema);

}

{

auto inner_dict_type = dictionary(int8(), int32(), /*ordered=*/true);

auto dict_type = dictionary(int16(), list(inner_dict_type));

Schema schema({field("f0", dict_type)});

CheckSchemaRoundtrip(schema);

}

{

auto dict_type1 = dictionary(int8(), utf8(), /*ordered=*/true);

auto dict_type2 = dictionary(int32(), fixed_size_binary(24));

auto dict_type3 = dictionary(int32(), binary());

auto dict_type4 = dictionary(int8(), decimal(19, 7));

auto struct_type1 = struct_({field("s1", dict_type1), field("s2", dict_type2)});

auto struct_type2 = struct_({field("s3", dict_type3), field("s4", dict_type4)});

Schema schema({field("f1", dictionary(int32(), struct_type1)),

field("f2", dictionary(int32(), struct_type2))});

CheckSchemaRoundtrip(schema);

}

auto field_metadata = key_value_metadata({{"key", "value"}});

auto schema_metadata = key_value_metadata({{"foo", "bar"}, {"bizz", "buzz"}});

auto f0 = field("f0", std::make_shared<Int8Type>());

auto f1 = field("f1", std::make_shared<Int16Type>(), false, field_metadata);

Schema schema({f0, f1}, schema_metadata);

CheckSchemaRoundtrip(schema);

// ARROW-9399

std::string root;

ASSERT_OK(GetTestResourceRoot(&root));

// schema_v6.arrow with currently non-existent MetadataVersion::V6

std::stringstream schema_v6_path;

schema_v6_path << root << "/forward-compatibility/schema_v6.arrow";

ASSERT_OK_AND_ASSIGN(auto schema_v6_file, io::ReadableFile::Open(schema_v6_path.str()));

DictionaryMemo placeholder_memo;

ASSERT_RAISES(Invalid, ReadSchema(schema_v6_file.get(), &placeholder_memo));

public:

// Register the extension types required to ensure roundtripping

ExtensionTypesMixin() : ext_guard_({uuid(), dict_extension_type(), complex128()}) {}

protected:

ExtensionTypeGuard ext_guard_;

public:

void SetUp() {

options_ = IpcWriteOptions::Defaults();

ASSERT_OK_AND_ASSIGN(temp_dir_, TemporaryDir::Make("ipc-test-"));

}

std::string TempFile(std::string_view file) {

return temp_dir_->path().Join(std::string(file)).ValueOrDie().ToString();

}

void DoSchemaRoundTrip(const Schema& schema, std::shared_ptr<Schema>* result) {

ASSERT_OK_AND_ASSIGN(std::shared_ptr<Buffer> serialized_schema,

SerializeSchema(schema, options_.memory_pool));

DictionaryMemo in_memo;

io::BufferReader buf_reader(serialized_schema);

ASSERT_OK_AND_ASSIGN(*result, ReadSchema(&buf_reader, &in_memo));

}

Result<std::shared_ptr<RecordBatch>> DoStandardRoundTrip(

const RecordBatch& batch, const IpcWriteOptions& options,

DictionaryMemo* dictionary_memo,

const IpcReadOptions& read_options = IpcReadOptions::Defaults()) {

ARROW_ASSIGN_OR_RAISE(std::shared_ptr<Buffer> serialized_batch,

SerializeRecordBatch(batch, options));

io::BufferReader buf_reader(serialized_batch);

return ReadRecordBatch(batch.schema(), dictionary_memo, read_options, &buf_reader);

}

Result<std::shared_ptr<RecordBatch>> DoLargeRoundTrip(const RecordBatch& batch,

bool zero_data) {

if (zero_data) {

RETURN_NOT_OK(ZeroMemoryMap(mmap_.get()));

}

RETURN_NOT_OK(mmap_->Seek(0));

auto options = options_;

options.allow_64bit = true;

ARROW_ASSIGN_OR_RAISE(auto file_writer,

MakeFileWriter(mmap_, batch.schema(), options));

RETURN_NOT_OK(file_writer->WriteRecordBatch(batch));

RETURN_NOT_OK(file_writer->Close());

ARROW_ASSIGN_OR_RAISE(int64_t offset, mmap_->Tell());

std::shared_ptr<RecordBatchFileReader> file_reader;

ARROW_ASSIGN_OR_RAISE(file_reader, RecordBatchFileReader::Open(mmap_.get(), offset));

return file_reader->ReadRecordBatch(0);

}

void CheckReadResult(const RecordBatch& result, const RecordBatch& expected) {

ASSERT_OK(result.ValidateFull());

EXPECT_EQ(expected.num_rows(), result.num_rows());

ASSERT_TRUE(expected.schema()->Equals(*result.schema()));

ASSERT_EQ(expected.num_columns(), result.num_columns())

<< expected.schema()->ToString() << " result: " << result.schema()->ToString();

CompareBatchColumnsDetailed(result, expected);

}

void CheckRoundtrip(const RecordBatch& batch,

IpcWriteOptions options = IpcWriteOptions::Defaults(),

IpcReadOptions read_options = IpcReadOptions::Defaults(),

int64_t buffer_size = 1 << 20) {

std::stringstream ss;

ss << "test-write-row-batch-" << g_file_number++;

ASSERT_OK_AND_ASSIGN(

mmap_, io::MemoryMapFixture::InitMemoryMap(buffer_size, TempFile(ss.str())));

std::shared_ptr<Schema> schema_result;

DoSchemaRoundTrip(*batch.schema(), &schema_result);

ASSERT_TRUE(batch.schema()->Equals(*schema_result));

DictionaryMemo dictionary_memo;

ASSERT_OK(::arrow::ipc::internal::CollectDictionaries(batch, &dictionary_memo));

ASSERT_OK_AND_ASSIGN(

auto result, DoStandardRoundTrip(batch, options, &dictionary_memo, read_options));

CheckReadResult(*result, batch);

ASSERT_OK_AND_ASSIGN(result, DoLargeRoundTrip(batch, /*zero_data=*/true));

CheckReadResult(*result, batch);

}

void CheckRoundtrip(const std::shared_ptr<Array>& array,

IpcWriteOptions options = IpcWriteOptions::Defaults(),

int64_t buffer_size = 1 << 20) {

auto f0 = arrow::field("f0", array->type());

std::vector<std::shared_ptr<Field>> fields = {f0};

auto schema = std::make_shared<Schema>(fields);

auto batch = RecordBatch::Make(schema, array->length(), {array});

CheckRoundtrip(*batch, options, IpcReadOptions::Defaults(), buffer_size);

}

protected:

std::shared_ptr<io::MemoryMappedFile> mmap_;

IpcWriteOptions options_;

std::unique_ptr<TemporaryDir> temp_dir_;

public:

void SetUp() override { IpcTestFixture::SetUp(); }

void TearDown() override { IpcTestFixture::TearDown(); }

public IpcTestFixture {

public:

void SetUp() override { IpcTestFixture::SetUp(); }

void TearDown() override { IpcTestFixture::TearDown(); }

void TestMetadataVersion(MetadataVersion expected_version) {

std::shared_ptr<RecordBatch> batch;

ASSERT_OK(MakeIntRecordBatch(&batch));

mmap_.reset(); // Ditch previous mmap view, to avoid errors on Windows

ASSERT_OK_AND_ASSIGN(mmap_,

io::MemoryMapFixture::InitMemoryMap(1 << 16, "test-metadata"));

int32_t metadata_length;

int64_t body_length;

const int64_t buffer_offset = 0;

ASSERT_OK(WriteRecordBatch(*batch, buffer_offset, mmap_.get(), &metadata_length,

&body_length, options_));

ASSERT_OK_AND_ASSIGN(std::unique_ptr<Message> message,

ReadMessage(0, metadata_length, mmap_.get()));

ASSERT_EQ(expected_version, message->metadata_version());

}

std::shared_ptr<RecordBatch> batch;

ASSERT_OK((*GetParam())(&batch)); // NOLINT clang-tidy gtest issue

for (const auto version : kMetadataVersions) {

options_.metadata_version = version;

CheckRoundtrip(*batch);

}

options_.metadata_version = MetadataVersion::V4;

TestMetadataVersion(MetadataVersion::V4);

options_.metadata_version = MetadataVersion::V5;

TestMetadataVersion(MetadataVersion::V5);

std::string data = "abc";

io::BufferReader reader(data);

ASSERT_RAISES(Invalid, ReadMessage(&reader));

// But no error on unsignaled EOS

io::BufferReader reader2("");

ASSERT_OK_AND_ASSIGN(auto message, ReadMessage(&reader2));

ASSERT_EQ(nullptr, message);

ASSERT_EQ(MetadataVersion::V1,

ipc::internal::GetMetadataVersion(flatbuf::MetadataVersion::V1));

ASSERT_EQ(MetadataVersion::V2,

ipc::internal::GetMetadataVersion(flatbuf::MetadataVersion::V2));

ASSERT_EQ(MetadataVersion::V3,

ipc::internal::GetMetadataVersion(flatbuf::MetadataVersion::V3));

ASSERT_EQ(MetadataVersion::V4,

ipc::internal::GetMetadataVersion(flatbuf::MetadataVersion::V4));

ASSERT_EQ(MetadataVersion::V5,

ipc::internal::GetMetadataVersion(flatbuf::MetadataVersion::V5));

ASSERT_EQ(MetadataVersion::V1,

ipc::internal::GetMetadataVersion(flatbuf::MetadataVersion::MIN));

ASSERT_EQ(MetadataVersion::V5,

ipc::internal::GetMetadataVersion(flatbuf::MetadataVersion::MAX));

std::shared_ptr<RecordBatch> batch;

ASSERT_OK((*GetParam())(&batch)); // NOLINT clang-tidy gtest issue

// Skip the zero-length case

if (batch->num_rows() < 2) {

return;

}

auto sliced_batch = batch->Slice(2, 10);

CheckRoundtrip(*sliced_batch);

std::shared_ptr<RecordBatch> batch;

ASSERT_OK((*GetParam())(&batch)); // NOLINT clang-tidy gtest issue

std::shared_ptr<RecordBatch> zero_length_batch;

if (batch->num_rows() > 2) {

zero_length_batch = batch->Slice(2, 0);

} else {

zero_length_batch = batch->Slice(0, 0);

}

CheckRoundtrip(*zero_length_batch);

// ARROW-544: check binary array

ASSERT_OK_AND_ASSIGN(auto value_offsets,

AllocateBuffer(sizeof(int32_t), options_.memory_pool));

*reinterpret_cast<int32_t*>(value_offsets->mutable_data()) = 0;

std::shared_ptr<Array> bin_array = std::make_shared<BinaryArray>(

0, std::move(value_offsets), std::make_shared<Buffer>(nullptr, 0),

std::make_shared<Buffer>(nullptr, 0));

// null value_offsets

std::shared_ptr<Array> bin_array2 = std::make_shared<BinaryArray>(0, nullptr, nullptr);

CheckRoundtrip(bin_array);

CheckRoundtrip(bin_array2);

auto storage_type = struct_({

field("i", int32()),

field("f", float32()),

field("s", utf8()),

});

auto storage = checked_pointer_cast<StructArray>(ArrayFromJSON(storage_type,

R"([

])"));

ASSERT_OK_AND_ASSIGN(

auto m01, StructArray::Make({storage->field(0), storage->field(1)},

{storage_type->field(0), storage_type->field(1)}));

ASSERT_OK_AND_ASSIGN(

auto m12, StructArray::Make({storage->field(1), storage->field(2)},

{storage_type->field(1), storage_type->field(2)}));

ASSERT_OK_AND_ASSIGN(

auto m20, StructArray::Make({storage->field(2), storage->field(0)},

{storage_type->field(2), storage_type->field(0)}));

auto ids = ArrayFromJSON(int8(), "[1, 12, 20, 1]");

ASSERT_OK_AND_ASSIGN(

auto sparse,

SparseUnionArray::Make(*ids, {m01, m12, m20}, {"m01", "m12", "m20"}, {01, 12, 20}));

auto expected = ArrayFromJSON(sparse_union(

{

field("m01", m01->type()),

field("m12", m12->type()),

field("m20", m20->type()),

},

{01, 12, 20}),

R"([

])");

AssertArraysEqual(*expected, *sparse, /*verbose=*/true);

DictionaryMemo ignored;

ASSERT_OK_AND_ASSIGN(

auto roundtripped_batch,

DoStandardRoundTrip(*RecordBatch::Make(schema({field("", sparse->type())}),

sparse->length(), {sparse}),

IpcWriteOptions::Defaults(), &ignored));

auto roundtripped =

checked_pointer_cast<SparseUnionArray>(roundtripped_batch->column(0));

AssertArraysEqual(*expected, *roundtripped, /*verbose=*/true);

auto roundtripped_m01 = checked_pointer_cast<StructArray>(roundtripped->field(0));

auto roundtripped_m12 = checked_pointer_cast<StructArray>(roundtripped->field(1));

auto roundtripped_m20 = checked_pointer_cast<StructArray>(roundtripped->field(2));

// The IPC writer does not take advantage of reusable buffers

ASSERT_NE(roundtripped_m01->field(0)->data()->buffers,

roundtripped_m20->field(1)->data()->buffers);

ASSERT_NE(roundtripped_m01->field(1)->data()->buffers,

roundtripped_m12->field(0)->data()->buffers);

ASSERT_NE(roundtripped_m12->field(1)->data()->buffers,

roundtripped_m20->field(0)->data()->buffers);

random::RandomArrayGenerator rg(/*seed=*/0);

// Generate both regular and dictionary encoded because the dictionary batch

// gets compressed also

int64_t length = 500;

int dict_size = 50;

std::shared_ptr<Array> dict = rg.String(dict_size, /*min_length=*/5, /*max_length=*/5,

/*null_probability=*/0);

std::shared_ptr<Array> indices = rg.Int32(length, /*min=*/0, /*max=*/dict_size - 1,

/*null_probability=*/0.1);

auto dict_type = dictionary(int32(), utf8());

auto dict_field = field("f1", dict_type);

ASSERT_OK_AND_ASSIGN(auto dict_array,

DictionaryArray::FromArrays(dict_type, indices, dict));

auto schema = ::arrow::schema({field("f0", utf8()), dict_field});

auto batch =

RecordBatch::Make(schema, length, {rg.String(500, 0, 10, 0.1), dict_array});

std::vector<Compression::type> codecs = {Compression::LZ4_FRAME, Compression::ZSTD};

for (auto codec : codecs) {

if (!util::Codec::IsAvailable(codec)) {

continue;

}

IpcWriteOptions write_options = IpcWriteOptions::Defaults();

ASSERT_OK_AND_ASSIGN(write_options.codec, util::Codec::Create(codec));

CheckRoundtrip(*batch, write_options);

// Check non-parallel read and write

IpcReadOptions read_options = IpcReadOptions::Defaults();

write_options.use_threads = false;

read_options.use_threads = false;

CheckRoundtrip(*batch, write_options, read_options);

}

std::vector<Compression::type> disallowed_codecs = {

Compression::BROTLI, Compression::BZ2, Compression::LZ4, Compression::GZIP,

Compression::SNAPPY};

for (auto codec : disallowed_codecs) {

if (!util::Codec::IsAvailable(codec)) {

continue;

}

IpcWriteOptions write_options = IpcWriteOptions::Defaults();

ASSERT_OK_AND_ASSIGN(write_options.codec, util::Codec::Create(codec));

ASSERT_RAISES(Invalid, SerializeRecordBatch(*batch, write_options));

}

// ARROW-6046

std::shared_ptr<Array> array;

ASSERT_OK(MakeRandomStringArray(500, false, default_memory_pool(), &array));

auto f0 = field("f0", array->type());

auto schema = ::arrow::schema({f0});

auto batch = RecordBatch::Make(schema, array->length(), {array});

auto sliced_batch = batch->Slice(0, 5);

ASSERT_OK_AND_ASSIGN(

mmap_, io::MemoryMapFixture::InitMemoryMap(/*buffer_size=*/1 << 20,

TempFile("test-truncate-offsets")));

DictionaryMemo dictionary_memo;

ASSERT_OK_AND_ASSIGN(

auto result,

DoStandardRoundTrip(*sliced_batch, IpcWriteOptions::Defaults(), &dictionary_memo));

ASSERT_EQ(6 * sizeof(int32_t), result->column(0)->data()->buffers[1]->size());

auto CheckArray = [this](const std::shared_ptr<Array>& array) {

auto f0 = field("f0", array->type());

auto schema = ::arrow::schema({f0});

auto batch = RecordBatch::Make(schema, array->length(), {array});

auto sliced_batch = batch->Slice(0, 5);

int64_t full_size;

int64_t sliced_size;

ASSERT_OK(GetRecordBatchSize(*batch, &full_size));

ASSERT_OK(GetRecordBatchSize(*sliced_batch, &sliced_size));

ASSERT_TRUE(sliced_size < full_size) << sliced_size << " " << full_size;

// make sure we can write and read it

this->CheckRoundtrip(*sliced_batch);

};

std::shared_ptr<Array> a0, a1;

auto pool = default_memory_pool();

// Integer

ASSERT_OK(MakeRandomInt32Array(500, false, pool, &a0));

CheckArray(a0);

// String / Binary

{

auto s = MakeRandomStringArray(500, false, pool, &a0);

ASSERT_TRUE(s.ok());

}

CheckArray(a0);

// Boolean

ASSERT_OK(MakeRandomBooleanArray(10000, false, &a0));

CheckArray(a0);

// List

ASSERT_OK(MakeRandomInt32Array(500, false, pool, &a0));

ASSERT_OK(MakeRandomListArray(a0, 200, false, pool, &a1));

CheckArray(a1);

// Struct

auto struct_type = struct_({field("f0", a0->type())});

std::vector<std::shared_ptr<Array>> struct_children = {a0};

a1 = std::make_shared<StructArray>(struct_type, a0->length(), struct_children);

CheckArray(a1);

// Sparse Union

auto union_type = sparse_union({field("f0", a0->type())}, {0});

std::vector<int32_t> type_ids(a0->length());

std::shared_ptr<Buffer> ids_buffer;

ASSERT_OK(CopyBufferFromVector(type_ids, default_memory_pool(), &ids_buffer));

a1 = std::make_shared<SparseUnionArray>(union_type, a0->length(), struct_children,

ids_buffer);

CheckArray(a1);

// Dense union

auto dense_union_type = dense_union({field("f0", a0->type())}, {0});

std::vector<int32_t> type_offsets;

for (int32_t i = 0; i < a0->length(); ++i) {

type_offsets.push_back(i);

}

std::shared_ptr<Buffer> offsets_buffer;

ASSERT_OK(CopyBufferFromVector(type_offsets, default_memory_pool(), &offsets_buffer));

a1 = std::make_shared<DenseUnionArray>(dense_union_type, a0->length(), struct_children,

ids_buffer, offsets_buffer);

CheckArray(a1);

// ARROW-7975

random::RandomArrayGenerator rg(/*seed=*/0);

int64_t length = 15;

auto arr = rg.String(length, 0, 10, 0.1);

auto batch = RecordBatch::Make(::arrow::schema({field("f0", utf8())}), length, {arr});

ASSERT_OK_AND_ASSIGN(

mmap_, io::MemoryMapFixture::InitMemoryMap(

/*buffer_size=*/1 << 20, TempFile("test-roundtrip-buffer-sizes")));

DictionaryMemo dictionary_memo;

ASSERT_OK_AND_ASSIGN(

auto result,

DoStandardRoundTrip(*batch, IpcWriteOptions::Defaults(), &dictionary_memo));

// Make sure that the validity bitmap is size 2 as expected

ASSERT_EQ(2, arr->data()->buffers[0]->size());

for (size_t i = 0; i < arr->data()->buffers.size(); ++i) {

ASSERT_EQ(arr->data()->buffers[i]->size(),

result->column(0)->data()->buffers[i]->size());

}

std::shared_ptr<RecordBatch> batch) {

io::MockOutputStream mock;

ipc::IpcPayload payload;

int32_t mock_metadata_length = -1;

int64_t mock_body_length = -1;

int64_t size = -1;

ASSERT_OK(WriteRecordBatch(*batch, 0, &mock, &mock_metadata_length, &mock_body_length,

options));

ASSERT_OK(GetRecordBatchPayload(*batch, options, &payload));

int64_t payload_size = GetPayloadSize(payload, options);

ASSERT_OK(GetRecordBatchSize(*batch, options, &size));

ASSERT_EQ(mock.GetExtentBytesWritten(), size);

ASSERT_EQ(mock.GetExtentBytesWritten(), payload_size);

std::shared_ptr<RecordBatch> batch;

ASSERT_OK(MakeIntRecordBatch(&batch));

TestGetRecordBatchSize(options_, batch);

ASSERT_OK(MakeListRecordBatch(&batch));

TestGetRecordBatchSize(options_, batch);

ASSERT_OK(MakeZeroLengthRecordBatch(&batch));

TestGetRecordBatchSize(options_, batch);

ASSERT_OK(MakeNonNullRecordBatch(&batch));

TestGetRecordBatchSize(options_, batch);

ASSERT_OK(MakeDeeplyNestedList(&batch));

TestGetRecordBatchSize(options_, batch);

public:

void SetUp() {

pool_ = default_memory_pool();

ASSERT_OK_AND_ASSIGN(temp_dir_, TemporaryDir::Make("ipc-recursion-limits-test-"));

}

std::string TempFile(std::string_view file) {

return temp_dir_->path().Join(std::string(file)).ValueOrDie().ToString();

}

void TearDown() { io::MemoryMapFixture::TearDown(); }

Status WriteToMmap(int recursion_level, bool override_level, int32_t* metadata_length,

int64_t* body_length, std::shared_ptr<RecordBatch>* batch,

std::shared_ptr<Schema>* schema) {

const int batch_length = 5;

auto type = int32();

std::shared_ptr<Array> array;

const bool include_nulls = true;

RETURN_NOT_OK(MakeRandomInt32Array(1000, include_nulls, pool_, &array));

for (int i = 0; i < recursion_level; ++i) {

type = list(type);

RETURN_NOT_OK(

MakeRandomListArray(array, batch_length, include_nulls, pool_, &array));

}

auto f0 = field("f0", type);

*schema = ::arrow::schema({f0});

*batch = RecordBatch::Make(*schema, batch_length, {array});

std::stringstream ss;

ss << "test-write-past-max-recursion-" << g_file_number++;

const int memory_map_size = 1 << 20;

ARROW_ASSIGN_OR_RAISE(

mmap_, io::MemoryMapFixture::InitMemoryMap(memory_map_size, TempFile(ss.str())));

auto options = IpcWriteOptions::Defaults();

if (override_level) {

options.max_recursion_depth = recursion_level + 1;

}

return WriteRecordBatch(**batch, 0, mmap_.get(), metadata_length, body_length,

options);

}

protected:

std::shared_ptr<io::MemoryMappedFile> mmap_;

std::unique_ptr<TemporaryDir> temp_dir_;

MemoryPool* pool_;

int32_t metadata_length = -1;

int64_t body_length = -1;

std::shared_ptr<Schema> schema;

std::shared_ptr<RecordBatch> batch;

ASSERT_RAISES(Invalid, WriteToMmap((1 << 8) + 1, false, &metadata_length, &body_length,

&batch, &schema));

int32_t metadata_length = -1;

int64_t body_length = -1;

std::shared_ptr<Schema> schema;

const int recursion_depth = 64;

std::shared_ptr<RecordBatch> batch;

ASSERT_OK(WriteToMmap(recursion_depth, true, &metadata_length, &body_length, &batch,

&schema));

ASSERT_OK_AND_ASSIGN(std::unique_ptr<Message> message,

ReadMessage(0, metadata_length, mmap_.get()));

io::BufferReader reader(message->body());

DictionaryMemo empty_memo;

ASSERT_RAISES(Invalid, ReadRecordBatch(*message->metadata(), schema, &empty_memo,

IpcReadOptions::Defaults(), &reader));

auto CheckDepth = [this](int recursion_depth, bool* it_works) {

int32_t metadata_length = -1;

int64_t body_length = -1;

std::shared_ptr<Schema> schema;

std::shared_ptr<RecordBatch> batch;

ASSERT_OK(WriteToMmap(recursion_depth, true, &metadata_length, &body_length, &batch,

&schema));

ASSERT_OK_AND_ASSIGN(std::unique_ptr<Message> message,

ReadMessage(0, metadata_length, mmap_.get()));

DictionaryMemo empty_memo;

auto options = IpcReadOptions::Defaults();

options.max_recursion_depth = recursion_depth + 1;

io::BufferReader reader(message->body());

std::shared_ptr<RecordBatch> result;

ASSERT_OK_AND_ASSIGN(result, ReadRecordBatch(*message->metadata(), schema,

&empty_memo, options, &reader));

*it_works = result->Equals(*batch);

};

bool it_works = false;

CheckDepth(100, &it_works);

ASSERT_TRUE(it_works);

#if !defined(ARROW_VALGRIND)