[yql] Fix DependsOn handling in horizontal join

YQL-14907

ref:9490a3a20e6471216f829b65d73c2e77291038cd

3 files changed, 24 insertions, 24 deletions

diff --git a/ydb/library/yql/core/common_opt/yql_co_flow1.cpp b/ydb/library/yql/core/common_opt/yql_co_flow1.cpp
index 99b5cbae475..d15804a3051 100644
--- a/ydb/library/yql/core/common_opt/yql_co_flow1.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_flow1.cpp
@@ -67,16 +67,8 @@ TExprNode::TPtr FuseFlatmaps(TCoFlatMapBase outerMap, TExprContext& ctx, TTypeAn
     }
 
     if (IsJustOrSingleAsList(innerBody.Ref())) {
-        auto placeHolder = ctx.NewArgument(outerMap.Pos(), "placeholder");
-
-        auto status = OptimizeExpr(outerLambda, outerLambda, [&placeHolder, outerLambdaArg](const TExprNode::TPtr& node, TExprContext& ctx) -> TExprNode::TPtr {
-            if (TCoDependsOn::Match(node.Get()) && node->Child(0) == outerLambdaArg) {
-                return ctx.ChangeChild(*node, 0, TExprNode::TPtr(placeHolder));
-            }
-            return node;
-        }, ctx, TOptimizeExprSettings(types));
-
-        if (status.Level == IGraphTransformer::TStatus::Error) {
+        auto [placeHolder, lambdaWithPlaceholder] = ReplaceDependsOn(outerLambda, ctx, types);
+        if (!placeHolder) {
             return {};
         }
 
@@ -99,7 +91,7 @@ TExprNode::TPtr FuseFlatmaps(TCoFlatMapBase outerMap, TExprContext& ctx, TTypeAn
                         .Args({"item"})
                         .template Body<TCoToFlow>()
                             .template Input<TExprApplier>()
-                                .Apply(TCoLambda(outerLambda))
+                                .Apply(TCoLambda(lambdaWithPlaceholder))
                                 .With(0, outerArgValue)
                                 .With(clonedInnerLambda.Args().Arg(0), "item")
                                 .With(TExprBase(placeHolder), "item")
@@ -169,23 +161,13 @@ TExprNode::TPtr FuseFlatmaps(TCoFlatMapBase outerMap, TExprContext& ctx, TTypeAn
         const auto clonedInnerLambda = TCoLambda(ctx.DeepCopyLambda(innerMap.Lambda().Ref()));
         const auto conditional = clonedInnerLambda.Body().template Cast<TCoConditionalValueBase>();
 
-        auto placeHolder = ctx.NewArgument(outerMap.Pos(), "placeholder");
-        auto outerLambda = outerMap.Lambda().Ptr();
-        auto outerLambdaArg = outerMap.Lambda().Args().Arg(0).Raw();
-
-        auto status = OptimizeExpr(outerLambda, outerLambda, [&placeHolder, outerLambdaArg](const TExprNode::TPtr& node, TExprContext& ctx) -> TExprNode::TPtr {
-            if (TCoDependsOn::Match(node.Get()) && node->Child(0) == outerLambdaArg) {
-                return ctx.ChangeChild(*node, 0, TExprNode::TPtr(placeHolder));
-            }
-            return node;
-        }, ctx, TOptimizeExprSettings(types));
-
-        if (status.Level == IGraphTransformer::TStatus::Error) {
+        auto [placeHolder, lambdaWithPlaceholder] = ReplaceDependsOn(outerMap.Lambda().Ptr(), ctx, types);
+        if (!placeHolder) {
             return {};
         }
 
         auto value = ctx.Builder(outerMap.Pos())
-            .Apply(outerLambda)
+            .Apply(lambdaWithPlaceholder)
             .With(0, conditional.Value().Ptr())
             .Seal()
             .Build();
diff --git a/ydb/library/yql/core/yql_opt_utils.cpp b/ydb/library/yql/core/yql_opt_utils.cpp
index fcaf1322f93..4ce61d228f6 100644
--- a/ydb/library/yql/core/yql_opt_utils.cpp
+++ b/ydb/library/yql/core/yql_opt_utils.cpp
@@ -1377,4 +1377,21 @@ IGraphTransformer::TStatus LocalUnorderedOptimize(TExprNode::TPtr input, TExprNo
 
 }
 
+std::pair<TExprNode::TPtr, TExprNode::TPtr> ReplaceDependsOn(TExprNode::TPtr lambda, TExprContext& ctx, TTypeAnnotationContext* typeCtx) {
+    auto placeHolder = ctx.NewArgument(lambda->Pos(), "placeholder");
+
+    auto status = OptimizeExpr(lambda, lambda, [&placeHolder, arg = &lambda->Head().Head()](const TExprNode::TPtr& node, TExprContext& ctx) -> TExprNode::TPtr {
+        if (TCoDependsOn::Match(node.Get()) && &node->Head() == arg) {
+            return ctx.ChangeChild(*node, 0, TExprNode::TPtr(placeHolder));
+        }
+        return node;
+    }, ctx, TOptimizeExprSettings{typeCtx});
+
+    if (status.Level == IGraphTransformer::TStatus::Error) {
+        return std::pair<TExprNode::TPtr, TExprNode::TPtr>{};
+    }
+
+    return {placeHolder, lambda};
+}
+
 }
diff --git a/ydb/library/yql/core/yql_opt_utils.h b/ydb/library/yql/core/yql_opt_utils.h
index 495f3c702a3..90501bfc949 100644
--- a/ydb/library/yql/core/yql_opt_utils.h
+++ b/ydb/library/yql/core/yql_opt_utils.h
@@ -100,4 +100,5 @@ bool WarnUnroderedSubquery(const TExprNode& unourderedSubquery, TExprContext& ct
 IGraphTransformer::TStatus LocalUnorderedOptimize(TExprNode::TPtr input, TExprNode::TPtr& output,
     const std::function<bool(const TExprNode*)>& stopTraverse, TExprContext& ctx, TTypeAnnotationContext* typeCtx);
 
+std::pair<TExprNode::TPtr, TExprNode::TPtr> ReplaceDependsOn(TExprNode::TPtr lambda, TExprContext& ctx, TTypeAnnotationContext* typeCtx);
 }