YQL-16302 Early auto map ToDict.

1 files changed, 26 insertions, 32 deletions

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 7c92bfae249..8496b920a0b 100644
--- a/ydb/library/yql/core/type_ann/type_ann_core.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_core.cpp
@@ -6586,37 +6586,35 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
             return IGraphTransformer::TStatus::Repeat;
         }
 
-        const TListExprType* listType;
-        bool isOptional;
-        if (input->Head().GetTypeAnn() && input->Head().GetTypeAnn()->GetKind() == ETypeAnnotationKind::Optional) {
-            auto itemType = input->Head().GetTypeAnn()->Cast<TOptionalExprType>()->GetItemType();
-            if (IsEmptyList(*itemType)) {
-                output = ctx.Expr.NewCallable(input->Pos(), "EmptyDict", {});
-                output = MakeConstMap(input->Pos(), input->HeadPtr(), output, ctx.Expr);
-                return IGraphTransformer::TStatus::Repeat;
-            }
-
-            if (!EnsureListType(input->Head().Pos(), *itemType, ctx.Expr)) {
-                return IGraphTransformer::TStatus::Error;
-            }
-
-            listType = itemType->Cast<TListExprType>();
-            isOptional = true;
-        } else {
-            if (IsEmptyList(input->Head())) {
-                output = ctx.Expr.NewCallable(input->Pos(), "EmptyDict", {});
-                return IGraphTransformer::TStatus::Repeat;
-            }
+        if (const auto inputType = input->Head().GetTypeAnn(); inputType && inputType->GetKind() == ETypeAnnotationKind::Optional
+            && inputType->Cast<TOptionalExprType>()->GetItemType()->GetKind() != ETypeAnnotationKind::Optional) {
+            output = ctx.Expr.Builder(input->Pos())
+                .Callable("Map")
+                    .Add(0, input->HeadPtr())
+                    .Lambda(1)
+                        .Param("list")
+                        .Callable(input->Content())
+                            .Arg(0, "list")
+                            .Add(1, input->ChildPtr(1))
+                            .Add(2, input->ChildPtr(2))
+                            .Add(3, input->ChildPtr(3))
+                        .Seal()
+                    .Seal()
+                .Seal().Build();
+            return IGraphTransformer::TStatus::Repeat;
+        }
 
-            if (!EnsureListType(input->Head(), ctx.Expr)) {
-                return IGraphTransformer::TStatus::Error;
-            }
+        if (IsEmptyList(input->Head())) {
+            output = ctx.Expr.NewCallable(input->Pos(), "EmptyDict", {});
+            return IGraphTransformer::TStatus::Repeat;
+        }
 
-            listType = input->Head().GetTypeAnn()->Cast<TListExprType>();
-            isOptional = false;
+        if (!EnsureListType(input->Head(), ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
         }
 
-        const TTypeAnnotationNode* itemType = listType->Cast<TListExprType>()->GetItemType();
+        const auto listType = input->Head().GetTypeAnn()->Cast<TListExprType>();
+        const auto itemType = listType->Cast<TListExprType>()->GetItemType();
         auto status = ConvertToLambda(input->ChildRef(1), ctx.Expr, 1);
         if (status.Level != IGraphTransformer::TStatus::Ok) {
             return status;
@@ -6690,10 +6688,6 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         }
 
         input->SetTypeAnn(dictType);
-        if (isOptional) {
-            input->SetTypeAnn(ctx.Expr.MakeType<TOptionalExprType>(input->GetTypeAnn()));
-        }
-
         return IGraphTransformer::TStatus::Ok;
     }
 
@@ -9147,7 +9141,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
                 if (item->ChildrenSize() == 4 && HasSetting(*item->Child(3), "autoName")) {
                     autoNameIndexes.push_back(i);
                 } else {
-                    addedInProjectionFields.emplace(fieldName);                    
+                    addedInProjectionFields.emplace(fieldName);
                     allItems.push_back(ctx.Expr.MakeType<TItemExprType>(fieldName, item->GetTypeAnn()));
                 }
             }