YQL-13710 row_number,lead,lag in pgtypes mode

ref:29114bab9e1dbcf29077d3e4e98ee6562ce4e8e3

11 files changed, 75 insertions, 15 deletions

diff --git a/ydb/library/yql/ast/yql_expr.cpp b/ydb/library/yql/ast/yql_expr.cpp
index cc380e3c5c0..351586131d9 100644
--- a/ydb/library/yql/ast/yql_expr.cpp
+++ b/ydb/library/yql/ast/yql_expr.cpp
@@ -2884,7 +2884,7 @@ const TString& TPgExprType::GetName() const {
 
 ui32 TPgExprType::GetFlags(ui32 typeId) {
     const auto& desc = NPg::LookupType(typeId);
-    ui32 ret = TypeHasManyValues;
+    ui32 ret = TypeHasManyValues | TypeHasOptional;
     if (!desc.SendFuncId || !desc.ReceiveFuncId) {
         ret |= TypeNonPersistable;
     }
diff --git a/ydb/library/yql/ast/yql_expr.h b/ydb/library/yql/ast/yql_expr.h
index 40c47f98959..fd7e2dbeb18 100644
--- a/ydb/library/yql/ast/yql_expr.h
+++ b/ydb/library/yql/ast/yql_expr.h
@@ -217,7 +217,7 @@ public:
 
     bool IsOptionalOrNull() const {
         auto kind = GetKind();
-        return kind == ETypeAnnotationKind::Optional || kind == ETypeAnnotationKind::Null;
+        return kind == ETypeAnnotationKind::Optional || kind == ETypeAnnotationKind::Null || kind == ETypeAnnotationKind::Pg;
     }
 
     bool IsSingleton() const {
diff --git a/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp b/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
index eee4c080b00..dbd69626d75 100644
--- a/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
@@ -921,7 +921,7 @@ TExprNode::TPtr BuildWindows(TPositionHandle pos, const TExprNode::TPtr& list, c
                     .Seal()
                     .Build();
 
-                if (!optCtx.Types->PgTypes && (node->Head().Content() == "row_number" || node->Head().Content() == "count")) {
+                if (!optCtx.Types->PgTypes && node->Head().Content() == "count" || node->Head().Content() == "row_number") {
                     ret = ctx.Builder(node->Pos())
                         .Callable("SafeCast")
                             .Add(0, ret)
@@ -930,6 +930,14 @@ TExprNode::TPtr BuildWindows(TPositionHandle pos, const TExprNode::TPtr& list, c
                         .Build();
                 }
 
+                if (optCtx.Types->PgTypes && node->Head().Content() == "row_number") {
+                    ret = ctx.Builder(node->Pos())
+                        .Callable("ToPg")
+                            .Add(0, ret)
+                        .Seal()
+                        .Build();
+                }
+
                 return ret;
             }
 
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 fd50cb8301e..21791599837 100644
--- a/ydb/library/yql/core/type_ann/type_ann_core.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_core.cpp
@@ -11243,7 +11243,6 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         Functions["SqlProjectItem"] = &SqlProjectItemWrapper;
         Functions["SqlProjectStarItem"] = &SqlProjectItemWrapper;
         Functions["PgStar"] = &PgStarWrapper;
-        Functions["PgWindowCall"] = &PgWindowCallWrapper;
         Functions["PgQualifiedStar"] = &PgQualifiedStarWrapper;
         Functions["PgColumnRef"] = &PgColumnRefWrapper;
         Functions["PgResultItem"] = &PgResultItemWrapper;
@@ -11395,6 +11394,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         ExtFunctions["PgAggWindowCall"] = &PgAggWrapper;
         ExtFunctions["PgCall"] = &PgCallWrapper;
         ExtFunctions["PgResolvedCall"] = &PgCallWrapper;
+        ExtFunctions["PgWindowCall"] = &PgWindowCallWrapper;
         ExtFunctions["PgResolvedCallCtx"] = &PgCallWrapper;
         ExtFunctions["PgOp"] = &PgOpWrapper;
         ExtFunctions["PgResolvedOp"] = &PgOpWrapper;
diff --git a/ydb/library/yql/core/type_ann/type_ann_list.cpp b/ydb/library/yql/core/type_ann/type_ann_list.cpp
index 1fe8d5d3eb6..d0a8e2b4ae6 100644
--- a/ydb/library/yql/core/type_ann/type_ann_list.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_list.cpp
@@ -1092,6 +1092,13 @@ namespace {
 
                 lambdaItemType = lambda->GetTypeAnn()->Cast<TOptionalExprType>()->GetItemType();
                 break;
+            case ETypeAnnotationKind::Pg:
+                if (input->Content().EndsWith("Warn")) {
+                    warn = true;
+                }
+
+                lambdaItemType = lambda->GetTypeAnn();
+                break;
             case ETypeAnnotationKind::Null: {
                 if (input->Content().EndsWith("Warn")) {
                     warn = true;
diff --git a/ydb/library/yql/core/type_ann/type_ann_pg.cpp b/ydb/library/yql/core/type_ann/type_ann_pg.cpp
index 94c01c68caa..53431940062 100644
--- a/ydb/library/yql/core/type_ann/type_ann_pg.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_pg.cpp
@@ -93,6 +93,18 @@ IGraphTransformer::TStatus PgCallWrapper(const TExprNode::TPtr& input, TExprNode
                 return IGraphTransformer::TStatus::Error;
             }
 
+            if (proc.Kind == NPg::EProcKind::Window) {
+                ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
+                    TStringBuilder() << "Window function " << name << " cannot be called directly with window specification"));
+                return IGraphTransformer::TStatus::Error;
+            }
+
+            if (proc.Kind == NPg::EProcKind::Aggregate) {
+                ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
+                    TStringBuilder() << "Aggregate function " << name << " cannot be called directly "));
+                return IGraphTransformer::TStatus::Error;
+            }
+
             auto result = ctx.Expr.MakeType<TPgExprType>(proc.ResultType);
             input->SetTypeAnn(result);
             return IGraphTransformer::TStatus::Ok;
@@ -345,7 +357,7 @@ IGraphTransformer::TStatus PgOpWrapper(const TExprNode::TPtr& input, TExprNode::
     }
 }
 
-IGraphTransformer::TStatus PgWindowCallWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+IGraphTransformer::TStatus PgWindowCallWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExtContext& ctx) {
     Y_UNUSED(output);
     if (!EnsureMinArgsCount(*input, 2, ctx.Expr)) {
         return IGraphTransformer::TStatus::Error;
@@ -370,9 +382,9 @@ IGraphTransformer::TStatus PgWindowCallWrapper(const TExprNode::TPtr& input, TEx
         }
 
         auto arg = input->Child(2)->GetTypeAnn();
-        if (arg->GetKind() == ETypeAnnotationKind::Null) {
-            input->SetTypeAnn(arg);
-        } else if (arg->GetKind() == ETypeAnnotationKind::Optional) {
+        if (arg->GetKind() == ETypeAnnotationKind::Null ||
+            arg->GetKind() == ETypeAnnotationKind::Optional ||
+            arg->GetKind() == ETypeAnnotationKind::Pg) {
             input->SetTypeAnn(arg);
         } else {
             input->SetTypeAnn(ctx.Expr.MakeType<TOptionalExprType>(arg));
@@ -384,8 +396,11 @@ IGraphTransformer::TStatus PgWindowCallWrapper(const TExprNode::TPtr& input, TEx
             return IGraphTransformer::TStatus::Error;
         }
 
-        auto result = ctx.Expr.MakeType<TOptionalExprType>(ctx.Expr.MakeType<TDataExprType>(EDataSlot::Int64));
-        input->SetTypeAnn(result);
+        if (ctx.Types.PgTypes) {
+            input->SetTypeAnn(ctx.Expr.MakeType<TPgExprType>(NPg::LookupType("int8").TypeId));
+        } else {
+            input->SetTypeAnn(ctx.Expr.MakeType<TOptionalExprType>(ctx.Expr.MakeType<TDataExprType>(EDataSlot::Int64)));
+        }
     } else {
         ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
             TStringBuilder() << "Unsupported function: " << name));
diff --git a/ydb/library/yql/core/type_ann/type_ann_pg.h b/ydb/library/yql/core/type_ann/type_ann_pg.h
index 9034916df84..7ebb538eea4 100644
--- a/ydb/library/yql/core/type_ann/type_ann_pg.h
+++ b/ydb/library/yql/core/type_ann/type_ann_pg.h
@@ -13,7 +13,7 @@ IGraphTransformer::TStatus PgCallWrapper(const TExprNode::TPtr& input, TExprNode
 IGraphTransformer::TStatus FromPgWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
 IGraphTransformer::TStatus ToPgWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
 IGraphTransformer::TStatus PgOpWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExtContext& ctx);
-IGraphTransformer::TStatus PgWindowCallWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
+IGraphTransformer::TStatus PgWindowCallWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExtContext& ctx);
 IGraphTransformer::TStatus PgAggWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExtContext& ctx);
 IGraphTransformer::TStatus PgQualifiedStarWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
 IGraphTransformer::TStatus PgColumnRefWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
diff --git a/ydb/library/yql/core/yql_expr_type_annotation.cpp b/ydb/library/yql/core/yql_expr_type_annotation.cpp
index ca825f41d45..4b05574b05e 100644
--- a/ydb/library/yql/core/yql_expr_type_annotation.cpp
+++ b/ydb/library/yql/core/yql_expr_type_annotation.cpp
@@ -3220,6 +3220,14 @@ bool EnsureNewSeqType(TPositionHandle position, const TTypeAnnotationNode& type,
                 return true;
             }
             break;
+        case ETypeAnnotationKind::Pg:
+            if constexpr (WithOptional) {
+                if (itemType) {
+                    *itemType = &type;
+                }
+                return true;
+            }
+            break;
         default: break;
     }
     ctx.AddError(TIssue(ctx.GetPosition(position), TStringBuilder() <<
@@ -3339,6 +3347,7 @@ bool EnsureCodeResourceType(const TExprNode& node, TExprContext& ctx) {
 const TTypeAnnotationNode* MakeSequenceType(ETypeAnnotationKind sequenceKind, const TTypeAnnotationNode& itemType, TExprContext& ctx) {
     switch (sequenceKind) {
         case ETypeAnnotationKind::Optional: return ctx.MakeType<TOptionalExprType>(&itemType);
+        case ETypeAnnotationKind::Pg:       return &itemType;
         case ETypeAnnotationKind::Flow:     return ctx.MakeType<TFlowExprType>(&itemType);
         case ETypeAnnotationKind::List:     return ctx.MakeType<TListExprType>(&itemType);
         case ETypeAnnotationKind::Stream:   return ctx.MakeType<TStreamExprType>(&itemType);
diff --git a/ydb/library/yql/core/yql_opt_window.cpp b/ydb/library/yql/core/yql_opt_window.cpp
index 68d8055036d..229b6d3fa51 100644
--- a/ydb/library/yql/core/yql_opt_window.cpp
+++ b/ydb/library/yql/core/yql_opt_window.cpp
@@ -115,7 +115,9 @@ TExprNode::TPtr AddOptionalIfNotAlreadyOptionalOrNull(const TExprNode::TPtr& lam
                 .Add(2, identity)
                 .Atom(3, "Null", TNodeFlags::Default)
                 .Add(4, identity)
-                .Lambda(5)
+                .Atom(5, "Pg", TNodeFlags::Default)
+                .Add(6, identity)
+                .Lambda(7)
                     .Param("result")
                     .Callable("Just")
                         .Arg(0, "result")
diff --git a/ydb/library/yql/minikql/mkql_node_builder.cpp b/ydb/library/yql/minikql/mkql_node_builder.cpp
index 34466b89446..5e2baa589cd 100644
--- a/ydb/library/yql/minikql/mkql_node_builder.cpp
+++ b/ydb/library/yql/minikql/mkql_node_builder.cpp
@@ -74,6 +74,11 @@ TType* UnpackOptional(TRuntimeNode data, bool& isOptional) {
 }
 
 TType* UnpackOptional(TType* type, bool& isOptional) {
+    if (type->IsPg()) {
+        isOptional = true;
+        return type;
+    }
+
     isOptional = type->IsOptional();
     if (!isOptional)
         return type;
diff --git a/ydb/library/yql/minikql/mkql_program_builder.cpp b/ydb/library/yql/minikql/mkql_program_builder.cpp
index c05f7617f2c..461a280a634 100644
--- a/ydb/library/yql/minikql/mkql_program_builder.cpp
+++ b/ydb/library/yql/minikql/mkql_program_builder.cpp
@@ -2063,6 +2063,10 @@ TType* TProgramBuilder::NewDecimalType(ui8 precision, ui8 scale) {
 }
 
 TType* TProgramBuilder::NewOptionalType(TType* itemType) {
+    if (itemType->IsPg()) {
+        return itemType;
+    }
+
     return TOptionalType::Create(itemType, Env);
 }
 
@@ -2614,7 +2618,7 @@ TRuntimeNode TProgramBuilder::IfPresent(TRuntimeNode optional, const TUnaryLambd
 
     MKQL_ENSURE(thenUnpacked->IsSameType(*elseUnpacked), "Different return types in branches.");
 
-    TCallableBuilder callableBuilder(Env, __func__, thenOpt || elseOpt ? NewOptionalType(thenUnpacked) : thenUnpacked);
+    TCallableBuilder callableBuilder(Env, __func__, (thenOpt || elseOpt) ? NewOptionalType(thenUnpacked) : thenUnpacked);
     callableBuilder.Add(optional);
     callableBuilder.Add(itemArg);
     callableBuilder.Add(then);
@@ -3240,10 +3244,20 @@ TRuntimeNode TProgramBuilder::BuildFlatMap(const std::string_view& callableName,
         const auto itemArg = Arg(AS_TYPE(TOptionalType, listType)->GetItemType());
         const auto newList = handler(itemArg);
         const auto type = newList.GetStaticType();
-        MKQL_ENSURE(type->IsList() || type->IsOptional() || type->IsStream() || type->IsFlow(), "Expected flow, list, stream or optional");
+        MKQL_ENSURE(type->IsList() || type->IsPg() || type->IsOptional() || type->IsStream() || type->IsFlow(), "Expected flow, list, stream or optional");
+        return IfPresent(list, [&](TRuntimeNode item) {
+            return handler(item);
+        }, (type->IsOptional() || type->IsPg()) ? NewEmptyOptional(type) : type->IsList() ? NewEmptyList(AS_TYPE(TListType, type)->GetItemType()) : EmptyIterator(type));
+    }
+
+    if (listType->IsPg()) {
+        const auto itemArg = Arg(listType);
+        const auto newList = handler(itemArg);
+        const auto type = newList.GetStaticType();
+        MKQL_ENSURE(type->IsList() || type->IsPg() || type->IsOptional() || type->IsStream() || type->IsFlow(), "Expected flow, list, stream or optional");
         return IfPresent(list, [&](TRuntimeNode item) {
             return handler(item);
-        }, type->IsOptional() ? NewEmptyOptional(type) : type->IsList() ? NewEmptyList(AS_TYPE(TListType, type)->GetItemType()) : EmptyIterator(type));
+        }, (type->IsOptional() || type->IsPg()) ? NewEmptyOptional(type) : type->IsList() ? NewEmptyList(AS_TYPE(TListType, type)->GetItemType()) : EmptyIterator(type));
     }
 
     const auto itemType = listType->IsFlow() ?