Fixed a bug with printing of actual stats (#3447)

2 files changed, 46 insertions, 23 deletions

diff --git a/ydb/core/kqp/opt/kqp_query_plan.cpp b/ydb/core/kqp/opt/kqp_query_plan.cpp
index 333b298b88..31d0748011 100644
--- a/ydb/core/kqp/opt/kqp_query_plan.cpp
+++ b/ydb/core/kqp/opt/kqp_query_plan.cpp
@@ -2119,8 +2119,6 @@ double ComputeCpuTimes(NJson::TJsonValue& plan) {
     }
 
     if (plan.GetMapSafe().contains("Stats") && plan.GetMapSafe().contains("Operators")) {
-        YQL_CLOG(TRACE, CoreDq) << "Found Operators";
-
         auto& ops = plan.GetMapSafe().at("Operators").GetArraySafe();
 
         const auto& stats = plan.GetMapSafe().at("Stats").GetMapSafe();
@@ -2144,6 +2142,45 @@ double ComputeCpuTimes(NJson::TJsonValue& plan) {
     return currCpuTime;
 }
 
+void ComputeTotalRows(NJson::TJsonValue& plan) {
+    
+    if (plan.GetMapSafe().contains("Plans")) {
+        for (auto& p : plan.GetMapSafe().at("Plans").GetArraySafe()) {
+            ComputeTotalRows(p);
+        }
+    }
+
+    if (plan.GetMapSafe().contains("Stats") && plan.GetMapSafe().contains("Operators")) {
+        auto& ops = plan.GetMapSafe().at("Operators").GetArraySafe();
+
+        const auto& stats = plan.GetMapSafe().at("Stats").GetMapSafe();
+
+        if (stats.contains("OutputRows")) {
+            auto outputRows = stats.at("OutputRows");
+            double nRows;
+            if (outputRows.IsMap()) {
+                nRows = outputRows.GetMapSafe().at("Sum").GetDouble();
+            } else {
+                nRows = outputRows.GetDouble();
+            }
+            ops[0]["A-Rows"] = nRows;
+        }
+    }
+}
+
+void RemoveStats(NJson::TJsonValue& plan) {
+
+    if (plan.GetMapSafe().contains("Plans")) {
+        for (auto& p : plan.GetMapSafe().at("Plans").GetArraySafe()) {
+            RemoveStats(p);
+        }
+    }
+
+    if (plan.GetMapSafe().contains("Stats")) {
+        plan.GetMapSafe().erase("Stats");
+    }
+}
+
 NJson::TJsonValue SimplifyQueryPlan(NJson::TJsonValue& plan) {
      static const THashSet<TString> redundantNodes = {
        "UnionAll",
@@ -2170,6 +2207,8 @@ NJson::TJsonValue SimplifyQueryPlan(NJson::TJsonValue& plan) {
     plan = ReconstructQueryPlanRec(plan, 0, planIndex, precomputes, nodeCounter);
     RemoveRedundantNodes(plan, redundantNodes);
     ComputeCpuTimes(plan);
+    ComputeTotalRows(plan);
+    RemoveStats(plan);
 
     return plan;
 }
diff --git a/ydb/public/lib/ydb_cli/common/format.cpp b/ydb/public/lib/ydb_cli/common/format.cpp
index eefed54abf..bd64efeb14 100644
--- a/ydb/public/lib/ydb_cli/common/format.cpp
+++ b/ydb/public/lib/ydb_cli/common/format.cpp
@@ -444,26 +444,6 @@ void TQueryPlanPrinter::PrintPrettyTableImpl(const NJson::TJsonValue& plan, TStr
     const auto& node = plan.GetMapSafe();
 
     auto& newRow = table.AddRow();
-    if (AnalyzeMode) {
-        TString cpuTime;
-        TString nRows;
-
-        if (node.contains("Stats")) {
-            const auto& stats = node.at("Stats").GetMapSafe();
-
-            if (stats.contains("OutputRows")) {
-                auto outputRows = stats.at("OutputRows");
-                if (outputRows.IsMap()) {
-                    nRows = FormatPrettyTableDouble(outputRows.GetMapSafe().at("Sum").GetString());
-                } else {
-                    nRows = FormatPrettyTableDouble(outputRows.GetString());
-                }
-            }
-        }
-
-        newRow.Column(1, std::move(cpuTime));
-        newRow.Column(2, std::move(nRows));
-    }
 
     NColorizer::TColors colors = NColorizer::AutoColors(Cout);
     TStringBuf color;
@@ -486,13 +466,16 @@ void TQueryPlanPrinter::PrintPrettyTableImpl(const NJson::TJsonValue& plan, TStr
         for (const auto& op : node.at("Operators").GetArraySafe()) {
             TVector<TString> info;
             TString aCpu;
+            TString aRows;
             TString eCost;
             TString eRows;
             TString eSize;
 
             for (const auto& [key, value] : op.GetMapSafe()) {
                 if (key == "A-Cpu") {
-                    aCpu = FormatPrettyTableDouble(value.GetString());
+                    aCpu = FormatPrettyTableDouble(std::to_string(value.GetDouble()));
+                } else if (key == "A-Rows") {
+                    aRows = FormatPrettyTableDouble(std::to_string(value.GetDouble()));
                 } else if (key == "E-Cost") {
                     eCost = FormatPrettyTableDouble(value.GetString());
                 } else if (key == "E-Rows") {
@@ -524,6 +507,7 @@ void TQueryPlanPrinter::PrintPrettyTableImpl(const NJson::TJsonValue& plan, TStr
             newRow.Column(0, std::move(operation));
             if (AnalyzeMode) {
                 newRow.Column(1, std::move(aCpu));
+                newRow.Column(2, std::move(aRows));
                 newRow.Column(3, std::move(eCost));
                 newRow.Column(4, std::move(eRows));
                 newRow.Column(5, std::move(eSize));