Avoid rewriting Coalesce on type annotation stage before its arguments are converted to common type. Fix EquiJoin flatten option processing

String should be returned here instead of Utf8: https://yql.yandex-team.ru/Operations/Y5o09slwvamzkR7aBqw7F6GOxIe1VAFgQQWxKRw4lEQ=

3 files changed, 17 insertions, 41 deletions

diff --git a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
index b5de86a3a8e..e13fe4c280b 100644
--- a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
@@ -248,60 +248,33 @@ TExprNode::TPtr ExpandFlattenEquiJoin(const TExprNode::TPtr& node, TExprContext&
 
     auto joins = node->Child(node->ChildrenSize() - 2);
     auto columnTypes = GetJoinColumnTypes(*joins, labels, ctx);
-    TMap<TString, std::pair<bool, TVector<TString>>> remap; // column -> isOptional, list of columns
+    TMap<TString, TVector<TString>> remap; // result column -> list of original columns
     for (auto it : labels.Inputs) {
         for (auto item : it.InputType->GetItems()) {
             TString fullName = it.FullName(item->GetName());
             auto type = columnTypes.FindPtr(fullName);
             if (type) {
-                auto columnName = it.ColumnName(fullName);
-                auto iter = remap.find(columnName);
-                if (iter != remap.end()) {
-                    if (iter->second.first) {
-                        // still optional
-                        if ((*type)->GetKind() == ETypeAnnotationKind::Optional) {
-                            iter->second.second.push_back(fullName);
-                        } else {
-                            iter->second.first = false;
-                            iter->second.second.clear();
-                            iter->second.second.push_back(fullName);
-                        }
-                    }
-                    continue;
-                }
-
-                remap.emplace(TString(columnName), std::make_pair((*type)->GetKind() == ETypeAnnotationKind::Optional,
-                    TVector<TString>(1, fullName)));
+                TString columnName(it.ColumnName(fullName));
+                remap[columnName].push_back(fullName);
             }
         }
     }
 
     auto lambdaArg = ctx.NewArgument(node->Pos(), "row");
     TExprNode::TListType remapItems;
-    for (auto& x : remap) {
-        TExprNode::TPtr value;
-        if (x.second.second.size() == 1) {
-            value = ctx.Builder(node->Pos())
+    for (auto& [resultName, sourceNames] : remap) {
+        TExprNode::TListType values;
+        for (auto& column : sourceNames) {
+            values.push_back(ctx.Builder(node->Pos())
                 .Callable("Member")
                     .Add(0, lambdaArg)
-                    .Atom(1, x.second.second.front())
+                    .Atom(1, column)
                 .Seal()
-                .Build();
-        } else {
-            TExprNode::TListType values;
-            for (auto& column : x.second.second) {
-                values.push_back(ctx.Builder(node->Pos())
-                    .Callable("Member")
-                        .Add(0, lambdaArg)
-                        .Atom(1, column)
-                    .Seal()
-                    .Build());
-            }
-
-            value = ctx.NewCallable(node->Pos(), "Coalesce", std::move(values));
+                .Build());
         }
 
-        remapItems.push_back(ctx.NewList(node->Pos(), { ctx.NewAtom(node->Pos(), x.first), value }));
+        TExprNode::TPtr coalesce = ctx.NewCallable(node->Pos(), "Coalesce", std::move(values));
+        remapItems.push_back(ctx.NewList(node->Pos(), { ctx.NewAtom(node->Pos(), resultName), coalesce }));
     }
 
     auto lambdaBody = ctx.NewCallable(node->Pos(), "AsStruct", std::move(remapItems));
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 ff05a762127..3a809f39715 100644
--- a/ydb/library/yql/core/type_ann/type_ann_core.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_core.cpp
@@ -4449,15 +4449,18 @@ namespace NTypeAnnImpl {
         }
 
         const auto leftType = input->Head().GetTypeAnn();
+        const auto rightType = input->Tail().GetTypeAnn();
+
         auto leftItemType = leftType;
         if (leftType->GetKind() == ETypeAnnotationKind::Optional) {
             leftItemType = leftType->Cast<TOptionalExprType>()->GetItemType();
-        } else if (leftType->GetKind() != ETypeAnnotationKind::Pg) {
+        } else if (leftType->GetKind() != ETypeAnnotationKind::Pg &&
+                   IsSameAnnotation(*RemoveOptionalType(leftType), *RemoveOptionalType(rightType)))
+        {
             output = input->HeadPtr();
             return IGraphTransformer::TStatus::Repeat;
         }
 
-        const auto rightType = input->Tail().GetTypeAnn();
         auto rightItemType = rightType;
         if (leftType->GetKind() != ETypeAnnotationKind::Optional &&
             rightType->GetKind() == ETypeAnnotationKind::Optional) {
diff --git a/ydb/library/yql/core/yql_join.cpp b/ydb/library/yql/core/yql_join.cpp
index b10fbd20276..ec3b7b49405 100644
--- a/ydb/library/yql/core/yql_join.cpp
+++ b/ydb/library/yql/core/yql_join.cpp
@@ -819,7 +819,7 @@ IGraphTransformer::TStatus EquiJoinAnnotation(
                 }
             }
 
-            if (isOptional) {
+            if (isOptional && commonType->GetKind() != ETypeAnnotationKind::Optional) {
                 commonType = ctx.MakeType<TOptionalExprType>(commonType);
             }