Added defaults args values in PgCall (#7421)

vitstn · web-flow · commit f64d67099796 · 2024-08-05T09:11:43.000+05:00
diff --git a/ydb/library/yql/core/type_ann/type_ann_pg.cpp b/ydb/library/yql/core/type_ann/type_ann_pg.cpp
@@ -307,9 +307,9 @@ IGraphTransformer::TStatus PgCallWrapper(const TExprNode::TPtr& input, TExprNode
     if (isResolved) {
         auto procId = FromString<ui32>(input->Child(1)->Content());
         const auto& proc = NPg::LookupProc(procId, argTypes);
-        if (proc.Name != name) {
+        if (!AsciiEqualsIgnoreCase(proc.Name, name)) {
             ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
-                TStringBuilder() << "Mismatch of resolved function name, expected: " << name << ", but got:" << proc.Name));
+                TStringBuilder() << "Mismatch of resolved function name, expected: " << name << ", but got: " << proc.Name));
             return IGraphTransformer::TStatus::Error;
         }
 
@@ -378,6 +378,27 @@ IGraphTransformer::TStatus PgCallWrapper(const TExprNode::TPtr& input, TExprNode
                         children[i+3] = WrapWithPgCast(std::move(children[i+3]), targetType, ctx.Expr);
                     }
                 }
