fix connections multiusage

3 files changed, 9 insertions, 4 deletions

diff --git a/ydb/core/kqp/opt/physical/kqp_opt_phy.cpp b/ydb/core/kqp/opt/physical/kqp_opt_phy.cpp
index 608f15c18c7..f992b8d8dde 100644
--- a/ydb/core/kqp/opt/physical/kqp_opt_phy.cpp
+++ b/ydb/core/kqp/opt/physical/kqp_opt_phy.cpp
@@ -103,8 +103,9 @@ protected:
         return output;
     }
 
-    TMaybeNode<TExprBase> BuildReadTableRangesStage(TExprBase node, TExprContext& ctx) {
-        TExprBase output = KqpBuildReadTableRangesStage(node, ctx, KqpCtx);
+    TMaybeNode<TExprBase> BuildReadTableRangesStage(TExprBase node, TExprContext& ctx, const TGetParents& getParents) {
+        auto parents = getParents();
+        TExprBase output = KqpBuildReadTableRangesStage(node, ctx, KqpCtx, *parents);
         DumpAppliedRule("BuildReadTableRangesStage", node.Ptr(), output.Ptr(), ctx);
         return output;
     }
diff --git a/ydb/core/kqp/opt/physical/kqp_opt_phy_build_stage.cpp b/ydb/core/kqp/opt/physical/kqp_opt_phy_build_stage.cpp
index 821ecbb2dab..a97c65bedbc 100644
--- a/ydb/core/kqp/opt/physical/kqp_opt_phy_build_stage.cpp
+++ b/ydb/core/kqp/opt/physical/kqp_opt_phy_build_stage.cpp
@@ -180,7 +180,7 @@ TExprBase KqpBuildReadTableStage(TExprBase node, TExprContext& ctx, const TKqpOp
 }
 
 TExprBase KqpBuildReadTableRangesStage(TExprBase node, TExprContext& ctx,
-    const TKqpOptimizeContext& kqpCtx)
+    const TKqpOptimizeContext& kqpCtx, const TParentsMap& parents)
 {
     if (!node.Maybe<TKqlReadTableRanges>()) {
         return node;
@@ -228,6 +228,10 @@ TExprBase KqpBuildReadTableRangesStage(TExprBase node, TExprContext& ctx,
                 if (!input.Maybe<TDqCnUnionAll>()) {
                     return node;
                 }
+                
+                if (!IsSingleConsumerConnection(input, parents, false)) {
+                    continue;
+                }
 
                 inputs.push_back(input);
                 stageInputs.push_back(
diff --git a/ydb/core/kqp/opt/physical/kqp_opt_phy_rules.h b/ydb/core/kqp/opt/physical/kqp_opt_phy_rules.h
index 54eb24f2273..9bbe1a746aa 100644
--- a/ydb/core/kqp/opt/physical/kqp_opt_phy_rules.h
+++ b/ydb/core/kqp/opt/physical/kqp_opt_phy_rules.h
@@ -15,7 +15,7 @@ NYql::NNodes::TExprBase KqpBuildReadTableStage(NYql::NNodes::TExprBase node, NYq
     const TKqpOptimizeContext& kqpCtx);
 
 NYql::NNodes::TExprBase KqpBuildReadTableRangesStage(NYql::NNodes::TExprBase node, NYql::TExprContext& ctx,
-    const TKqpOptimizeContext& kqpCtx);
+    const TKqpOptimizeContext& kqpCtx, const NYql::TParentsMap& parents);
 
 NYql::NNodes::TExprBase KqpBuildLookupTableStage(NYql::NNodes::TExprBase node, NYql::TExprContext& ctx);