Fix scope calculation for native cbo output + tests

4 files changed, 103 insertions, 18 deletions

diff --git a/ydb/library/yql/core/cbo/cbo_optimizer.cpp b/ydb/library/yql/core/cbo/cbo_optimizer.cpp
index f3963ade7ad..c0717782728 100644
--- a/ydb/library/yql/core/cbo/cbo_optimizer.cpp
+++ b/ydb/library/yql/core/cbo/cbo_optimizer.cpp
@@ -113,13 +113,15 @@ void PrettyPrintRel(TStringBuilder& b, const IOptimizer::TInput* input, const au
 
 } // namespace
 
-TString IOptimizer::TOutput::ToString() const {
+TString IOptimizer::TOutput::ToString(bool printCost) const {
     TStringBuilder b;
-    char buf[1024];
-    snprintf(buf, sizeof(buf), "%.2lf", Rows);
-    b << "Rows: " << buf << "\n";
-    snprintf(buf, sizeof(buf), "%.2lf", TotalCost);
-    b << "TotalCost: " << buf << "\n";
+    if (printCost) {
+        char buf[1024];
+        snprintf(buf, sizeof(buf), "%.2lf", Rows);
+        b << "Rows: " << buf << "\n";
+        snprintf(buf, sizeof(buf), "%.2lf", TotalCost);
+        b << "TotalCost: " << buf << "\n";
+    }
     b << "{\n";
     if (!Nodes.empty()) {
         PrettyPrintNode(0, b, *this, 0);
diff --git a/ydb/library/yql/core/cbo/cbo_optimizer.h b/ydb/library/yql/core/cbo/cbo_optimizer.h
index 047a7050f50..ea428a8fd28 100644
--- a/ydb/library/yql/core/cbo/cbo_optimizer.h
+++ b/ydb/library/yql/core/cbo/cbo_optimizer.h
@@ -65,7 +65,7 @@ struct IOptimizer {
         double Rows = 0;
         double TotalCost = 0;
 
-        TString ToString() const;
+        TString ToString(bool printCost = true) const;
     };
 
     virtual ~IOptimizer() = default;
diff --git a/ydb/library/yql/dq/opt/dq_cbo_ut.cpp b/ydb/library/yql/dq/opt/dq_cbo_ut.cpp
index 8d977162528..f8d0c2d5ad9 100644
--- a/ydb/library/yql/dq/opt/dq_cbo_ut.cpp
+++ b/ydb/library/yql/dq/opt/dq_cbo_ut.cpp
@@ -36,6 +36,89 @@ Y_UNIT_TEST(Empty) {
     std::unique_ptr<IOptimizer> optimizer = std::unique_ptr<IOptimizer>(NDq::MakeNativeOptimizer(input, {}));
 }
 
+Y_UNIT_TEST(JoinSearch2Rels) {
+    IOptimizer::TRel rel1 = {100000, 1000000, {{'a'}}};
+    IOptimizer::TRel rel2 = {1000000, 9000009, {{'b'}}};
+    IOptimizer::TInput input = {{rel1, rel2}, {}, {}, {}};
+
+    input.EqClasses.emplace_back(IOptimizer::TEq {
+        {{1, 1}, {2, 1}}
+    });
+
+    auto log = [](const TString& str) {
+        Cerr << str << "\n";
+    };
+
+    auto optimizer = std::unique_ptr<IOptimizer>(MakeNativeOptimizer(input, log));
+
+    auto res = optimizer->JoinSearch();
+    UNIT_ASSERT(res.Rows > 0);
+    UNIT_ASSERT(res.TotalCost > 0);
+    auto resStr = res.ToString(false);
+    Cerr << resStr;
+    TString expected = R"__({
+ Inner Join
+ Rels: [2,1]
+ Op: b = a
+ {
+  Node
+  Rels: [2]
+ }
+ {
+  Node
+  Rels: [1]
+ }
+}
+)__";
+    UNIT_ASSERT_STRINGS_EQUAL(expected, resStr);
+}
+
+Y_UNIT_TEST(JoinSearch3Rels) {
+    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 log = [](const TString& str) {
+        Cerr << str << "\n";
+    };
+
+    auto optimizer = std::unique_ptr<IOptimizer>(MakeNativeOptimizer(input, log));
+    auto res = optimizer->JoinSearch();
+    UNIT_ASSERT(res.Rows > 0);
+    UNIT_ASSERT(res.TotalCost > 0);
+    auto resStr = res.ToString(false);
+    Cerr << resStr;
+    TString expected = R"__({
+ Inner Join
+ Rels: [1,3,2]
+ Op: a = b
+ {
+  Inner Join
+  Rels: [1,3]
+  Op: a = c
+  {
+   Node
+   Rels: [1]
+  }
+  {
+   Node
+   Rels: [3]
+  }
+ }
+ {
+  Node
+  Rels: [2]
+ }
+}
+)__";
+	UNIT_ASSERT_STRINGS_EQUAL(expected, resStr);
+}
+
 Y_UNIT_TEST(RelCollector) {
     TExprContext ctx;
     auto pos = ctx.AppendPosition({});
diff --git a/ydb/library/yql/dq/opt/dq_opt_join_cost_based.cpp b/ydb/library/yql/dq/opt/dq_opt_join_cost_based.cpp
index 078a036a195..b9d74d18772 100644
--- a/ydb/library/yql/dq/opt/dq_opt_join_cost_based.cpp
+++ b/ydb/library/yql/dq/opt/dq_opt_join_cost_based.cpp
@@ -996,7 +996,8 @@ public:
 
         TOutput output;
         output.Input = &Input;
-        BuildOutput(&output, result.get());
+        TVector<int> scope;
+        BuildOutput(&output, result.get(), scope);
         output.Rows = result->Stats->Nrows;
         output.TotalCost = *result->Stats->Cost;
         if (Log) {
@@ -1006,24 +1007,26 @@ public:
     }
 
 private:
-    int BuildOutput(TOutput* output, IBaseOptimizerNode* node) {
+    int BuildOutput(TOutput* output, IBaseOptimizerNode* node, TVector<int>& scope) {
         int index = (int)output->Nodes.size();
         TJoinNode r = output->Nodes.emplace_back();
         switch (node->Kind) {
         case EOptimizerNodeKind::RelNodeType: {
             // leaf
             TRelOptimizerNode* n = static_cast<TRelOptimizerNode*>(node);
-            r.Rels.emplace_back(FromString<int>(n->Label));
+            int relId = FromString<int>(n->Label);
+            r.Rels.emplace_back(relId);
+            scope.emplace_back(relId);
             break;
         }
         case EOptimizerNodeKind::JoinNodeType: {
             // node
             r.Mode = IOptimizer::EJoinType::Inner;
             TJoinOptimizerNode* n = static_cast<TJoinOptimizerNode*>(node);
-            r.Outer = BuildOutput(output, n->LeftArg.get());
-            r.Inner = BuildOutput(output, n->RightArg.get());
+            int index = scope.size();
+            r.Outer = BuildOutput(output, n->LeftArg.get(), scope);
+            r.Inner = BuildOutput(output, n->RightArg.get(), scope);
 
-            std::set<int> rels;
             for (auto& [col1, col2] : n->JoinConditions) {
                 int relId1 = FromString<int>(col1.RelName);
                 int colId1 = FromString<int>(col1.AttributeName);
@@ -1032,13 +1035,10 @@ private:
 
                 r.LeftVars.emplace_back(std::make_tuple(relId1, colId1));
                 r.RightVars.emplace_back(std::make_tuple(relId2, colId2));
-
-                rels.emplace(relId1);
-                rels.emplace(relId2);
             }
 
-            r.Rels.reserve(rels.size());
-            r.Rels.insert(r.Rels.end(), rels.begin(), rels.end());
+            r.Rels.reserve(scope.size());
+            r.Rels.insert(r.Rels.end(), scope.begin() + index, scope.end());
             break;
         }
         default: