From ebf0e885cff4d926ddb1bd07b900ac081dce24a9 Mon Sep 17 00:00:00 2001
From: Ivan Chelyubeev <ijon@ydb.tech>
Date: Mon, 19 Feb 2024 01:31:14 +0300
Subject: [PATCH] ydb workload kv init: add --max-partitions arg

---
 ydb/library/workload/kv_workload.cpp | 20 +++++++++++---------
 ydb/library/workload/kv_workload.h   |  6 +++---
 2 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/ydb/library/workload/kv_workload.cpp b/ydb/library/workload/kv_workload.cpp
index c98857470fac..fc16af69ab41 100644
--- a/ydb/library/workload/kv_workload.cpp
+++ b/ydb/library/workload/kv_workload.cpp
@@ -24,12 +24,12 @@ using TRow = TKvWorkloadGenerator::TRow;
 // Note: there is no mechanism to update row values for now so all keys should be different
 struct TRowsVerifyer {
     TRowsVerifyer(size_t capacity = 10000)
-        : Capacity(capacity) 
+        : Capacity(capacity)
     {
         Rows.reserve(Capacity);
     }
 
-    void TryInsert(TRow row) {        
+    void TryInsert(TRow row) {
         std::unique_lock<std::mutex> lock(Mutex, std::try_to_lock);
         if (!lock.owns_lock()) {
             return;
@@ -43,7 +43,7 @@ struct TRowsVerifyer {
         Rows[RandomNumber<size_t>(Rows.size())] = row;
     }
 
-    std::optional<TRow> GetRandom() {        
+    std::optional<TRow> GetRandom() {
         std::unique_lock<std::mutex> lock(Mutex, std::try_to_lock);
         if (!lock.owns_lock()) {
             return { };
@@ -75,7 +75,7 @@ void AddResultSet(const NYdb::TResultSet& resultSet, TVector<TRow>& rows) {
     NYdb::TResultSetParser parser(resultSet);
     while (parser.TryNextRow()) {
         TRow row;
-        
+
         for (size_t col = 0; col < parser.ColumnsCount(); col++) {
             auto& valueParser = parser.ColumnParser(col);
             if (valueParser.GetPrimitiveType() == NYdb::EPrimitiveType::Uint64) {
@@ -101,7 +101,7 @@ void VerifyRows(const TRow& checkRow, const TVector<TRow>& readRows, TString mes
 
     if (readRows.size() > 1) {
         Cerr << "Expected to have " << checkRow.ToString()
-            << " but got " << readRows.size() << " rows " 
+            << " but got " << readRows.size() << " rows "
             << message
             << Endl;
 
@@ -388,9 +388,9 @@ TQueryInfoList TKvWorkloadGenerator::Mixed() {
             if (Params.MixedDoReadRows) doReadRows = true;
             if (Params.MixedDoSelect) doSelect = true;
         } else {
-            if (RandomNumber<ui32>(2) == 0) 
-                doReadRows = true; 
-            else 
+            if (RandomNumber<ui32>(2) == 0)
+                doReadRows = true;
+            else
                 doSelect = true;
         }
         Y_ABORT_UNLESS(doReadRows ^ doSelect);
@@ -417,7 +417,7 @@ TQueryInfoList TKvWorkloadGenerator::Mixed() {
             }
 
             return selectQuery;
-        } 
+        }
         if (doReadRows) {
             auto readRowsQuery = ReadRows(std::move(rows));
 
@@ -495,6 +495,8 @@ void TKvWorkloadParams::ConfigureOpts(NLastGetopt::TOpts& opts, const ECommandTy
             .DefaultValue((ui64)KvWorkloadConstants::INIT_ROW_COUNT).StoreResult(&InitRowCount);
         opts.AddLongOption("min-partitions", "Minimum partitions for tables.")
             .DefaultValue((ui64)KvWorkloadConstants::MIN_PARTITIONS).StoreResult(&MinPartitions);
+        opts.AddLongOption("max-partitions", "Maximum partitions for tables.")
+            .DefaultValue((ui64)KvWorkloadConstants::MAX_PARTITIONS).StoreResult(&MaxPartitions);
         opts.AddLongOption("partition-size", "Maximum partition size in megabytes (AUTO_PARTITIONING_PARTITION_SIZE_MB).")
             .DefaultValue((ui64)KvWorkloadConstants::PARTITION_SIZE_MB).StoreResult(&PartitionSizeMb);
         opts.AddLongOption("auto-partition", "Enable auto partitioning by load.")
diff --git a/ydb/library/workload/kv_workload.h b/ydb/library/workload/kv_workload.h
index 7cf2791cbefa..ae571063bb87 100644
--- a/ydb/library/workload/kv_workload.h
+++ b/ydb/library/workload/kv_workload.h
@@ -32,7 +32,7 @@ class TKvWorkloadParams : public TWorkloadParams {
     THolder<IWorkloadQueryGenerator> CreateGenerator() const override;
     TString GetWorkloadName() const override;
     ui64 MinPartitions = KvWorkloadConstants::MIN_PARTITIONS;
-    const ui64 MaxPartitions = KvWorkloadConstants::MAX_PARTITIONS;
+    ui64 MaxPartitions = KvWorkloadConstants::MAX_PARTITIONS;
     ui64 PartitionSizeMb = KvWorkloadConstants::PARTITION_SIZE_MB;
     ui64 InitRowCount = KvWorkloadConstants::INIT_ROW_COUNT;
     ui64 MaxFirstKey = KvWorkloadConstants::MAX_FIRST_KEY;
@@ -42,11 +42,11 @@ class TKvWorkloadParams : public TWorkloadParams {
     ui64 KeyColumnsCnt = KvWorkloadConstants::KEY_COLUMNS_CNT;
     ui64 RowsCnt = KvWorkloadConstants::ROWS_CNT;
     bool PartitionsByLoad = KvWorkloadConstants::PARTITIONS_BY_LOAD;
-    
+
     ui64 MixedChangePartitionsSize = KvWorkloadConstants::MIXED_CHANGE_PARTITIONS_SIZE;
     ui64 MixedDoReadRows = KvWorkloadConstants::MIXED_DO_READ_ROWS;
     ui64 MixedDoSelect = KvWorkloadConstants::MIXED_DO_SELECT;
-    
+
     const std::string TableName = "kv_test";
 };