Function TestWriteOffsetIndex

cpp/src/parquet/page_index_test.cc:422–470

}

void TestWriteOffsetIndex(bool write_size_stats) {
  /// Create offset index via the OffsetIndexBuilder interface.
  auto builder = OffsetIndexBuilder::Make();
  const size_t num_pages = 5;
  const std::vector<int64_t> offsets = {100, 200, 300, 400, 500};
  const std::vector<int32_t> page_sizes = {1024, 2048, 3072, 4096, 8192};
  const std::vector<int64_t> first_row_indices = {0, 10000, 20000, 30000, 40000};
  const std::vector<int64_t> unencoded_byte_array_lengths = {1111, 2222, 0, 3333, 4444};
  for (size_t i = 0; i < num_pages; ++i) {
    auto unencoded_byte_array_length =
        write_size_stats ? std::make_optional(unencoded_byte_array_lengths[i])
                         : std::nullopt;
    builder->AddPage(offsets[i], page_sizes[i], first_row_indices[i],
                     unencoded_byte_array_length);
  }
  const int64_t final_position = 4096;
  builder->Finish(final_position);

  std::vector<std::unique_ptr<OffsetIndex>> offset_indexes;
  /// 1st element is the offset index just built.
  offset_indexes.emplace_back(builder->Build());
  /// 2nd element is the offset index restored by serialize-then-deserialize round trip.
  auto sink = CreateOutputStream();
  builder->WriteTo(sink.get());
  PARQUET_ASSIGN_OR_THROW(auto buffer, sink->Finish());
  offset_indexes.emplace_back(OffsetIndex::Make(buffer->data(),
                                                static_cast<uint32_t>(buffer->size()),
                                                default_reader_properties()));

  /// Verify the data of the offset index.
  for (const auto& offset_index : offset_indexes) {
    ASSERT_EQ(num_pages, offset_index->page_locations().size());
    if (write_size_stats) {
      ASSERT_EQ(num_pages, offset_index->unencoded_byte_array_data_bytes().size());
    } else {
      ASSERT_TRUE(offset_index->unencoded_byte_array_data_bytes().empty());
    }
    for (size_t i = 0; i < num_pages; ++i) {
      const auto& page_location = offset_index->page_locations().at(i);
      ASSERT_EQ(offsets[i] + final_position, page_location.offset);
      ASSERT_EQ(page_sizes[i], page_location.compressed_page_size);
      ASSERT_EQ(first_row_indices[i], page_location.first_row_index);
      if (write_size_stats) {
        ASSERT_EQ(unencoded_byte_array_lengths[i],
                  offset_index->unencoded_byte_array_data_bytes()[i]);
      }
    }
  }
}

TEST(PageIndex, WriteOffsetIndexWithoutSizeStats) {
  TestWriteOffsetIndex(/*write_size_stats=*/false);