Merge be8b231 into 9588de1

ydb-platform · Mar 1, 2024 · a37e003 · a37e003
2 parents 9588de1 + be8b231
commit a37e003
Show file tree

Hide file tree

Showing 2 changed files with 75 additions and 21 deletions.
diff --git a/ydb/library/yql/providers/yt/provider/ut/yql_yt_cbo_ut.cpp b/ydb/library/yql/providers/yt/provider/ut/yql_yt_cbo_ut.cpp
@@ -2,6 +2,7 @@
 
 #include <ydb/library/yql/providers/yt/provider/yql_yt_join_impl.h>
 #include <ydb/library/yql/core/cbo/cbo_optimizer_new.h>
+#include <ydb/library/yql/dq/opt/dq_opt_log.h>
 
 namespace NYql {
 
@@ -71,6 +72,21 @@ Y_UNIT_TEST(OrderJoinsDoesNothingWhenCBODisabled) {
     UNIT_ASSERT_VALUES_EQUAL(tree, optimizedTree);
 }
 
+Y_UNIT_TEST(NonReordable) {
+    auto stat = std::make_shared<TOptimizerStatistics>();
+    auto left = std::make_shared<TRelOptimizerNode>("a", stat);
+    auto right = std::make_shared<TRelOptimizerNode>("a", stat);
+
+    std::set<std::pair<NDq::TJoinColumn, NDq::TJoinColumn>> joinConditions;
+    joinConditions.insert({NDq::TJoinColumn{"a", "b"}, NDq::TJoinColumn{"a","c"}});
+    auto root = std::make_shared<TJoinOptimizerNode>(
+        left, right, joinConditions, EJoinKind::InnerJoin, EJoinAlgoType::GraceJoin, true);
+    TDummyProviderContext optCtx;
+    std::unique_ptr<IOptimizerNew> opt = std::unique_ptr<IOptimizerNew>(NDq::MakeNativeOptimizerNew(optCtx, 1024));
+    auto result = opt->JoinSearch(root);
+    UNIT_ASSERT(root == result);
+}
+
 Y_UNIT_TEST(BuildOptimizerTree2Tables) {
     TExprContext exprCtx;
     auto tree = MakeOp({"c", "c_nationkey"}, {"n", "n_nationkey"}, {"c", "n"}, exprCtx);
@@ -134,6 +150,22 @@ Y_UNIT_TEST(BuildYtJoinTree2Tables) {
     UNIT_ASSERT(AreSimilarTrees(joinTree, tree));
 }
 
+Y_UNIT_TEST(BuildYtJoinTree2TablesForceMergeJoib) {
+    TExprContext exprCtx;
+    auto tree = MakeOp({"c", "c_nationkey"}, {"n", "n_nationkey"}, {"c", "n"}, exprCtx);
+    tree->Left = MakeLeaf({"c"}, {"c"},  100000, 12333, exprCtx);
+    tree->Right = MakeLeaf({"n"}, {"n"}, 1000, 1233, exprCtx);
+    tree->LinkSettings.ForceSortedMerge = true;
+
+    std::shared_ptr<IBaseOptimizerNode> resultTree;
+    std::shared_ptr<IProviderContext> resultCtx;
+    BuildOptimizerJoinTree(resultTree, resultCtx, tree);
+
+    auto joinTree = BuildYtJoinTree(resultTree, exprCtx, {});
+
+    UNIT_ASSERT(joinTree == tree);
+}
+
 Y_UNIT_TEST(BuildYtJoinTree2TablesComplexLabel) {
     TExprContext exprCtx;
     auto tree = MakeOp({"c", "c_nationkey"}, {"n", "n_nationkey"}, {"c", "n", "e"}, exprCtx);

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
@@ -134,6 +134,21 @@ class TYtRelOptimizerNode: public TRelOptimizerNode {
     TYtJoinNodeLeaf* OriginalLeaf;
 };
 
+class TYtJoinOptimizerNode: public TJoinOptimizerNode {
+public:
+    TYtJoinOptimizerNode(const std::shared_ptr<IBaseOptimizerNode>& left,
+        const std::shared_ptr<IBaseOptimizerNode>& right,
+        const std::set<std::pair<NDq::TJoinColumn, NDq::TJoinColumn>>& joinConditions,
+        const EJoinKind joinType,
+        const EJoinAlgoType joinAlgo,
+        TYtJoinNodeOp* originalOp)
+        : TJoinOptimizerNode(left, right, joinConditions, joinType, joinAlgo, originalOp != nullptr)
+        , OriginalOp(originalOp)
+    { }
+
+    TYtJoinNodeOp* OriginalOp; // Only for nonReorderable
+};
+
 class TOptimizerTreeBuilder
 {
 public:
@@ -175,9 +190,10 @@ class TOptimizerTreeBuilder
             NDq::TJoinColumn rcol{TString(rtable), TString(rcolumn)};
             joinConditions.insert({lcol, rcol});
         }
+        bool nonReorderable = op->LinkSettings.ForceSortedMerge;
 
-        return std::make_shared<TJoinOptimizerNode>(
-            left, right, joinConditions, joinKind, EJoinAlgoType::GraceJoin
+        return std::make_shared<TYtJoinOptimizerNode>(
+            left, right, joinConditions, joinKind, EJoinAlgoType::GraceJoin, nonReorderable ? op : nullptr
             );
     }
 
@@ -229,27 +245,33 @@ TYtJoinNode::TPtr BuildYtJoinTree(std::shared_ptr<IBaseOptimizerNode> node, TVec
         scope.insert(scope.end(), leaf->Scope.begin(), leaf->Scope.end());
         return leaf;
     } else if (node->Kind == JoinNodeType) {
-        auto ret = MakeIntrusive<TYtJoinNodeOp>();
         auto* op = static_cast<TJoinOptimizerNode*>(node.get());
-        ret->JoinKind = ctx.NewAtom(pos, ConvertToJoinString(op->JoinType));
-        TVector<TExprNodePtr> leftLabel, rightLabel;
-        leftLabel.reserve(op->JoinConditions.size() * 2);
-        rightLabel.reserve(op->JoinConditions.size() * 2);
-        for (auto& [left, right] : op->JoinConditions) {
-            leftLabel.emplace_back(ctx.NewAtom(pos, left.RelName));
-            leftLabel.emplace_back(ctx.NewAtom(pos, left.AttributeName));
-
-            rightLabel.emplace_back(ctx.NewAtom(pos, right.RelName));
-            rightLabel.emplace_back(ctx.NewAtom(pos, right.AttributeName));
+        auto* ytop = dynamic_cast<TYtJoinOptimizerNode*>(op);
+        TYtJoinNodeOp::TPtr ret;
+        if (ytop && !ytop->IsReorderable) {
+            ret = ytop->OriginalOp;
+        } else {
+            ret = MakeIntrusive<TYtJoinNodeOp>();
+            ret->JoinKind = ctx.NewAtom(pos, ConvertToJoinString(op->JoinType));
+            TVector<TExprNodePtr> leftLabel, rightLabel;
+            leftLabel.reserve(op->JoinConditions.size() * 2);
+            rightLabel.reserve(op->JoinConditions.size() * 2);
+            for (auto& [left, right] : op->JoinConditions) {
+                leftLabel.emplace_back(ctx.NewAtom(pos, left.RelName));
+                leftLabel.emplace_back(ctx.NewAtom(pos, left.AttributeName));
+
+                rightLabel.emplace_back(ctx.NewAtom(pos, right.RelName));
+                rightLabel.emplace_back(ctx.NewAtom(pos, right.AttributeName));
+            }
+            ret->LeftLabel = Build<TCoAtomList>(ctx, pos)
+                .Add(leftLabel)
+                .Done()
+                .Ptr();
+            ret->RightLabel = Build<TCoAtomList>(ctx, pos)
+                .Add(rightLabel)
+                .Done()
+                .Ptr();
         }
-        ret->LeftLabel = Build<TCoAtomList>(ctx, pos)
-            .Add(leftLabel)
-            .Done()
-            .Ptr();
-        ret->RightLabel = Build<TCoAtomList>(ctx, pos)
-            .Add(rightLabel)
-            .Done()
-            .Ptr();
         int index = scope.size();
         ret->Left = BuildYtJoinTree(op->LeftArg, scope, ctx, pos);
         ret->Right = BuildYtJoinTree(op->RightArg, scope, ctx, pos);