Repository: apache/arrow
File: cpp/src/parquet/column_io_benchmark.cc (C++)

Code before:
                                                     ->disable_dictionary()
                                                     ->build();

  auto metadata = ColumnChunkMetaDataBuilder::Make(
      properties, schema.get(), reinterpret_cast<uint8_t*>(&thrift_metadata));

  auto stream = CreateOutputStream();
  std::shared_ptr<Int64Writer> writer = BuildWriter(
      state.range(0), stream, metadata.get(), schema.get(), properties.get(), codec);
  writer->WriteBatch(int64_values.length(), definition_levels.data(),
                     repetition_levels.data(), int64_values.raw_values());
  writer->Close();

  PARQUET_ASSIGN_OR_THROW(auto src, stream->Finish());
  int64_t stream_size = src->size();
  int64_t data_size = int64_values.length() * sizeof(int64_t);

  std::vector<int64_t> values_out(state.range(1));
  std::vector<int16_t> definition_levels_out(state.range(1));
  std::vector<int16_t> repetition_levels_out(state.range(1));
  while (state.KeepRunning()) {
    std::shared_ptr<Int64Reader> reader =
        BuildReader(src, state.range(1), codec, schema.get());
    int64_t values_read = 0;
    for (int64_t i = 0; i < int64_values.length(); i += values_read) {
      reader->ReadBatch(values_out.size(), definition_levels_out.data(),
                        repetition_levels_out.data(), values_out.data(), &values_read);
    }
  }
  SetBytesProcessed(state, repetition);
  state.counters["compression_ratio"] = static_cast<double>(data_size) / stream_size;
}

template <Repetition::type repetition,
          Compression::type codec = Compression::UNCOMPRESSED,
          Encoding::type encoding = Encoding::PLAIN>
static void BM_ReadInt64Column(::benchmark::State& state) {
  BM_ReadInt64Column(state, repetition, codec, encoding);
}

void ReadColumnSetArgs(::benchmark::internal::Benchmark* bench) {
  // Small column, tiny reads
  bench->Args({1024, 16});
  // Small column, full read
  bench->Args({1024, 1024});
  // Midsize column, midsize reads
  bench->Args({65536, 1024});
}

BENCHMARK_TEMPLATE(BM_ReadInt64Column, Repetition::REQUIRED)->Apply(ReadColumnSetArgs);


Review comment (performance):
This is the fix: we were instantiating the ColumnReader with the read batch size (`state.range(1)`) instead of the total number of values (`state.range(0)`).

Code after:
                                                     .compression(codec)
                                                     ->encoding(encoding)
                                                     ->disable_dictionary()
                                                     ->build();

  format::ColumnChunk thrift_metadata;
  auto metadata = ColumnChunkMetaDataBuilder::Make(
      properties, schema.get(), reinterpret_cast<uint8_t*>(&thrift_metadata));

  auto stream = CreateOutputStream();
  std::shared_ptr<Int64Writer> writer = BuildWriter(
      kNumValues, stream, metadata.get(), schema.get(), properties.get(), codec);
  writer->WriteBatch(int64_values.length(), definition_levels.data(),
                     repetition_levels.data(), int64_values.raw_values());
  writer->Close();

  PARQUET_ASSIGN_OR_THROW(auto src, stream->Finish());
  int64_t stream_size = src->size();
  int64_t data_size = int64_values.length() * sizeof(int64_t);

  std::vector<int64_t> values_out(kBatchSize);
  std::vector<int16_t> definition_levels_out(kBatchSize);
  std::vector<int16_t> repetition_levels_out(kBatchSize);
  while (state.KeepRunning()) {
    std::shared_ptr<Int64Reader> reader =
        BuildReader(src, kNumValues, codec, schema.get());
    int64_t values_read = 0;
    for (int64_t i = 0; i < kNumValues; i += values_read) {
      reader->ReadBatch(kBatchSize, definition_levels_out.data(),
                        repetition_levels_out.data(), values_out.data(), &values_read);
      ARROW_CHECK_NE(values_read, 0) << "Unexpected end of column";
    }
  }
  SetBytesProcessed(state, repetition);
  state.counters["compression_ratio"] = static_cast<double>(data_size) / stream_size;
}

template <Repetition::type repetition,
          Compression::type codec = Compression::UNCOMPRESSED,
          Encoding::type encoding = Encoding::PLAIN>
static void BM_ReadInt64Column(::benchmark::State& state) {
  BM_ReadInt64Column(state, repetition, codec, encoding);
}

void ReadColumnSetArgs(::benchmark::internal::Benchmark* bench) {
  // Small column, tiny reads
  bench->Args({1024, 16});
  // Small column, full read
  bench->Args({1024, 1024});
  // Midsize column, midsize reads
  bench->Args({65536, 1024});