Repository: apache/arrow
File: cpp/src/arrow/ipc/read_write_test.cc (C++)

Code before:
  std::shared_ptr<RecordBatch> batch;
  for (int i = 0; i < num_batches; i++) {
    ARROW_EXPECT_OK(MakeBooleanInt32Int64Batch(num_rows, &batch));
    ARROW_EXPECT_OK(writer->WriteRecordBatch(*batch));
  }

  ARROW_EXPECT_OK(writer->Close());
  EXPECT_OK_AND_ASSIGN(auto buffer, sink->Finish());
  return buffer;
}

void GetReadRecordBatchReadRanges(
    uint32_t num_rows, const std::vector<int>& included_fields,
    const std::vector<int64_t>& expected_body_read_lengths) {
  auto buffer = MakeBooleanInt32Int64File(num_rows, /*num_batches=*/1);

  io::BufferReader buffer_reader(buffer);
  std::unique_ptr<io::TrackedRandomAccessFile> tracked =
      io::TrackedRandomAccessFile::Make(&buffer_reader);

  auto read_options = IpcReadOptions::Defaults();
  // if empty, return all fields
  read_options.included_fields = included_fields;
  ASSERT_OK_AND_ASSIGN(auto reader,
                       RecordBatchFileReader::Open(tracked.get(), read_options));
  ASSERT_OK_AND_ASSIGN(auto out_batch, reader->ReadRecordBatch(0));

  ASSERT_EQ(out_batch->num_rows(), num_rows);
  ASSERT_EQ(out_batch->num_columns(),
            included_fields.empty() ? 3 : included_fields.size());

  auto read_ranges = tracked->get_read_ranges();

  const int32_t magic_size = static_cast<int>(strlen(ipc::internal::kArrowMagicBytes));
  // read magic and footer length IO
  auto file_end_size = magic_size + sizeof(int32_t);
  auto footer_length_offset = buffer->size() - file_end_size;
  auto footer_length = bit_util::FromLittleEndian(
      util::SafeLoadAs<int32_t>(buffer->data() + footer_length_offset));

  // read magic and footer length IO
  EXPECT_EQ(read_ranges[0].length, file_end_size);
  // read footer IO
  EXPECT_EQ(read_ranges[1].length, footer_length);

  // there are 3 read IOs before reading body:
  // 1) read magic and footer length IO
  // 2) read footer IO
  // 3) read record batch metadata IO
  if (included_fields.empty()) {
    EXPECT_EQ(read_ranges.size(), 3);

Review comment (bug):
Ok, can we keep a less strict assertion to check that we can at least access the two first elements without segfaulting?
Something like:
```c++
  // there are at least 2 read IOs before reading body:
  // 1) read magic and footer length IO
  // 2) read footer IO
  EXPECT_GE(read_ranges.size(), 2);
```

Code after:
Status MakeBooleanInt32Int64Batch(const int length, std::shared_ptr<RecordBatch>* out) {
  auto schema_ = MakeBooleanInt32Int64Schema();
  std::shared_ptr<Array> a0, a1, a2;
  RETURN_NOT_OK(MakeRandomBooleanArray(length, false, &a0));
  RETURN_NOT_OK(MakeRandomInt32Array(length, false, arrow::default_memory_pool(), &a1));
  RETURN_NOT_OK(MakeRandomInt64Array(length, false, arrow::default_memory_pool(), &a2));
  *out = RecordBatch::Make(std::move(schema_), length, {a0, a1, a2});
  return Status::OK();
}

std::shared_ptr<Buffer> MakeBooleanInt32Int64File(int num_rows, int num_batches) {
  auto schema_ = MakeBooleanInt32Int64Schema();
  EXPECT_OK_AND_ASSIGN(auto sink, io::BufferOutputStream::Create(0));
  EXPECT_OK_AND_ASSIGN(auto writer, MakeFileWriter(sink.get(), schema_));

  std::shared_ptr<RecordBatch> batch;
  for (int i = 0; i < num_batches; i++) {
    ARROW_EXPECT_OK(MakeBooleanInt32Int64Batch(num_rows, &batch));
    ARROW_EXPECT_OK(writer->WriteRecordBatch(*batch));
  }

  ARROW_EXPECT_OK(writer->Close());
  EXPECT_OK_AND_ASSIGN(auto buffer, sink->Finish());
  return buffer;
}

void GetReadRecordBatchReadRanges(
    uint32_t num_rows, const std::vector<int>& included_fields,
    const std::vector<int64_t>& expected_body_read_lengths) {
  auto buffer = MakeBooleanInt32Int64File(num_rows, /*num_batches=*/1);

  io::BufferReader buffer_reader(buffer);
  std::unique_ptr<io::TrackedRandomAccessFile> tracked =
      io::TrackedRandomAccessFile::Make(&buffer_reader);

  auto read_options = IpcReadOptions::Defaults();
  // if empty, return all fields
  read_options.included_fields = included_fields;
  ASSERT_OK_AND_ASSIGN(auto reader,
                       RecordBatchFileReader::Open(tracked.get(), read_options));
  ASSERT_OK_AND_ASSIGN(auto out_batch, reader->ReadRecordBatch(0));

  ASSERT_EQ(out_batch->num_rows(), num_rows);
  ASSERT_EQ(out_batch->num_columns(),
            included_fields.empty() ? 3 : included_fields.size());

  auto read_ranges = tracked->get_read_ranges();

  const int32_t magic_size = static_cast<int>(strlen(ipc::internal::kArrowMagicBytes));
  // read magic and footer length IO
  auto file_end_size = magic_size + sizeof(int32_t);