Support of functions with lang=sql (#8233)

Author: vitstn

ydb/library/yql/core/peephole_opt/yql_opt_peephole_physical.cpp

@@ -2679,6 +2679,16 @@ TExprNode::TPtr ExpandVersion(const TExprNode::TPtr& node, TExprContext& ctx) {
         .Seal().Build();
 }
 
+TExprNode::TPtr ExpandRightOverCons(const TExprNode::TPtr& node, TExprContext& ctx) {
+    Y_UNUSED(ctx);
+    if (node->Head().IsCallable(ConsName)) {
+        YQL_CLOG(DEBUG, CorePeepHole) << "Expand Right! over Cons!";
+        return node->Head().TailPtr();
+    }
+
+    return node;
+}
+
 TExprNode::TPtr ExpandPartitionsByKeys(const TExprNode::TPtr& node, TExprContext& ctx) {
     YQL_CLOG(DEBUG, CorePeepHole) << "Expand " << node->Content();
     const bool isStream = node->Head().GetTypeAnn()->GetKind() == ETypeAnnotationKind::Flow ||
@@ -8351,6 +8361,7 @@ struct TPeepHoleRules {
         {"JsonExists", &ExpandJsonExists},
         {"EmptyIterator", &DropDependsOnFromEmptyIterator},
         {"Version", &ExpandVersion},
+        {RightName, &ExpandRightOverCons},
     };
 
     const TExtPeepHoleOptimizerMap CommonStageExtRules = {

ydb/library/yql/core/type_ann/type_ann_pg.cpp

@@ -368,21 +368,22 @@ IGraphTransformer::TStatus PgCallWrapper(const TExprNode::TPtr& input, TExprNode
             const auto procOrType = NPg::LookupProcWithCasts(TString(name), argTypes);
             auto children = input->ChildrenList();
             if (const auto* procPtr = std::get_if<const NPg::TProcDesc*>(&procOrType)) {
-                auto idNode = ctx.Expr.NewAtom(input->Pos(), ToString((*procPtr)->ProcId));
+                const auto& proc = *(*procPtr);
+                auto idNode = ctx.Expr.NewAtom(input->Pos(), ToString(proc.ProcId));
                 children.insert(children.begin() + 1, idNode);
 
-                const auto& fargTypes = (*procPtr)->ArgTypes;
+                const auto& fargTypes = proc.ArgTypes;
                 for (size_t i = 0; i < argTypes.size(); ++i) {
-                    auto targetType = (i >= fargTypes.size()) ? (*procPtr)->VariadicType : fargTypes[i];
+                    auto targetType = (i >= fargTypes.size()) ? proc.VariadicType : fargTypes[i];
                     if (IsCastRequired(argTypes[i], targetType)) {
                         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());
+                    YQL_ENSURE(fargTypes.size() - argTypes.size() <= proc.DefaultArgs.size());
                     for (size_t i = argTypes.size(); i < fargTypes.size(); ++i) {
-                        const auto& value = (*procPtr)->DefaultArgs[i + (*procPtr)->DefaultArgs.size() - fargTypes.size()];
+                        const auto& value = proc.DefaultArgs[i + proc.DefaultArgs.size() - fargTypes.size()];
                         TExprNode::TPtr defNode;
                         if (!value) {
                             defNode = ctx.Expr.NewCallable(input->Pos(), "Null", {});
@@ -399,6 +400,27 @@ IGraphTransformer::TStatus PgCallWrapper(const TExprNode::TPtr& input, TExprNode
                         children.insert(children.end(), defNode);
                     }
                 }
+
+                if (proc.Lang == NPg::LangSQL) {
+                    if (!proc.ExprNode) {
+                        throw yexception() << "Function " << proc.Name << " has no implementation";
+                    }
+
+                    // substibute by lambda
+                    YQL_ENSURE(proc.ExprNode->IsLambda());
+                    YQL_ENSURE(proc.ExprNode->Head().ChildrenSize() == fargTypes.size());
+                    TNodeOnNodeOwnedMap deepClones;
+                    YQL_ENSURE(NPg::GetSqlLanguageParser());
+                    auto lambda = ctx.Expr.DeepCopy(*proc.ExprNode, NPg::GetSqlLanguageParser()->GetContext(), deepClones, true, false);
+                    TNodeOnNodeOwnedMap replaces;
+                    for (ui32 i = 0; i < fargTypes.size(); ++i) {
+                        replaces[lambda->Head().Child(i)] = children[i + 3];
+                    }
+
+                    output = ctx.Expr.ReplaceNodes(lambda->TailPtr(), replaces);
+                    return IGraphTransformer::TStatus::Repeat;
+                }
+
                 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);

ydb/library/yql/parser/pg_catalog/catalog.cpp

@@ -476,10 +476,12 @@ class TProcsParser : public TParser {
         } else if (key == "prosrc") {
             LastProc.Src = value;
         } else if (key == "prolang") {
-            if (value != "c") {
-                IsSupported = false;
-            } else {
+            if (value == "sql") {
+                LastProc.Lang = LangSQL;
+            } else if (value == "c") {
                 LastProc.Lang = LangC;
+            } else {
+                IsSupported = false;
             }
         } else if (key == "proargtypes") {
             TVector<TString> strArgs;
@@ -2272,6 +2274,8 @@ struct TCatalog : public IExtensionSqlBuilder {
 
         bool AllowAllFunctions = false;
         THashSet<TString> AllowedProcs;
+
+        bool SystemFunctionInit = false;
     };
 
     mutable TMaybe<TFile> ExportFile;
@@ -3598,6 +3602,49 @@ void AllowFunction(const TString& name) {
     }
 }
 
+struct TSqlLanguageParserHolder {
+    std::unique_ptr<ISqlLanguageParser> Parser;
+};
+
+void SetSqlLanguageParser(std::unique_ptr<ISqlLanguageParser> parser) {
+    Singleton<TSqlLanguageParserHolder>()->Parser = std::move(parser);
+}
+
+ISqlLanguageParser* GetSqlLanguageParser() {
+    return Singleton<TSqlLanguageParserHolder>()->Parser.get();
+}
+
+void LoadSystemFunctions(ISystemFunctionsParser& parser) {
+    auto& catalog = TCatalog::MutableInstance();
+    with_lock (catalog.ExtensionsGuard) {
+        Y_ENSURE(!catalog.State->SystemFunctionInit);
+        TString data;
+        Y_ENSURE(NResource::FindExact("system_functions.sql", &data));
+        TVector<TProcDesc> procs;
+        parser.Parse(data, procs);
+        for (const auto& p : procs) {
+            auto procIdPtr = catalog.State->ProcByName.FindPtr(p.Name);
+            Y_ENSURE(procIdPtr);
+            TProcDesc* foundProc = nullptr;
+            for (auto procId : *procIdPtr) {
+                auto procPtr = catalog.State->Procs.FindPtr(procId);
+                Y_ENSURE(procPtr);
+                if (procPtr->ArgTypes == p.ArgTypes && procPtr->VariadicType == p.VariadicType) {
+                    foundProc = procPtr;
+                    break;
+                }
+            }
+
+            Y_ENSURE(foundProc);
+            foundProc->DefaultArgs = p.DefaultArgs;
+            foundProc->Src = p.Src;
+            foundProc->ExprNode = p.ExprNode;
+        }
+
+        catalog.State->SystemFunctionInit = true;
+    }
+}
+
 void RegisterExtensions(const TVector<TExtensionDesc>& extensions, bool typesOnly,
     IExtensionSqlParser& parser, IExtensionLoader* loader) {
     if (extensions.size() > MaximumExtensionsCount) {

ydb/library/yql/parser/pg_catalog/catalog.h

@@ -8,6 +8,13 @@
 #include <variant>
 #include <functional>
 
+struct Node;
+
+namespace NYql {
+class TExprNode;
+struct TExprContext;
+}
+
 namespace NYql::NPg {
 
 constexpr ui32 UnknownOid = 705;
@@ -75,6 +82,7 @@ struct TProcDesc {
     ui32 VariadicArgType = 0;
     TString VariadicArgName;
     TVector<TMaybe<TString>> DefaultArgs;
+    TExprNode* ExprNode = nullptr;
     ui32 ExtensionIndex = 0;
 };
 
@@ -436,6 +444,26 @@ class IExtensionLoader {
     virtual void Load(ui32 extensionIndex, const TString& name, const TString& path) = 0;
 };
 
+class ISystemFunctionsParser {
+public:
+    virtual ~ISystemFunctionsParser() = default;
+    virtual void Parse(const TString& sql, TVector<TProcDesc>& procs) const = 0;
+};
+
+class ISqlLanguageParser {
+public:
+    virtual ~ISqlLanguageParser() = default;
+    virtual void Parse(const TString& sql, TProcDesc& proc) = 0;
+    virtual void ParseNode(const Node* stmt, TProcDesc& proc) = 0;
+    virtual void Freeze() = 0;
+    virtual TExprContext& GetContext() = 0;
+};
+
+void SetSqlLanguageParser(std::unique_ptr<ISqlLanguageParser> parser);
+ISqlLanguageParser* GetSqlLanguageParser();
+
+void LoadSystemFunctions(ISystemFunctionsParser& parser);
+
 // either RegisterExtensions or ImportExtensions should be called at most once, see ClearExtensions as well
 void RegisterExtensions(const TVector<TExtensionDesc>& extensions, bool typesOnly,
     IExtensionSqlParser& parser, IExtensionLoader* loader);

ydb/library/yql/parser/pg_catalog/postgis_procs.h

@@ -90,6 +90,7 @@
 "_postgis_join_selectivity",
 "_postgis_selectivity",
 "_postgis_stats",
+"postgis_transform_geometry",
 "spheroid_in",
 "spheroid_out",
 "st_3dclosestpoint",
@@ -106,6 +107,7 @@
 "st_3dshortestline",
 "st_addmeasure",
 "st_addpoint",
+"st_affine",
 "st_angle",
 "st_area",
 "st_area2d",
@@ -114,6 +116,7 @@
 "st_asewkb",
 "st_asewkt",
 "st_asgeojson",
+"_st_asgml",
 "st_asgml",
 "st_askml",
 "st_aslatlontext",
@@ -122,6 +125,7 @@
 "st_assvg",
 "st_astext",
 "st_astwkb",
+"_st_asx3d",
 "st_azimuth",
 "_st_bestsrid",
 "st_boundary",
@@ -153,15 +157,15 @@
 "st_distance",
 "st_distancecpa",
 "st_distancespheroid",
-"st_dump",
-"st_dumppoints",
-"st_dumpsegments",
+"_st_distancetree",
+"_st_distanceuncached",
 "st_dwithin",
 "st_endpoint",
 "st_envelope",
 "st_equals",
 "st_estimatedextent",
 "st_expand",
+"st_exteriorring",
 "st_filterbym",
 "st_flipcoordinates",
 "st_force2d",
@@ -184,6 +188,7 @@
 "st_geomfromewkb",
 "st_geomfromewkt",
 "st_geomfromgeojson",
+"_st_geomfromgml",
 "st_geomfromgml",
 "st_geomfromkml",
 "st_geomfrommarc21",
@@ -225,13 +230,15 @@
 "st_locatealong",
 "st_locatebetween",
 "st_locatebetweenelevations",
+"_st_longestline",
 "st_m",
 "st_makebox2d",
 "st_makeenvelope",
 "st_makepoint",
 "st_makepointm",
 "st_makepolygon",
 "st_makevalid",
+"_st_maxdistance",
 "st_maximuminscribedcircle",
 "st_memsize",
 "st_minimumboundingcircle",

ydb/library/yql/parser/pg_catalog/used_procs.h

@@ -38,6 +38,7 @@
 "any_out",
 "anyrange_in",
 "anyrange_out",
+"anytextcat",
 "any_value_transfn",
 "area",
 "array_agg_array_combine",
@@ -702,6 +703,7 @@
 "int8or",
 "int8out",
 "int8pl",
+"int8pl_inet",
 "int8recv",
 "int8send",
 "int8shl",
@@ -710,6 +712,7 @@
 "int8um",
 "int8up",
 "int8xor",
+"integer_pl_date",
 "inter_lb",
 "internal_in",
 "internal_out",
@@ -734,6 +737,11 @@
 "interval_ne",
 "interval_out",
 "interval_pl",
+"interval_pl_date",
+"interval_pl_time",
+"interval_pl_timestamp",
+"interval_pl_timestamptz",
+"interval_pl_timetz",
 "interval_recv",
 "interval_send",
 "interval_smaller",
@@ -1015,6 +1023,7 @@
 "numeric_mul",
 "numeric_ne",
 "numeric_out",
+"numeric_pl_pg_lsn",
 "numeric_poly_avg",
 "numeric_poly_combine",
 "numeric_poly_deserialize",
@@ -1074,9 +1083,11 @@
 "ordered_set_transition_multi",
 "overlaps",
 "overlay",
+"parse_ident",
 "path",
 "path_add",
 "path_add_pt",
+"path_contain_pt",
 "path_distance",
 "path_div_pt",
 "path_in",
@@ -1312,6 +1323,7 @@
 "sha384",
 "sha512",
 "sign",
+"similar_to_escape",
 "sin",
 "sind",
 "sinh",
@@ -1336,6 +1348,7 @@
 "tand",
 "tanh",
 "text",
+"textanycat",
 "textcat",
 "texteq",
 "texteqname",
@@ -1381,6 +1394,7 @@
 "tidsmaller",
 "time",
 "time_cmp",
+"timedate_pl",
 "time_eq",
 "time_ge",
 "time_gt",
@@ -1465,6 +1479,7 @@
 "timetypmodout",
 "timetz",
 "timetz_cmp",
+"timetzdate_pl",
 "timetz_eq",
 "timetz_ge",
 "timetz_gt",

ydb/library/yql/parser/pg_catalog/ya.make

@@ -12,6 +12,7 @@ RESOURCE(../pg_wrapper/postgresql/src/include/catalog/pg_amop.dat pg_amop.dat)
 RESOURCE(../pg_wrapper/postgresql/src/include/catalog/pg_am.dat pg_am.dat)
 RESOURCE(../pg_wrapper/postgresql/src/include/catalog/pg_conversion.dat pg_conversion.dat)
 RESOURCE(../pg_wrapper/postgresql/src/include/catalog/pg_language.dat pg_language.dat)
+RESOURCE(../pg_wrapper/postgresql/src/backend/catalog/system_functions.sql system_functions.sql)
 
 SRCS(
     catalog.cpp

ydb/library/yql/parser/pg_wrapper/interface/parser.h

@@ -14,5 +14,7 @@ namespace NSQLTranslationPG {
 NYql::TAstParseResult PGToYql(const TString& query, const NSQLTranslation::TTranslationSettings& settings, NYql::TStmtParseInfo* stmtParseInfo = nullptr);
 TVector<NYql::TAstParseResult> PGToYqlStatements(const TString& query, const NSQLTranslation::TTranslationSettings& settings, TVector<NYql::TStmtParseInfo>* stmtParseInfo = nullptr);
 std::unique_ptr<NYql::NPg::IExtensionSqlParser> CreateExtensionSqlParser();
+std::unique_ptr<NYql::NPg::ISystemFunctionsParser> CreateSystemFunctionsParser();
+std::unique_ptr<NYql::NPg::ISqlLanguageParser> CreateSqlLanguageParser();
 
 } // NSQLTranslationPG