Compact tpc-* queries and don't use pragma. (#6650)

ydb-platform · Jul 13, 2024 · 1daaf66 · 1daaf66
1 parent 0602b3c
commit 1daaf66
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 20 deletions.
diff --git a/ydb/library/workload/tpc_base/tpc_base.cpp b/ydb/library/workload/tpc_base/tpc_base.cpp
@@ -5,6 +5,7 @@
 #include <library/cpp/resource/resource.h>
 #include <util/stream/file.h>
 #include <util/string/split.h>
+#include <util/string/strip.h>
 
 namespace NYdbWorkload {
 
@@ -17,32 +18,55 @@ void TTpcBaseWorkloadGenerator::PatchQuery(TString& query) const {
     TString header;
     switch (Params.GetFloatMode()) {
     case TTpcBaseWorkloadParams::EFloatMode::FLOAT:
-        header = NResource::Find("consts.yql");
+        header = FilterHeader(NResource::Find("consts.yql"), query);
         break;
     case TTpcBaseWorkloadParams::EFloatMode::DECIMAL:
-        header = NResource::Find("consts_decimal.yql");
+    case TTpcBaseWorkloadParams::EFloatMode::DECIMAL_YDB:
+        header = FilterHeader(NResource::Find("consts_decimal.yql"), query);
         break;
-    case TTpcBaseWorkloadParams::EFloatMode::DECIMAL_YDB: {
-            header = NResource::Find("consts_decimal.yql");
-            header.to_lower();
-            const TStringBuf dec("decimal(");
-            auto p = header.find(dec);
-            while (p != TString::npos) {
-                p += dec.length();
-                const auto q = header.find(')', p);
-                TVector<ui32> decParams;
-                StringSplitter(header.cbegin() + p, q - p).SplitBySet(", ").SkipEmpty().Limit(2).ParseInto(&decParams);
-                TStringBuilder newDecParams;
-                newDecParams
-                    << Max(decParams[0], NKikimr::NScheme::DECIMAL_PRECISION)
-                    << "," << Max(decParams[1], NKikimr::NScheme::DECIMAL_SCALE);
-                header.replace(p, q - p, newDecParams);
-                p = header.find(dec, q);
+    }
+    PatchHeader(header);
+    SubstGlobal(query, "{% include 'header.sql.jinja' %}", header);
+}
+
+TString TTpcBaseWorkloadGenerator::FilterHeader(TStringBuf header, const TString& query) const {
+    TStringBuilder result;
+    for(TStringBuf line; header.ReadLine(line);) {
+        const auto pos = line.find('=');
+        if (pos == line.npos) {
+            continue;
+        }
+        const auto name = StripString(line.SubString(0, pos));
+        for(auto posInQ = query.find(name); posInQ != query.npos; posInQ = query.find(name, posInQ)) {
+            posInQ += name.length();
+            if (posInQ >= query.length() || !IsAsciiAlnum(query[posInQ]) && query[posInQ] != '_') {
+                result << line << Endl;
+                break;
             }
         }
-        break;
     }
-    SubstGlobal(query, "{% include 'header.sql.jinja' %}", header);
+    return result;
+}
+
+void TTpcBaseWorkloadGenerator::PatchHeader(TString& header) const {
+    if (Params.GetFloatMode() != TTpcBaseWorkloadParams::EFloatMode::DECIMAL_YDB) {
+        return;
+    }
+    header.to_lower();
+    const TStringBuf dec("decimal(");
+    auto p = header.find(dec);
+    while (p != TString::npos) {
+        p += dec.length();
+        const auto q = header.find(')', p);
+        TVector<ui32> decParams;
+        StringSplitter(header.cbegin() + p, q - p).SplitBySet(", ").SkipEmpty().Limit(2).ParseInto(&decParams);
+        TStringBuilder newDecParams;
+        newDecParams
+            << Max(decParams[0], NKikimr::NScheme::DECIMAL_PRECISION)
+            << "," << Max(decParams[1], NKikimr::NScheme::DECIMAL_SCALE);
+        header.replace(p, q - p, newDecParams);
+        p = header.find(dec, q);
+    }
 }
 
 void TTpcBaseWorkloadParams::ConfigureOpts(NLastGetopt::TOpts& opts, const ECommandType commandType, int workloadType) {

diff --git a/ydb/library/workload/tpc_base/tpc_base.h b/ydb/library/workload/tpc_base/tpc_base.h
@@ -25,6 +25,8 @@ class TTpcBaseWorkloadGenerator: public TWorkloadGeneratorBase {
 
 private:
     const TTpcBaseWorkloadParams& Params;
+    TString FilterHeader(TStringBuf header, const TString& query) const;
+    void PatchHeader(TString& header) const;
 };
 
 } // namespace NYdbWorkload