add explain plan to benchmarks (#9167)

Author: iddqdex

ydb/public/lib/ydb_cli/commands/benchmark_utils.cpp

@@ -1,6 +1,7 @@
 #include "benchmark_utils.h"
 
 #include <util/string/split.h>
+#include <util/string/builder.h>
 #include <util/stream/file.h>
 #include <util/folder/pathsplit.h>
 #include <util/folder/path.h>
@@ -291,9 +292,10 @@ class TCSVResultScanner: public IQueryResultScanner, public TQueryResultInfo {
     }
 };
 
-TQueryBenchmarkResult Execute(const TString& query, NTable::TTableClient& client) {
+TQueryBenchmarkResult ExecuteImpl(const TString& query, NTable::TTableClient& client, bool explainOnly) {
     TStreamExecScanQuerySettings settings;
     settings.CollectQueryStats(ECollectQueryStatsMode::Full);
+    settings.Explain(explainOnly);
     auto it = client.StreamExecuteScanQuery(query, settings).GetValueSync();
     ThrowOnError(it);
 
@@ -316,9 +318,18 @@ TQueryBenchmarkResult Execute(const TString& query, NTable::TTableClient& client
     }
 }
 
-TQueryBenchmarkResult Execute(const TString& query, NQuery::TQueryClient& client) {
+TQueryBenchmarkResult Execute(const TString& query, NTable::TTableClient& client) {
+    return ExecuteImpl(query, client, false);
+}
+
+TQueryBenchmarkResult Explain(const TString& query, NTable::TTableClient& client) {
+    return ExecuteImpl(query, client, true);
+}
+
+TQueryBenchmarkResult ExecuteImpl(const TString& query, NQuery::TQueryClient& client, bool explainOnly) {
     NQuery::TExecuteQuerySettings settings;
     settings.StatsMode(NQuery::EStatsMode::Full);
+    settings.ExecMode(explainOnly ? NQuery::EExecMode::Explain : NQuery::EExecMode::Execute);
     auto it = client.StreamExecuteQuery(
         query,
         NYdb::NQuery::TTxControl::BeginTx().CommitTx(),
@@ -344,6 +355,14 @@ TQueryBenchmarkResult Execute(const TString& query, NQuery::TQueryClient& client
     }
 }
 
+TQueryBenchmarkResult Execute(const TString& query, NQuery::TQueryClient& client) {
+    return ExecuteImpl(query, client, false);
+}
+
+TQueryBenchmarkResult Explain(const TString& query, NQuery::TQueryClient& client) {
+    return ExecuteImpl(query, client, true);
+}
+
 NJson::TJsonValue GetQueryLabels(ui32 queryId) {
     NJson::TJsonValue labels(NJson::JSON_MAP);
     labels.InsertValue("query", Sprintf("Query%02u", queryId));

ydb/public/lib/ydb_cli/commands/benchmark_utils.h

@@ -91,6 +91,8 @@ void ThrowOnError(const TStatus& status);
 bool HasCharsInString(const TString& str);
 TQueryBenchmarkResult Execute(const TString & query, NTable::TTableClient & client);
 TQueryBenchmarkResult Execute(const TString & query, NQuery::TQueryClient & client);
+TQueryBenchmarkResult Explain(const TString & query, NTable::TTableClient & client);
+TQueryBenchmarkResult Explain(const TString & query, NQuery::TQueryClient & client);
 NJson::TJsonValue GetQueryLabels(ui32 queryId);
 NJson::TJsonValue GetSensorValue(TStringBuf sensor, TDuration& value, ui32 queryId);
 NJson::TJsonValue GetSensorValue(TStringBuf sensor, double value, ui32 queryId);

ydb/public/lib/ydb_cli/commands/ydb_benchmark.cpp

@@ -331,6 +331,16 @@ bool TWorkloadCommandBenchmark::RunBench(TClient& client, NYdbWorkload::IWorkloa
         ui32 failsCount = 0;
         ui32 diffsCount = 0;
         std::optional<TString> prevResult;
+        if (PlanFileName) {
+            TQueryBenchmarkResult res = TQueryBenchmarkResult::Error("undefined", "undefined", "undefined");
+            try {
+                res = Explain(query, client);
+            } catch (...) {
+                res = TQueryBenchmarkResult::Error(CurrentExceptionMessage(), "", "");
+            }
+            SavePlans(res, queryN, "explain");
+        }
+
         for (ui32 i = 0; i < IterationsCount; ++i) {
             auto t1 = TInstant::Now();
             TQueryBenchmarkResult res = TQueryBenchmarkResult::Error("undefined", "undefined", "undefined");
@@ -371,32 +381,7 @@ bool TWorkloadCommandBenchmark::RunBench(TClient& client, NYdbWorkload::IWorkloa
                 Cerr << query << Endl << Endl;
                 Sleep(TDuration::Seconds(1));
             }
-            if (PlanFileName) {
-                TFsPath(PlanFileName).Parent().MkDirs();
-                const TString planFName =  TStringBuilder() << PlanFileName << "." << i << ".";
-                if (res.GetQueryPlan()) {
-                    {
-                        TFileOutput out(planFName + "table");
-                        TQueryPlanPrinter queryPlanPrinter(EDataFormat::PrettyTable, true, out, 120);
-                        queryPlanPrinter.Print(res.GetQueryPlan());
-                    }
-                    {
-                        TFileOutput out(planFName + "json");
-                        TQueryPlanPrinter queryPlanPrinter(EDataFormat::JsonBase64, true, out, 120);
-                        queryPlanPrinter.Print(res.GetQueryPlan());
-                    }
-                    {
-                        TPlanVisualizer pv;
-                        pv.LoadPlans(res.GetQueryPlan());
-                        TFileOutput out(planFName + "svg");
-                        out << pv.PrintSvgSafe();
-                    }
-                }
-                if (res.GetPlanAst()) {
-                    TFileOutput out(planFName + "ast");
-                    out << res.GetPlanAst();
-                }
-            }
+            SavePlans(res, queryN, ToString(i));
         }
 
         auto [inserted, success] = queryRuns.emplace(queryN, TTestInfo(std::move(clientTimings), std::move(serverTimings)));
@@ -463,6 +448,36 @@ bool TWorkloadCommandBenchmark::RunBench(TClient& client, NYdbWorkload::IWorkloa
     return !someFailQueries;
 }
 
+void TWorkloadCommandBenchmark::SavePlans(const BenchmarkUtils::TQueryBenchmarkResult& res, ui32 queryNum, const TStringBuf name) const {
+    if (!PlanFileName) {
+        return;
+    }
+    TFsPath(PlanFileName).Parent().MkDirs();
+    const TString planFName =  TStringBuilder() << PlanFileName << "." << queryNum << "." << name << ".";
+    if (res.GetQueryPlan()) {
+        {
+            TFileOutput out(planFName + "table");
+            TQueryPlanPrinter queryPlanPrinter(EDataFormat::PrettyTable, true, out, 120);
+            queryPlanPrinter.Print(res.GetQueryPlan());
+        }
+        {
+            TFileOutput out(planFName + "json");
+            TQueryPlanPrinter queryPlanPrinter(EDataFormat::JsonBase64, true, out, 120);
+            queryPlanPrinter.Print(res.GetQueryPlan());
+        }
+        {
+            TPlanVisualizer pv;
+            pv.LoadPlans(res.GetQueryPlan());
+            TFileOutput out(planFName + "svg");
+            out << pv.PrintSvgSafe();
+        }
+    }
+    if (res.GetPlanAst()) {
+        TFileOutput out(planFName + "ast");
+        out << res.GetPlanAst();
+    }
+}
+
 int TWorkloadCommandBenchmark::DoRun(NYdbWorkload::IWorkloadQueryGenerator& workloadGen, TConfig& /*config*/) {
     if (QueryExecuterType == "scan") {
         return !RunBench(*TableClient, workloadGen);

ydb/public/lib/ydb_cli/commands/ydb_benchmark.h

@@ -2,6 +2,10 @@
 
 namespace NYdb::NConsoleClient {
 
+namespace BenchmarkUtils {
+    class TQueryBenchmarkResult;
+}
+
 class TWorkloadCommandBenchmark final: public TWorkloadCommandBase {
 public:
     TWorkloadCommandBenchmark(NYdbWorkload::TWorkloadParams& params, const NYdbWorkload::IWorkloadQueryGenerator::TWorkloadType& workload);
@@ -16,6 +20,7 @@ class TWorkloadCommandBenchmark final: public TWorkloadCommandBase {
 
     template <typename TClient>
     bool RunBench(TClient& client, NYdbWorkload::IWorkloadQueryGenerator& workloadGen);
+    void SavePlans(const BenchmarkUtils::TQueryBenchmarkResult& res, ui32 queryNum, const TStringBuf name) const;
 
 private:
     TString QueryExecuterType;

ydb/tests/olap/lib/ydb_cli.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+from typing import Optional
 import yatest.common
 import json
 import os
@@ -25,7 +26,7 @@ def get_cli_command() -> list[str]:
         else:
             return [cli]
 
-    class QueuePlan:
+    class QueryPlan:
         def __init__(self, plan: dict | None = None, table: str | None = None, ast: str | None = None, svg: str | None = None) -> None:
             self.plan = plan
             self.table = table
@@ -34,9 +35,9 @@ def __init__(self, plan: dict | None = None, table: str | None = None, ast: str
 
     class WorkloadRunResult:
         def __init__(
-            self, stats: dict[str, dict[str, any]] = {}, query_out: str = None, stdout: str = None, stderr: str = None,
-            error_message: str | None = None, plans: list[YdbCliHelper.QueuePlan] | None = None,
-            errors_by_iter: dict[int, str] | None = None
+            self, stats: dict[str, dict[str, any]] = {}, query_out: Optional[str] = None, stdout: Optional[str] = None, stderr: Optional[str] = None,
+            error_message: Optional[str] = None, plans: Optional[list[YdbCliHelper.QueryPlan]] = None,
+            errors_by_iter: Optional[dict[int, str]] = None, explain_plan: Optional[YdbCliHelper.QueryPlan] = None
         ) -> None:
             self.stats = stats
             self.query_out = query_out if str != '' else None
@@ -45,6 +46,7 @@ def __init__(
             self.success = error_message is None
             self.error_message = '' if self.success else error_message
             self.plans = plans
+            self.explain_plan = explain_plan
             self.errors_by_iter = errors_by_iter
 
     @staticmethod
@@ -76,6 +78,23 @@ def _try_extract_error_message(stderr: str) -> str:
                 else:
                     result[iter] = stderr[begin_pos:end_pos].strip()
 
+        def _load_plans(plan_path: str, name: str) -> YdbCliHelper.QueryPlan:
+            result = YdbCliHelper.QueryPlan()
+            pp = f'{plan_path}.{query_num}.{name}'
+            if (os.path.exists(f'{pp}.json')):
+                with open(f'{pp}.json') as f:
+                    result.plan = json.load(f)
+            if (os.path.exists(f'{pp}.table')):
+                with open(f'{pp}.table') as f:
+                    result.table = f.read()
+            if (os.path.exists(f'{pp}.ast')):
+                with open(f'{pp}.ast') as f:
+                    result.ast = f.read()
+            if (os.path.exists(f'{pp}.svg')):
+                with open(f'{pp}.svg') as f:
+                    result.svg = f.read()
+            return result
+
         errors_by_iter = {}
         try:
             wait_error = YdbCluster.wait_ydb_alive(300, path)
@@ -129,29 +148,16 @@ def _try_extract_error_message(stderr: str) -> str:
             if (os.path.exists(qout_path)):
                 with open(qout_path, 'r') as r:
                     qout = r.read()
-            plans = []
-            for i in range(iterations):
-                plans.append(YdbCliHelper.QueuePlan())
-                pp = f'{plan_path}.{i}'
-                if (os.path.exists(f'{pp}.json')):
-                    with open(f'{pp}.json') as f:
-                        plans[i].plan = json.load(f)
-                if (os.path.exists(f'{pp}.table')):
-                    with open(f'{pp}.table') as f:
-                        plans[i].table = f.read()
-                if (os.path.exists(f'{pp}.ast')):
-                    with open(f'{pp}.ast') as f:
-                        plans[i].ast = f.read()
-                if (os.path.exists(f'{pp}.svg')):
-                    with open(f'{pp}.svg') as f:
-                        plans[i].svg = f.read()
+            plans = [_load_plans(plan_path, str(i)) for i in range(iterations)]
+            explain_plan = _load_plans(plan_path, 'explain')
 
             return YdbCliHelper.WorkloadRunResult(
                 stats=stats,
                 query_out=qout,
                 plans=plans,
-                stdout=exec.stdout.decode('utf-8'),
-                stderr=exec.stderr.decode('utf-8'),
+                explain_plan=explain_plan,
+                stdout=exec.stdout.decode('utf-8', 'ignore'),
+                stderr=exec.stderr.decode('utf-8', 'ignore'),
                 error_message=err,
                 errors_by_iter=errors_by_iter
             )

ydb/tests/olap/load/conftest.py

@@ -29,6 +29,16 @@ def _get_duraton(stats, field):
             result = stats.get(field)
             return float(result) / 1e3 if result is not None else None
 
+        def _attach_plans(plan: YdbCliHelper.QueryPlan) -> None:
+            if plan.plan is not None:
+                allure.attach(json.dumps(plan.plan), 'Plan json', attachment_type=allure.attachment_type.JSON)
+            if plan.table is not None:
+                allure.attach(plan.table, 'Plan table', attachment_type=allure.attachment_type.TEXT)
+            if plan.ast is not None:
+                allure.attach(plan.ast, 'Plan ast', attachment_type=allure.attachment_type.TEXT)
+            if plan.svg is not None:
+                allure.attach(plan.svg, 'Plan svg', attachment_type=allure.attachment_type.SVG)
+
         test = f'Query{query_num:02d}'
         allure_test_description(self.suite, test, refference_set=self.refference)
         allure_listener = next(filter(lambda x: isinstance(x, AllureListener), plugin_manager.get_plugin_manager().get_plugins()))
@@ -47,19 +57,16 @@ def _get_duraton(stats, field):
             stats = {}
         if result.query_out is not None:
             allure.attach(result.query_out, 'Query output', attachment_type=allure.attachment_type.TEXT)
+
+        if result.explain_plan is not None:
+            with allure.step('Explain'):
+                _attach_plans(result.explain_plan)
+
         if result.plans is not None:
             for i in range(self.iterations):
                 try:
                     with allure.step(f'Iteration {i}'):
-                        plan = result.plans[i]
-                        if plan.plan is not None:
-                            allure.attach(json.dumps(plan.plan), 'Plan json', attachment_type=allure.attachment_type.JSON)
-                        if plan.table is not None:
-                            allure.attach(plan.table, 'Plan table', attachment_type=allure.attachment_type.TEXT)
-                        if plan.ast is not None:
-                            allure.attach(plan.ast, 'Plan ast', attachment_type=allure.attachment_type.TEXT)
-                        if plan.svg is not None:
-                            allure.attach(plan.svg, 'Plan svg', attachment_type=allure.attachment_type.SVG)
+                        _attach_plans(result.plans[i])
                         if i in result.errors_by_iter:
                             pytest.fail(result.errors_by_iter[i])
                 except BaseException: