Fix FlatMap over Optional with Pg lambda

commit_hash:a52a24f4c4ad29082428362e7072c8d9cf8d2feb

2 files changed, 6 insertions, 2 deletions

diff --git a/yql/essentials/core/peephole_opt/yql_opt_peephole_physical.cpp b/yql/essentials/core/peephole_opt/yql_opt_peephole_physical.cpp
index 07e5880c23..6131a31b15 100644
--- a/yql/essentials/core/peephole_opt/yql_opt_peephole_physical.cpp
+++ b/yql/essentials/core/peephole_opt/yql_opt_peephole_physical.cpp
@@ -2469,7 +2469,10 @@ TExprNode::TPtr ExpandFlatMap(const TExprNode::TPtr& node, TExprContext& ctx) {
         }
     }
 
-    if (ETypeAnnotationKind::Optional == node->GetTypeAnn()->GetKind()) {
+    if (ETypeAnnotationKind::Optional == node->GetTypeAnn()->GetKind() ||
+        ETypeAnnotationKind::Pg == node->GetTypeAnn()->GetKind())
+    {
+        YQL_ENSURE(node->Head().GetTypeAnn()->GetKind() == ETypeAnnotationKind::Optional);
         YQL_CLOG(DEBUG, CorePeepHole) << node->Content() << " over Optional";
         return ctx.Builder(node->Pos())
             .Callable("IfPresent")
diff --git a/yql/essentials/core/type_ann/type_ann_list.cpp b/yql/essentials/core/type_ann/type_ann_list.cpp
index 674260bf35..b8c6782267 100644
--- a/yql/essentials/core/type_ann/type_ann_list.cpp
+++ b/yql/essentials/core/type_ann/type_ann_list.cpp
@@ -1224,6 +1224,7 @@ namespace {
         // lambda L F L S L
         // lambda S F L S S
         // lambda O F L S O
+        // lambda P F L S P
         // lambda F F F - F
 
         bool warn = false;
@@ -1300,7 +1301,7 @@ namespace {
             return IGraphTransformer::TStatus::Repeat;
         }
 
-        input->SetTypeAnn(MakeSequenceType(resultKind, *lambdaItemType, ctx.Expr));
+        input->SetTypeAnn(resultKind == ETypeAnnotationKind::Pg ? lambdaItemType : MakeSequenceType(resultKind, *lambdaItemType, ctx.Expr));
         return IGraphTransformer::TStatus::Ok;
     }