KIKIMR-19451: add simple query plan format

feat(cli): simplify query plan

4 files changed, 239 insertions, 4 deletions

diff --git a/ydb/public/lib/ydb_cli/commands/ydb_service_table.cpp b/ydb/public/lib/ydb_cli/commands/ydb_service_table.cpp
index 818bc2620c9..fe5dfbabedd 100644
--- a/ydb/public/lib/ydb_cli/commands/ydb_service_table.cpp
+++ b/ydb/public/lib/ydb_cli/commands/ydb_service_table.cpp
@@ -693,6 +693,7 @@ void TCommandExplain::Config(TConfig& config) {
 
     AddFormats(config, {
             EOutputFormat::Pretty,
+            EOutputFormat::PrettyTable,
             EOutputFormat::JsonUnicode,
             EOutputFormat::JsonBase64
     });
diff --git a/ydb/public/lib/ydb_cli/common/format.cpp b/ydb/public/lib/ydb_cli/common/format.cpp
index 5ab81206435..56e39258fc1 100644
--- a/ydb/public/lib/ydb_cli/common/format.cpp
+++ b/ydb/public/lib/ydb_cli/common/format.cpp
@@ -1,5 +1,4 @@
 #include "format.h"
-#include "pretty_table.h"
 
 #include <util/string/vector.h>
 #include <library/cpp/json/json_prettifier.h>
@@ -30,6 +29,7 @@ namespace {
 
     THashMap<EOutputFormat, TString> FormatDescriptions = {
         { EOutputFormat::Pretty, "Human readable output" },
+        { EOutputFormat::PrettyTable, "Human readable table output" },
         { EOutputFormat::Json, "Output in json format" },
         { EOutputFormat::JsonUnicode, "Output in json format, binary strings are encoded with unicode characters. "
                                       "Every row is a separate json on a separate line." },
@@ -232,7 +232,8 @@ void TCommandWithFormat::ParseMessagingFormats() {
 void TQueryPlanPrinter::Print(const TString& plan) {
     switch (Format) {
         case EOutputFormat::Default:
-        case EOutputFormat::Pretty: {
+        case EOutputFormat::Pretty:
+        case EOutputFormat::PrettyTable: {
             NJson::TJsonValue planJson;
             NJson::ReadJsonTree(plan, &planJson, true);
 
@@ -246,10 +247,19 @@ void TQueryPlanPrinter::Print(const TString& plan) {
                 for (size_t queryId = 0; queryId < queries.size(); ++queryId) {
                     const auto& query = queries[queryId];
                     Output << "Query " << queryId << ":" << Endl;
-                    PrintPretty(query);
+
+                    if (Format == EOutputFormat::PrettyTable) {
+                        PrintPrettyTable(query);
+                    } else {
+                        PrintPretty(query);
+                    }
                 }
             } else {
-                PrintPretty(planJson);
+                if (Format == EOutputFormat::PrettyTable) {
+                    PrintPrettyTable(planJson);
+                } else {
+                    PrintPretty(planJson);
+                }
             }
 
             break;
@@ -350,6 +360,220 @@ void TQueryPlanPrinter::PrintPrettyImpl(const NJson::TJsonValue& plan, TVector<T
     }
 }
 
+void TQueryPlanPrinter::PrintPrettyTable(const NJson::TJsonValue& plan) {
+    static const TVector<TString> explainColumnNames = {"Operation", "E-Cost", "E-Rows"};
+    static const TVector<TString> explainAnalyzeColumnNames = {"Operation", "DurationMs", "Rows"};
+
+    if (plan.GetMapSafe().contains("Plan")) {
+        auto queryPlan = plan.GetMapSafe().at("Plan");
+        SimplifyQueryPlan(queryPlan);
+
+        TPrettyTable table(AnalyzeMode ? explainAnalyzeColumnNames : explainColumnNames,
+            TPrettyTableConfig().WithoutRowDelimiters());
+
+        Y_ENSURE(queryPlan.GetMapSafe().contains("Plans"));
+
+        TString offset;
+        for (auto subplan : queryPlan.GetMapSafe().at("Plans").GetArraySafe()) {
+            PrintPrettyTableImpl(subplan, offset, table);
+        }
+
+        Output << table;
+    } else { /* old format plan */
+        PrintJson(plan.GetStringRobust());
+    }
+}
+
+void TQueryPlanPrinter::PrintPrettyTableImpl(const NJson::TJsonValue& plan, TString& offset, TPrettyTable& table) {
+    const auto& node = plan.GetMapSafe();
+
+    auto& newRow = table.AddRow();
+    if (AnalyzeMode) {
+        TString duration;
+        TString nRows;
+
+        if (node.contains("Stats")) {
+            const auto& stats = node.at("Stats").GetMapSafe();
+
+            if (stats.contains("TotalOutputRows")) {
+                nRows = JsonToString(stats.at("TotalOutputRows"));
+            }
+
+            if (stats.contains("TotalDurationMs")) {
+                duration = JsonToString(stats.at("TotalDurationMs"));
+            }
+        }
+
+        newRow.Column(1, std::move(duration));
+        newRow.Column(2, std::move(nRows));
+    }
+
+
+    if (node.contains("Operators")) {
+        for (const auto& op : node.at("Operators").GetArraySafe()) {
+            TVector<TString> info;
+            TString eCost;
+            TString eRows;
+
+            for (const auto& [key, value] : op.GetMapSafe()) {
+                if (key == "E-Cost") {
+                    eCost = JsonToString(value);
+                } else if (key == "E-Rows") {
+                    eRows = JsonToString(value);
+                } else if (key != "Name") {
+                    info.emplace_back(TStringBuilder() << key << ": " << JsonToString(value));
+                }
+            }
+
+            TStringBuilder operation;
+            if (info.empty()) {
+                operation << offset << " -> " << op.GetMapSafe().at("Name").GetString();
+            } else {
+                operation << offset << " -> " << op.GetMapSafe().at("Name").GetString()
+                     << " (" << JoinStrings(info, ", ") << ")";
+            }
+
+            auto& row = table.AddRow();
+            row.Column(0, std::move(operation));
+
+            if (!AnalyzeMode) {
+                row.Column(1, std::move(eCost));
+                row.Column(2, std::move(eRows));
+            }
+        }
+    } else {
+        TStringBuilder operation;
+        operation << offset << " -> " << node.at("Node Type").GetString();
+        newRow.Column(0, std::move(operation));
+    }
+
+    if (node.contains("Plans")) {
+        auto& plans = node.at("Plans").GetArraySafe();
+        for (auto subplan : plans) {
+            offset += "  ";
+            PrintPrettyTableImpl(subplan, offset, table);
+            offset.resize(offset.size() - 2);
+        }
+    }
+}
+
+TVector<NJson::TJsonValue> TQueryPlanPrinter::RemoveRedundantNodes(NJson::TJsonValue& plan, const THashSet<TString>& redundantNodes) {
+    auto& planMap = plan.GetMapSafe();
+
+    TVector<NJson::TJsonValue> children;
+    if (planMap.contains("Plans")) {
+        for (auto& child : planMap.at("Plans").GetArraySafe()) {
+            auto newChildren = RemoveRedundantNodes(child, redundantNodes);
+            children.insert(children.end(), newChildren.begin(), newChildren.end());
+        }
+    }
+
+    planMap.erase("Plans");
+    if (!children.empty()) {
+        auto& plans = planMap["Plans"];
+        for (auto& child :  children) {
+            plans.AppendValue(child);
+        }
+    }
+
+    const auto typeName = planMap.at("Node Type").GetStringSafe();
+    if (redundantNodes.contains(typeName)) {
+        return children;
+    }
+
+    return {plan};
+}
+
+THashMap<TString, NJson::TJsonValue> TQueryPlanPrinter::ExtractPrecomputes(NJson::TJsonValue& plan) {
+    auto& planMap = plan.GetMapSafe();
+
+    if (!planMap.contains("Plans")) {
+        return {};
+    }
+
+    THashMap<TString, NJson::TJsonValue> precomputes;
+    TVector<NJson::TJsonValue> results;
+    for (auto& child : planMap.at("Plans").GetArraySafe()) {
+        if (child.GetMapSafe().contains("Subplan Name")) {
+            const auto& precomputeName = child.GetMapSafe().at("Subplan Name").GetStringSafe();
+
+            auto pos = precomputeName.find("precompute");
+            if (pos != TString::npos) {
+                precomputes[precomputeName.substr(pos)] = std::move(child);
+            }
+        } else {
+            results.push_back(std::move(child));
+        }
+    }
+
+    planMap.erase("Plans");
+    if (!results.empty()) {
+        auto& plans = planMap["Plans"];
+        for (auto& result : results) {
+            plans.AppendValue(std::move(result));
+        }
+    }
+
+    return precomputes;
+}
+
+void TQueryPlanPrinter::ResolvePrecomputeLinks(NJson::TJsonValue& plan, const THashMap<TString, NJson::TJsonValue>& precomputes) {
+    auto& planMap = plan.GetMapSafe();
+
+    if (planMap.contains("CTE Name")) {
+        const auto& precomputeName = planMap.at("CTE Name").GetStringSafe();
+
+        auto precomputeIt = precomputes.find(precomputeName);
+        if (precomputeIt != precomputes.end()) {
+            auto& precompute = precomputeIt->second;
+
+            if (precompute.GetMapSafe().contains("Plans")) {
+                auto& plans = planMap["Plans"];
+                for (auto& child : precompute.GetMapSafe().at("Plans").GetArraySafe()) {
+                    plans.AppendValue(child);
+                }
+            }
+
+            if (planMap.contains("Operators")) {
+                // delete precompute link from operator block
+                for (auto &op : planMap.at("Operators").GetArraySafe()) {
+                    if (op.GetMapSafe().contains("Iterator")) {
+                        op.GetMapSafe().erase("Iterator");
+                    } else if (op.GetMapSafe().contains("Input")) {
+                        op.GetMapSafe().erase("Input");
+                    }
+                }
+            }
+
+        }
+
+        // delete precompute link
+        planMap.erase("CTE Name");
+    }
+
+    if (planMap.contains("Plans")) {
+        for (auto& child : planMap.at("Plans").GetArraySafe()) {
+            ResolvePrecomputeLinks(child, precomputes);
+        }
+    }
+}
+
+void TQueryPlanPrinter::SimplifyQueryPlan(NJson::TJsonValue& plan) {
+    static const THashSet<TString> redundantNodes = {
+        "UnionAll",
+        "Broadcast",
+        "Map",
+        "HashShuffle",
+        "Merge",
+        "Collect",
+        "Stage"
+    };
+
+    RemoveRedundantNodes(plan, redundantNodes);
+    auto precomputes = ExtractPrecomputes(plan);
+    ResolvePrecomputeLinks(plan, precomputes);
+}
+
 TString TQueryPlanPrinter::JsonToString(const NJson::TJsonValue& jsonValue) {
     return jsonValue.GetStringRobust();
 }
diff --git a/ydb/public/lib/ydb_cli/common/format.h b/ydb/public/lib/ydb_cli/common/format.h
index 7782d36fe77..18fd898d1d4 100644
--- a/ydb/public/lib/ydb_cli/common/format.h
+++ b/ydb/public/lib/ydb_cli/common/format.h
@@ -2,6 +2,7 @@
 
 #include "command.h"
 #include "formats.h"
+#include "pretty_table.h"
 
 #include <ydb/public/lib/json_value/ydb_json_value.h>
 #include <ydb/public/sdk/cpp/client/ydb_result/result.h>
@@ -94,9 +95,17 @@ public:
 private:
     void PrintPretty(const NJson::TJsonValue& plan);
     void PrintPrettyImpl(const NJson::TJsonValue& plan, TVector<TString>& offsets);
+    void PrintPrettyTable(const NJson::TJsonValue& plan);
+    void PrintPrettyTableImpl(const NJson::TJsonValue& plan, TString& offset, TPrettyTable& table);
     void PrintJson(const TString& plan);
     TString JsonToString(const NJson::TJsonValue& jsonValue);
 
+    void SimplifyQueryPlan(NJson::TJsonValue& plan);
+    TVector<NJson::TJsonValue> RemoveRedundantNodes(NJson::TJsonValue& plan, const THashSet<TString>& redundantNodes);
+    THashMap<TString, NJson::TJsonValue> ExtractPrecomputes(NJson::TJsonValue& planJson);
+    void ResolvePrecomputeLinks(NJson::TJsonValue& planJson, const THashMap<TString, NJson::TJsonValue>& precomputes);
+
+private:
     EOutputFormat Format;
     bool AnalyzeMode;
     IOutputStream& Output;
diff --git a/ydb/public/lib/ydb_cli/common/formats.h b/ydb/public/lib/ydb_cli/common/formats.h
index 25bb08421c2..d7aef93699c 100644
--- a/ydb/public/lib/ydb_cli/common/formats.h
+++ b/ydb/public/lib/ydb_cli/common/formats.h
@@ -7,6 +7,7 @@ namespace NConsoleClient {
 enum class EOutputFormat {
     Default /* "default" */,
     Pretty /* "pretty" */,
+    PrettyTable /* "pretty-table" */,
     Json /* "json" */,
     JsonUnicode /* "json-unicode" */,
     JsonUnicodeArray /* "json-unicode-array" */,