YQL-14728 settings in PgCall and related functions

ref:fff2e4ec3422f27aa928c8875ef3854a7768b259

5 files changed, 156 insertions, 75 deletions

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 47377d76dad..2b664b130be 100644
--- a/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
@@ -52,10 +52,10 @@ TExprNode::TPtr JoinColumns(TPositionHandle pos, const TExprNode::TPtr& list1, c
                 .Add(0, list2)
                 .Atom(1, "b")
             .Seal()
-            .List(2)

-                .Atom(0, leftJoinColumns ? "Left" : "Cross")

-                .Atom(1, "a")

-                .Atom(2, "b")

+            .List(2)
+                .Atom(0, leftJoinColumns ? "Left" : "Cross")
+                .Atom(1, "a")
+                .Atom(2, "b")
                 .List(3)
                     .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder & {
                         if (leftJoinColumns) {
@@ -67,7 +67,7 @@ TExprNode::TPtr JoinColumns(TPositionHandle pos, const TExprNode::TPtr& list1, c
 
                         return parent;
                     })
-                .Seal()

+                .Seal()
                 .List(4)
                     .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder & {
                         if (leftJoinColumns) {
@@ -79,12 +79,12 @@ TExprNode::TPtr JoinColumns(TPositionHandle pos, const TExprNode::TPtr& list1, c
 
                         return parent;
                     })
-                .Seal()

-                .List(5)

-                .Seal()

+                .Seal()
+                .List(5)
+                .Seal()
             .Seal()
             .List(3)
-            .Seal()

+            .Seal()
         .Seal()
         .Build();
 
@@ -266,15 +266,15 @@ std::pair<TExprNode::TPtr, TExprNode::TPtr> RewriteSubLinks(TPositionHandle pos,
 
                 auto select = ExpandPgSelectSublink(node->TailPtr(), ctx, optCtx, uniqueOuterList);
 
-                auto exportsPtr = optCtx.Types->Modules->GetModule("/lib/yql/aggregate.yql");

+                auto exportsPtr = optCtx.Types->Modules->GetModule("/lib/yql/aggregate.yql");
                 YQL_ENSURE(exportsPtr);
                 const auto& exports = exportsPtr->Symbols();
 
-                auto selectTypeNode = ctx.Builder(node->Pos())

-                    .Callable("TypeOf")

-                        .Add(0, select)

-                    .Seal()

-                    .Build();

+                auto selectTypeNode = ctx.Builder(node->Pos())
+                    .Callable("TypeOf")
+                        .Add(0, select)
+                    .Seal()
+                    .Build();
 
                 TExprNode::TPtr countAllTraits;
                 TExprNode::TPtr someTraits;
@@ -282,7 +282,7 @@ std::pair<TExprNode::TPtr, TExprNode::TPtr> RewriteSubLinks(TPositionHandle pos,
                 for (ui32 factoryIndex = 0; factoryIndex < 3; ++factoryIndex)
                 {
                     TStringBuf name;
-                    switch (factoryIndex) {

+                    switch (factoryIndex) {
                     case 0:
                         if (linkType != "exists" && linkType != "expr") {
                             continue;
@@ -296,23 +296,23 @@ std::pair<TExprNode::TPtr, TExprNode::TPtr> RewriteSubLinks(TPositionHandle pos,
                         }
 
                         name = "some_traits_factory";
-                        break;

+                        break;
                     case 2:
                         if (linkType != "any" && linkType != "all") {
                             continue;
                         }
 
                         name = "count_if_traits_factory";
-                        break;

-                    }

+                        break;
+                    }
 
-                    const auto ex = exports.find(name);

-                    YQL_ENSURE(exports.cend() != ex);

+                    const auto ex = exports.find(name);
+                    YQL_ENSURE(exports.cend() != ex);
                     auto lambda = ctx.DeepCopy(*ex->second, exportsPtr->ExprCtx(), deepClones, true, false);
-                    auto arg = ctx.NewArgument(node->Pos(), "row");

+                    auto arg = ctx.NewArgument(node->Pos(), "row");
                     auto arguments = ctx.NewArguments(node->Pos(), { arg });
                     TExprNode::TPtr root;
-                    switch (factoryIndex) {

+                    switch (factoryIndex) {
                     case 0:
                         root = arg;
                         break;
@@ -340,29 +340,29 @@ std::pair<TExprNode::TPtr, TExprNode::TPtr> RewriteSubLinks(TPositionHandle pos,
                         root = ctx.NewCallable(node->Pos(), "FromPg", { testRoot });
                         break;
                     }
-                    }

-

-                    auto extractor = ctx.NewLambda(node->Pos(), std::move(arguments), std::move(root));

-

-                    auto traits = ctx.ReplaceNodes(lambda->TailPtr(), {

-                        {lambda->Head().Child(0), selectTypeNode},

-                        {lambda->Head().Child(1), extractor}

-                    });

-

-                    ctx.Step.Repeat(TExprStep::ExpandApplyForLambdas);

-                    auto status = ExpandApply(traits, traits, ctx);

+                    }
+
+                    auto extractor = ctx.NewLambda(node->Pos(), std::move(arguments), std::move(root));
+
+                    auto traits = ctx.ReplaceNodes(lambda->TailPtr(), {
+                        {lambda->Head().Child(0), selectTypeNode},
+                        {lambda->Head().Child(1), extractor}
+                    });
+
+                    ctx.Step.Repeat(TExprStep::ExpandApplyForLambdas);
+                    auto status = ExpandApply(traits, traits, ctx);
                     YQL_ENSURE(status != IGraphTransformer::TStatus::Error);
-                    switch (factoryIndex) {

+                    switch (factoryIndex) {
                     case 0:
                         countAllTraits = traits;
                         break;
                     case 1:
                         someTraits = traits;
-                        break;

+                        break;
                     case 2:
                         countIfTraits = traits;
-                        break;

-                    }

+                        break;
+                    }
                 }
 
                 auto columnName = "_yql_sublink_" + ToString(sublinkColumnIndex++);
@@ -1057,7 +1057,7 @@ std::tuple<TAggs, TNodeMap<ui32>> GatherAggregations(const TExprNode::TPtr& proj
             return false;
         }
 
-        if (node->IsCallable("PgAgg") || node->IsCallable("PgAggAll")) {
+        if (node->IsCallable("PgAgg")) {
             aggId[node.Get()] = aggs.size();
             aggs.push_back({ node, projectionLambda->Head().HeadPtr() });
         }
@@ -1072,7 +1072,7 @@ std::tuple<TAggs, TNodeMap<ui32>> GatherAggregations(const TExprNode::TPtr& proj
                 return false;
             }
 
-            if (node->IsCallable("PgAgg") || node->IsCallable("PgAggAll")) {
+            if (node->IsCallable("PgAgg")) {
                 aggId[node.Get()] = aggs.size();
                 aggs.push_back({ node, havingLambda->Head().HeadPtr() });
             }
@@ -1091,7 +1091,7 @@ TExprNode::TPtr BuildAggregationTraits(TPositionHandle pos, bool onWindow,
     auto arguments = ctx.NewArguments(pos, { arg });
     auto func = agg.first->Head().Content();
     TExprNode::TListType aggFuncArgs;
-    for (ui32 j = onWindow ? 2 : 1; j < agg.first->ChildrenSize(); ++j) {
+    for (ui32 j = onWindow ? 3 : 2; j < agg.first->ChildrenSize(); ++j) {
         aggFuncArgs.push_back(ctx.ReplaceNode(agg.first->ChildPtr(j), *agg.second, arg));
     }
 
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 a734b1dca19..d5064eee1b5 100644
--- a/ydb/library/yql/core/type_ann/type_ann_pg.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_pg.cpp
@@ -52,7 +52,7 @@ IGraphTransformer::TStatus PgStarWrapper(const TExprNode::TPtr& input, TExprNode
 
 IGraphTransformer::TStatus PgCallWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExtContext& ctx) {
     bool isResolved = input->Content().StartsWith("PgResolvedCall");
-    if (!EnsureMinArgsCount(*input, isResolved ? 2 : 1, ctx.Expr)) {
+    if (!EnsureMinArgsCount(*input, isResolved ? 3 : 2, ctx.Expr)) {
         return IGraphTransformer::TStatus::Error;
     }
 
@@ -68,9 +68,29 @@ IGraphTransformer::TStatus PgCallWrapper(const TExprNode::TPtr& input, TExprNode
         }
     }
 
+    if (!EnsureTuple(*input->Child(isResolved ? 2 : 1), ctx.Expr)) {
+        return IGraphTransformer::TStatus::Error;
+    }
+
+    for (const auto& setting : input->Child(isResolved ? 2 : 1)->Children()) {
+        if (!EnsureTupleMinSize(*setting, 1, ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        if (!EnsureAtom(setting->Head(), ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        auto content = setting->Head().Content();
+        ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
+            TStringBuilder() << "Unexpected setting " << content << " in function " << name));
+
+        return IGraphTransformer::TStatus::Error;
+    }
+
     TVector<ui32> argTypes;
     bool needRetype = false;
-    for (ui32 i = isResolved ? 2 : 1; i < input->ChildrenSize(); ++i) {
+    for (ui32 i = isResolved ? 3 : 2; i < input->ChildrenSize(); ++i) {
         auto type = input->Child(i)->GetTypeAnn();
         ui32 argType;
         bool convertToPg;
@@ -305,7 +325,7 @@ IGraphTransformer::TStatus PgOpWrapper(const TExprNode::TPtr& input, TExprNode::
 
 IGraphTransformer::TStatus PgWindowCallWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExtContext& ctx) {
     Y_UNUSED(output);
-    if (!EnsureMinArgsCount(*input, 2, ctx.Expr)) {
+    if (!EnsureMinArgsCount(*input, 3, ctx.Expr)) {
         return IGraphTransformer::TStatus::Error;
     }
 
@@ -320,23 +340,40 @@ IGraphTransformer::TStatus PgWindowCallWrapper(const TExprNode::TPtr& input, TEx
         return IGraphTransformer::TStatus::Error;
     }
 
+    if (!EnsureTuple(*input->Child(2), ctx.Expr)) {
+        return IGraphTransformer::TStatus::Error;
+    }
+
+    for (const auto& setting : input->Child(2)->Children()) {
+        if (!EnsureTupleMinSize(*setting, 1, ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        if (!EnsureAtom(setting->Head(), ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        auto content = setting->Head().Content();
+        ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
+            TStringBuilder() << "Unexpected setting " << content << " in window function " << name));
+        return IGraphTransformer::TStatus::Error;
+    }
+
     if (name == "lead" || name == "lag") {
-        if (input->ChildrenSize() != 3) {
+        if (input->ChildrenSize() != 4) {
             ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
                 TStringBuilder() << "Expected one argument in " << name << " function"));
             return IGraphTransformer::TStatus::Error;
         }
 
-        auto arg = input->Child(2)->GetTypeAnn();
-        if (arg->GetKind() == ETypeAnnotationKind::Null ||
-            arg->GetKind() == ETypeAnnotationKind::Optional ||
-            arg->GetKind() == ETypeAnnotationKind::Pg) {
+        auto arg = input->Child(3)->GetTypeAnn();
+        if (arg->IsOptionalOrNull()) {
             input->SetTypeAnn(arg);
         } else {
             input->SetTypeAnn(ctx.Expr.MakeType<TOptionalExprType>(arg));
         }
     } else if (name == "row_number") {
-        if (input->ChildrenSize() != 2) {
+        if (input->ChildrenSize() != 3) {
             ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
                 "Expected no arguments in row_number function"));
             return IGraphTransformer::TStatus::Error;
@@ -355,7 +392,7 @@ IGraphTransformer::TStatus PgWindowCallWrapper(const TExprNode::TPtr& input, TEx
 IGraphTransformer::TStatus PgAggWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExtContext& ctx) {
     Y_UNUSED(output);
     bool overWindow = (input->Content() == "PgAggWindowCall");
-    if (!EnsureMinArgsCount(*input, overWindow ? 2 : 1, ctx.Expr)) {
+    if (!EnsureMinArgsCount(*input, overWindow ? 3 : 2, ctx.Expr)) {
         return IGraphTransformer::TStatus::Error;
     }
 
@@ -372,9 +409,28 @@ IGraphTransformer::TStatus PgAggWrapper(const TExprNode::TPtr& input, TExprNode:
         }
     }
 
+    if (!EnsureTuple(*input->Child(overWindow ? 2 : 1), ctx.Expr)) {
+        return IGraphTransformer::TStatus::Error;
+    }
+
+    for (const auto& setting : input->Child(overWindow ? 2 : 1)->Children()) {
+        if (!EnsureTupleMinSize(*setting, 1, ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        if (!EnsureAtom(setting->Head(), ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        auto content = setting->Head().Content();
+        ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
+            TStringBuilder() << "Unexpected setting " << content << " in aggregate function " << name));
+        return IGraphTransformer::TStatus::Error;
+    }
+
     TVector<ui32> argTypes;
     bool needRetype = false;
-    for (ui32 i = overWindow ? 2 : 1; i < input->ChildrenSize(); ++i) {
+    for (ui32 i = overWindow ? 3 : 2; i < input->ChildrenSize(); ++i) {
         auto type = input->Child(i)->GetTypeAnn();
         ui32 argType;
         bool convertToPg;
@@ -1012,7 +1068,9 @@ IGraphTransformer::TStatus PgAggregationTraitsWrapper(const TExprNode::TPtr& inp
             .Callable("PgResolvedCallCtx")
                 .Atom(0, NPg::LookupProc(aggDesc.FinalFuncId).Name)
                 .Atom(1, ToString(aggDesc.FinalFuncId))
-                .Arg(2, "state")
+                .List(2)
+                .Seal()
+                .Arg(3, "state")
             .Seal()
             .Seal()
             .Build();
@@ -1052,8 +1110,10 @@ IGraphTransformer::TStatus PgAggregationTraitsWrapper(const TExprNode::TPtr& inp
                 .Callable("PgResolvedCallCtx")
                     .Atom(0, transFuncDesc.Name)
                     .Atom(1, ToString(aggDesc.TransFuncId))
-                    .Add(2, initValue)
-                    .Apply(3, lambda)
+                    .List(2)
+                    .Seal()
+                    .Add(3, initValue)
+                    .Apply(4, lambda)
                         .With(0, "row")
                     .Seal()
                 .Seal()
@@ -1068,11 +1128,13 @@ IGraphTransformer::TStatus PgAggregationTraitsWrapper(const TExprNode::TPtr& inp
                     .Callable(0, "PgResolvedCallCtx")
                         .Atom(0, transFuncDesc.Name)
                         .Atom(1, ToString(aggDesc.TransFuncId))
-                        .Callable(2, "Coalesce")
+                        .List(2)
+                        .Seal()
+                        .Callable(3, "Coalesce")
                             .Arg(0, "state")
                             .Add(1, initValue)
                         .Seal()
-                        .Apply(3, lambda)
+                        .Apply(4, lambda)
                             .With(0, "row")
                         .Seal()
                     .Seal()
@@ -1118,8 +1180,10 @@ IGraphTransformer::TStatus PgAggregationTraitsWrapper(const TExprNode::TPtr& inp
                         .Callable(0, "PgResolvedCallCtx")
                             .Atom(0, transFuncDesc.Name)
                             .Atom(1, ToString(aggDesc.TransFuncId))
-                            .Arg(2, "state")
-                            .Apply(3, lambda)
+                            .List(2)
+                            .Seal()
+                            .Arg(3, "state")
+                            .Apply(4, lambda)
                                 .With(0, "row")
                             .Seal()
                         .Seal()
@@ -1161,7 +1225,9 @@ IGraphTransformer::TStatus PgAggregationTraitsWrapper(const TExprNode::TPtr& inp
                 .Callable("PgResolvedCallCtx")
                     .Atom(0, serializeFuncDesc.Name)
                     .Atom(1, ToString(aggDesc.SerializeFuncId))
-                    .Arg(2, "state")
+                    .List(2)
+                    .Seal()
+                    .Arg(3, "state")
                 .Seal()
             .Seal()
             .Build();
@@ -1175,8 +1241,10 @@ IGraphTransformer::TStatus PgAggregationTraitsWrapper(const TExprNode::TPtr& inp
                 .Callable("PgResolvedCallCtx")
                     .Atom(0, deserializeFuncDesc.Name)
                     .Atom(1, ToString(aggDesc.DeserializeFuncId))
-                    .Arg(2, "state")
-                    .Callable(3, "PgInternal0")
+                    .List(2)
+                    .Seal()
+                    .Arg(3, "state")
+                    .Callable(4, "PgInternal0")
                     .Seal()
                 .Seal()
             .Seal()
@@ -1198,8 +1266,10 @@ IGraphTransformer::TStatus PgAggregationTraitsWrapper(const TExprNode::TPtr& inp
                             .Callable(0, "PgResolvedCallCtx")
                                 .Atom(0, combineFuncDesc.Name)
                                 .Atom(1, ToString(aggDesc.CombineFuncId))
-                                .Arg(2, "state1")
-                                .Arg(3, "state2")
+                                .List(2)
+                                .Seal()
+                                .Arg(3, "state1")
+                                .Arg(4, "state2")
                             .Seal()
                             .Arg(1, "state1")
                         .Seal()
@@ -1215,8 +1285,10 @@ IGraphTransformer::TStatus PgAggregationTraitsWrapper(const TExprNode::TPtr& inp
                     .Callable("PgResolvedCallCtx")
                         .Atom(0, combineFuncDesc.Name)
                         .Atom(1, ToString(aggDesc.CombineFuncId))
-                        .Arg(2, "state1")
-                        .Arg(3, "state2")
+                        .List(2)
+                        .Seal()
+                        .Arg(3, "state1")
+                        .Arg(4, "state2")
                     .Seal()
                 .Seal()
                 .Build();
@@ -3053,7 +3125,9 @@ IGraphTransformer::TStatus PgTypeModWrapper(const TExprNode::TPtr& input, TExprN
     output = ctx.Expr.Builder(input->Pos())
         .Callable("PgCall")
             .Atom(0, NPg::LookupProc(typeDesc.TypeModInFuncId, { 0 }).Name)
-            .Add(1, arr)
+            .List(1)
+            .Seal()
+            .Add(2, arr)
         .Seal()
         .Build();
 
diff --git a/ydb/library/yql/providers/common/mkql/yql_provider_mkql.cpp b/ydb/library/yql/providers/common/mkql/yql_provider_mkql.cpp
index 08bdb52f2fc..710b6dbe585 100644
--- a/ydb/library/yql/providers/common/mkql/yql_provider_mkql.cpp
+++ b/ydb/library/yql/providers/common/mkql/yql_provider_mkql.cpp
@@ -2257,8 +2257,8 @@ TMkqlCommonCallableCompiler::TShared::TShared() {
         auto name = node.Head().Content();
         auto id = FromString<ui32>(node.Child(1)->Content());
         std::vector<TRuntimeNode> args;
-        args.reserve(node.ChildrenSize() - 2);
-        for (ui32 i = 2; i < node.ChildrenSize(); ++i) {
+        args.reserve(node.ChildrenSize() - 3);
+        for (ui32 i = 3; i < node.ChildrenSize(); ++i) {
             args.push_back(MkqlBuildExpr(*node.Child(i), ctx));
         }
 
diff --git a/ydb/library/yql/sql/pg/pg_sql.cpp b/ydb/library/yql/sql/pg/pg_sql.cpp
index e64560a9369..b75ed39af09 100644
--- a/ydb/library/yql/sql/pg/pg_sql.cpp
+++ b/ydb/library/yql/sql/pg/pg_sql.cpp
@@ -1343,11 +1343,6 @@ public:
             return nullptr;
         }
 
-        if (value->agg_distinct) {
-            AddError("FuncCall: unsupported agg_distinct");
-            return nullptr;
-        }
-
         if (value->func_variadic) {
             AddError("FuncCall: unsupported func_variadic");
             return nullptr;
@@ -1436,6 +1431,17 @@ public:
             args.push_back(window);
         }
 
+        TVector<TAstNode*> callSettings;
+        if (value->agg_distinct) {
+            if (!isAggregateFunc) {
+                AddError("FuncCall: agg_distinct must be set only for aggregate functions");
+                return nullptr;
+            }
+
+            callSettings.push_back(QL(QA("distinct")));
+        }
+
+        args.push_back(QVL(callSettings.data(), callSettings.size()));
         if (value->agg_star) {
             if (name != "count") {
                 AddError("FuncCall: * is expected only in count function");
@@ -1553,7 +1559,7 @@ public:
                         args.push_back(L(A("PgConst"), QAX(s), L(A("PgType"), QA("cstring"))));
                     }
 
-                    typeMod = L(A("PgCall"), QA(procDesc.Name), VL(args.data(), args.size()));
+                    typeMod = L(A("PgCall"), QA(procDesc.Name), QL(), VL(args.data(), args.size()));
                 }
 
                 return L(A("PgCast"), input, L(A("PgType"), QAX(finalType)), typeMod);
diff --git a/ydb/library/yql/sql/v1/builtin.cpp b/ydb/library/yql/sql/v1/builtin.cpp
index d26811c75d7..78139688c7d 100644
--- a/ydb/library/yql/sql/v1/builtin.cpp
+++ b/ydb/library/yql/sql/v1/builtin.cpp
@@ -739,6 +739,7 @@ public:
         }
 
         Args[0] = BuildQuotedAtom(Args[0]->GetPos(), Args[0]->GetLiteralValue());
+        Args.insert(Args.begin() + 1, Q(Y()));
         return TCallNode::DoInit(ctx, src);
     }