YQL-14686 fix.

ref:e9691da4c4c3e9a017b36afd8d6bb9d17d893140

3 files changed, 21 insertions, 13 deletions

diff --git a/ydb/core/kqp/opt/physical/kqp_opt_phy.cpp b/ydb/core/kqp/opt/physical/kqp_opt_phy.cpp
index db89edff15..fb4814170e 100644
--- a/ydb/core/kqp/opt/physical/kqp_opt_phy.cpp
+++ b/ydb/core/kqp/opt/physical/kqp_opt_phy.cpp
@@ -60,6 +60,7 @@ public:
         AddHandler(0, &TCoUnorderedBase::Match, HNDL(DropUnordered));
         AddHandler(0, &TDqStage::Match, HNDL(FloatUpStage));
         AddHandler(0, &TCoHasItems::Match, HNDL(BuildHasItems));
+        AddHandler(0, &TCoHead::Match, HNDL(BuildScalarPrecompute<false>));
         AddHandler(0, &TCoToOptional::Match, HNDL(BuildScalarPrecompute<false>));
         AddHandler(0, &TCoAsList::Match, HNDL(PropagatePrecomuteScalarRowset<false>));
         AddHandler(0, &TCoTake::Match, HNDL(PropagatePrecomuteTake<false>));
@@ -78,6 +79,7 @@ public:
         AddHandler(1, &TDqJoin::Match, HNDL(BuildJoin<true>));
         AddHandler(1, &TCoLMap::Match, HNDL(PushLMapToStage<true>));
         AddHandler(1, &TCoOrderedLMap::Match, HNDL(PushOrderedLMapToStage<true>));
+        AddHandler(1, &TCoHead::Match, HNDL(BuildScalarPrecompute<true>));
         AddHandler(1, &TCoToOptional::Match, HNDL(BuildScalarPrecompute<true>));
         AddHandler(1, &TCoAsList::Match, HNDL(PropagatePrecomuteScalarRowset<true>));
         AddHandler(1, &TCoTake::Match, HNDL(PropagatePrecomuteTake<true>));
diff --git a/ydb/library/yql/core/type_ann/type_ann_core.cpp b/ydb/library/yql/core/type_ann/type_ann_core.cpp
index 3ef30edece..8d85fe04c7 100644
--- a/ydb/library/yql/core/type_ann/type_ann_core.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_core.cpp
@@ -3116,11 +3116,6 @@ namespace NTypeAnnImpl {
         return IGraphTransformer::TStatus::Ok;
     }
 
-    IGraphTransformer::TStatus HeadWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
-        output = ctx.Expr.RenameNode(*input, "ToOptional");
-        return IGraphTransformer::TStatus::Repeat;
-    }
-
     IGraphTransformer::TStatus ToFlowWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
         if (!EnsureMinArgsCount(*input, 1, ctx.Expr)) {
             return IGraphTransformer::TStatus::Error;
@@ -10992,7 +10987,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         Functions["ToList"] = &ToListWrapper;
         Functions["ToOptional"] = &ToOptionalWrapper;
         Functions["Iterable"] = &IterableWrapper;
-        Functions["Head"] = &HeadWrapper;
+        Functions["Head"] = &ToOptionalWrapper;
         Functions["Last"] = &ToOptionalWrapper;
         Functions["AsTagged"] = &AsTaggedWrapper;
         Functions["Untag"] = &UntagWrapper;
diff --git a/ydb/library/yql/dq/opt/dq_opt_phy.cpp b/ydb/library/yql/dq/opt/dq_opt_phy.cpp
index a44f8bcfa2..063fdba022 100644
--- a/ydb/library/yql/dq/opt/dq_opt_phy.cpp
+++ b/ydb/library/yql/dq/opt/dq_opt_phy.cpp
@@ -567,6 +567,16 @@ TExprBase DqBuildAggregationResultStage(TExprBase node, TExprContext& ctx, IOpti
                 return false;
             }
 
+            if (expr.Maybe<TCoHead>().Input().Maybe<TDqCnUnionAll>()) {
+                if (connection && (connection != expr.Cast<TCoHead>().Input().Ptr())) {
+                    dependsOnManyConnections = true;
+                    return false;
+                }
+
+                connection = expr.Cast<TCoHead>().Input().Ptr();
+                return false;
+            }
+
             if (expr.Maybe<TDqPhyPrecompute>().IsValid()) {
                 auto precompute = expr.Cast<TDqPhyPrecompute>();
                 auto maybeConnection = precompute.Connection().Maybe<TDqCnValue>();
@@ -1560,17 +1570,18 @@ TExprBase DqBuildHasItems(TExprBase node, TExprContext& ctx, IOptimizationContex
 TExprBase DqBuildScalarPrecompute(TExprBase node, TExprContext& ctx, IOptimizationContext& optCtx,
     const TParentsMap& parentsMap, bool allowStageMultiUsage)
 {
-    if (!node.Maybe<TCoToOptional>()) {
-        return node;
+    TMaybeNode<TDqCnUnionAll> maybeUnionAll;
+    if (const auto maybeToOptional = node.Maybe<TCoToOptional>()) {
+        maybeUnionAll = maybeToOptional.Cast().List().Maybe<TDqCnUnionAll>();
+    } else if (const auto maybeHead = node.Maybe<TCoHead>()) {
+        maybeUnionAll = maybeHead.Cast().Input().Maybe<TDqCnUnionAll>();
     }
 
-    auto toOptional = node.Cast<TCoToOptional>();
-
-    if (!toOptional.List().Maybe<TDqCnUnionAll>()) {
+    if (!maybeUnionAll) {
         return node;
     }
 
-    auto unionAll = toOptional.List().Cast<TDqCnUnionAll>();
+    const auto unionAll = maybeUnionAll.Cast();
     if (!IsSingleConsumerConnection(unionAll, parentsMap, allowStageMultiUsage)) {
         return node;
     }
@@ -1606,7 +1617,7 @@ TExprBase DqBuildScalarPrecompute(TExprBase node, TExprContext& ctx, IOptimizati
                 .Build()
             .Build()
         .State<TCoNothing>()
-            .OptionalType(ExpandType(node.Pos(), *toOptional.Ptr()->GetTypeAnn(), ctx))
+            .OptionalType(ExpandType(node.Pos(), *node.Ref().GetTypeAnn(), ctx))
             .Build()
         .SwitchHandler()
             .Args({"item", "state"})