+
+                if (argTypes.size() < fargTypes.size()) {
+                    YQL_ENSURE(fargTypes.size() - argTypes.size() <= (*procPtr)->DefaultArgs.size());
+                    for (size_t i = argTypes.size(); i < fargTypes.size(); ++i) {
+                        const auto& value = (*procPtr)->DefaultArgs[i + (*procPtr)->DefaultArgs.size() - fargTypes.size()];
+                        TExprNode::TPtr defNode;
+                        if (!value) {
+                            defNode = ctx.Expr.NewCallable(input->Pos(), "Null", {});
+                        } else {
+                            defNode = ctx.Expr.Builder(input->Pos())
+                                .Callable("PgConst")
+                                    .Atom(0, *value)
+                                    .Callable(1, "PgType")
+                                        .Atom(0, NPg::LookupType(fargTypes[i]).Name)
+                                    .Seal()
+                                .Seal()
+                                .Build();
+                        }
+                        children.insert(children.end(), defNode);
+                    }
+                }
                 output = ctx.Expr.NewCallable(input->Pos(), "PgResolvedCall", std::move(children));
             } else if (const auto* typePtr = std::get_if<const NPg::TTypeDesc*>(&procOrType)) {
                 output = WrapWithPgCast(std::move(children[2]), (*typePtr)->TypeId, ctx.Expr);
diff --git a/ydb/library/yql/minikql/mkql_type_builder.cpp b/ydb/library/yql/minikql/mkql_type_builder.cpp
@@ -26,14 +26,23 @@ namespace {
 static const TString UdfName("UDF");
 
 class TPgTypeIndex {
-    static constexpr ui32 MaxOid = 15000;
-    using TUdfTypes = std::array<NYql::NUdf::TPgTypeDescription, MaxOid>;
+    using TUdfTypes = TVector<NYql::NUdf::TPgTypeDescription>;
     TUdfTypes Types;
 
 public:
     TPgTypeIndex() {
+        Rebuild();
+    }
+
+    void Rebuild() {
+        Types.clear();
+        ui32 maxTypeId = 0;
+        NYql::NPg::EnumTypes([&](ui32 typeId, const NYql::NPg::TTypeDesc&) {
+            maxTypeId = Max(maxTypeId, typeId);
+        });
+
+        Types.resize(maxTypeId + 1);
         NYql::NPg::EnumTypes([&](ui32 typeId, const NYql::NPg::TTypeDesc& t) {
-            Y_ABORT_UNLESS(typeId < Types.size());
             auto& e = Types[typeId];
             e.Name = t.Name;
             e.TypeId = t.TypeId;
@@ -2680,5 +2689,9 @@ TType* TTypeBuilder::NewVariantType(TType* underlyingType) const {
     return TVariantType::Create(underlyingType, Env);
 }
 
+void RebuildTypeIndex() {
+    HugeSingleton<TPgTypeIndex>()->Rebuild();
+}
+
 } // namespace NMiniKQL
 } // namespace Nkikimr
diff --git a/ydb/library/yql/minikql/mkql_type_builder.h b/ydb/library/yql/minikql/mkql_type_builder.h
@@ -324,6 +324,7 @@ class TTypeBuilder : public TMoveOnly {
     bool UseNullType = true;
 };
 
+void RebuildTypeIndex();
 
 } // namespace NMiniKQL
 } // namespace Nkikimr
diff --git a/ydb/library/yql/parser/pg_catalog/catalog.cpp b/ydb/library/yql/parser/pg_catalog/catalog.cpp
@@ -2529,13 +2529,14 @@ bool IsExactMatch(const TVector<ui32>& procArgTypes, ui32 procVariadicType, cons
     return true;
 }
 
-ui64 CalcProcScore(const TVector<ui32>& procArgTypes, ui32 procVariadicType, const TVector<ui32>& argTypeIds, const TCatalog& catalog) {
+ui64 CalcProcScore(const TVector<ui32>& procArgTypes, ui32 procVariadicType, ui32 procDefArgs, const TVector<ui32>& argTypeIds, const TCatalog& catalog) {
     ui64 result = 0UL;
     if (!procVariadicType) {
         ++result;
     }
 
-    if (argTypeIds.size() < procArgTypes.size()) {
+    Y_ENSURE(procArgTypes.size() >= procDefArgs);
+    if (argTypeIds.size() < procArgTypes.size() - procDefArgs) {
         return ArgTypeMismatch;
     }
 
@@ -2804,7 +2805,7 @@ std::variant<const TProcDesc*, const TTypeDesc*> LookupProcWithCasts(const TStri
         }
 
         // https://www.postgresql.org/docs/14/typeconv-func.html, steps 4.a, 4.c, 4.d
-        auto score = NPrivate::CalcProcScore(d->ArgTypes, d->VariadicType, argTypeIds, catalog);
+        auto score = NPrivate::CalcProcScore(d->ArgTypes, d->VariadicType, d->DefaultArgs.size(), argTypeIds, catalog);
         if (bestScore < score) {
             bestScore = score;
 
@@ -3169,7 +3170,7 @@ const TAggregateDesc& LookupAggregation(const TString& name, const TVector<ui32>
         }
 
         // https://www.postgresql.org/docs/14/typeconv-func.html, steps 4.a, 4.c, 4.d
-        auto score = NPrivate::CalcProcScore(d->ArgTypes, 0, argTypeIds, catalog);
+        auto score = NPrivate::CalcProcScore(d->ArgTypes, 0, 0, argTypeIds, catalog);
 
         if (bestScore < score) {
             bestScore = score;
diff --git a/ydb/library/yql/parser/pg_catalog/catalog.h b/ydb/library/yql/parser/pg_catalog/catalog.h
@@ -72,6 +72,7 @@ struct TProcDesc {
     ui32 VariadicType = 0;
     ui32 VariadicArgType = 0;
     TString VariadicArgName;
+    TVector<TMaybe<TString>> DefaultArgs;
     ui32 ExtensionIndex = 0;
 };
 
diff --git a/ydb/library/yql/sql/pg/pg_sql.cpp b/ydb/library/yql/sql/pg/pg_sql.cpp
@@ -36,6 +36,7 @@ extern "C" {
 #include <ydb/library/yql/parser/pg_wrapper/postgresql/src/backend/catalog/pg_type_d.h>
 #include <ydb/library/yql/parser/pg_catalog/catalog.h>
 #include <ydb/library/yql/providers/common/provider/yql_provider_names.h>
+#include <ydb/library/yql/minikql/mkql_type_builder.h>
 #include <ydb/library/yql/core/issue/yql_issue.h>
 #include <ydb/library/yql/core/yql_callable_names.h>
 #include <ydb/library/yql/parser/pg_catalog/catalog.h>
@@ -5423,15 +5424,30 @@ class TExtensionHandler : public IPGParseEvents {
         }
 
         bool hasArgNames = false;
-        ui32 defArgsCount = 0;
         for (int i = 0; i < ListLength(value->parameters); ++i) {
             auto node = LIST_CAST_NTH(FunctionParameter, value->parameters, i);
             hasArgNames = hasArgNames || (node->name != nullptr);
             if (node->mode == FUNC_PARAM_IN || node->mode == FUNC_PARAM_DEFAULT) {
                 if (node->defexpr) {
-                    ++defArgsCount;
+                    desc.DefaultArgs.emplace_back();
+                    auto& value = desc.DefaultArgs.back();
+                    auto expr = node->defexpr;
+                    if (NodeTag(expr) == T_TypeCast) {
+                        expr = CAST_NODE(TypeCast, expr)->arg;
+                    }
+
+                    if (NodeTag(expr) != T_A_Const) {
+                        return false;
+                    }
+
+                    auto pgConst = GetValueNType(CAST_NODE(A_Const, expr));
+                    if (!pgConst) {
+                        return false;
+                    }
+
+                    value = pgConst->value;
                 } else {
-                    Y_ENSURE(!defArgsCount);
+                    Y_ENSURE(desc.DefaultArgs.empty());
                 }
 
                 desc.InputArgNames.push_back(node->name ? node->name : "");
@@ -5470,17 +5486,6 @@ class TExtensionHandler : public IPGParseEvents {
         }
 
         Builder.CreateProc(desc);
-        if (defArgsCount) {
-            Y_ENSURE(!desc.VariadicType);
-            for (ui32 i = 0; i < defArgsCount; ++i) {
-                desc.ArgTypes.pop_back();
-                if (!desc.InputArgNames.empty()) {
-                    desc.InputArgNames.pop_back();
-                }
-
-                Builder.CreateProc(desc);
-            }
-        }
         return true;
     }
 
@@ -5584,6 +5589,8 @@ class TExtensionSqlParser : public NYql::NPg::IExtensionSqlParser {
         for (const auto& sql : sqls) {
             NYql::PGParse(sql, handler);
         }
+
+        NKikimr::NMiniKQL::RebuildTypeIndex();
     }
 };
 
diff --git a/ydb/library/yql/sql/pg/ya.make b/ydb/library/yql/sql/pg/ya.make
@@ -8,6 +8,7 @@ PEERDIR(
     ydb/library/yql/ast
     ydb/library/yql/core
     ydb/library/yql/parser/pg_catalog
+    ydb/library/yql/minikql
     ydb/library/yql/sql/settings
     ydb/public/api/protos
 )

Original file line number	Diff line number	Diff line change
`@@ -2529,13 +2529,14 @@ bool IsExactMatch(const TVector<ui32>& procArgTypes, ui32 procVariadicType, cons`
`2529`	`2529`	`return true;`
`2530`	`2530`	`}`
`2531`	`2531`
`2532`		`-ui64 CalcProcScore(const TVector<ui32>& procArgTypes, ui32 procVariadicType, const TVector<ui32>& argTypeIds, const TCatalog& catalog) {`
	`2532`	`+ui64 CalcProcScore(const TVector<ui32>& procArgTypes, ui32 procVariadicType, ui32 procDefArgs, const TVector<ui32>& argTypeIds, const TCatalog& catalog) {`
`2533`	`2533`	`ui64 result = 0UL;`
`2534`	`2534`	`if (!procVariadicType) {`
`2535`	`2535`	`++result;`
`2536`	`2536`	`}`
`2537`	`2537`
`2538`		`- if (argTypeIds.size() < procArgTypes.size()) {`
	`2538`	`+ Y_ENSURE(procArgTypes.size() >= procDefArgs);`
	`2539`	`+ if (argTypeIds.size() < procArgTypes.size() - procDefArgs) {`
`2539`	`2540`	`return ArgTypeMismatch;`
`2540`	`2541`	`}`
`2541`	`2542`
`@@ -2804,7 +2805,7 @@ std::variant<const TProcDesc, const TTypeDesc> LookupProcWithCasts(const TStri`
`2804`	`2805`	`}`
`2805`	`2806`
`2806`	`2807`	`// https://www.postgresql.org/docs/14/typeconv-func.html, steps 4.a, 4.c, 4.d`
`2807`		`- auto score = NPrivate::CalcProcScore(d->ArgTypes, d->VariadicType, argTypeIds, catalog);`
	`2808`	`+ auto score = NPrivate::CalcProcScore(d->ArgTypes, d->VariadicType, d->DefaultArgs.size(), argTypeIds, catalog);`
`2808`	`2809`	`if (bestScore < score) {`
`2809`	`2810`	`bestScore = score;`
`2810`	`2811`
`@@ -3169,7 +3170,7 @@ const TAggregateDesc& LookupAggregation(const TString& name, const TVector<ui32>`
`3169`	`3170`	`}`
`3170`	`3171`
`3171`	`3172`	`// https://www.postgresql.org/docs/14/typeconv-func.html, steps 4.a, 4.c, 4.d`
`3172`		`- auto score = NPrivate::CalcProcScore(d->ArgTypes, 0, argTypeIds, catalog);`
	`3173`	`+ auto score = NPrivate::CalcProcScore(d->ArgTypes, 0, 0, argTypeIds, catalog);`
`3173`	`3174`
`3174`	`3175`	`if (bestScore < score) {`
`3175`	`3176`	`bestScore = score;`
Original file line number	Diff line number	Diff line change
`@@ -8,6 +8,7 @@ PEERDIR(`
`8`	`8`	`ydb/library/yql/ast`
`9`	`9`	`ydb/library/yql/core`
`10`	`10`	`ydb/library/yql/parser/pg_catalog`
	`11`	`+ ydb/library/yql/minikql`
`11`	`12`	`ydb/library/yql/sql/settings`
`12`	`13`	`ydb/public/api/protos`
`13`	`14`	`)`