pitrou commented on code in PR #45085: URL: https://github.com/apache/arrow/pull/45085#discussion_r1907071817
########## cpp/src/parquet/arrow/size_stats_benchmark.cc: ########## @@ -0,0 +1,165 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "benchmark/benchmark.h" + +#include <cstdint> +#include <numeric> + +#include "parquet/arrow/writer.h" +#include "parquet/file_reader.h" +#include "parquet/metadata.h" +#include "parquet/platform.h" +#include "parquet/properties.h" + +#include "arrow/array.h" +#include "arrow/io/buffered.h" +#include "arrow/io/memory.h" +#include "arrow/table.h" +#include "arrow/testing/gtest_util.h" +#include "arrow/testing/random.h" + +namespace parquet::benchmark { + +// This should result in multiple pages for most primitive types +constexpr int64_t kBenchmarkSize = 1024 * 1024; +// Use a skewed null probability to reduce levels encoding overhead +constexpr double kNullProbability = 0.95; + +int64_t GetTotalBytes(const std::shared_ptr<::arrow::ArrayData>& data) { + if (data == nullptr) { + return 0; + } + int64_t total_bytes = + std::accumulate(data->buffers.cbegin(), data->buffers.cend(), int64_t{0}, + [](int64_t acc, const auto& buffer) { + return acc + (buffer != nullptr ? buffer->size() : int64_t{0}); + }); + total_bytes += std::accumulate( + data->child_data.cbegin(), data->child_data.cend(), int64_t{0}, + [](int64_t acc, const auto& child) { return acc + GetTotalBytes(child); }); + total_bytes += GetTotalBytes(data->dictionary); + return total_bytes; +} + +int64_t GetTotalBytes(const std::shared_ptr<::arrow::Table>& table) { + int64_t total_bytes = 0; + for (const auto& column : table->columns()) { + for (const auto& chunk : column->chunks()) { + total_bytes += GetTotalBytes(chunk->data()); + } + } + return total_bytes; +} + +int64_t GetTotalPageIndexSize(const std::shared_ptr<::parquet::FileMetaData>& metadata) { + int64_t total_page_index_size = 0; + for (int i = 0; i < metadata->num_row_groups(); ++i) { + auto row_group = metadata->RowGroup(i); + for (int j = 0; j < row_group->num_columns(); ++j) { + auto column = row_group->ColumnChunk(j); + total_page_index_size += + column->GetColumnIndexLocation().value_or(parquet::IndexLocation{0, 0}).length; + } + } + return total_page_index_size; +} + +void WriteColumn(::benchmark::State& state, const std::shared_ptr<::arrow::Table>& table, + SizeStatisticsLevel stats_level) { + // Use the fastest possible encoding and compression settings, to better exhibit + // the size statistics overhead. + auto properties = WriterProperties::Builder() + .enable_statistics() + ->enable_write_page_index() + ->disable_dictionary() + ->encoding(Encoding::PLAIN) + ->set_size_statistics_level(stats_level) + ->build(); + + for (auto _ : state) { + auto output = parquet::CreateOutputStream(); + ARROW_EXPECT_OK(::parquet::arrow::WriteTable( + *table, ::arrow::default_memory_pool(), + std::static_pointer_cast<::arrow::io::OutputStream>(output), + DEFAULT_MAX_ROW_GROUP_LENGTH, properties)); + + if (state.counters.find("page_index_size") == state.counters.end()) { + state.PauseTiming(); + auto metadata = parquet::ReadMetaData( + std::make_shared<::arrow::io::BufferReader>(output->Finish().ValueOrDie())); + state.counters["output_size"] = static_cast<double>(output->Tell().ValueOrDie()); + state.counters["page_index_size"] = + static_cast<double>(GetTotalPageIndexSize(metadata)); + state.ResumeTiming(); + } + } + + state.SetItemsProcessed(state.iterations() * kBenchmarkSize); + state.SetBytesProcessed(state.iterations() * GetTotalBytes(table)); +} + +template <SizeStatisticsLevel level, typename ArrowType> +void BM_WritePrimitiveColumn(::benchmark::State& state) { + ::arrow::random::RandomArrayGenerator generator(/*seed=*/42); + auto type = std::make_shared<ArrowType>(); + auto array = generator.ArrayOf(type, kBenchmarkSize, kNullProbability); + auto table = ::arrow::Table::Make( + ::arrow::schema({::arrow::field("column", type, kNullProbability > 0)}), {array}); + WriteColumn(state, table, level); +} + +template <SizeStatisticsLevel level, typename ArrowType> +void BM_WriteListColumn(::benchmark::State& state) { + ::arrow::random::RandomArrayGenerator generator(/*seed=*/42); + auto element_type = std::make_shared<ArrowType>(); + auto element_array = generator.ArrayOf(element_type, kBenchmarkSize, kNullProbability); + auto list_type = ::arrow::list(element_type); + auto list_array = generator.List(*element_array, kBenchmarkSize / 10, kNullProbability); Review Comment: Hmm, I think we made a mistake here. This generates a list array with length ~10 and very long lists. Is it desired @wgtmac ? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: github-unsubscr...@arrow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
