Support multi-keys in CBO

4 files changed, 59 insertions, 56 deletions

diff --git a/ydb/library/yql/core/cbo/cbo_optimizer.cpp b/ydb/library/yql/core/cbo/cbo_optimizer.cpp
index 13d0a48b1a..4c8b04ad5f 100644
--- a/ydb/library/yql/core/cbo/cbo_optimizer.cpp
+++ b/ydb/library/yql/core/cbo/cbo_optimizer.cpp
@@ -39,6 +39,15 @@ void PrettyPrintVar(TStringBuilder& b, const IOptimizer::TInput* input, IOptimiz
     }
 }
 
+void PrettyPrintVars(TStringBuilder& b, const IOptimizer::TInput* input, const std::vector<IOptimizer::TVarId>& vars) {
+    for (ui32 j = 0; j < vars.size(); ++j) {
+        PrettyPrintVar(b, input, vars[j]);
+        if (j != vars.size() - 1) {
+            b << ",";
+        }
+    } 
+}
+
 void PrettyPrintNode(int level, TStringBuilder& b, const IOptimizer::TOutput& output, int id) {
     TStringBuf prefix = Prefix(level);
     const auto& node = output.Nodes[id];
@@ -65,16 +74,11 @@ void PrettyPrintNode(int level, TStringBuilder& b, const IOptimizer::TOutput& ou
     }
 
     {
-        auto isEmpty = [](IOptimizer::TVarId id) -> bool {
-            auto& [a, b] = id;
-            return a<=0 || b<=0;
-        };
-
-        if (!isEmpty(node.LeftVar) && !isEmpty(node.RightVar)) {
+        if (!node.LeftVars.empty() && !node.RightVars.empty()) {
             b << prefix << " Op: ";
-            PrettyPrintVar(b, output.Input, node.LeftVar);
+            PrettyPrintVars(b, output.Input, node.LeftVars);
             b << " = ";
-            PrettyPrintVar(b, output.Input, node.RightVar);
+            PrettyPrintVars(b, output.Input, node.RightVars);
             b << "\n";
         }
     }
@@ -130,12 +134,7 @@ TString IOptimizer::TInput::ToString() const {
     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 << ",";
-            }
-        }
+        PrettyPrintVars(b, this, EqClasses[i].Vars);
         b << "]";
         if (i != EqClasses.size() - 1) {
             b << ",";
diff --git a/ydb/library/yql/core/cbo/cbo_optimizer.h b/ydb/library/yql/core/cbo/cbo_optimizer.h
index 56702caf3a..12afdec5c5 100644
--- a/ydb/library/yql/core/cbo/cbo_optimizer.h
+++ b/ydb/library/yql/core/cbo/cbo_optimizer.h
@@ -51,9 +51,9 @@ struct IOptimizer {
     struct TJoinNode {
         EJoinType Mode = EJoinType::Unknown;
         EJoinStrategy Strategy = EJoinStrategy::Unknown;
-        // only a = b supported yet
-        TVarId LeftVar = {};
-        TVarId RightVar = {};
+        // only a = b && c = d ...  supported yet
+        std::vector<TVarId> LeftVars = {};
+        std::vector<TVarId> RightVars = {};
         std::vector<int> Rels = {};
         int Outer = -1; // index in Nodes
         int Inner = -1; // index in Nodes
diff --git a/ydb/library/yql/providers/yt/provider/yql_yt_join_reorder.cpp b/ydb/library/yql/providers/yt/provider/yql_yt_join_reorder.cpp
index 3e4d700ea7..fef2232ea1 100644
--- a/ydb/library/yql/providers/yt/provider/yql_yt_join_reorder.cpp
+++ b/ydb/library/yql/providers/yt/provider/yql_yt_join_reorder.cpp
@@ -301,14 +301,15 @@ private:
         }
     }
 
-    TExprNode::TPtr MakeLabel(IOptimizer::TVarId var) const {
-        auto [relId, varId] = var;
-        auto [table, column] = Var2TableCol[relId - 1][varId - 1];
-
-        TVector<TExprNodePtr> label = {
-            Ctx.NewAtom(Root->JoinKind->Pos(), table),
-            Ctx.NewAtom(Root->JoinKind->Pos(), column)
-        };
+    TExprNode::TPtr MakeLabel(const std::vector<IOptimizer::TVarId>& vars) const {
+        TVector<TExprNodePtr> label; label.reserve(vars.size() * 2);
+ 
+        for (auto [relId, varId] : vars) {
+            auto [table, column] = Var2TableCol[relId - 1][varId - 1];
+
+            label.emplace_back(Ctx.NewAtom(Root->JoinKind->Pos(), table));
+            label.emplace_back(Ctx.NewAtom(Root->JoinKind->Pos(), column));
+        }
 
         return Build<TCoAtomList>(Ctx, Root->JoinKind->Pos())
                 .Add(label)
@@ -344,8 +345,8 @@ private:
                 break;
             }
             ret->JoinKind = Ctx.NewAtom(Root->JoinKind->Pos(), joinKind);
-            ret->LeftLabel = MakeLabel(node->LeftVar);
-            ret->RightLabel = MakeLabel(node->RightVar);
+            ret->LeftLabel = MakeLabel(node->LeftVars);
+            ret->RightLabel = MakeLabel(node->RightVars);
             int index = scope.size();
             ret->Left = Convert(node->Outer, scope);
             ret->Right = Convert(node->Inner, scope);
diff --git a/ydb/library/yql/sql/pg/optimizer.cpp b/ydb/library/yql/sql/pg/optimizer.cpp
index 1279ff9ed5..8a3745c056 100644
--- a/ydb/library/yql/sql/pg/optimizer.cpp
+++ b/ydb/library/yql/sql/pg/optimizer.cpp
@@ -210,39 +210,42 @@ int TPgOptimizer::MakeOutputJoin(TOutput& output, Path* path) {
             break;
         }
 
-        YQL_ENSURE(list_length(jpath->joinrestrictinfo) == 1, "Unsupported joinrestrictinfo len");
-        RestrictInfo* rinfo = (RestrictInfo*)jpath->joinrestrictinfo->elements[0].ptr_value;
-        Var* left;
-        Var* right;
-
-        if (jpath->jointype == JOIN_INNER) {
-            YQL_ENSURE(rinfo->left_em->em_expr->type == T_Var, "Unsupported left em type");
-            YQL_ENSURE(rinfo->right_em->em_expr->type == T_Var, "Unsupported right em type");
-
-            left = (Var*)rinfo->left_em->em_expr;
-            right = (Var*)rinfo->right_em->em_expr;
-        } else if (jpath->jointype == JOIN_LEFT || jpath->jointype == JOIN_RIGHT) {
-            YQL_ENSURE(rinfo->clause->type == T_OpExpr);
-            OpExpr* expr = (OpExpr*)rinfo->clause;
-            YQL_ENSURE(list_length(expr->args) == 2);
-            Expr* a1 = (Expr*)list_nth(expr->args, 0);
-            Expr* a2 = (Expr*)list_nth(expr->args, 1);
-            YQL_ENSURE(a1->type == T_Var, "Unsupported left arg type");
-            YQL_ENSURE(a2->type == T_Var, "Unsupported right arg type");
-
-            left = (Var*)a1;
-            right = (Var*)a2;
-        }
+       YQL_ENSURE(list_length(jpath->joinrestrictinfo) >= 1, "Unsupported joinrestrictinfo len");
+
+        for (int i = 0; i < list_length(jpath->joinrestrictinfo); i++) {
+            RestrictInfo* rinfo = (RestrictInfo*)jpath->joinrestrictinfo->elements[i].ptr_value;
+            Var* left;
+            Var* right;
+
+            if (jpath->jointype == JOIN_INNER) {
+                YQL_ENSURE(rinfo->left_em->em_expr->type == T_Var, "Unsupported left em type");
+                YQL_ENSURE(rinfo->right_em->em_expr->type == T_Var, "Unsupported right em type");
+
+                left = (Var*)rinfo->left_em->em_expr;
+                right = (Var*)rinfo->right_em->em_expr;
+            } else if (jpath->jointype == JOIN_LEFT || jpath->jointype == JOIN_RIGHT) {
+                YQL_ENSURE(rinfo->clause->type == T_OpExpr);
+                OpExpr* expr = (OpExpr*)rinfo->clause;
+                YQL_ENSURE(list_length(expr->args) == 2);
+                Expr* a1 = (Expr*)list_nth(expr->args, 0);
+                Expr* a2 = (Expr*)list_nth(expr->args, 1);
+                YQL_ENSURE(a1->type == T_Var, "Unsupported left arg type");
+                YQL_ENSURE(a2->type == T_Var, "Unsupported right arg type");
+
+                left = (Var*)a1;
+                right = (Var*)a2;
+            }
 
-        node.LeftVar = std::make_tuple(left->varno, left->varattno);
-        node.RightVar = std::make_tuple(right->varno, right->varattno);
+            node.LeftVars.emplace_back(std::make_tuple(left->varno, left->varattno));
+            node.RightVars.emplace_back(std::make_tuple(right->varno, right->varattno));
+
+            if (!bms_is_member(left->varno, jpath->outerjoinpath->parent->relids)) {
+                std::swap(node.LeftVars.back(), node.RightVars.back());
+            }
+        }
 
         node.Inner = MakeOutputJoin(output, jpath->innerjoinpath);
         node.Outer = MakeOutputJoin(output, jpath->outerjoinpath);
-
-        if (!bms_is_member(left->varno, jpath->outerjoinpath->parent->relids)) {
-            std::swap(node.LeftVar, node.RightVar);
-        }
     }
 
     output.Nodes[id] = node;