Do not perform FuseEquiJoins opt when called from peephole

commit_hash:9909fd2bda625151a5763f049f0fdf99e9246b8d

7 files changed, 27 insertions, 20 deletions

diff --git a/yql/essentials/core/common_opt/yql_co.h b/yql/essentials/core/common_opt/yql_co.h
index 027fb26f0a2..451db6ccad4 100644
--- a/yql/essentials/core/common_opt/yql_co.h
+++ b/yql/essentials/core/common_opt/yql_co.h
@@ -10,6 +10,7 @@ namespace NYql {
 struct TOptimizeContext {
     TTypeAnnotationContext* Types = nullptr;
     TParentsMap* ParentsMap = nullptr;
+    bool ForPeephole = false;
 
     const TExprNode* GetParentIfSingle(const TExprNode& node) const {
         YQL_ENSURE(ParentsMap);
diff --git a/yql/essentials/core/common_opt/yql_co_flow2.cpp b/yql/essentials/core/common_opt/yql_co_flow2.cpp
index 8ebb5243ffa..2bde8dc2aad 100644
--- a/yql/essentials/core/common_opt/yql_co_flow2.cpp
+++ b/yql/essentials/core/common_opt/yql_co_flow2.cpp
@@ -2000,15 +2000,18 @@ void RegisterCoFlowCallables2(TCallableOptimizerMap& map) {
     };
 
     map["EquiJoin"] = [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) {
-        ui32 inputsCount = node->ChildrenSize() - 2;
-        for (ui32 i = 0; i < inputsCount; ++i) {
-            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);
-                if (ret != node) {
-                    YQL_CLOG(DEBUG, Core) << "FuseEquiJoins";
-                    return ret;
+        if (!optCtx.ForPeephole) {
+            // Peephole splits EquiJoin to pairs, so we don't perform FuseEquiJoin here
+            ui32 inputsCount = node->ChildrenSize() - 2;
+            for (ui32 i = 0; i < inputsCount; ++i) {
+                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);
+                    if (ret != node) {
+                        YQL_CLOG(DEBUG, Core) << "FuseEquiJoins";
+                        return ret;
+                    }
                 }
             }
         }
diff --git a/yql/essentials/core/common_opt/yql_co_transformer.cpp b/yql/essentials/core/common_opt/yql_co_transformer.cpp
index bde50bb6077..f249c80e0bd 100644
--- a/yql/essentials/core/common_opt/yql_co_transformer.cpp
+++ b/yql/essentials/core/common_opt/yql_co_transformer.cpp
@@ -22,9 +22,10 @@ namespace {
 
 class TCommonOptTransformer final : public TSyncTransformerBase {
 public:
-    TCommonOptTransformer(TTypeAnnotationContext* typeCtx, bool final)
+    TCommonOptTransformer(TTypeAnnotationContext* typeCtx, bool final, bool forPeephole)
         : TypeCtx(typeCtx)
         , Final(final)
+        , ForPeephole(forPeephole)
     {}
 
     IGraphTransformer::TStatus DoTransform(TExprNode::TPtr input, TExprNode::TPtr& output, TExprContext& ctx) final;
@@ -47,17 +48,18 @@ private:
     THashSet<TIssue> AddedErrors;
     TTypeAnnotationContext* TypeCtx;
     const bool Final;
+    const bool ForPeephole;
     bool CheckMissingWorld = false;
 };
 
 }
 
-TAutoPtr<IGraphTransformer> CreateCommonOptTransformer(TTypeAnnotationContext* typeCtx) {
-    return new TCommonOptTransformer(typeCtx, false);
+TAutoPtr<IGraphTransformer> CreateCommonOptTransformer(bool forPeephole, TTypeAnnotationContext* typeCtx) {
+    return new TCommonOptTransformer(typeCtx, false, forPeephole);
 }
 
 TAutoPtr<IGraphTransformer> CreateCommonOptFinalTransformer(TTypeAnnotationContext* typeCtx) {
-    return new TCommonOptTransformer(typeCtx, true);
+    return new TCommonOptTransformer(typeCtx, true, false);
 }
 
 IGraphTransformer::TStatus TCommonOptTransformer::DoTransform(TExprNode::TPtr input, TExprNode::TPtr& output, TExprContext& ctx) {
@@ -148,6 +150,7 @@ IGraphTransformer::TStatus TCommonOptTransformer::DoTransform(
     TParentsMap parentsMap;
     TOptimizeContext optCtx;
     optCtx.Types = TypeCtx;
+    optCtx.ForPeephole = ForPeephole;
     if (withParents) {
         GatherParents(*input, parentsMap);
         optCtx.ParentsMap = &parentsMap;
diff --git a/yql/essentials/core/common_opt/yql_co_transformer.h b/yql/essentials/core/common_opt/yql_co_transformer.h
index 7ff5875244f..8480e945425 100644
--- a/yql/essentials/core/common_opt/yql_co_transformer.h
+++ b/yql/essentials/core/common_opt/yql_co_transformer.h
@@ -8,7 +8,7 @@
 
 namespace NYql {
 
-TAutoPtr<IGraphTransformer> CreateCommonOptTransformer(TTypeAnnotationContext* typeCtx);
+TAutoPtr<IGraphTransformer> CreateCommonOptTransformer(bool forPeephole, TTypeAnnotationContext* typeCtx);
 TAutoPtr<IGraphTransformer> CreateCommonOptFinalTransformer(TTypeAnnotationContext* typeCtx);
 
 }
diff --git a/yql/essentials/core/peephole_opt/yql_opt_peephole_physical.cpp b/yql/essentials/core/peephole_opt/yql_opt_peephole_physical.cpp
index 28ee1ead8f5..ef78839b26d 100644
--- a/yql/essentials/core/peephole_opt/yql_opt_peephole_physical.cpp
+++ b/yql/essentials/core/peephole_opt/yql_opt_peephole_physical.cpp
@@ -9108,7 +9108,7 @@ THolder<IGraphTransformer> CreatePeepHoleCommonStageTransformer(TTypeAnnotationC
 
     auto issueCode = TIssuesIds::CORE_EXEC;
 
-    pipeline.AddCommonOptimization(issueCode);
+    pipeline.AddCommonOptimization(true, issueCode);
     pipeline.Add(CreateFunctorTransformer(
             [&types](const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) {
                 return PeepHoleCommonStage(input, output, ctx, types,
diff --git a/yql/essentials/core/services/yql_transform_pipeline.cpp b/yql/essentials/core/services/yql_transform_pipeline.cpp
index 2dd8f2c2787..c6bb45a2189 100644
--- a/yql/essentials/core/services/yql_transform_pipeline.cpp
+++ b/yql/essentials/core/services/yql_transform_pipeline.cpp
@@ -137,13 +137,13 @@ TTransformationPipeline& TTransformationPipeline::AddPostTypeAnnotation(bool for
     return *this;
 }
 
-TTransformationPipeline& TTransformationPipeline::AddCommonOptimization(EYqlIssueCode issueCode) {
+TTransformationPipeline& TTransformationPipeline::AddCommonOptimization(bool forPeephole, EYqlIssueCode issueCode) {
     // auto instantCallableTransformer =
     //    CreateExtCallableTypeAnnotationTransformer(*TypeAnnotationContext_, true);
     // TypeAnnotationContext_->CustomInstantTypeTransformer =
     //     CreateTypeAnnotationTransformer(instantCallableTransformer, *TypeAnnotationContext_);
     Transformers_.push_back(TTransformStage(
-        CreateCommonOptTransformer(TypeAnnotationContext_.Get()),
+        CreateCommonOptTransformer(forPeephole, TypeAnnotationContext_.Get()),
         "CommonOptimization",
         issueCode));
     return *this;
@@ -158,7 +158,7 @@ TTransformationPipeline& TTransformationPipeline::AddFinalCommonOptimization(EYq
 }
 
 TTransformationPipeline& TTransformationPipeline::AddOptimization(bool checkWorld, bool withFinalOptimization, EYqlIssueCode issueCode) {
-    AddCommonOptimization(issueCode);
+    AddCommonOptimization(false, issueCode);
     Transformers_.push_back(TTransformStage(
         CreateChoiceGraphTransformer(
             [&typesCtx = std::as_const(*TypeAnnotationContext_)](const TExprNode::TPtr&, TExprContext&) {
@@ -199,7 +199,7 @@ TTransformationPipeline& TTransformationPipeline::AddOptimization(bool checkWorl
 }
 
 TTransformationPipeline& TTransformationPipeline::AddLineageOptimization(TMaybe<TString>& lineageOut, EYqlIssueCode issueCode) {
-    AddCommonOptimization(issueCode);
+    AddCommonOptimization(false, issueCode);
     Transformers_.push_back(TTransformStage(
         CreateFunctorTransformer(
             [typeCtx = TypeAnnotationContext_, &lineageOut](const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx) {
diff --git a/yql/essentials/core/services/yql_transform_pipeline.h b/yql/essentials/core/services/yql_transform_pipeline.h
index e9326746faf..83250336db0 100644
--- a/yql/essentials/core/services/yql_transform_pipeline.h
+++ b/yql/essentials/core/services/yql_transform_pipeline.h
@@ -32,7 +32,7 @@ public:
     TTransformationPipeline& AddIOAnnotation(bool withEpochsTransformer = true, EYqlIssueCode issueCode = TIssuesIds::CORE_PRE_TYPE_ANN);
     TTransformationPipeline& AddTypeAnnotation(EYqlIssueCode issueCode = TIssuesIds::CORE_TYPE_ANN, bool twoStages = false);
     TTransformationPipeline& AddPostTypeAnnotation(bool forSubGraph = false, bool disableConstraintCheck = false, EYqlIssueCode issueCode = TIssuesIds::CORE_POST_TYPE_ANN);
-    TTransformationPipeline& AddCommonOptimization(EYqlIssueCode issueCode = TIssuesIds::CORE_OPTIMIZATION);
+    TTransformationPipeline& AddCommonOptimization(bool forPeephole = false, EYqlIssueCode issueCode = TIssuesIds::CORE_OPTIMIZATION);
     TTransformationPipeline& AddFinalCommonOptimization(EYqlIssueCode issueCode = TIssuesIds::CORE_OPTIMIZATION);
     TTransformationPipeline& AddOptimization(bool checkWorld = true, bool withFinalOptimization = true, EYqlIssueCode issueCode = TIssuesIds::CORE_OPTIMIZATION);
     TTransformationPipeline& AddLineageOptimization(TMaybe<TString>& lineageOut, EYqlIssueCode issueCode = TIssuesIds::CORE_OPTIMIZATION);