change partitioning in tpch (#17416)

Co-authored-by: Nikolay Perfilov <pnv1@yandex-team.ru>

Author: iddqdex

ydb/apps/ydb/CHANGELOG.md

@@ -1,3 +1,4 @@
+* Added `--scale` option to `ydb workload tpch init` and `ydb workload tpcds init` commands. Sets the percentage of the benchmark's data size and workload to use, relative to full scale.
 * Added "--no-discovery" option. It allows to skip discovery and use user provided endpoint to connect to YDB cluster.
 * Added `--retries` to `ydb workload <clickbenh|tpch|tpcds> run` command.
 * Added `--partition-size` param to `ydb workload <clickbench/tpcds/tpch> init`.

ydb/library/workload/benchmark_base/workload.cpp

@@ -35,6 +35,22 @@ const TString TWorkloadGeneratorBase::CsvFormatString = [] () {
     return settings.SerializeAsString();
 } ();
 
+namespace {
+
+    TString KeysList(const NJson::TJsonValue& table, const TString& key) {
+        TVector<TStringBuf> keysV;
+        for (const auto& k: table[key].GetArray()) {
+            keysV.emplace_back(k.GetString());
+        }
+        return JoinSeq(", ", keysV);
+    }
+
+}
+
+ui32 TWorkloadGeneratorBase::GetDefaultPartitionsCount(const TString& /*tableName*/) const {
+    return 64;
+}
+
 void TWorkloadGeneratorBase::GenerateDDLForTable(IOutputStream& result, const NJson::TJsonValue& table, bool single) const {
     auto specialTypes = GetSpecialDataTypes();
     specialTypes["string_type"] = Params.GetStringType();
@@ -65,11 +81,7 @@ void TWorkloadGeneratorBase::GenerateDDLForTable(IOutputStream& result, const NJ
         }
     }
     result << JoinSeq(",\n", columns);
-    TVector<TStringBuf> keysV;
-    for (const auto& k: table["primary_key"].GetArray()) {
-        keysV.emplace_back(k.GetString());
-    }
-    const TString keys = JoinSeq(", ", keysV);
+    const auto keys = KeysList(table, "primary_key");
     if (Params.GetStoreType() == TWorkloadBaseParams::EStoreType::ExternalS3) {
         result << Endl;
     } else {
@@ -78,7 +90,7 @@ void TWorkloadGeneratorBase::GenerateDDLForTable(IOutputStream& result, const NJ
     result << ")" << Endl;
 
     if (Params.GetStoreType() == TWorkloadBaseParams::EStoreType::Column) {
-        result << "PARTITION BY HASH (" << keys << ")" << Endl;
+        result << "PARTITION BY HASH (" <<  (table.Has("partition_by") ? KeysList(table, "partition_by") : keys) << ")" << Endl;
     }
 
     result << "WITH (" << Endl;
@@ -89,12 +101,12 @@ void TWorkloadGeneratorBase::GenerateDDLForTable(IOutputStream& result, const NJ
         break;
     case TWorkloadBaseParams::EStoreType::Column:
         result << "    STORE = COLUMN," << Endl;
-        result << "    AUTO_PARTITIONING_MIN_PARTITIONS_COUNT = " << table["partitioning"].GetUIntegerSafe(64) << Endl;
+        result << "    AUTO_PARTITIONING_MIN_PARTITIONS_COUNT = " << table["partitioning"].GetUIntegerSafe(GetDefaultPartitionsCount(tableName)) << Endl;
         break;
     case TWorkloadBaseParams::EStoreType::Row:
         result << "    STORE = ROW," << Endl;
         result << "    AUTO_PARTITIONING_PARTITION_SIZE_MB = " << Params.GetPartitionSizeMb() << ", " << Endl;
-        result << "    AUTO_PARTITIONING_MIN_PARTITIONS_COUNT = " << table["partitioning"].GetUIntegerSafe(64) << Endl;
+        result << "    AUTO_PARTITIONING_MIN_PARTITIONS_COUNT = " << table["partitioning"].GetUIntegerSafe(GetDefaultPartitionsCount(tableName)) << Endl;
     }
     result << ");" << Endl;
 }

ydb/library/workload/benchmark_base/workload.h

@@ -53,6 +53,7 @@ class TWorkloadGeneratorBase : public IWorkloadQueryGenerator {
     virtual TString GetTablesYaml() const = 0;
     virtual TSpecialDataTypes GetSpecialDataTypes() const = 0;
     NJson::TJsonValue GetTablesJson() const;
+    virtual ui32 GetDefaultPartitionsCount(const TString& tableName) const;
 
     THolder<TGeneratorStateProcessor> StateProcessor;
 private:

ydb/library/workload/tpc_base/tpc_base.cpp

@@ -167,7 +167,7 @@ void TTpcBaseWorkloadParams::ConfigureOpts(NLastGetopt::TOpts& opts, const EComm
             .StoreResult(&ExternalQueriesDir);
         opts.AddLongOption( "syntax", "Query syntax [" + GetEnumAllNames<EQuerySyntax>() + "].")
             .StoreResult(&Syntax).DefaultValue(Syntax);
-        opts.AddLongOption("scale", "scale in percents")
+        opts.AddLongOption("scale", "Sets the percentage of the benchmark's data size and workload to use, relative to full scale.")
             .DefaultValue(Scale).StoreResult(&Scale);
         opts.AddLongOption("float-mode", "Float mode. Can be float, decimal or decimal_ydb. If set to 'float' - float will be used, 'decimal' means that decimal will be used with canonical size and 'decimal_ydb' means that all floats will be converted to decimal(22,9) because YDB supports only this type.")
             .StoreResult(&FloatMode).DefaultValue(FloatMode);
@@ -177,6 +177,8 @@ void TTpcBaseWorkloadParams::ConfigureOpts(NLastGetopt::TOpts& opts, const EComm
     case TWorkloadParams::ECommandType::Init:
         opts.AddLongOption("float-mode", "Float mode. Can be float, decimal or decimal_ydb. If set to 'float' - float will be used, 'decimal' means that decimal will be used with canonical size and 'decimal_ydb' means that all floats will be converted to decimal(22,9) because YDB supports only this type.")
             .StoreResult(&FloatMode).DefaultValue(FloatMode);
+        opts.AddLongOption("scale", "Sets the percentage of the benchmark's data size and workload to use, relative to full scale.")
+            .DefaultValue(Scale).StoreResult(&Scale);
         break;
     default:
         break;

ydb/library/workload/tpcds/data_generator.cpp

@@ -19,7 +19,7 @@ TTpcdsWorkloadDataInitializerGenerator::TTpcdsWorkloadDataInitializerGenerator(c
 
 void TTpcdsWorkloadDataInitializerGenerator::ConfigureOpts(NLastGetopt::TOpts& opts) {
     TWorkloadDataInitializerBase::ConfigureOpts(opts);
-    opts.AddLongOption("scale", "scale in percents")
+    opts.AddLongOption("scale", "Sets the percentage of the benchmark's data size and workload to use, relative to full scale.")
         .DefaultValue(Scale).StoreResult(&Scale);
     opts.AddLongOption("tables", "Commaseparated list of tables for generate. Empty means all tables.\n"
             "Enabled tables: " + JoinSeq(", ", TBulkDataGenerator::TFactory::GetRegisteredKeys()))

ydb/library/workload/tpch/data_generator.cpp

@@ -11,7 +11,7 @@ TTpchWorkloadDataInitializerGenerator::TTpchWorkloadDataInitializerGenerator(con
 
 void TTpchWorkloadDataInitializerGenerator::ConfigureOpts(NLastGetopt::TOpts& opts) {
     TWorkloadDataInitializerBase::ConfigureOpts(opts);
-    opts.AddLongOption("scale", "scale in percents")
+    opts.AddLongOption("scale", "Sets the percentage of the benchmark's data size and workload to use, relative to full scale.")
         .DefaultValue(Scale).StoreResult(&Scale);
     opts.AddLongOption("tables", "Commaseparated list of tables for generate. Empty means all tables.\n"
             "Enabled tables: " + JoinSeq(", ", TBulkDataGenerator::TFactory::GetRegisteredKeys()))

ydb/library/workload/tpch/tpch.cpp

@@ -29,6 +29,10 @@ TWorkloadGeneratorBase::TSpecialDataTypes TTpchWorkloadGenerator::GetSpecialData
     }
 }
 
+ui32 TTpchWorkloadGenerator::GetDefaultPartitionsCount(const TString& /*tableName*/) const {
+    return Params.GetScale() <= 10 ? 64 : 256;
+}
+
 
 THolder<IWorkloadQueryGenerator> TTpchWorkloadParams::CreateGenerator() const {
     return MakeHolder<TTpchWorkloadGenerator>(*this);

ydb/library/workload/tpch/tpch.h

@@ -19,6 +19,7 @@ class TTpchWorkloadGenerator final: public TTpcBaseWorkloadGenerator {
 protected:
     TString GetTablesYaml() const override;
     TWorkloadGeneratorBase::TSpecialDataTypes GetSpecialDataTypes() const override;
+    ui32 GetDefaultPartitionsCount(const TString& tableName) const override;
 
 private:
     const TTpchWorkloadParams& Params;

ydb/library/workload/tpch/tpch_schema.yaml

@@ -81,6 +81,8 @@ tables:
     primary_key:
       - l_orderkey
       - l_linenumber
+    partition_by:
+      - l_orderkey
 
   - name: nation
     columns:

ydb/tests/functional/benchmarks_init/canondata/result.json

@@ -32,6 +32,9 @@
     "test_init.TestClickbenchInit.test_s1_s3": {
         "uri": "file://test_init.TestClickbenchInit.test_s1_s3/s1_s3"
     },
+    "test_init.TestTpcdsInit.test_s100_column": {
+        "uri": "file://test_init.TestTpcdsInit.test_s100_column/s100_column"
+    },
     "test_init.TestTpcdsInit.test_s1_column": {
         "uri": "file://test_init.TestTpcdsInit.test_s1_column/s1_column"
     },
@@ -47,6 +50,9 @@
     "test_init.TestTpcdsInit.test_s1_s3": {
         "uri": "file://test_init.TestTpcdsInit.test_s1_s3/s1_s3"
     },
+    "test_init.TestTpchInit.test_s100_column": {
+        "uri": "file://test_init.TestTpchInit.test_s100_column/s100_column"
+    },
     "test_init.TestTpchInit.test_s1_column": {
         "uri": "file://test_init.TestTpchInit.test_s1_column/s1_column"
     },