Can print CBO input

3 files changed, 71 insertions, 5 deletions

diff --git a/ydb/library/yql/core/cbo/cbo_optimizer.cpp b/ydb/library/yql/core/cbo/cbo_optimizer.cpp
index 3a755b45445..d14f5a6f874 100644
--- a/ydb/library/yql/core/cbo/cbo_optimizer.cpp
+++ b/ydb/library/yql/core/cbo/cbo_optimizer.cpp
@@ -23,10 +23,10 @@ TStringBuf Prefix(int level) {
         : Prefixes.back();
 }
 
-void PrettyPrintVar(TStringBuilder& b, const IOptimizer::TOutput& output, IOptimizer::TVarId varId) {
+void PrettyPrintVar(TStringBuilder& b, const IOptimizer::TInput* input, IOptimizer::TVarId varId) {
     const auto& [relno, varno] = varId;
-    auto varName = output.Input
-        ? output.Input->Rels[relno-1].TargetVars[varno-1].Name
+    auto varName = input
+        ? input->Rels[relno-1].TargetVars[varno-1].Name
         : '\0';
     if (!varName) {
         b << "(" << relno << "," << varno << ")";
@@ -66,9 +66,9 @@ void PrettyPrintNode(int level, TStringBuilder& b, const IOptimizer::TOutput& ou
 
         if (!isEmpty(node.LeftVar) && !isEmpty(node.RightVar)) {
             b << prefix << " Op: ";
-            PrettyPrintVar(b, output, node.LeftVar);
+            PrettyPrintVar(b, output.Input, node.LeftVar);
             b << " = ";
-            PrettyPrintVar(b, output, node.RightVar);
+            PrettyPrintVar(b, output.Input, node.RightVar);
             b << "\n";
         }
     }
@@ -85,6 +85,22 @@ void PrettyPrintNode(int level, TStringBuilder& b, const IOptimizer::TOutput& ou
     }
 }
 
+void PrettyPrintRel(TStringBuilder& b, const IOptimizer::TInput* input, const auto& relId) {
+    const auto& rel = input->Rels[relId - 1];
+    b << "{";
+    b << "rows: " << rel.Rows << ",";
+    b << "cost: " << rel.TotalCost << ",";
+    b << "vars: [";
+    for (ui32 i = 0; i < rel.TargetVars.size(); i++) {
+        PrettyPrintVar(b, input, {relId, i + 1});
+        if (i != rel.TargetVars.size() - 1) {
+            b << ", ";
+        }
+    }
+    b << "]";
+    b << "}";
+}
+
 } // namespace
 
 TString IOptimizer::TOutput::ToString() const {
@@ -95,4 +111,32 @@ TString IOptimizer::TOutput::ToString() const {
     return b;
 }
 
+TString IOptimizer::TInput::ToString() {
+    TStringBuilder b;
+    b << "Rels: [";
+    for (ui32 i = 0; i < Rels.size(); ++i) {
+        PrettyPrintRel(b, this, i + 1);
+        if (i != Rels.size() - 1) {
+            b << ",\n";
+        }
+    }
+    b << "]\n";
+    b << "EqClasses: [";
+    for (ui32 i = 0; i < EqClasses.size(); ++i) {
+        b << "[";
+        for (ui32 j = 0; j < EqClasses[i].Vars.size(); ++j) {
+            PrettyPrintVar(b, this, EqClasses[i].Vars[j]);
+            if (j != EqClasses[i].Vars.size() - 1) {
+                b << ",";
+            }
+        }
+        b << "]";
+        if (i != EqClasses.size() - 1) {
+            b << ",";
+        }
+    }
+    b << "]\n";
+    return b;
+}
+
 } // namespace NYql
diff --git a/ydb/library/yql/core/cbo/cbo_optimizer.h b/ydb/library/yql/core/cbo/cbo_optimizer.h
index 982dd1790a4..ad6db62ab89 100644
--- a/ydb/library/yql/core/cbo/cbo_optimizer.h
+++ b/ydb/library/yql/core/cbo/cbo_optimizer.h
@@ -28,6 +28,8 @@ struct IOptimizer {
     struct TInput {
         std::vector<TRel> Rels;
         std::vector<TEq> EqClasses;
+
+        TString ToString();
     };
 
     enum class EJoinType {
diff --git a/ydb/library/yql/sql/pg/optimizer_ut.cpp b/ydb/library/yql/sql/pg/optimizer_ut.cpp
index 0b814718c22..2cff43bb215 100644
--- a/ydb/library/yql/sql/pg/optimizer_ut.cpp
+++ b/ydb/library/yql/sql/pg/optimizer_ut.cpp
@@ -97,4 +97,24 @@ Y_UNIT_TEST(PgJoinSearch3Rels) {
     UNIT_ASSERT_STRINGS_EQUAL(expected, resStr);
 }
 
+Y_UNIT_TEST(InputToString) {
+    IOptimizer::TRel rel1 = {100000, 1000000, {{'a'}}};
+    IOptimizer::TRel rel2 = {1000000, 9000009, {{'b'}}};
+    IOptimizer::TRel rel3 = {10000, 9009, {{'c'}}};
+    IOptimizer::TInput input = {{rel1, rel2, rel3}};
+
+    input.EqClasses.emplace_back(IOptimizer::TEq {
+        {{1, 1}, {2, 1}, {3, 1}}
+    });
+
+    auto str = input.ToString();
+
+    TString expected = R"__(Rels: [{rows: 100000,cost: 1000000,vars: [a]},
+{rows: 1000000,cost: 9000009,vars: [b]},
+{rows: 10000,cost: 9009,vars: [c]}]
+EqClasses: [[a,b,c]]
+)__";
+    UNIT_ASSERT_STRINGS_EQUAL(expected, str);
+}
+
 } // Y_UNIT_TEST_SUITE(PgOptimizer)