Add tests and fixes for generic cbo

7 files changed, 93 insertions, 11 deletions

diff --git a/ydb/library/yql/dq/opt/dq_cbo_ut.cpp b/ydb/library/yql/dq/opt/dq_cbo_ut.cpp
index aa640297e0d..8d977162528 100644
--- a/ydb/library/yql/dq/opt/dq_cbo_ut.cpp
+++ b/ydb/library/yql/dq/opt/dq_cbo_ut.cpp
@@ -2,6 +2,7 @@
 #include <library/cpp/testing/hook/hook.h>
 #include <ydb/library/yql/core/yql_type_annotation.h>
 #include <ydb/library/yql/providers/common/provider/yql_provider.h>
+#include <ydb/library/yql/parser/pg_wrapper/interface/optimizer.h>
 
 #include "dq_opt_log.h"
 #include "dq_opt_join.h"
@@ -10,6 +11,24 @@ using namespace NYql;
 using namespace NNodes;
 using namespace NYql::NDq;
 
+namespace {
+
+TExprNode::TPtr MakeLabel(TExprContext& ctx, const std::vector<TStringBuf>& vars) {
+    TVector<TExprNodePtr> label; label.reserve(vars.size());
+
+    auto pos = ctx.AppendPosition({});
+    for (auto var : vars) {
+        label.emplace_back(ctx.NewAtom(pos, var));
+    }
+
+    return Build<TCoAtomList>(ctx, pos)
+            .Add(label)
+            .Done()
+        .Ptr();
+}
+
+} // namespace
+
 Y_UNIT_TEST_SUITE(DQCBO) {
 
 Y_UNIT_TEST(Empty) {
@@ -68,5 +87,62 @@ Y_UNIT_TEST(RelCollectorBrokenEquiJoin) {
     UNIT_ASSERT(DqCollectJoinRelationsWithStats(typeCtx, equiJoin, [&](auto, auto) {}) == false);
 }
 
+void _DqOptimizeEquiJoinWithCosts(const std::function<IOptimizer*(IOptimizer::TInput&&)>& optFactory) {
+    TTypeAnnotationContext typeCtx;
+    TExprContext ctx;
+    auto pos = ctx.AppendPosition({});
+    TVector<TExprBase> joinArgs;
+    TVector<TExprBase> tables;
+    tables.emplace_back(Build<TCoEquiJoinInput>(ctx, pos).List(Build<TCoAtomList>(ctx, pos).Done().Ptr()).Scope(ctx.NewAtom(pos, "orders")).Done());
+    tables.emplace_back(Build<TCoEquiJoinInput>(ctx, pos).List(Build<TCoAtomList>(ctx, pos).Done().Ptr()).Scope(ctx.NewAtom(pos, "customer")).Done());
+
+    auto settings = Build<TCoAtomList>(ctx, pos).Done().Ptr();
+
+    auto joinTree = Build<TCoEquiJoinTuple>(ctx, pos)
+        .Type(ctx.NewAtom(pos, "Inner"))
+        .LeftScope(ctx.NewAtom(pos, "orders"))
+        .RightScope(ctx.NewAtom(pos, "customer"))
+        .LeftKeys(MakeLabel(ctx, {"orders", "a"}))
+        .RightKeys(MakeLabel(ctx, {"customer", "b"}))
+        .Options(settings)
+        .Done().Ptr();
+
+    joinArgs.insert(joinArgs.end(), tables.begin(), tables.end());
+    joinArgs.emplace_back(joinTree);
+    joinArgs.emplace_back(settings);
+
+    typeCtx.StatisticsMap[tables[0].Ptr()->Child(0)] = std::make_shared<TOptimizerStatistics>(1, 1, 1);
+    typeCtx.StatisticsMap[tables[1].Ptr()->Child(0)] = std::make_shared<TOptimizerStatistics>(1, 1, 1);
+
+    TCoEquiJoin equiJoin = Build<TCoEquiJoin>(ctx, pos)
+        .Add(joinArgs)
+        .Done();
+
+    auto res = DqOptimizeEquiJoinWithCosts(equiJoin, ctx, typeCtx, optFactory, true);
+    UNIT_ASSERT(equiJoin.Ptr() != res.Ptr());
+    UNIT_ASSERT(equiJoin.Ptr()->ChildrenSize() == res.Ptr()->ChildrenSize());
+    UNIT_ASSERT(equiJoin.Maybe<TCoEquiJoin>());
+}
+
+Y_UNIT_TEST(DqOptimizeEquiJoinWithCostsNative) {
+    std::function<void(const TString&)> log = [&](auto str) {
+        Cerr << str;
+    };
+    std::function<IOptimizer*(IOptimizer::TInput&&)> optFactory = [&](auto input) {
+        return MakeNativeOptimizer(input, log);
+    };
+    _DqOptimizeEquiJoinWithCosts(optFactory);
+}
+
+Y_UNIT_TEST(DqOptimizeEquiJoinWithCostsPG) {
+    std::function<void(const TString&)> log = [&](auto str) {
+        Cerr << str;
+    };
+    std::function<IOptimizer*(IOptimizer::TInput&&)> optFactory = [&](auto input) {
+        return MakePgOptimizer(input, log);
+    };
+    _DqOptimizeEquiJoinWithCosts(optFactory);
+}
+
 } // DQCBO
 
diff --git a/ydb/library/yql/dq/opt/dq_opt_join_cost_based_generic.cpp b/ydb/library/yql/dq/opt/dq_opt_join_cost_based_generic.cpp
index 95459586414..ee4e100a1ca 100644
--- a/ydb/library/yql/dq/opt/dq_opt_join_cost_based_generic.cpp
+++ b/ydb/library/yql/dq/opt/dq_opt_join_cost_based_generic.cpp
@@ -14,10 +14,10 @@ namespace {
 struct TState {
 	IOptimizer::TInput Input;
     IOptimizer::TOutput Result;
-    std::vector<TStringBuf> Tables; // relId -> table
-    std::vector<THashMap<TStringBuf, int>> VarIds; // relId -> varsIds
-    THashMap<TStringBuf, std::vector<int>> Table2RelIds;
-	std::vector<std::vector<std::tuple<TStringBuf, TStringBuf>>> Var2TableCol; // relId, varId -> table, col
+    std::vector<TString> Tables; // relId -> table
+    std::vector<THashMap<TString, int>> VarIds; // relId -> varsIds
+    THashMap<TString, std::vector<int>> Table2RelIds;
+    std::vector<std::vector<std::tuple<TStringBuf, TStringBuf>>> Var2TableCol; // relId, varId -> table, col
     TPositionHandle Pos;
 
     TState(const TCoEquiJoin& join)
@@ -38,9 +38,10 @@ struct TState {
         return varId;
     }
 
-    void CollectRel(TStringBuf label, auto stat) {
+    void CollectRel(const TString& label, auto stat) {
         Input.Rels.emplace_back();
-		Var2TableCol.emplace_back();
+        Var2TableCol.emplace_back();
+        VarIds.emplace_back();
         int relId = Input.Rels.size();
         Input.Rels.back().Rows = stat->Nrows;
         Input.Rels.back().TotalCost = *stat->Cost;
diff --git a/ydb/library/yql/dq/opt/ut/CMakeLists.darwin-x86_64.txt b/ydb/library/yql/dq/opt/ut/CMakeLists.darwin-x86_64.txt
index ccb809ca9eb..692c17ce3d0 100644
--- a/ydb/library/yql/dq/opt/ut/CMakeLists.darwin-x86_64.txt
+++ b/ydb/library/yql/dq/opt/ut/CMakeLists.darwin-x86_64.txt
@@ -20,7 +20,8 @@ target_link_libraries(ydb-library-yql-dq-opt-ut PUBLIC
   library-cpp-cpuid_check
   cpp-testing-unittest_main
   yql-dq-opt
-  yql-sql-pg_dummy
+  yql-sql-pg
+  yql-parser-pg_wrapper
   udf-service-stub
 )
 target_link_options(ydb-library-yql-dq-opt-ut PRIVATE
diff --git a/ydb/library/yql/dq/opt/ut/CMakeLists.linux-aarch64.txt b/ydb/library/yql/dq/opt/ut/CMakeLists.linux-aarch64.txt
index efd7e7e6dcb..07b6d1728ab 100644
--- a/ydb/library/yql/dq/opt/ut/CMakeLists.linux-aarch64.txt
+++ b/ydb/library/yql/dq/opt/ut/CMakeLists.linux-aarch64.txt
@@ -20,7 +20,8 @@ target_link_libraries(ydb-library-yql-dq-opt-ut PUBLIC
   yutil
   cpp-testing-unittest_main
   yql-dq-opt
-  yql-sql-pg_dummy
+  yql-sql-pg
+  yql-parser-pg_wrapper
   udf-service-stub
 )
 target_link_options(ydb-library-yql-dq-opt-ut PRIVATE
diff --git a/ydb/library/yql/dq/opt/ut/CMakeLists.linux-x86_64.txt b/ydb/library/yql/dq/opt/ut/CMakeLists.linux-x86_64.txt
index 11f64a2b30a..93b94a42a7c 100644
--- a/ydb/library/yql/dq/opt/ut/CMakeLists.linux-x86_64.txt
+++ b/ydb/library/yql/dq/opt/ut/CMakeLists.linux-x86_64.txt
@@ -21,7 +21,8 @@ target_link_libraries(ydb-library-yql-dq-opt-ut PUBLIC
   library-cpp-cpuid_check
   cpp-testing-unittest_main
   yql-dq-opt
-  yql-sql-pg_dummy
+  yql-sql-pg
+  yql-parser-pg_wrapper
   udf-service-stub
 )
 target_link_options(ydb-library-yql-dq-opt-ut PRIVATE
diff --git a/ydb/library/yql/dq/opt/ut/CMakeLists.windows-x86_64.txt b/ydb/library/yql/dq/opt/ut/CMakeLists.windows-x86_64.txt
index a3566121c53..b024a91e985 100644
--- a/ydb/library/yql/dq/opt/ut/CMakeLists.windows-x86_64.txt
+++ b/ydb/library/yql/dq/opt/ut/CMakeLists.windows-x86_64.txt
@@ -20,7 +20,8 @@ target_link_libraries(ydb-library-yql-dq-opt-ut PUBLIC
   library-cpp-cpuid_check
   cpp-testing-unittest_main
   yql-dq-opt
-  yql-sql-pg_dummy
+  yql-sql-pg
+  yql-parser-pg_wrapper
   udf-service-stub
 )
 target_sources(ydb-library-yql-dq-opt-ut PRIVATE
diff --git a/ydb/library/yql/dq/opt/ut/ya.make b/ydb/library/yql/dq/opt/ut/ya.make
index 3805d49de6e..e240034eaef 100644
--- a/ydb/library/yql/dq/opt/ut/ya.make
+++ b/ydb/library/yql/dq/opt/ut/ya.make
@@ -6,7 +6,8 @@ SRCS(
 
 PEERDIR(
     ydb/library/yql/dq/opt
-    ydb/library/yql/sql/pg_dummy
+    ydb/library/yql/sql/pg
+    ydb/library/yql/parser/pg_wrapper
     ydb/library/yql/public/udf/service/stub
 )