YT-24419 Allow multiple CBO runs on the same YTEquiJoin

commit_hash:9df70eb109c8492f701ed6f47250369dfec1adb4

6 files changed, 63 insertions, 46 deletions

diff --git a/yql/essentials/core/yql_cost_function.h b/yql/essentials/core/yql_cost_function.h
index 1831bdcc220..dbee46ad0c0 100644
--- a/yql/essentials/core/yql_cost_function.h
+++ b/yql/essentials/core/yql_cost_function.h
@@ -34,9 +34,10 @@ namespace NDq {
  * attribute name, used in join conditions
 */
 struct TJoinColumn {
-    TString RelName{};
+    TString RelName{};  // TODO: this should be a list. now list of relations is separated by comma in this string
     TString AttributeName{};
     TString AttributeNameWithAliases{};
+    std::optional<TString> OriginalRelName;
     std::optional<ui32> EquivalenceClass{};
     bool IsConstant = false;
 
diff --git a/yt/yql/providers/yt/provider/yql_yt_cbo_helpers.cpp b/yt/yql/providers/yt/provider/yql_yt_cbo_helpers.cpp
index f745ea8601f..b2e81b8106a 100644
--- a/yt/yql/providers/yt/provider/yql_yt_cbo_helpers.cpp
+++ b/yt/yql/providers/yt/provider/yql_yt_cbo_helpers.cpp
@@ -103,11 +103,13 @@ IGraphTransformer::TStatus CollectCboStatsLeaf(
     TVector<TYtPathInfo::TPtr> tables;
     if (maxChunkCountExtendedStats) {
         TVector<TString> requestedColumnList;
-        auto columnsPos = relJoinColumns.find(JoinLeafLabel(leaf.Label));
-        if (columnsPos != relJoinColumns.end()) {
-            requestedColumnList.assign(columnsPos->second.begin(), columnsPos->second.end());
+        auto labels = JoinLeafLabels(leaf.Label);
+        for (const auto& relName : labels) {
+            auto columnsPos = relJoinColumns.find(relName);
+            if (columnsPos != relJoinColumns.end()) {
+                std::copy(columnsPos->second.begin(), columnsPos->second.end(), std::back_inserter(requestedColumnList));
+            }
         }
-
         THashSet<TString> memSizeColumns(requestedColumnList.begin(), requestedColumnList.end());
         TVector<IYtGateway::TPathStatReq> pathStatReqs;
 
@@ -278,4 +280,16 @@ IGraphTransformer::TStatus CollectCboStats(const TString& cluster, TYtJoinNodeOp
     return CollectCboStatsNode(relJoinColumns, cluster, op, state, ctx);
 }
 
+TVector<TString> JoinLeafLabels(TExprNode::TPtr label) {
+    if (label->ChildrenSize() == 0) {
+        return TVector<TString>{TString(label->Content())};
+    }
+    TVector<TString> result;
+    for (ui32 i = 0; i < label->ChildrenSize(); ++i) {
+        result.push_back(TString(label->Child(i)->Content()));
+    }
+    return result;
+}
+
+
 }
diff --git a/yt/yql/providers/yt/provider/yql_yt_cbo_helpers.h b/yt/yql/providers/yt/provider/yql_yt_cbo_helpers.h
index 1fb9f5ad3bf..0d05beea7af 100644
--- a/yt/yql/providers/yt/provider/yql_yt_cbo_helpers.h
+++ b/yt/yql/providers/yt/provider/yql_yt_cbo_helpers.h
@@ -9,4 +9,6 @@ IGraphTransformer::TStatus CollectCboStats(const TString& cluster, TYtJoinNodeOp
 
 IGraphTransformer::TStatus PopulateJoinStrategySizeInfo(TRelSizeInfo& outLeft, TRelSizeInfo& outRight, const TYtState::TPtr& state, TString cluster, TExprContext& ctx, TYtJoinNodeOp* op);
 
+TVector<TString> JoinLeafLabels(TExprNode::TPtr label);
+
 }  // namespace NYql
diff --git a/yt/yql/providers/yt/provider/yql_yt_join_impl.cpp b/yt/yql/providers/yt/provider/yql_yt_join_impl.cpp
index fae1c0f3f72..c74ab89154d 100644
--- a/yt/yql/providers/yt/provider/yql_yt_join_impl.cpp
+++ b/yt/yql/providers/yt/provider/yql_yt_join_impl.cpp
@@ -5206,19 +5206,4 @@ TMaybeNode<TExprBase> ExportYtEquiJoin(TYtEquiJoin equiJoin, const TYtJoinNodeOp
     return TExprBase(ctx.ChangeChildren(join.Ref(), std::move(children)));
 }
 
-TString JoinLeafLabel(TExprNode::TPtr label) {
-    if (label->ChildrenSize() == 0) {
-        return TString(label->Content());
-    }
-    TString result;
-    for (ui32 i = 0; i < label->ChildrenSize(); ++i) {
-        result += label->Child(i)->Content();
-        if (i+1 != label->ChildrenSize()) {
-            result += ",";
-        }
-    }
-
-    return result;
-}
-
 }
diff --git a/yt/yql/providers/yt/provider/yql_yt_join_impl.h b/yt/yql/providers/yt/provider/yql_yt_join_impl.h
index 51927a91038..a36685582bf 100644
--- a/yt/yql/providers/yt/provider/yql_yt_join_impl.h
+++ b/yt/yql/providers/yt/provider/yql_yt_join_impl.h
@@ -73,7 +73,6 @@ IGraphTransformer::TStatus RewriteYtEquiJoinLeaves(TYtEquiJoin equiJoin, TYtJoin
 IGraphTransformer::TStatus RewriteYtEquiJoin(TYtEquiJoin equiJoin, TYtJoinNodeOp& op, const TYtState::TPtr& state, TExprContext& ctx);
 TMaybeNode<TExprBase> ExportYtEquiJoin(TYtEquiJoin equiJoin, const TYtJoinNodeOp& op, TExprContext& ctx, const TYtState::TPtr& state);
 TYtJoinNodeOp::TPtr OrderJoins(TYtJoinNodeOp::TPtr op, const TYtState::TPtr& state, const TString& cluster,  TExprContext& ctx, bool debug = false);
-TString JoinLeafLabel(TExprNode::TPtr label);
 
 struct IBaseOptimizerNode;
 struct IProviderContext;
diff --git a/yt/yql/providers/yt/provider/yql_yt_join_reorder.cpp b/yt/yql/providers/yt/provider/yql_yt_join_reorder.cpp
index 805d9796978..5b9c326df92 100644
--- a/yt/yql/providers/yt/provider/yql_yt_join_reorder.cpp
+++ b/yt/yql/providers/yt/provider/yql_yt_join_reorder.cpp
@@ -11,6 +11,7 @@
 #include <yt/yql/providers/yt/opt/yql_yt_join.h>
 #include <yt/yql/providers/yt/provider/yql_yt_provider_context.h>
 #include <yql/essentials/utils/log/log.h>
+#include <util/string/vector.h>
 
 #include <yt/cpp/mapreduce/common/helpers.h>
 
@@ -212,15 +213,21 @@ public:
     }
 
 private:
-    TVector<TString> GetJoinColumns(const TString& label) {
-        auto pos = RelJoinColumns.find(label);
-        if (pos == RelJoinColumns.end()) {
-            return TVector<TString>{};
-        }
+    TVector<TString> GetJoinColumns(const TVector<TString>& labels) {
+        TVector<TString> result;
+        for (const auto& label : labels) {
+            auto pos = RelJoinColumns.find(label);
+            if (pos == RelJoinColumns.end()) {
+                YQL_CLOG(ERROR, ProviderYt) << "Did not find any join columns for label " << label;
+                return TVector<TString>{};
+            }
 
-        return TVector<TString>(pos->second.begin(), pos->second.end());
+            std::copy(pos->second.begin(), pos->second.end(), std::back_inserter(result));
+        }
+        return result;
     }
 
+
     std::shared_ptr<IBaseOptimizerNode> ProcessNode(TYtJoinNode::TPtr node, TRelSizeInfo sizeInfo) {
         if (auto* op = dynamic_cast<TYtJoinNodeOp*>(node.Get())) {
             return OnOp(op);
@@ -255,6 +262,20 @@ private:
 
         auto left = ProcessNode(op->Left, leftSizeInfo);
         auto right = ProcessNode(op->Right, rightSizeInfo);
+
+        for (auto& joinColumn : leftKeys) {
+            if (MultiLabelIndex_.count(joinColumn.RelName) > 0) {
+                joinColumn.OriginalRelName = joinColumn.RelName;
+                joinColumn.RelName = JoinStrings(MultiLabels_[MultiLabelIndex_[joinColumn.RelName]], ",");
+            }
+        }
+        for (auto& joinColumn : rightKeys) {
+            if (MultiLabelIndex_.count(joinColumn.RelName) > 0) {
+                joinColumn.OriginalRelName = joinColumn.RelName;
+                joinColumn.RelName = JoinStrings(MultiLabels_[MultiLabelIndex_[joinColumn.RelName]], ",");
+            }
+        }
+
         bool nonReorderable = op->LinkSettings.ForceSortedMerge;
         LinkSettings.HasForceSortedMerge = LinkSettings.HasForceSortedMerge || op->LinkSettings.ForceSortedMerge;
         LinkSettings.HasHints = LinkSettings.HasHints || !op->LinkSettings.LeftHints.empty() || !op->LinkSettings.RightHints.empty();
@@ -265,11 +286,12 @@ private:
     }
 
     std::shared_ptr<IBaseOptimizerNode> OnLeaf(TYtJoinNodeLeaf* leaf, TRelSizeInfo sizeInfo) {
-        TString label = JoinLeafLabel(leaf->Label);
+        auto labels = JoinLeafLabels(leaf->Label);
+        TString label = JoinStrings(labels, ",");
 
         const TMaybe<ui64> maxChunkCountExtendedStats = State->Configuration->ExtendedStatsMaxChunkCount.Get();
 
-        TVector<TString> keyList = GetJoinColumns(label);
+        TVector<TString> keyList = GetJoinColumns(labels);
 
         TYtSection section{leaf->Section};
         auto stat = std::make_shared<TOptimizerStatistics>();
@@ -331,6 +353,13 @@ private:
         });
 
         stat->Specific = std::move(providerStats);
+
+        if (labels.size() != 1) {
+            MultiLabels_.push_back(labels);
+            for (const auto& label : labels) {
+                MultiLabelIndex_[label] = std::ssize(MultiLabels_) - 1;
+            }
+        }
         return std::make_shared<TYtRelOptimizerNode>(std::move(label), std::move(*stat), leaf);
     }
 
@@ -424,21 +453,6 @@ private:
         entry.first->second.insert(rcolumn);
     }
 
-    TString JoinLeafLabel(TExprNode::TPtr label) {
-        if (label->ChildrenSize() == 0) {
-            return TString(label->Content());
-        }
-        TString result;
-        for (ui32 i = 0; i < label->ChildrenSize(); ++i) {
-            result += label->Child(i)->Content();
-            if (i+1 != label->ChildrenSize()) {
-                result += ",";
-            }
-        }
-
-        return result;
-    }
-
     TYtState::TPtr State;
     const TString Cluster;
     std::shared_ptr<IBaseOptimizerNode>& Tree;
@@ -447,6 +461,8 @@ private:
     TYtJoinNodeOp::TPtr InputTree;
     TExprContext& Ctx;
     TVector<TYtProviderRelInfo> ProviderRelInfo_;
+    TVector<TVector<TString>> MultiLabels_;
+    THashMap<TString, int> MultiLabelIndex_;
 };
 
 TYtJoinNode::TPtr BuildYtJoinTree(std::shared_ptr<IBaseOptimizerNode> node, TVector<TString>& scope, TExprContext& ctx, TPositionHandle pos) {
@@ -468,11 +484,11 @@ TYtJoinNode::TPtr BuildYtJoinTree(std::shared_ptr<IBaseOptimizerNode> node, TVec
             leftLabel.reserve(op->LeftJoinKeys.size() * 2);
             rightLabel.reserve(op->RightJoinKeys.size() * 2);
             for (auto& left : op->LeftJoinKeys) {
-                leftLabel.emplace_back(ctx.NewAtom(pos, left.RelName));
+                leftLabel.emplace_back(ctx.NewAtom(pos, left.OriginalRelName ? *left.OriginalRelName : left.RelName));
                 leftLabel.emplace_back(ctx.NewAtom(pos, left.AttributeName));
             }
             for (auto& right : op->RightJoinKeys) {
-                rightLabel.emplace_back(ctx.NewAtom(pos, right.RelName));
+                rightLabel.emplace_back(ctx.NewAtom(pos, right.OriginalRelName ? *right.OriginalRelName : right.RelName));
                 rightLabel.emplace_back(ctx.NewAtom(pos, right.AttributeName));
             }
             ret->LeftLabel = Build<TCoAtomList>(ctx, pos)