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

Code before:
  } else {
    EXPECT_EQ(read_ranges.size(), 3 + expected_body_read_lengths.size());

    // read record batch metadata.  The exact size is tricky to determine but it doesn't
    // matter for this test and it should be smaller than the footer.
    EXPECT_LE(read_ranges[2].length, footer_length);
    for (uint32_t i = 0; i < expected_body_read_lengths.size(); i++) {
      EXPECT_EQ(read_ranges[3 + i].length, expected_body_read_lengths[i]);
    }
  }
}

void GetReadRecordBatchReadRanges(
    const std::vector<int>& included_fields,
    const std::vector<int64_t>& expected_body_read_lengths) {
  return GetReadRecordBatchReadRanges(5, included_fields, expected_body_read_lengths);
}

TEST(TestRecordBatchFileReaderIo, LoadAllFieldsShouldReadTheEntireBody) {
  // read the entire record batch body in single read
  // the batch has 5 * bool + 5 * int32 + 5 * int64
  // ==>
  // + 5 bool:  5 bits      (aligned to  8 bytes)
  // + 5 int32: 5 * 4 bytes (aligned to 24 bytes)
  // + 5 int64: 5 * 8 bytes (aligned to 40 bytes)
  GetReadRecordBatchReadRanges({}, {8 + 24 + 40});
}

TEST(TestRecordBatchFileReaderIo, ReadSingleFieldAtTheStart) {
  // read only the bool field
  // + 5 bool:  5 bits (1 byte)
  GetReadRecordBatchReadRanges({0}, {1});
}

TEST(TestRecordBatchFileReaderIo, ReadSingleFieldInTheMiddle) {
  // read only the int32 field
  // + 5 int32: 5 * 4 bytes
  GetReadRecordBatchReadRanges({1}, {20});
}

TEST(TestRecordBatchFileReaderIo, ReadSingleFieldInTheEnd) {
  // read only the int64 field
  // + 5 int64: 5 * 8 bytes
  GetReadRecordBatchReadRanges({2}, {40});
}

TEST(TestRecordBatchFileReaderIo, SkipTheFieldInTheMiddle) {
  // read the bool field and the int64 field
  // two IOs for body are expected, first for reading bool and the second for reading
  // int64
  // + 5 bool:  5 bits (1 byte)

Review comment (suggestion):
Add a small comment explaining why? Also can take inspiration from the other code path and add:
```
  EXPECT_LE(read_ranges[2].length, total_body + footer_length);
```

Code after:
  // 1) read magic and footer length IO
  // 2) footer IO
  EXPECT_GE(read_ranges.size(), 2);

  // 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);

  if (included_fields.empty()) {
    // When no fields are explicitly included, the reader optimizes by
    // reading metadata and the entire body in a single IO.
    // Thus, there are exactly 3 read IOs in total:
    // 1) magic and footer length
    // 2) footer
    // 3) record batch metadata + body
    EXPECT_EQ(read_ranges.size(), 3);

    int64_t total_body = 0;
    for (auto len : expected_body_read_lengths) total_body += len;

    // In the optimized path (included_fields is empty), the 3rd read operation
    // fetches both the message metadata (flatbuffer) and the entire message body
    // in one contiguous block. Therefore, its length must at least exceed the
    // total body length by the size of the metadata.
    EXPECT_GT(read_ranges[2].length, total_body);
    EXPECT_LE(read_ranges[2].length, total_body + footer_length);
  } else {
    // When fields are filtered, we see 3 initial reads followed by N body reads
    // (one for each field/buffer range):
    // 1) magic and footer length
    // 2) footer
    // 3) record batch metadata
    // 4) individual body buffer reads
    EXPECT_EQ(read_ranges.size(), 3 + expected_body_read_lengths.size());

    // read record batch metadata.  The exact size is tricky to determine but it doesn't
    // matter for this test and it should be smaller than the footer.
    EXPECT_LE(read_ranges[2].length, footer_length);
    for (uint32_t i = 0; i < expected_body_read_lengths.size(); i++) {
      EXPECT_EQ(read_ranges[3 + i].length, expected_body_read_lengths[i]);
    }
  }
}

void GetReadRecordBatchReadRanges(
    const std::vector<int>& included_fields,
    const std::vector<int64_t>& expected_body_read_lengths) {
  return GetReadRecordBatchReadRanges(5, included_fields, expected_body_read_lengths);
}