PR from branch users/deniskhalikov/fuse_join_tree

Add fuse EquiJoins with label list.
commit_hash:8f7cbce0cd1bd2bd8e9146fdf4b64805519d4af2

3 files changed, 128 insertions, 39 deletions

diff --git a/yql/essentials/core/common_opt/yql_co_flow2.cpp b/yql/essentials/core/common_opt/yql_co_flow2.cpp
index 2bde8dc2aad..ee0a9074f18 100644
--- a/yql/essentials/core/common_opt/yql_co_flow2.cpp
+++ b/yql/essentials/core/common_opt/yql_co_flow2.cpp
@@ -261,13 +261,13 @@ TExprNode::TPtr RenameJoinTree(TExprNode::TPtr joinTree, const THashMap<TString,
     return ret;
 }
 
-TExprNode::TPtr ReassembleJoinEquality(TExprNode::TPtr columns, const TStringBuf& upstreamLabel,
+TExprNode::TPtr ReassembleJoinEquality(TExprNode::TPtr columns, const THashSet<TStringBuf>& upstreamLabels,
     const THashMap<TString, TString>& upstreamTablesRename,
     const THashMap<TString, TString>& upstreamColumnsBackRename, TExprContext& ctx)
 {
     TExprNode::TListType newChildren(columns->ChildrenList());
     for (ui32 i = 0; i < columns->ChildrenSize(); i += 2) {
-        if (columns->Child(i)->Content() != upstreamLabel) {
+        if (!upstreamLabels.contains(columns->Child(i)->Content())) {
             continue;
         }
 
@@ -280,14 +280,16 @@ TExprNode::TPtr ReassembleJoinEquality(TExprNode::TPtr columns, const TStringBuf
                 upstreamTablesRename, ctx);
             newChildren[i + 1] = ctx.NewAtom(columns->Pos(), part2);
         } else {
-            TStringBuf part1;
-            TStringBuf part2;
-            SplitTableName(column->Content(), part1, part2);
+            TStringBuf part1 = columns->Child(i)->Content();
+            TStringBuf part2 = columns->Child(i + 1)->Content();
+
+            if (TString(column->Content()).find(".") != TString::npos) {
+                SplitTableName(column->Content(), part1, part2);
+            }
+
             newChildren[i] = RenameJoinTable(columns->Pos(), ctx.NewAtom(columns->Pos(), part1),
                 upstreamTablesRename, ctx);
             newChildren[i + 1] = ctx.NewAtom(columns->Pos(), part2);
-
-            return nullptr;
         }
     }
 
@@ -295,13 +297,13 @@ TExprNode::TPtr ReassembleJoinEquality(TExprNode::TPtr columns, const TStringBuf
     return ret;
 }
 
-TExprNode::TPtr FuseJoinTree(TExprNode::TPtr downstreamJoinTree, TExprNode::TPtr upstreamJoinTree, const TStringBuf& upstreamLabel,
+TExprNode::TPtr FuseJoinTree(TExprNode::TPtr downstreamJoinTree, TExprNode::TPtr upstreamJoinTree, const THashSet<TStringBuf>& upstreamLabels,
     const THashMap<TString, TString>& upstreamTablesRename, const THashMap<TString, TString>& upstreamColumnsBackRename,
     TExprContext& ctx)
 {
     TExprNode::TPtr left;
     if (downstreamJoinTree->Child(1)->IsAtom()) {
-        if (downstreamJoinTree->Child(1)->Content() != upstreamLabel) {
+        if (!upstreamLabels.contains(downstreamJoinTree->Child(1)->Content())) {
             left = downstreamJoinTree->Child(1);
         }
         else {
@@ -309,7 +311,7 @@ TExprNode::TPtr FuseJoinTree(TExprNode::TPtr downstreamJoinTree, TExprNode::TPtr
         }
     }
     else {
-        left = FuseJoinTree(downstreamJoinTree->Child(1), upstreamJoinTree, upstreamLabel, upstreamTablesRename,
+        left = FuseJoinTree(downstreamJoinTree->Child(1), upstreamJoinTree, upstreamLabels, upstreamTablesRename,
             upstreamColumnsBackRename, ctx);
         if (!left) {
             return nullptr;
@@ -318,14 +320,14 @@ TExprNode::TPtr FuseJoinTree(TExprNode::TPtr downstreamJoinTree, TExprNode::TPtr
 
     TExprNode::TPtr right;
     if (downstreamJoinTree->Child(2)->IsAtom()) {
-        if (downstreamJoinTree->Child(2)->Content() != upstreamLabel) {
+        if (!upstreamLabels.contains(downstreamJoinTree->Child(2)->Content())) {
             right = downstreamJoinTree->Child(2);
         }
         else {
             right = RenameJoinTree(upstreamJoinTree, upstreamTablesRename, ctx);
         }
     } else {
-        right = FuseJoinTree(downstreamJoinTree->Child(2), upstreamJoinTree, upstreamLabel, upstreamTablesRename,
+        right = FuseJoinTree(downstreamJoinTree->Child(2), upstreamJoinTree, upstreamLabels, upstreamTablesRename,
             upstreamColumnsBackRename, ctx);
         if (!right) {
             return nullptr;
@@ -335,9 +337,9 @@ TExprNode::TPtr FuseJoinTree(TExprNode::TPtr downstreamJoinTree, TExprNode::TPtr
     TExprNode::TListType newChildren(downstreamJoinTree->ChildrenList());
     newChildren[1] = left;
     newChildren[2] = right;
-    newChildren[3] = ReassembleJoinEquality(downstreamJoinTree->Child(3), upstreamLabel, upstreamTablesRename,
+    newChildren[3] = ReassembleJoinEquality(downstreamJoinTree->Child(3), upstreamLabels, upstreamTablesRename,
         upstreamColumnsBackRename, ctx);
-    newChildren[4] = ReassembleJoinEquality(downstreamJoinTree->Child(4), upstreamLabel, upstreamTablesRename,
+    newChildren[4] = ReassembleJoinEquality(downstreamJoinTree->Child(4), upstreamLabels, upstreamTablesRename,
         upstreamColumnsBackRename, ctx);
     if (!newChildren[3] || !newChildren[4]) {
         return nullptr;
@@ -347,18 +349,37 @@ TExprNode::TPtr FuseJoinTree(TExprNode::TPtr downstreamJoinTree, TExprNode::TPtr
     return ret;
 }
 
-TExprNode::TPtr FuseEquiJoins(const TExprNode::TPtr& node, ui32 upstreamIndex, TExprContext& ctx) {
+bool IsSuitableToFuseInputMultiLabels(TOptimizeContext &optCtx) {
+    YQL_ENSURE(optCtx.Types);
+    static const char optName[] = "FuseEquiJoinsInputMultiLabels";
+    return IsOptimizerEnabled<optName>(*optCtx.Types);
+}
+
+TExprNode::TPtr FuseEquiJoins(const TExprNode::TPtr& node, ui32 upstreamIndex, TExprContext& ctx, TOptimizeContext &optCtx) {
     ui32 downstreamInputs = node->ChildrenSize() - 2;
     auto upstreamList = node->Child(upstreamIndex)->Child(0);
     auto upstreamLabel = node->Child(upstreamIndex)->Child(1);
+    THashSet<TStringBuf> upstreamLabelsAssociatedByInputIndex;
     THashSet<TStringBuf> downstreamLabels;
     for (ui32 i = 0; i < downstreamInputs; ++i) {
         auto label = node->Child(i)->Child(1);
-        if (!label->IsAtom()) {
-            return node;
+        if (auto list = TMaybeNode<TCoAtomList>(label)) {
+            if (!IsSuitableToFuseInputMultiLabels(optCtx)) {
+                return node;
+            }
+            for (auto labelAtom : list.Cast()) {
+                auto label = labelAtom.Value();
+                downstreamLabels.insert(label);
+                if (upstreamIndex == i) {
+                    upstreamLabelsAssociatedByInputIndex.insert(label);
+                }
+            }
+        } else {
+            if (upstreamIndex == i) {
+                upstreamLabelsAssociatedByInputIndex.insert(label->Content());
+            }
+            downstreamLabels.insert(label->Content());
         }
-
-        downstreamLabels.insert(label->Content());
     }
 
     THashMap<TString, TString> upstreamTablesRename; // rename of conflicted upstream tables
@@ -381,7 +402,18 @@ TExprNode::TPtr FuseEquiJoins(const TExprNode::TPtr& node, ui32 upstreamIndex, T
             return node;
         }
 
-        if (downstreamLabels.contains(label->Content())) {
+        if (upstreamLabelsAssociatedByInputIndex.size() == 1 && downstreamLabels.contains(label->Content()) ||
+            // In case multiple labels input, we are not renaming labels associated with upstream input index.
+            // For example:
+            // (let ej1 = (EquiJoin '(input1, 'a), '(input2, 'b), upstreamJoinTree, '()))
+            // (let ej2 = (EquiJoin '(ej1, '('a 'b)), '(input3, 'c), downstreamJoinTree, '())))
+            // Upstream labels: [a, b];
+            // Downstream labels: [a, b, c];
+            // Not renaming [a, b] because their associated with input index.
+            // As result we should get:
+            // (let ejFused = (EquiJoin '(input1, 'a), '(input2, 'b), '(input3, 'c), fusedJoinTree, '()))
+            (upstreamLabelsAssociatedByInputIndex.size() > 1 && downstreamLabels.contains(label->Content()) &&
+             !upstreamLabelsAssociatedByInputIndex.contains(label->Content()))) {
             // fix conflict for labels
             for (ui32 suffix = 1;; ++suffix) {
                 auto newName = TString::Join(label->Content(), "_", ToString(suffix));
@@ -481,26 +513,35 @@ TExprNode::TPtr FuseEquiJoins(const TExprNode::TPtr& node, ui32 upstreamIndex, T
         }
     }
 
-    for (auto& x : upstreamColumnsRename) {
-        for (auto& y : x.second) {
-            TStringBuf part1;
-            TStringBuf part2;
-            SplitTableName(x.first, part1, part2);
-            if (auto renamed = upstreamTablesRename.FindPtr(part1)) {
-                part1 = *renamed;
-            }
+   for (auto& x : upstreamColumnsRename) {
+       for (auto& y : x.second) {
+           TStringBuf part1;
+           TStringBuf part2;
+           SplitTableName(x.first, part1, part2);
+           TStringBuf labelName = upstreamLabel->Content();
+           if (upstreamLabelsAssociatedByInputIndex.size() > 1) {
+               if (upstreamLabelsAssociatedByInputIndex.contains(part1)) {
+                   continue;
+               } else {
+                   labelName = part1;
+               }
+           }
 
-            settingsChildren.push_back(ctx.Builder(node->Pos())
-                .List()
-                .Atom(0, "rename")
-                .Atom(1, TString::Join(part1, ".", part2))
-                .Atom(2, TString::Join(upstreamLabel->Content(), ".", y))
-                .Seal()
-                .Build());
-        }
-    }
+           if (auto renamed = upstreamTablesRename.FindPtr(part1)) {
+               part1 = *renamed;
+           }
 
-    auto joinTree = FuseJoinTree(downstreamJoinTree, upstreamJoinTree, upstreamLabel->Content(),
+           settingsChildren.push_back(ctx.Builder(node->Pos())
+               .List()
+                   .Atom(0, "rename")
+                   .Atom(1, TString::Join(part1, ".", part2))
+                   .Atom(2, TString::Join(labelName, ".", y))
+               .Seal()
+               .Build());
+       }
+   }
+
+    auto joinTree = FuseJoinTree(downstreamJoinTree, upstreamJoinTree, upstreamLabelsAssociatedByInputIndex,
         upstreamTablesRename, upstreamColumnsBackRename, ctx);
     if (!joinTree) {
         return node;
@@ -514,6 +555,7 @@ TExprNode::TPtr FuseEquiJoins(const TExprNode::TPtr& node, ui32 upstreamIndex, T
     return ret;
 }
 
+
 bool IsRenamingOrPassthroughFlatMap(const TCoFlatMapBase& flatMap, THashMap<TStringBuf, TStringBuf>& renames,
     THashSet<TStringBuf>& outputMembers, bool& isIdentity)
 {
@@ -2007,7 +2049,7 @@ void RegisterCoFlowCallables2(TCallableOptimizerMap& map) {
                 if (node->Child(i)->Child(0)->IsCallable("EquiJoin") &&
                     optCtx.IsSingleUsage(*node->Child(i)) &&
                     optCtx.IsSingleUsage(*node->Child(i)->Child(0))) {
-                    auto ret = FuseEquiJoins(node, i, ctx);
+                    auto ret = FuseEquiJoins(node, i, ctx, optCtx);
                     if (ret != node) {
                         YQL_CLOG(DEBUG, Core) << "FuseEquiJoins";
                         return ret;
diff --git a/yql/essentials/tests/s-expressions/minirun/part5/canondata/result.json b/yql/essentials/tests/s-expressions/minirun/part5/canondata/result.json
index 98fcc4f04ff..0bf1144eecf 100644
--- a/yql/essentials/tests/s-expressions/minirun/part5/canondata/result.json
+++ b/yql/essentials/tests/s-expressions/minirun/part5/canondata/result.json
@@ -128,6 +128,20 @@
             "uri": "https://{canondata_backend}/1942525/ebc309da64241c6f5f13249c90361edaa566a9cc/resource.tar.gz#test.test_EquiJoin-EquiConvertToCommonTypeAlias-default.txt-Results_/results.txt"
         }
     ],
+    "test.test[EquiJoin-FuseEquiJoins-default.txt-Debug]": [
+        {
+            "checksum": "e4789e7f286c2d891f84c3f5a9016dbd",
+            "size": 981,
+            "uri": "https://{canondata_backend}/1942173/e0468c17e455c4e7740e0b19765633d5a5aca7af/resource.tar.gz#test.test_EquiJoin-FuseEquiJoins-default.txt-Debug_/opt.yql"
+        }
+    ],
+    "test.test[EquiJoin-FuseEquiJoins-default.txt-Results]": [
+        {
+            "checksum": "d4c82ff9c9a0c83a3f0b79b0472c76af",
+            "size": 2466,
+            "uri": "https://{canondata_backend}/1942173/e0468c17e455c4e7740e0b19765633d5a5aca7af/resource.tar.gz#test.test_EquiJoin-FuseEquiJoins-default.txt-Results_/results.txt"
+        }
+    ],
     "test.test[EquiJoin-JoinInMemOpt1Opt2-default.txt-Debug]": [
         {
             "checksum": "6fab425c7eb72691262bf6b82c6f60fb",
diff --git a/yql/essentials/tests/s-expressions/suites/EquiJoin/FuseEquiJoins.yqls b/yql/essentials/tests/s-expressions/suites/EquiJoin/FuseEquiJoins.yqls
new file mode 100644
index 00000000000..6333e6f3432
--- /dev/null
+++ b/yql/essentials/tests/s-expressions/suites/EquiJoin/FuseEquiJoins.yqls
@@ -0,0 +1,33 @@
+(
+(let flag (Configure! world (DataSource 'config) 'OptimizerFlags 'FuseEquiJoinsInputMultiLabels))
+(let mr_source (DataSource 'yt 'plato))
+(let t1 (AsList 
+  (AsStruct '('key1 (Int32 '1)) '('value1 (String 'A))) 
+  (AsStruct '('key1 (Int32 '2)) '('value1 (String 'B)))
+))
+
+(let t2 (AsList 
+  (AsStruct '('key2 (Int32 '3)) '('value2 (String 'C))) 
+  (AsStruct '('key2 (Int32 '4)) '('value2 (String 'D)))
+))
+
+(let t3 (AsList 
+  (AsStruct '('key3 (Int32 '4)) '('value3 (String 'E))) 
+  (AsStruct '('key3 (Int32 '5)) '('value3 (String 'F)))
+))
+
+(let t4 (AsList 
+  (AsStruct '('key4 (Int32 '6)) '('value4 (String 'J))) 
+  (AsStruct '('key4 (Int32 '7)) '('value4 (String 'H)))
+))
+
+(let ej1 (EquiJoin '(t1 'a) '((FlatMap t2 (lambda '(row) (OptionalIf (Coalesce (> (Member row 'key2) (Int32 '1)) (Bool 'false)) row))) 'b) '('Inner 'a 'b '('a 'key1) '('b 'key2) '()) '()))
+(let ej2 (EquiJoin '(ej1 '('a 'b)) '(t3 'c) '('Inner 'a 'c '('a 'key1) '('c 'key3) '()) '()))
+(let ej3 (EquiJoin '(ej2 '('a 'b 'c)) '(t4 'd) '('Inner 'c 'd '('c 'key3) '('d 'key4) '()) '()))
+
+(let res_sink (DataSink 'result))
+(let world (Write! flag res_sink (Key) ej3 '('('type))))
+
+(let world (Commit! world res_sink))
+(return world)
+)