add RowsCnt parameter and insert command to kv workload

6 files changed, 170 insertions, 50 deletions

diff --git a/ydb/core/blobstorage/testload/test_load_kqp.cpp b/ydb/core/blobstorage/testload/test_load_kqp.cpp
index 1efe502282d..e5c9ab36395 100644
--- a/ydb/core/blobstorage/testload/test_load_kqp.cpp
+++ b/ydb/core/blobstorage/testload/test_load_kqp.cpp
@@ -367,6 +367,7 @@ public:
             params.MaxFirstKey = cmd.GetKv().GetMaxFirstKey();
             params.StringLen = cmd.GetKv().GetStringLen();
             params.ColumnsCnt = cmd.GetKv().GetColumnsCnt();
+            params.RowsCnt = cmd.GetKv().GetRowsCnt();
             params.MinPartitions = UniformPartitionsCount;
             params.DbPath = WorkingDir;
             WorkloadQueryGen = factory.GetWorkloadQueryGenerator(NYdbWorkload::EWorkload::KV, &params);
@@ -411,6 +412,7 @@ public:
                 << "MinPartitions: " << params->MinPartitions << " "
                 << "StringLen: " << params->StringLen << " "
                 << "ColumnsCnt: " << params->ColumnsCnt << " "
+                << "RowsCnt: " << params->RowsCnt << " "
                 << "DbPath: " << params->DbPath);
         }
 
