| /* | |
| * 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_suite.h" | |
| #include <atomic> | |
| #include <cstdint> | |
| #include <iostream> | |
| #include <map> | |
| #include <memory> | |
| #include <mutex> | |
| #include <string> | |
| #include <tuple> | |
| #include <utility> | |
| #include <vector> | |
| #include "arrow/api.h" | |
| #include "arrow/c/bridge.h" | |
| #include "arrow/io/api.h" | |
| #include "benchmark/benchmark_helpers.h" | |
| #include "benchmark/cli_option_parsing.h" | |
| #include "paimon/api.h" | |
| #include "paimon/catalog/catalog.h" | |
| #include "paimon/testing/utils/testharness.h" | |
| #if __has_include("parquet/arrow/reader.h") | |
| #include "parquet/arrow/reader.h" | |
| #include "parquet/file_reader.h" | |
| #define PAIMON_BENCHMARK_HAS_PARQUET_READER 1 | |
| #else | |
| #define PAIMON_BENCHMARK_HAS_PARQUET_READER 0 | |
| #endif | |
| namespace paimon::benchmark { | |
| namespace { | |
| constexpr int64_t kSourceBatchMaxRows = 4096; | |
| constexpr int32_t kRowToBatchThreadNumber = 3; | |
| struct BenchmarkCliOptions { | |
| std::string source_data_file; | |
| std::string source_table_path; | |
| std::vector<std::string> pk_columns; | |
| std::vector<std::pair<std::string, std::string>> extra_options; | |
| }; | |
| struct SourceDataSpec { | |
| std::string format; | |
| std::string path; | |
| }; | |
| BenchmarkCliOptions& MutableBenchmarkCliOptions() { | |
| static BenchmarkCliOptions options; | |
| return options; | |
| } | |
| const BenchmarkCliOptions& GetBenchmarkCliOptions() { | |
| return MutableBenchmarkCliOptions(); | |
| } | |
| Status ParsePaimonBenchmarkCliArgsImpl(int32_t* argc, char** argv) { | |
| auto& options = MutableBenchmarkCliOptions(); | |
| options = BenchmarkCliOptions{}; | |
| const int32_t parsed_argc = *argc; | |
| int32_t write_index = 1; | |
| for (int32_t arg_index = 1; arg_index < parsed_argc; ++arg_index) { | |
| const std::string arg(argv[arg_index]); | |
| PAIMON_ASSIGN_OR_RAISE(bool is_parsed, | |
| paimon::benchmark::ParseStringOptionArg( | |
| parsed_argc, argv, arg, "--paimon_source_data_file", &arg_index, | |
| &options.source_data_file)); | |
| if (is_parsed) { | |
| continue; | |
| } | |
| PAIMON_ASSIGN_OR_RAISE(is_parsed, paimon::benchmark::ParseStringOptionArg( | |
| parsed_argc, argv, arg, "--paimon_source_table_path", | |
| &arg_index, &options.source_table_path)); | |
| if (is_parsed) { | |
| continue; | |
| } | |
| PAIMON_ASSIGN_OR_RAISE(is_parsed, paimon::benchmark::ParseCommaSeparatedOptionArg( | |
| parsed_argc, argv, arg, "--paimon_pk_columns", | |
| &arg_index, &options.pk_columns)); | |
| if (is_parsed) { | |
| continue; | |
| } | |
| PAIMON_ASSIGN_OR_RAISE(is_parsed, paimon::benchmark::ParseDelimitedRepeatableOptionArg( | |
| parsed_argc, argv, arg, "--paimon_option", &arg_index, | |
| &options.extra_options)); | |
| if (is_parsed) { | |
| continue; | |
| } | |
| argv[write_index++] = argv[arg_index]; | |
| } | |
| *argc = write_index; | |
| argv[write_index] = nullptr; | |
| return Status::OK(); | |
| } | |
| bool HasHelpFlagImpl(int32_t argc, char** argv) { | |
| for (int32_t arg_index = 1; arg_index < argc; ++arg_index) { | |
| const std::string arg(argv[arg_index]); | |
| if (arg == "-h" || arg == "--help" || arg == "--help=true") { | |
| return true; | |
| } | |
| } | |
| return false; | |
| } | |
| void PrintPaimonBenchmarkCliHelpImpl() { | |
| std::cout | |
| << "Paimon benchmark custom options:\n" | |
| << " --paimon_source_data_file=<path>\n" | |
| << " Required. External source data file used to build benchmark data.\n" | |
| << " Currently supports Parquet source files.\n" | |
| << " Also supports: --paimon_source_data_file <path>\n" | |
| << " --paimon_source_table_path=<path>\n" | |
| << " Optional for BM_Read and BM_MOR_Read. If set, read directly from existing\n" | |
| << " table path and skip source file loading and pre-write stage.\n" | |
| << " Also supports: --paimon_source_table_path <path>\n" | |
| << " --paimon_pk_columns=<col1,col2,...>\n" | |
| << " Required by BM_PK_Write and BM_MOR_Read.\n" | |
| << " Also supports: --paimon_pk_columns <col1,col2,...>\n" | |
| << " --paimon_option=<key1>:<value1>;<key2>:<value2>\n" | |
| << " Optional and repeatable. Pass through table options as-is.\n" | |
| << " Default table file format is parquet; use file.format:<format> to override.\n" | |
| << " Also supports: --paimon_option <key1>:<value1>;<key2>:<value2>\n" | |
| << " Note: use quotes in shell, e.g. \"--paimon_option k1:v1;k2:v2\".\n" | |
| << "\n" | |
| << "Example:\n" | |
| << " paimon-read-write-benchmark --paimon_source_data_file /path/data.parquet \\\n" | |
| << " --paimon_pk_columns=id --paimon_option \"read.batch-size:8192\" \\\n" | |
| << " --benchmark_filter=BM_Read\n" | |
| << std::endl; | |
| } | |
| Result<std::unique_ptr<paimon::test::UniqueTestDirectory>> CreateBenchmarkWorkspace() { | |
| auto workspace = paimon::test::UniqueTestDirectory::Create(); | |
| if (workspace == nullptr) { | |
| return Status::Invalid("failed to create benchmark workspace"); | |
| } | |
| return workspace; | |
| } | |
| uint64_t NextTableId() { | |
| static std::atomic<uint64_t> id{0}; | |
| return ++id; | |
| } | |
| std::string RequirePath(const std::string& root_path, const std::string& db_name, | |
| const std::string& table_name) { | |
| return root_path + "/" + db_name + ".db/" + table_name; | |
| } | |
| template <typename T> | |
| Result<T> AddContext(paimon::Result<T>&& result, const std::string& context) { | |
| if (!result.ok()) { | |
| const Status status = result.status(); | |
| return status.WithMessage(context, ": ", status.message()); | |
| } | |
| return std::move(result).value(); | |
| } | |
| Status AddContext(const paimon::Status& status, const std::string& context) { | |
| if (!status.ok()) { | |
| return status.WithMessage(context, ": ", status.message()); | |
| } | |
| return Status::OK(); | |
| } | |
| void SkipWithMessage(::benchmark::State& state, const std::string& message) { | |
| state.SkipWithError(message); | |
| } | |
| std::string GetConfiguredFileFormat() { | |
| std::string file_format = "parquet"; | |
| for (const auto& kv : GetBenchmarkCliOptions().extra_options) { | |
| if (kv.first == paimon::Options::FILE_FORMAT) { | |
| file_format = kv.second; | |
| } | |
| } | |
| return file_format; | |
| } | |
| bool IsFileFormatSupported(const std::string& format) { | |
| if (format == "parquet") { | |
| return true; | |
| } | |
| if (format == "orc") { | |
| #ifdef PAIMON_ENABLE_ORC | |
| return true; | |
| #else | |
| return false; | |
| #endif | |
| } | |
| return false; | |
| } | |
| void ApplyExtraOptions(std::map<std::string, std::string>* options) { | |
| for (const auto& kv : GetBenchmarkCliOptions().extra_options) { | |
| (*options)[kv.first] = kv.second; | |
| } | |
| } | |
| std::map<std::string, std::string> BuildOptions(const std::string& file_format) { | |
| std::map<std::string, std::string> options = { | |
| {paimon::Options::FILE_FORMAT, file_format}, | |
| }; | |
| ApplyExtraOptions(&options); | |
| return options; | |
| } | |
| std::map<std::string, std::string> BuildPkOptions(const std::string& file_format) { | |
| auto options = BuildOptions(file_format); | |
| options[paimon::Options::BUCKET] = "1"; | |
| options[paimon::Options::MERGE_ENGINE] = "deduplicate"; | |
| return options; | |
| } | |
| std::string GetSourceDataFilePath() { | |
| return GetBenchmarkCliOptions().source_data_file; | |
| } | |
| std::string GetSourceTablePath() { | |
| return GetBenchmarkCliOptions().source_table_path; | |
| } | |
| const std::vector<std::string>& GetPkColumns() { | |
| return GetBenchmarkCliOptions().pk_columns; | |
| } | |
| SourceDataSpec GetSourceDataSpec() { | |
| const std::string source_data_file_path = GetSourceDataFilePath(); | |
| if (!source_data_file_path.empty()) { | |
| return {"parquet", source_data_file_path}; | |
| } | |
| return {"", ""}; | |
| } | |
| int64_t GetSourceBatchMaxRows() { | |
| return kSourceBatchMaxRows; | |
| } | |
| int32_t GetRowToBatchThreadNumber() { | |
| return kRowToBatchThreadNumber; | |
| } | |
| bool SupportsParquetSourceDataMode() { | |
| #if PAIMON_BENCHMARK_HAS_PARQUET_READER | |
| return true; | |
| #else | |
| return false; | |
| #endif | |
| } | |
| bool SupportsSourceDataMode(const std::string& source_format) { | |
| if (source_format == "parquet") { | |
| return SupportsParquetSourceDataMode(); | |
| } | |
| return false; | |
| } | |
| struct SourceDataMetadata { | |
| std::shared_ptr<arrow::Schema> schema; | |
| int64_t total_rows = 0; | |
| std::string format; | |
| std::string path; | |
| }; | |
| #if PAIMON_BENCHMARK_HAS_PARQUET_READER | |
| Result<std::unique_ptr<parquet::arrow::FileReader>> OpenParquetSourceReader( | |
| const std::string& path) { | |
| auto input = arrow::io::ReadableFile::Open(path); | |
| if (!input.ok()) { | |
| return Status::Invalid("open Parquet source failed: ", path, ", ", | |
| input.status().ToString()); | |
| } | |
| std::unique_ptr<parquet::arrow::FileReader> parquet_reader; | |
| const auto open_status = parquet::arrow::OpenFile( | |
| input.ValueUnsafe(), arrow::default_memory_pool(), &parquet_reader); | |
| if (!open_status.ok()) { | |
| return Status::Invalid("create Parquet reader failed: ", open_status.ToString()); | |
| } | |
| parquet_reader->set_batch_size(GetSourceBatchMaxRows()); | |
| return parquet_reader; | |
| } | |
| #endif | |
| Result<SourceDataMetadata> LoadParquetSourceMetadata(const std::string& path) { | |
| #if !PAIMON_BENCHMARK_HAS_PARQUET_READER | |
| return Status::Invalid( | |
| "Parquet source data mode requires parquet::arrow reader support in this build"); | |
| #else | |
| static SourceDataMetadata cache; | |
| static std::mutex cache_mutex; | |
| std::lock_guard<std::mutex> lock(cache_mutex); | |
| if (cache.path == path && cache.format == "parquet") { | |
| return cache; | |
| } | |
| PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<parquet::arrow::FileReader> parquet_reader, | |
| OpenParquetSourceReader(path)); | |
| std::shared_ptr<arrow::Schema> schema; | |
| const auto schema_status = parquet_reader->GetSchema(&schema); | |
| if (!schema_status.ok()) { | |
| return Status::Invalid("read Parquet source schema failed: ", schema_status.ToString()); | |
| } | |
| const int64_t total_rows = parquet_reader->parquet_reader()->metadata()->num_rows(); | |
| if (total_rows <= 0) { | |
| return Status::Invalid("Parquet source is empty: ", path); | |
| } | |
| cache.schema = std::move(schema); | |
| cache.total_rows = total_rows; | |
| cache.format = "parquet"; | |
| cache.path = path; | |
| return cache; | |
| #endif | |
| } | |
| Result<SourceDataMetadata> LoadSourceDataMetadata(const SourceDataSpec& source_spec) { | |
| if (source_spec.format == "parquet") { | |
| return LoadParquetSourceMetadata(source_spec.path); | |
| } | |
| return Status::Invalid("unknown source format: ", source_spec.format); | |
| } | |
| std::shared_ptr<arrow::StructArray> BuildStructArrayFromRecordBatch( | |
| const std::shared_ptr<arrow::RecordBatch>& batch) { | |
| return std::make_shared<arrow::StructArray>(arrow::struct_(batch->schema()->fields()), | |
| batch->num_rows(), batch->columns()); | |
| } | |
| Result<std::unique_ptr<paimon::RecordBatch>> MakeRecordBatch( | |
| const std::shared_ptr<arrow::StructArray>& arr) { | |
| ArrowArray c_array; | |
| if (!arrow::ExportArray(*arr, &c_array).ok()) { | |
| return Status::Invalid("failed to export arrow array"); | |
| } | |
| paimon::RecordBatchBuilder builder(&c_array); | |
| return AddContext(builder.Finish(), "build paimon record batch"); | |
| } | |
| Status EnsureTable(const std::string& root_path, const std::string& db_name, | |
| const std::string& table_name, const std::map<std::string, std::string>& options, | |
| const std::shared_ptr<arrow::Schema>& schema, | |
| const std::vector<std::string>& primary_keys = {}) { | |
| PAIMON_ASSIGN_OR_RAISE( | |
| std::unique_ptr<paimon::Catalog> catalog, | |
| AddContext(paimon::Catalog::Create(root_path, options), "create catalog")); | |
| PAIMON_RETURN_NOT_OK( | |
| AddContext(catalog->CreateDatabase(db_name, options, true), "create database")); | |
| ArrowSchema c_schema; | |
| if (!arrow::ExportSchema(*schema, &c_schema).ok()) { | |
| return Status::Invalid("failed to export table schema"); | |
| } | |
| PAIMON_RETURN_NOT_OK( | |
| AddContext(catalog->CreateTable(paimon::Identifier(db_name, table_name), &c_schema, | |
| /*partition_keys=*/{}, primary_keys, options, | |
| /*ignore_if_exists=*/false), | |
| "create table")); | |
| return Status::OK(); | |
| } | |
| Status WriteSourceDataToWriter(paimon::FileStoreWrite* writer, const SourceDataSpec& source_spec) { | |
| if (source_spec.format != "parquet") { | |
| return Status::Invalid("unknown source format: ", source_spec.format); | |
| } | |
| #if !PAIMON_BENCHMARK_HAS_PARQUET_READER | |
| return Status::Invalid( | |
| "Parquet source data mode requires parquet::arrow reader support in this build"); | |
| #else | |
| PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<parquet::arrow::FileReader> parquet_reader, | |
| OpenParquetSourceReader(source_spec.path)); | |
| std::unique_ptr<arrow::RecordBatchReader> batch_reader; | |
| const auto reader_status = parquet_reader->GetRecordBatchReader(&batch_reader); | |
| if (!reader_status.ok()) { | |
| return Status::Invalid("create Parquet source batch reader failed: ", | |
| reader_status.ToString()); | |
| } | |
| int64_t written_rows = 0; | |
| while (true) { | |
| std::shared_ptr<arrow::RecordBatch> record_batch; | |
| const auto read_status = batch_reader->ReadNext(&record_batch); | |
| if (!read_status.ok()) { | |
| return Status::Invalid("read Parquet source batch failed: ", read_status.ToString()); | |
| } | |
| if (record_batch == nullptr) { | |
| break; | |
| } | |
| if (record_batch->num_rows() <= 0) { | |
| continue; | |
| } | |
| auto struct_array = BuildStructArrayFromRecordBatch(record_batch); | |
| PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<paimon::RecordBatch> batch, | |
| MakeRecordBatch(struct_array)); | |
| PAIMON_RETURN_NOT_OK(AddContext(writer->Write(std::move(batch)), "write batch")); | |
| written_rows += record_batch->num_rows(); | |
| } | |
| if (written_rows <= 0) { | |
| return Status::Invalid("source file has no non-empty data batches: ", source_spec.path); | |
| } | |
| return Status::OK(); | |
| #endif | |
| } | |
| Status WriteAndCommit(const std::string& table_path, | |
| const std::map<std::string, std::string>& options, | |
| const SourceDataSpec& source_spec) { | |
| paimon::WriteContextBuilder write_builder(table_path, "benchmark-writer"); | |
| PAIMON_ASSIGN_OR_RAISE( | |
| std::unique_ptr<paimon::WriteContext> write_ctx, | |
| AddContext(write_builder.SetOptions(options).Finish(), "create write context")); | |
| PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<paimon::FileStoreWrite> writer, | |
| AddContext(paimon::FileStoreWrite::Create(std::move(write_ctx)), | |
| "create file store writer")); | |
| PAIMON_RETURN_NOT_OK(WriteSourceDataToWriter(writer.get(), source_spec)); | |
| PAIMON_ASSIGN_OR_RAISE(std::vector<std::shared_ptr<paimon::CommitMessage>> messages, | |
| AddContext(writer->PrepareCommit(), "prepare commit")); | |
| paimon::CommitContextBuilder commit_builder(table_path, "benchmark-writer"); | |
| PAIMON_ASSIGN_OR_RAISE( | |
| std::unique_ptr<paimon::CommitContext> commit_ctx, | |
| AddContext(commit_builder.SetOptions(options).Finish(), "create commit context")); | |
| PAIMON_ASSIGN_OR_RAISE( | |
| std::unique_ptr<paimon::FileStoreCommit> committer, | |
| AddContext(paimon::FileStoreCommit::Create(std::move(commit_ctx)), "create committer")); | |
| PAIMON_RETURN_NOT_OK(AddContext(committer->Commit(messages), "commit write")); | |
| return Status::OK(); | |
| } | |
| struct SharedReadTableCache { | |
| std::string key; | |
| std::unique_ptr<paimon::test::UniqueTestDirectory> workspace; | |
| std::string table_path; | |
| int64_t total_rows = 0; | |
| }; | |
| struct SharedMorReadTableCache { | |
| std::string key; | |
| std::unique_ptr<paimon::test::UniqueTestDirectory> workspace; | |
| std::string table_path; | |
| int64_t total_rows = 0; | |
| }; | |
| std::string BuildReadTableCacheKey(const std::string& file_format, | |
| const SourceDataSpec& source_spec) { | |
| return file_format + "|" + source_spec.format + "|" + source_spec.path + "|" + | |
| std::to_string(GetSourceBatchMaxRows()); | |
| } | |
| std::string JoinColumns(const std::vector<std::string>& columns) { | |
| std::string joined; | |
| for (size_t i = 0; i < columns.size(); ++i) { | |
| if (i > 0) { | |
| joined.append(","); | |
| } | |
| joined.append(columns[i]); | |
| } | |
| return joined; | |
| } | |
| Result<const SharedMorReadTableCache*> GetOrCreateSharedMorReadTable( | |
| const std::string& file_format, const SourceDataSpec& source_spec) { | |
| static SharedMorReadTableCache cache; | |
| static std::mutex cache_mutex; | |
| const std::vector<std::string>& pk_columns = GetPkColumns(); | |
| const std::string cache_key = | |
| BuildReadTableCacheKey(file_format, source_spec) + "|pk=" + JoinColumns(pk_columns); | |
| std::lock_guard<std::mutex> lock(cache_mutex); | |
| if (cache.workspace != nullptr && cache.key == cache_key) { | |
| std::cout << "[benchmark][mor-read] reuse_output_table_path=" << cache.table_path | |
| << std::endl; | |
| return &cache; | |
| } | |
| auto options = BuildPkOptions(file_format); | |
| PAIMON_ASSIGN_OR_RAISE(const SourceDataMetadata source_metadata, | |
| LoadSourceDataMetadata(source_spec)); | |
| PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<paimon::test::UniqueTestDirectory> workspace, | |
| CreateBenchmarkWorkspace()); | |
| const std::string db_name = "bench_db"; | |
| const std::string table_name = "mor_read_shared_" + std::to_string(NextTableId()); | |
| PAIMON_RETURN_NOT_OK(EnsureTable(workspace->Str(), db_name, table_name, options, | |
| source_metadata.schema, | |
| /*primary_keys=*/pk_columns)); | |
| const std::string table_path = RequirePath(workspace->Str(), db_name, table_name); | |
| std::cout << "[benchmark][mor-read] create_shared_output_table_path=" << table_path | |
| << std::endl; | |
| PAIMON_RETURN_NOT_OK(WriteAndCommit(table_path, options, source_spec)); | |
| cache.key = cache_key; | |
| cache.workspace = std::move(workspace); | |
| cache.table_path = table_path; | |
| cache.total_rows = source_metadata.total_rows; | |
| return &cache; | |
| } | |
| Result<const SharedReadTableCache*> GetOrCreateSharedReadTable(const std::string& file_format, | |
| const SourceDataSpec& source_spec) { | |
| static SharedReadTableCache cache; | |
| static std::mutex cache_mutex; | |
| const std::string cache_key = BuildReadTableCacheKey(file_format, source_spec); | |
| std::lock_guard<std::mutex> lock(cache_mutex); | |
| if (cache.workspace != nullptr && cache.key == cache_key) { | |
| std::cout << "[benchmark][read] reuse_output_table_path=" << cache.table_path << std::endl; | |
| return &cache; | |
| } | |
| auto options = BuildOptions(file_format); | |
| PAIMON_ASSIGN_OR_RAISE(const SourceDataMetadata source_metadata, | |
| LoadSourceDataMetadata(source_spec)); | |
| PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<paimon::test::UniqueTestDirectory> workspace, | |
| CreateBenchmarkWorkspace()); | |
| const std::string db_name = "bench_db"; | |
| const std::string table_name = "read_shared_" + std::to_string(NextTableId()); | |
| PAIMON_RETURN_NOT_OK( | |
| EnsureTable(workspace->Str(), db_name, table_name, options, source_metadata.schema)); | |
| const std::string table_path = RequirePath(workspace->Str(), db_name, table_name); | |
| std::cout << "[benchmark][read] create_shared_output_table_path=" << table_path << std::endl; | |
| PAIMON_RETURN_NOT_OK(WriteAndCommit(table_path, options, source_spec)); | |
| cache.key = cache_key; | |
| cache.workspace = std::move(workspace); | |
| cache.table_path = table_path; | |
| cache.total_rows = source_metadata.total_rows; | |
| return &cache; | |
| } | |
| Result<int64_t> ReadRows(const std::string& table_path, | |
| const std::map<std::string, std::string>& options, | |
| int32_t prefetch_parallel_num) { | |
| paimon::ScanContextBuilder scan_builder(table_path); | |
| PAIMON_ASSIGN_OR_RAISE( | |
| std::unique_ptr<paimon::ScanContext> scan_ctx, | |
| AddContext(scan_builder.SetOptions(options).Finish(), "create scan context")); | |
| PAIMON_ASSIGN_OR_RAISE( | |
| std::unique_ptr<paimon::TableScan> scanner, | |
| AddContext(paimon::TableScan::Create(std::move(scan_ctx)), "create scanner")); | |
| PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<paimon::Plan> plan, | |
| AddContext(scanner->CreatePlan(), "create plan")); | |
| paimon::ReadContextBuilder read_builder(table_path); | |
| constexpr int32_t kPrefetchBatchCount = 600; | |
| read_builder.SetOptions(options) | |
| .EnablePrefetch(true) | |
| .SetPrefetchBatchCount(kPrefetchBatchCount) | |
| .SetPrefetchMaxParallelNum(prefetch_parallel_num) | |
| .EnableMultiThreadRowToBatch(GetRowToBatchThreadNumber() > 1) | |
| .SetRowToBatchThreadNumber(GetRowToBatchThreadNumber()); | |
| PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<paimon::ReadContext> read_ctx, | |
| AddContext(read_builder.Finish(), "create read context")); | |
| PAIMON_ASSIGN_OR_RAISE( | |
| std::unique_ptr<paimon::TableRead> reader, | |
| AddContext(paimon::TableRead::Create(std::move(read_ctx)), "create table reader")); | |
| PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<paimon::BatchReader> batch_reader, | |
| AddContext(reader->CreateReader(plan->Splits()), "create batch reader")); | |
| int64_t total_rows = 0; | |
| while (true) { | |
| PAIMON_ASSIGN_OR_RAISE(paimon::BatchReader::ReadBatch batch, | |
| AddContext(batch_reader->NextBatch(), "read next batch")); | |
| if (paimon::BatchReader::IsEofBatch(batch)) { | |
| break; | |
| } | |
| auto& [array, schema] = batch; | |
| auto imported = arrow::ImportArray(array.get(), schema.get()); | |
| if (!imported.ok()) { | |
| return Status::Invalid("import c data array failed: ", imported.status().ToString()); | |
| } | |
| total_rows += imported.ValueUnsafe()->length(); | |
| } | |
| return total_rows; | |
| } | |
| struct PreparedSourceData { | |
| std::shared_ptr<arrow::Schema> schema; | |
| int64_t total_rows = 0; | |
| }; | |
| bool TryGetSourceSpec(::benchmark::State& state, SourceDataSpec* source_spec) { | |
| (void)state; | |
| *source_spec = GetSourceDataSpec(); | |
| return true; | |
| } | |
| bool TryPrepareSourceData(::benchmark::State& state, const SourceDataSpec& source_spec, | |
| PreparedSourceData* prepared) { | |
| auto source_metadata = LoadSourceDataMetadata(source_spec); | |
| if (!source_metadata.ok()) { | |
| SkipWithMessage(state, source_metadata.status().ToString()); | |
| return false; | |
| } | |
| prepared->schema = source_metadata.value().schema; | |
| prepared->total_rows = source_metadata.value().total_rows; | |
| return true; | |
| } | |
| } // namespace | |
| Status ParsePaimonBenchmarkCliArgs(int* argc, char** argv) { | |
| auto parsed_argc = static_cast<int32_t>(*argc); | |
| PAIMON_RETURN_NOT_OK(ParsePaimonBenchmarkCliArgsImpl(&parsed_argc, argv)); | |
| *argc = static_cast<int>(parsed_argc); | |
| return Status::OK(); | |
| } | |
| bool HasHelpFlag(int32_t argc, char** argv) { | |
| return HasHelpFlagImpl(argc, argv); | |
| } | |
| void PrintPaimonBenchmarkCliHelp() { | |
| PrintPaimonBenchmarkCliHelpImpl(); | |
| } | |
| void RunBMWrite(::benchmark::State& state) { | |
| const std::string file_format = GetConfiguredFileFormat(); | |
| SourceDataSpec source_spec; | |
| if (!TryGetSourceSpec(state, &source_spec)) { | |
| return; | |
| } | |
| if (!BenchmarkHelpers::ValidateSourcePresenceOrSkip( | |
| state, source_spec.path, "--paimon_source_data_file is required", &SkipWithMessage)) { | |
| return; | |
| } | |
| if (!BenchmarkHelpers::ValidateSourceSupportOrSkip(state, source_spec.format, | |
| SupportsSourceDataMode(source_spec.format), | |
| &SkipWithMessage)) { | |
| return; | |
| } | |
| if (!BenchmarkHelpers::ValidateFileFormatOrSkip( | |
| state, file_format, IsFileFormatSupported(file_format), &SkipWithMessage)) { | |
| return; | |
| } | |
| auto options = BuildOptions(file_format); | |
| PreparedSourceData prepared; | |
| if (!TryPrepareSourceData(state, source_spec, &prepared)) { | |
| return; | |
| } | |
| auto workspace = CreateBenchmarkWorkspace(); | |
| if (!workspace.ok()) { | |
| SkipWithMessage(state, workspace.status().ToString()); | |
| return; | |
| } | |
| for (auto _ : state) { | |
| const std::string db_name = "bench_db"; | |
| const std::string table_name = "write_" + std::to_string(NextTableId()); | |
| const Status ensure_status = | |
| EnsureTable(workspace.value()->Str(), db_name, table_name, options, prepared.schema); | |
| if (!ensure_status.ok()) { | |
| SkipWithMessage(state, ensure_status.ToString()); | |
| return; | |
| } | |
| const std::string table_path = RequirePath(workspace.value()->Str(), db_name, table_name); | |
| std::cout << "[benchmark][write] output_table_path=" << table_path << std::endl; | |
| const Status write_status = WriteAndCommit(table_path, options, source_spec); | |
| if (!write_status.ok()) { | |
| SkipWithMessage(state, write_status.ToString()); | |
| return; | |
| } | |
| } | |
| state.SetItemsProcessed(state.iterations() * prepared.total_rows); | |
| } | |
| void RunBMRead(::benchmark::State& state) { | |
| const auto prefetch_parallel_num = static_cast<int32_t>(state.range(0)); | |
| const std::string file_format = GetConfiguredFileFormat(); | |
| const std::string source_table_path = GetSourceTablePath(); | |
| SourceDataSpec source_spec; | |
| if (!TryGetSourceSpec(state, &source_spec)) { | |
| return; | |
| } | |
| if (!BenchmarkHelpers::ValidateFileFormatOrSkip( | |
| state, file_format, IsFileFormatSupported(file_format), &SkipWithMessage)) { | |
| return; | |
| } | |
| if (!BenchmarkHelpers::ValidatePrefetchParallelOrSkip(state, prefetch_parallel_num, | |
| &SkipWithMessage)) { | |
| return; | |
| } | |
| auto options = BuildOptions(file_format); | |
| auto source_table_read_result = BenchmarkHelpers::TryRunSourceTableReadMode( | |
| state, "read", source_table_path, | |
| [&]() { return ReadRows(source_table_path, options, prefetch_parallel_num); }); | |
| if (!source_table_read_result.ok()) { | |
| SkipWithMessage(state, source_table_read_result.status().ToString()); | |
| return; | |
| } | |
| if (source_table_read_result.value()) { | |
| return; | |
| } | |
| if (!BenchmarkHelpers::ValidateSourcePresenceOrSkip( | |
| state, source_spec.path, | |
| "--paimon_source_data_file is required when --paimon_source_table_path is not set", | |
| &SkipWithMessage)) { | |
| return; | |
| } | |
| if (!BenchmarkHelpers::ValidateSourceSupportOrSkip(state, source_spec.format, | |
| SupportsSourceDataMode(source_spec.format), | |
| &SkipWithMessage)) { | |
| return; | |
| } | |
| auto shared_table = GetOrCreateSharedReadTable(file_format, source_spec); | |
| if (!shared_table.ok()) { | |
| SkipWithMessage(state, shared_table.status().ToString()); | |
| return; | |
| } | |
| auto rows_read = BenchmarkHelpers::RunReadIterations(state, [&]() { | |
| return ReadRows(shared_table.value()->table_path, options, prefetch_parallel_num); | |
| }); | |
| if (!rows_read.ok()) { | |
| SkipWithMessage(state, rows_read.status().ToString()); | |
| return; | |
| } | |
| state.SetItemsProcessed(state.iterations() * rows_read.value()); | |
| } | |
| void RunBMPkWrite(::benchmark::State& state) { | |
| const std::string file_format = GetConfiguredFileFormat(); | |
| SourceDataSpec source_spec; | |
| if (!TryGetSourceSpec(state, &source_spec)) { | |
| return; | |
| } | |
| if (!BenchmarkHelpers::ValidateSourcePresenceOrSkip( | |
| state, source_spec.path, "--paimon_source_data_file is required", &SkipWithMessage)) { | |
| return; | |
| } | |
| if (!BenchmarkHelpers::ValidateSourceSupportOrSkip(state, source_spec.format, | |
| SupportsSourceDataMode(source_spec.format), | |
| &SkipWithMessage)) { | |
| return; | |
| } | |
| if (!BenchmarkHelpers::ValidateFileFormatOrSkip( | |
| state, file_format, IsFileFormatSupported(file_format), &SkipWithMessage)) { | |
| return; | |
| } | |
| const std::vector<std::string>& pk_columns = GetPkColumns(); | |
| if (pk_columns.empty()) { | |
| SkipWithMessage(state, "--paimon_pk_columns is required for BM_PK_Write"); | |
| return; | |
| } | |
| auto options = BuildPkOptions(file_format); | |
| PreparedSourceData prepared; | |
| if (!TryPrepareSourceData(state, source_spec, &prepared)) { | |
| return; | |
| } | |
| auto workspace = CreateBenchmarkWorkspace(); | |
| if (!workspace.ok()) { | |
| SkipWithMessage(state, workspace.status().ToString()); | |
| return; | |
| } | |
| for (auto _ : state) { | |
| const std::string db_name = "bench_db"; | |
| const std::string table_name = "pk_write_" + std::to_string(NextTableId()); | |
| const Status ensure_status = | |
| EnsureTable(workspace.value()->Str(), db_name, table_name, options, prepared.schema, | |
| /*primary_keys=*/pk_columns); | |
| if (!ensure_status.ok()) { | |
| SkipWithMessage(state, ensure_status.ToString()); | |
| return; | |
| } | |
| const std::string table_path = RequirePath(workspace.value()->Str(), db_name, table_name); | |
| std::cout << "[benchmark][pk-write] output_table_path=" << table_path << std::endl; | |
| const Status write_status = WriteAndCommit(table_path, options, source_spec); | |
| if (!write_status.ok()) { | |
| SkipWithMessage(state, write_status.ToString()); | |
| return; | |
| } | |
| } | |
| state.SetItemsProcessed(state.iterations() * prepared.total_rows); | |
| } | |
| void RunBMMorRead(::benchmark::State& state) { | |
| const auto prefetch_parallel_num = static_cast<int32_t>(state.range(0)); | |
| const std::string file_format = GetConfiguredFileFormat(); | |
| const std::string source_table_path = GetSourceTablePath(); | |
| SourceDataSpec source_spec; | |
| if (!TryGetSourceSpec(state, &source_spec)) { | |
| return; | |
| } | |
| if (!BenchmarkHelpers::ValidateFileFormatOrSkip( | |
| state, file_format, IsFileFormatSupported(file_format), &SkipWithMessage)) { | |
| return; | |
| } | |
| if (!BenchmarkHelpers::ValidatePrefetchParallelOrSkip(state, prefetch_parallel_num, | |
| &SkipWithMessage)) { | |
| return; | |
| } | |
| const auto source_table_read_options = BuildOptions(file_format); | |
| auto source_table_read_result = | |
| BenchmarkHelpers::TryRunSourceTableReadMode(state, "mor-read", source_table_path, [&]() { | |
| return ReadRows(source_table_path, source_table_read_options, prefetch_parallel_num); | |
| }); | |
| if (!source_table_read_result.ok()) { | |
| SkipWithMessage(state, source_table_read_result.status().ToString()); | |
| return; | |
| } | |
| if (source_table_read_result.value()) { | |
| return; | |
| } | |
| if (!BenchmarkHelpers::ValidateSourcePresenceOrSkip( | |
| state, source_spec.path, | |
| "--paimon_source_data_file is required when --paimon_source_table_path is not set", | |
| &SkipWithMessage)) { | |
| return; | |
| } | |
| if (!BenchmarkHelpers::ValidateSourceSupportOrSkip(state, source_spec.format, | |
| SupportsSourceDataMode(source_spec.format), | |
| &SkipWithMessage)) { | |
| return; | |
| } | |
| if (GetPkColumns().empty()) { | |
| SkipWithMessage(state, "--paimon_pk_columns is required for BM_MOR_Read"); | |
| return; | |
| } | |
| auto options = BuildPkOptions(file_format); | |
| auto shared_table = GetOrCreateSharedMorReadTable(file_format, source_spec); | |
| if (!shared_table.ok()) { | |
| SkipWithMessage(state, shared_table.status().ToString()); | |
| return; | |
| } | |
| auto rows_read = BenchmarkHelpers::RunReadIterations(state, [&]() { | |
| return ReadRows(shared_table.value()->table_path, options, prefetch_parallel_num); | |
| }); | |
| if (!rows_read.ok()) { | |
| SkipWithMessage(state, rows_read.status().ToString()); | |
| return; | |
| } | |
| state.SetItemsProcessed(state.iterations() * rows_read.value()); | |
| } | |
| } // namespace paimon::benchmark |