Merge 68b1397 into c7db92a

ydb-platform · Jan 10, 2025 · 71458b5 · 71458b5
2 parents c7db92a + 68b1397
commit 71458b5
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 9 deletions.
diff --git a/ydb/library/workload/kv/kv.cpp b/ydb/library/workload/kv/kv.cpp
@@ -151,10 +151,14 @@ std::string TKvWorkloadGenerator::GetDDLQueries() const {
 
     for (size_t i = 0; i < Params.ColumnsCnt; ++i) {
         if (i < Params.IntColumnsCnt) {
-            ss << "c" << i << " Uint64, ";
+            ss << "c" << i << " Uint64";
         } else {
-            ss << "c" << i << " String, ";
+            ss << "c" << i << " String";
         }
+        if (i < Params.KeyColumnsCnt && Params.GetStoreType() == TKvWorkloadParams::EStoreType::Column) {
+            ss << " NOT NULL";
+        }
+        ss << ", ";
     }
 
     ss << "PRIMARY KEY(";
@@ -166,13 +170,23 @@ std::string TKvWorkloadGenerator::GetDDLQueries() const {
     }
     ss << ")) WITH (";
 
-    if (Params.PartitionsByLoad) {
-        ss << "AUTO_PARTITIONING_BY_LOAD = ENABLED, ";
+    switch (Params.GetStoreType()) {
+        case TKvWorkloadParams::EStoreType::Row:
+            ss << "STORE = ROW, ";
+            if (Params.PartitionsByLoad) {
+                ss << "AUTO_PARTITIONING_BY_LOAD = ENABLED, ";
+            }
+            ss << "UNIFORM_PARTITIONS = " << Params.MinPartitions << ", ";
+            ss << "AUTO_PARTITIONING_MAX_PARTITIONS_COUNT = " << Max(Params.MinPartitions, Params.MaxPartitions) << ", ";
+            ss << "AUTO_PARTITIONING_PARTITION_SIZE_MB = " << Params.PartitionSizeMb << ", ";
+            break;
+        case TKvWorkloadParams::EStoreType::Column:
+            ss << "STORE = COLUMN, ";
+            break;
+        default:
+            throw yexception() << "Unsupported store type: " << Params.GetStoreType();
     }
-    ss << "AUTO_PARTITIONING_MIN_PARTITIONS_COUNT = " << Params.MinPartitions << ", ";
-    ss << "AUTO_PARTITIONING_PARTITION_SIZE_MB = " << Params.PartitionSizeMb << ", ";
-    ss << "UNIFORM_PARTITIONS = " << Params.MinPartitions << ", ";
-    ss << "AUTO_PARTITIONING_MAX_PARTITIONS_COUNT = " << Max(Params.MinPartitions, Params.MaxPartitions) << ")";
+    ss << "AUTO_PARTITIONING_MIN_PARTITIONS_COUNT = " << Params.MinPartitions << ")";
 
     return ss.str();
 }
@@ -517,6 +531,17 @@ void TKvWorkloadParams::ConfigureOpts(NLastGetopt::TOpts& opts, const ECommandTy
             .DefaultValue((ui64)KvWorkloadConstants::KEY_COLUMNS_CNT).StoreResult(&KeyColumnsCnt);
         opts.AddLongOption("rows", "Number of rows")
             .DefaultValue((ui64)KvWorkloadConstants::ROWS_CNT).StoreResult(&RowsCnt);
+        opts.AddLongOption("store", "Storage type."
+                " Options: row, column\n"
+                "  row - use row-based storage engine;\n"
+                "  column - use column-based storage engine.")
+            .DefaultValue(StoreType)
+            .Handler1T<TStringBuf>([this](TStringBuf arg) {
+                const auto l = to_lower(TString(arg));
+                if (!TryFromString(arg, StoreType)) {
+                    throw yexception() << "Ivalid store type: " << arg;
+                }
+            });
         break;
     case TWorkloadParams::ECommandType::Run:
         opts.AddLongOption("max-first-key", "Maximum value of a first primary key")

diff --git a/ydb/library/workload/kv/kv.h b/ydb/library/workload/kv/kv.h
@@ -30,6 +30,11 @@ enum KvWorkloadConstants : ui64 {
 
 class TKvWorkloadParams : public TWorkloadParams {
 public:
+    enum class EStoreType {
+        Row     /* "row"    */,
+        Column  /* "column" */,
+    };
+
     void ConfigureOpts(NLastGetopt::TOpts& opts, const ECommandType commandType, int workloadType) override;
     THolder<IWorkloadQueryGenerator> CreateGenerator() const override;
     TString GetWorkloadName() const override;
@@ -52,6 +57,7 @@ class TKvWorkloadParams : public TWorkloadParams {
     const std::string TableName = "kv_test";
 
     bool StaleRO = KvWorkloadConstants::STALE_RO;
+    YDB_READONLY(EStoreType, StoreType, EStoreType::Row);
 };
 
 class TKvWorkloadGenerator final: public TWorkloadQueryGeneratorBase<TKvWorkloadParams> {

diff --git a/ydb/tests/workloads/kv/tests/test_workload.py b/ydb/tests/workloads/kv/tests/test_workload.py
@@ -34,7 +34,8 @@ def test(self):
                 "--init-upserts", "0",
                 "--cols", "5",
                 "--int-cols", "2",
-                "--key-cols", "3"
+                "--key-cols", "3",
+                "--store", "column",
             ],
             wait=True
         )