diff --git a/ydb/core/protos/blobstorage.proto b/ydb/core/protos/blobstorage.proto
index e8d154d4873..302b1e158ff 100644
--- a/ydb/core/protos/blobstorage.proto
+++ b/ydb/core/protos/blobstorage.proto
@@ -1578,9 +1578,10 @@ message TEvTestLoadRequest {
         message TKvWorkload {
             optional uint64 InitRowCount = 1 [default = 1000];
             optional bool PartitionsByLoad = 2 [default = true];
-            optional uint64 MaxFirstKey = 3 [default = 5000];
+            optional uint64 MaxFirstKey = 3 [default = 18446744073709551615];
             optional uint64 StringLen = 4 [default = 8];
             optional uint64 ColumnsCnt = 5 [default = 2];
+            optional uint64 RowsCnt = 6 [default = 1];
         }
         optional uint64 Tag = 1;
         optional uint32 DurationSeconds = 2;
diff --git a/ydb/library/workload/kv_workload.cpp b/ydb/library/workload/kv_workload.cpp
index 49786e1b92e..e29562ddea3 100644
--- a/ydb/library/workload/kv_workload.cpp
+++ b/ydb/library/workload/kv_workload.cpp
@@ -11,6 +11,12 @@
 #include <random>
 #include <sstream>
 
+template <>
+void Out<NYdbWorkload::KvWorkloadConstants>(IOutputStream& out, NYdbWorkload::KvWorkloadConstants constant)
+{
+    out << static_cast<ui64>(constant);
+}
+
 namespace NYdbWorkload {
 
 TKvWorkloadGenerator::TKvWorkloadGenerator(const TKvWorkloadParams* params)
@@ -56,6 +62,8 @@ TQueryInfoList TKvWorkloadGenerator::GetWorkload(int type) {
     switch (static_cast<EType>(type)) {
         case EType::UpsertRandom:
             return UpsertRandom();
+        case EType::InsertRandom:
+            return InsertRandom();
         case EType::SelectRandom:
             return SelectRandom();
         default:
@@ -63,21 +71,27 @@ TQueryInfoList TKvWorkloadGenerator::GetWorkload(int type) {
     }
 }
 
-TQueryInfoList TKvWorkloadGenerator::UpsertRandom() {
+
+TQueryInfoList TKvWorkloadGenerator::AddOperation(TString operation) {
     std::stringstream ss;
 
     NYdb::TParamsBuilder paramsBuilder;
 
     ss << "--!syntax_v1\n";
-    ss << "DECLARE $c0 AS Uint64;\n";
-    paramsBuilder.AddParam("$c0").Uint64(KeyUniformDistGen(Gen)).Build();
 
-    for (size_t col = 1; col < Params.ColumnsCnt; ++col) {
-        ss << "DECLARE $c" << col << " AS String;\n";
-        paramsBuilder.AddParam("$c" + std::to_string(col)).String(BigString).Build();
+    for (size_t row = 0; row < Params.RowsCnt; ++row) {
+        TString pkname = "$r" + std::to_string(row);
+        ss << "DECLARE " << pkname << " AS Uint64;\n";
+        paramsBuilder.AddParam(pkname).Uint64(KeyUniformDistGen(Gen)).Build();
+
+        for (size_t col = 1; col < Params.ColumnsCnt; ++col) {
+            TString cname = "$c" + std::to_string(row) + std::to_string(col);
+            ss << "DECLARE " << cname << " AS String;\n";
+            paramsBuilder.AddParam(cname).String(BigString).Build();
+        }
     }
 
-    ss << "UPSERT INTO `kv_test`(";
+    ss << operation << " INTO `kv_test`(";
 
     for (size_t col = 0; col < Params.ColumnsCnt; ++col) {
         ss << "c" << col;
@@ -86,27 +100,50 @@ TQueryInfoList TKvWorkloadGenerator::UpsertRandom() {
         }
     }
 
-    ss << ") VALUES (";
+    ss << ") VALUES ";
 
-    for (size_t col = 0; col < Params.ColumnsCnt; ++col) {
-        ss << "$c" << col;
-        if (col + 1 < Params.ColumnsCnt) {
+    for (size_t row = 0; row < Params.RowsCnt; ++row) {
+
+        ss << "(";
+        ss << "$r" << row;
+
+        for (size_t col = 1; col < Params.ColumnsCnt; ++col) {
+            ss << ", $c" << row << col;
+        }
+
+        ss << ")";
+
+        if (row + 1 < Params.RowsCnt) {
             ss << ", ";
         }
-    }
 
-    ss << ")";
+    }
 
     auto params = paramsBuilder.Build();
 
     return TQueryInfoList(1, TQueryInfo(ss.str(), std::move(params)));
 }
 
+TQueryInfoList TKvWorkloadGenerator::UpsertRandom() {
+    return AddOperation("UPSERT");
+}
+
+TQueryInfoList TKvWorkloadGenerator::InsertRandom() {
+    return AddOperation("INSERT");
+}
+
 TQueryInfoList TKvWorkloadGenerator::SelectRandom() {
     std::stringstream ss;
 
+    NYdb::TParamsBuilder paramsBuilder;
+
     ss << "--!syntax_v1\n";
-    ss << "DECLARE $c0 AS Uint64;\n";
+
+    for (size_t row = 0; row < Params.RowsCnt; ++row) {
+        ss << "DECLARE $r" << row << " AS Uint64;\n";
+        paramsBuilder.AddParam("$r" + std::to_string(row)).Uint64(KeyUniformDistGen(Gen)).Build();
+    }
+
     ss << "SELECT ";
     for (size_t col = 0; col < Params.ColumnsCnt; ++col) {
         ss << "c" << col;
@@ -116,16 +153,15 @@ TQueryInfoList TKvWorkloadGenerator::SelectRandom() {
         ss << " ";
     }
 
-    ss << "FROM `kv_test` WHERE c0 = $c0";
-
-    ui64 x = KeyUniformDistGen(Gen);
+    ss << "FROM `kv_test` WHERE ";
+    for (size_t row = 0; row < Params.RowsCnt; ++row) {
+        ss << "c0 = $r" << row;
+        if (row + 1 < Params.RowsCnt) {
+            ss << " OR ";
+        }
+    }
 
-    NYdb::TParamsBuilder paramsBuilder;
-    auto params = paramsBuilder
-        .AddParam("$c0")
-            .Uint64(x)
-            .Build()
-        .Build();
+    auto params = paramsBuilder.Build();
 
     return TQueryInfoList(1, TQueryInfo(ss.str(), std::move(params)));
 }
diff --git a/ydb/library/workload/kv_workload.h b/ydb/library/workload/kv_workload.h
index b5e492cfddb..3a53317ff8a 100644
--- a/ydb/library/workload/kv_workload.h
+++ b/ydb/library/workload/kv_workload.h
@@ -7,13 +7,24 @@
 
 namespace NYdbWorkload {
 
+enum KvWorkloadConstants : ui64 {
+    MIN_PARTITIONS = 40,
+    INIT_ROW_COUNT = 1000,
+    MAX_FIRST_KEY = Max<ui64>(),
+    STRING_LEN = 8,
+    COLUMNS_CNT = 2,
+    ROWS_CNT = 1,
+    PARTITIONS_BY_LOAD = true
+};
+
 struct TKvWorkloadParams : public TWorkloadParams {
-    ui64 MinPartitions = 1;
-    ui64 InitRowCount = 1000;
-    ui64 MaxFirstKey = 5000;
-    ui64 StringLen = 8;
-    ui64 ColumnsCnt = 2;
-    bool PartitionsByLoad = true;
+    ui64 MinPartitions = KvWorkloadConstants::MIN_PARTITIONS;
+    ui64 InitRowCount = KvWorkloadConstants::INIT_ROW_COUNT;
+    ui64 MaxFirstKey = KvWorkloadConstants::MAX_FIRST_KEY;
+    ui64 StringLen = KvWorkloadConstants::STRING_LEN;
+    ui64 ColumnsCnt = KvWorkloadConstants::COLUMNS_CNT;
+    ui64 RowsCnt = KvWorkloadConstants::ROWS_CNT;
+    bool PartitionsByLoad = KvWorkloadConstants::PARTITIONS_BY_LOAD;
 };
 
 class TKvWorkloadGenerator : public IWorkloadQueryGenerator {
@@ -40,16 +51,19 @@ public:
 
     enum class EType {
         UpsertRandom,
+        InsertRandom,
         SelectRandom
     };
 
 private:
     TQueryInfoList UpsertRandom();
+    TQueryInfoList InsertRandom();
     TQueryInfoList SelectRandom();
 
     TKvWorkloadGenerator(const TKvWorkloadParams* params);
 
     TQueryInfo FillKvData() const;
+    TQueryInfoList AddOperation(TString operation);
 
     std::string DbPath;
     TKvWorkloadParams Params;
diff --git a/ydb/public/lib/ydb_cli/commands/kv_workload.cpp b/ydb/public/lib/ydb_cli/commands/kv_workload.cpp
index c8bff4f05c1..1a48bca2054 100644
--- a/ydb/public/lib/ydb_cli/commands/kv_workload.cpp
+++ b/ydb/public/lib/ydb_cli/commands/kv_workload.cpp
@@ -16,12 +16,13 @@ TCommandKv::TCommandKv()
 
 TCommandKvInit::TCommandKvInit()
     : TWorkloadCommand("init", {}, "Create and initialize tables for workload")
-    , InitRowCount(1000)
-    , MinPartitions(1)
-    , MaxFirstKey(5000)
-    , StringLen(8)
-    , ColumnsCnt(2)
-    , PartitionsByLoad(true)
+    , InitRowCount(NYdbWorkload::KvWorkloadConstants::INIT_ROW_COUNT)
+    , MinPartitions(NYdbWorkload::KvWorkloadConstants::MIN_PARTITIONS)
+    , MaxFirstKey(NYdbWorkload::KvWorkloadConstants::MAX_FIRST_KEY)
+    , StringLen(NYdbWorkload::KvWorkloadConstants::STRING_LEN)
+    , ColumnsCnt(NYdbWorkload::KvWorkloadConstants::COLUMNS_CNT)
+    , RowsCnt(NYdbWorkload::KvWorkloadConstants::ROWS_CNT)
+    , PartitionsByLoad(NYdbWorkload::KvWorkloadConstants::PARTITIONS_BY_LOAD)
 {}
 
 void TCommandKvInit::Config(TConfig& config) {
@@ -29,18 +30,20 @@ void TCommandKvInit::Config(TConfig& config) {
 
     config.SetFreeArgsNum(0);
 
-    config.Opts->AddLongOption("rows-cnt", "count of rows need to Insert while table initialization")
-        .DefaultValue(1000).StoreResult(&InitRowCount);
+    config.Opts->AddLongOption("init-upserts", "count of upserts need to create while table initialization")
+        .DefaultValue(NYdbWorkload::KvWorkloadConstants::INIT_ROW_COUNT).StoreResult(&InitRowCount);
     config.Opts->AddLongOption("min-partitions", "Minimum partitions for tables.")
-        .DefaultValue(40).StoreResult(&MinPartitions);
+        .DefaultValue(NYdbWorkload::KvWorkloadConstants::MIN_PARTITIONS).StoreResult(&MinPartitions);
     config.Opts->AddLongOption("auto-partition", "Enable auto partitioning by load.")
-        .DefaultValue(true).StoreResult(&PartitionsByLoad);
+        .DefaultValue(NYdbWorkload::KvWorkloadConstants::PARTITIONS_BY_LOAD).StoreResult(&PartitionsByLoad);
     config.Opts->AddLongOption("max-first-key", "maximum value of first primary key")
-        .DefaultValue(5000).StoreResult(&MaxFirstKey);
+        .DefaultValue(NYdbWorkload::KvWorkloadConstants::MAX_FIRST_KEY).StoreResult(&MaxFirstKey);
     config.Opts->AddLongOption("len", "String len")
-        .DefaultValue(8).StoreResult(&StringLen);
+        .DefaultValue(NYdbWorkload::KvWorkloadConstants::STRING_LEN).StoreResult(&StringLen);
     config.Opts->AddLongOption("cols", "Number of columns")
-        .DefaultValue(2).StoreResult(&ColumnsCnt);
+        .DefaultValue(NYdbWorkload::KvWorkloadConstants::COLUMNS_CNT).StoreResult(&ColumnsCnt);
+    config.Opts->AddLongOption("rows", "Number of rows to upsert")
+        .DefaultValue(NYdbWorkload::KvWorkloadConstants::ROWS_CNT).StoreResult(&RowsCnt);
 }
 
 void TCommandKvInit::Parse(TConfig& config) {
@@ -58,6 +61,7 @@ int TCommandKvInit::Run(TConfig& config) {
     params.MaxFirstKey = MaxFirstKey;
     params.StringLen = StringLen;
     params.ColumnsCnt = ColumnsCnt;
+    params.RowsCnt = RowsCnt;
 
     NYdbWorkload::TWorkloadFactory factory;
     auto workloadGen = factory.GetWorkloadQueryGenerator(NYdbWorkload::EWorkload::KV, &params);
@@ -94,11 +98,12 @@ TCommandKvRun::TCommandKvRun()
     : TClientCommandTree("run", {}, "Run YDB KV workload")
 {
     AddCommand(std::make_unique<TCommandKvRunUpsertRandom>());
+    AddCommand(std::make_unique<TCommandKvRunInsertRandom>());
     AddCommand(std::make_unique<TCommandKvRunSelectRandom>());
 }
 
 TCommandKvRunUpsertRandom::TCommandKvRunUpsertRandom()
-    : TWorkloadCommand("upsert", {}, "upsert random pairs (a, b) into table")
+    : TWorkloadCommand("upsert", {}, "upsert random rows into table")
 {}
 
 void TCommandKvRunUpsertRandom::Config(TConfig& config) {
@@ -106,11 +111,13 @@ void TCommandKvRunUpsertRandom::Config(TConfig& config) {
     config.SetFreeArgsNum(0);
 
     config.Opts->AddLongOption("max-first-key", "maximum value of first primary key")
-        .DefaultValue(5000).StoreResult(&MaxFirstKey);
+        .DefaultValue(NYdbWorkload::KvWorkloadConstants::MAX_FIRST_KEY).StoreResult(&MaxFirstKey);
     config.Opts->AddLongOption("len", "String len")
-        .DefaultValue(8).StoreResult(&StringLen);
+        .DefaultValue(NYdbWorkload::KvWorkloadConstants::STRING_LEN).StoreResult(&StringLen);
     config.Opts->AddLongOption("cols", "Number of columns")
-        .DefaultValue(2).StoreResult(&ColumnsCnt);
+        .DefaultValue(NYdbWorkload::KvWorkloadConstants::COLUMNS_CNT).StoreResult(&ColumnsCnt);
+    config.Opts->AddLongOption("rows", "Number of rows to upsert")
+        .DefaultValue(NYdbWorkload::KvWorkloadConstants::ROWS_CNT).StoreResult(&RowsCnt);
 }
 
 void TCommandKvRunUpsertRandom::Parse(TConfig& config) {
@@ -125,6 +132,7 @@ int TCommandKvRunUpsertRandom::Run(TConfig& config) {
     params.MaxFirstKey = MaxFirstKey;
     params.StringLen = StringLen;
     params.ColumnsCnt = ColumnsCnt;
+    params.RowsCnt = RowsCnt;
 
     NYdbWorkload::TWorkloadFactory factory;
     auto workloadGen = factory.GetWorkloadQueryGenerator(NYdbWorkload::EWorkload::KV, &params);
@@ -132,8 +140,46 @@ int TCommandKvRunUpsertRandom::Run(TConfig& config) {
     return RunWorkload(workloadGen, static_cast<int>(NYdbWorkload::TKvWorkloadGenerator::EType::UpsertRandom));
 }
 
+TCommandKvRunInsertRandom::TCommandKvRunInsertRandom()
+    : TWorkloadCommand("insert", {}, "insert random rows into table")
+{}
+
+void TCommandKvRunInsertRandom::Config(TConfig& config) {
+    TWorkloadCommand::Config(config);
+    config.SetFreeArgsNum(0);
+
+    config.Opts->AddLongOption("max-first-key", "maximum value of first primary key")
+        .DefaultValue(NYdbWorkload::KvWorkloadConstants::MAX_FIRST_KEY).StoreResult(&MaxFirstKey);
+    config.Opts->AddLongOption("len", "String len")
+        .DefaultValue(NYdbWorkload::KvWorkloadConstants::STRING_LEN).StoreResult(&StringLen);
+    config.Opts->AddLongOption("cols", "Number of columns")
+        .DefaultValue(NYdbWorkload::KvWorkloadConstants::COLUMNS_CNT).StoreResult(&ColumnsCnt);
+    config.Opts->AddLongOption("rows", "Number of rows to insert")
+        .DefaultValue(NYdbWorkload::KvWorkloadConstants::ROWS_CNT).StoreResult(&RowsCnt);
+}
+
+void TCommandKvRunInsertRandom::Parse(TConfig& config) {
+    TClientCommand::Parse(config);
+}
+
+int TCommandKvRunInsertRandom::Run(TConfig& config) {
+    PrepareForRun(config);
+
+    NYdbWorkload::TKvWorkloadParams params;
+    params.DbPath = config.Database;
+    params.MaxFirstKey = MaxFirstKey;
+    params.StringLen = StringLen;
+    params.ColumnsCnt = ColumnsCnt;
+    params.RowsCnt = RowsCnt;
+
+    NYdbWorkload::TWorkloadFactory factory;
+    auto workloadGen = factory.GetWorkloadQueryGenerator(NYdbWorkload::EWorkload::KV, &params);
+
+    return RunWorkload(workloadGen, static_cast<int>(NYdbWorkload::TKvWorkloadGenerator::EType::InsertRandom));
+}
+
 TCommandKvRunSelectRandom::TCommandKvRunSelectRandom()
-    : TWorkloadCommand("select", {}, "select row by exactly matching of a")
+    : TWorkloadCommand("select", {}, "select rows by exactly matching of a")
 {}
 
 void TCommandKvRunSelectRandom::Config(TConfig& config) {
@@ -141,9 +187,11 @@ void TCommandKvRunSelectRandom::Config(TConfig& config) {
     config.SetFreeArgsNum(0);
 
     config.Opts->AddLongOption("max-first-key", "maximum value of first primary key")
-        .DefaultValue(5000).StoreResult(&MaxFirstKey);
+        .DefaultValue(NYdbWorkload::KvWorkloadConstants::MAX_FIRST_KEY).StoreResult(&MaxFirstKey);
     config.Opts->AddLongOption("cols", "Number of columns")
-        .DefaultValue(2).StoreResult(&ColumnsCnt);
+        .DefaultValue(NYdbWorkload::KvWorkloadConstants::COLUMNS_CNT).StoreResult(&ColumnsCnt);
+    config.Opts->AddLongOption("rows", "Number of rows to select")
+        .DefaultValue(NYdbWorkload::KvWorkloadConstants::ROWS_CNT).StoreResult(&RowsCnt);
 }
 
 void TCommandKvRunSelectRandom::Parse(TConfig& config) {
@@ -157,6 +205,7 @@ int TCommandKvRunSelectRandom::Run(TConfig& config) {
     params.DbPath = config.Database;
     params.MaxFirstKey = MaxFirstKey;
     params.ColumnsCnt = ColumnsCnt;
+    params.RowsCnt = RowsCnt;
 
     NYdbWorkload::TWorkloadFactory factory;
     auto workloadGen = factory.GetWorkloadQueryGenerator(NYdbWorkload::EWorkload::KV, &params);
diff --git a/ydb/public/lib/ydb_cli/commands/kv_workload.h b/ydb/public/lib/ydb_cli/commands/kv_workload.h
index b3e55598312..bb3de13fc5c 100644
--- a/ydb/public/lib/ydb_cli/commands/kv_workload.h
+++ b/ydb/public/lib/ydb_cli/commands/kv_workload.h
@@ -23,6 +23,7 @@ private:
     ui64 MaxFirstKey;
     ui64 StringLen;
     ui64 ColumnsCnt;
+    ui64 RowsCnt;
     bool PartitionsByLoad;
 };
 
@@ -50,6 +51,22 @@ private:
     ui64 MaxFirstKey;
     ui64 StringLen;
     ui64 ColumnsCnt;
+    ui64 RowsCnt;
+
+};
+
+class TCommandKvRunInsertRandom : public TWorkloadCommand {
+public:
+    TCommandKvRunInsertRandom();
+    virtual void Config(TConfig& config) override;
+    virtual void Parse(TConfig& config) override;
+    virtual int Run(TConfig& config) override;
+
+private:
+    ui64 MaxFirstKey;
+    ui64 StringLen;
+    ui64 ColumnsCnt;
+    ui64 RowsCnt;
 
 };
 
@@ -63,6 +80,7 @@ public:
 private:
     ui64 MaxFirstKey;
     ui64 ColumnsCnt;
+    ui64 RowsCnt;
 
 };