YQL-13710 range functions (scalar)

ref:0d6bd7840402417bcacf17debadf1656fc9024f6

3 files changed, 150 insertions, 53 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 a3824aa844b..7997f4fc318 100644
--- a/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
@@ -153,21 +153,29 @@ void FillInputIndices(const TExprNode::TPtr& from, TUsedColumns& usedColumns, TO
         for (ui32 inputIndex = 0; inputIndex < from->Tail().ChildrenSize(); ++inputIndex) {
             const auto& read = from->Tail().Child(inputIndex)->Head();
             const auto& columns = from->Tail().Child(inputIndex)->Tail();
-            if (columns.ChildrenSize() > 0) {
-                auto readOrder = optCtx.Types->LookupColumnOrder(read);
-                YQL_ENSURE(*readOrder);
-                for (ui32 i = 0; i < columns.ChildrenSize(); ++i) {
-                    if (columns.Child(i)->Content() == x.first) {
-                        foundColumn = true;
-                        x.second.second = (*readOrder)[i];
-                        break;
+            if (read.IsCallable("PgResolvedCall")) {
+                const auto& alias = from->Tail().Child(inputIndex)->Child(1)->Content();
+                Y_ENSURE(!alias.empty());
+                Y_ENSURE(columns.ChildrenSize() == 0 || columns.ChildrenSize() == 1);
+                auto memberName = (columns.ChildrenSize() == 1) ? columns.Head().Content() : alias;
+                foundColumn = (memberName == x.first);
+            } else {
+                if (columns.ChildrenSize() > 0) {
+                    auto readOrder = optCtx.Types->LookupColumnOrder(read);
+                    YQL_ENSURE(readOrder);
+                    for (ui32 i = 0; i < columns.ChildrenSize(); ++i) {
+                        if (columns.Child(i)->Content() == x.first) {
+                            foundColumn = true;
+                            x.second.second = (*readOrder)[i];
+                            break;
+                        }
                     }
+                } else {
+                    auto type = read.GetTypeAnn()->
+                        Cast<TListExprType>()->GetItemType()->Cast<TStructExprType>();
+                    auto pos = type->FindItem(x.first);
+                    foundColumn = pos.Defined();
                 }
-            } else {
-                auto type = read.GetTypeAnn()->
-                    Cast<TListExprType>()->GetItemType()->Cast<TStructExprType>();
-                auto pos = type->FindItem(x.first);
-                foundColumn = pos.Defined();
             }
 
             if (foundColumn) {
@@ -183,9 +191,28 @@ void FillInputIndices(const TExprNode::TPtr& from, TUsedColumns& usedColumns, TO
 TExprNode::TListType BuildCleanedColumns(TPositionHandle pos, const TExprNode::TPtr& from, const TUsedColumns& usedColumns, TExprContext& ctx) {
     TExprNode::TListType cleanedInputs;
     for (ui32 i = 0; i < from->Tail().ChildrenSize(); ++i) {
+        auto list = from->Tail().Child(i)->HeadPtr();
+        if (list->IsCallable("PgResolvedCall")) {
+            const auto& alias = from->Tail().Child(i)->Child(1)->Content();
+            const auto& columns = from->Tail().Child(i)->Tail();
+            Y_ENSURE(!alias.empty());
+            Y_ENSURE(columns.ChildrenSize() == 0 || columns.ChildrenSize() == 1);
+            auto memberName = (columns.ChildrenSize() == 1) ? columns.Head().Content() : alias;
+            list = ctx.Builder(pos)
+                .Callable("AsList")
+                    .Callable(0, "AsStruct")
+                        .List(0)
+                            .Atom(0, memberName)
+                            .Add(1, list)
+                        .Seal()
+                    .Seal()
+                .Seal()
+                .Build();
+        }
+
         auto cleaned = ctx.Builder(pos)
             .Callable("Map")
-                .Add(0, from->Tail().Child(i)->HeadPtr())
+                .Add(0, list)
                 .Lambda(1)
                     .Param("row")
                     .Callable("AsStruct")
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 242ca5db1e7..5a8db570335 100644
--- a/ydb/library/yql/core/type_ann/type_ann_pg.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_pg.cpp
@@ -23,26 +23,6 @@ IGraphTransformer::TStatus PgStarWrapper(const TExprNode::TPtr& input, TExprNode
     return IGraphTransformer::TStatus::Ok;
 }
 
-struct TPgFuncDesc {
-    ui32 MinArgs;
-    ui32 MaxArgs;
-    EDataSlot ReturnType;
-    TVector<EDataSlot> DataTypes;
-};
-
-class TPgFuncMap {
-public:
-    static const TPgFuncMap& Instance() {
-        return *Singleton<TPgFuncMap>();
-    }
-
-    THashMap<TString, TPgFuncDesc> Funcs;
-
-    TPgFuncMap() {
-        Funcs["substring"] = { 3, 3, EDataSlot::Utf8, { EDataSlot::Utf8, EDataSlot::Int32, EDataSlot::Int32 } };
-    }
-};
-
 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)) {
@@ -1779,18 +1759,42 @@ IGraphTransformer::TStatus PgSetItemWrapper(const TExprNode::TPtr& input, TExprN
                         }
                     }
 
+                    auto alias = TString(p->Child(1)->Content());
                     auto columnOrder = ctx.Types.LookupColumnOrder(p->Head());
-                    if (!EnsureListType(p->Head(), ctx.Expr)) {
-                        return IGraphTransformer::TStatus::Error;
-                    }
+                    const bool isRangeFunction = p->Head().IsCallable("PgResolvedCall");
+                    const TStructExprType* inputStructType = nullptr;
+                    if (isRangeFunction) {
+                        if (alias.empty()) {
+                            ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(option->Head().Pos()),
+                                "Empty alias for range function is not allowed"));
+                            return IGraphTransformer::TStatus::Error;
+                        }
 
-                    auto inputRowType = p->Head().GetTypeAnn()->Cast<TListExprType>()->GetItemType();
-                    if (!EnsureStructType(p->Head().Pos(), *inputRowType, ctx.Expr)) {
-                        return IGraphTransformer::TStatus::Error;
+                        if (p->Child(2)->ChildrenSize() > 0 && p->Child(2)->ChildrenSize() != 1) {
+                            ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(option->Head().Pos()),
+                                TStringBuilder() << "Expected exactly one column name for range function, but got: " << p->Child(2)->ChildrenSize()));
+                            return IGraphTransformer::TStatus::Error;
+                        }
+
+                        auto memberName = (p->Child(2)->ChildrenSize() == 1) ? p->Child(2)->Head().Content() : alias;
+                        TVector<const TItemExprType*> items;
+                        auto itemType = p->Head().GetTypeAnn();
+                        items.push_back(ctx.Expr.MakeType<TItemExprType>(memberName, itemType));
+                        inputStructType = ctx.Expr.MakeType<TStructExprType>(items);
+                        columnOrder = TColumnOrder({ TString(memberName) });
+                    } else {
+                        if (!EnsureListType(p->Head(), ctx.Expr)) {
+                            return IGraphTransformer::TStatus::Error;
+                        }
+
+                        auto inputRowType = p->Head().GetTypeAnn()->Cast<TListExprType>()->GetItemType();
+                        if (!EnsureStructType(p->Head().Pos(), *inputRowType, ctx.Expr)) {
+                            return IGraphTransformer::TStatus::Error;
+                        }
+
+                        inputStructType = inputRowType->Cast<TStructExprType>();
                     }
 
-                    auto inputStructType = inputRowType->Cast<TStructExprType>();
-                    auto alias = TString(p->Child(1)->Content());
                     if (!alias.empty()) {
                         if (!possibleAliases.insert(alias).second) {
                             ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(option->Head().Pos()),
diff --git a/ydb/library/yql/sql/pg/pg_sql.cpp b/ydb/library/yql/sql/pg/pg_sql.cpp
index dc778e894f5..c6822ddfc30 100644
--- a/ydb/library/yql/sql/pg/pg_sql.cpp
+++ b/ydb/library/yql/sql/pg/pg_sql.cpp
@@ -107,8 +107,17 @@ const Node* ListNodeNth(const List* list, int index) {
 
 class TConverter : public IPGParseEvents {
 public:
-    using TFromDesc = std::tuple<TAstNode*, TString, TVector<TString>, bool>;
-    using TInsertDesc = std::tuple<TAstNode*,TAstNode*>;
+    struct TFromDesc {
+        TAstNode* Source = nullptr;
+        TString Alias;
+        TVector<TString> ColNames;
+        bool InjectRead = false;
+    };
+
+    struct TInsertDesc {
+        TAstNode* Sink = nullptr;
+        TAstNode* Key = nullptr;
+    };
 
     struct TExprSettings {
         bool AllowColumns = false;
@@ -247,7 +256,7 @@ public:
                 auto node = ListNodeNth(x->fromClause, i);
                 if (NodeTag(node) != T_JoinExpr) {
                     auto p = ParseFromClause(node);
-                    if (!std::get<0>(p)) {
+                    if (!p.Source) {
                         return nullptr;
                     }
 
@@ -263,7 +272,7 @@ public:
                         if (NodeTag(top.first) != T_JoinExpr) {
                             // leaf
                             auto p = ParseFromClause(top.first);
-                            if (!std::get<0>(p)) {
+                            if (!p.Source) {
                                 return nullptr;
                             }
 
@@ -704,7 +713,7 @@ public:
         }
 
         auto insertDesc = ParseWriteRangeVar(value->relation);
-        if (!std::get<0>(insertDesc)) {
+        if (!insertDesc.Sink) {
             return nullptr;
         }
 
@@ -715,7 +724,7 @@ public:
 
         auto writeOptions = QL(QL(QA("mode"), QA("append")));
         Statements.push_back(L(A("let"), A("world"), L(A("Write!"),
-            A("world"), std::get<0>(insertDesc), std::get<1>(insertDesc), select, writeOptions)));
+            A("world"), insertDesc.Sink, insertDesc.Key, select, writeOptions)));
 
         return Statements.back();
     }
@@ -726,6 +735,8 @@ public:
             return ParseRangeVar(CAST_NODE(RangeVar, node));
         case T_RangeSubselect:
             return ParseRangeSubselect(CAST_NODE(RangeSubselect, node));
+        case T_RangeFunction:
+            return ParseRangeFunction(CAST_NODE(RangeFunction, node));
         default:
             NodeNotImplementedImpl<SelectStmt>(node);
             return {};
@@ -733,21 +744,21 @@ public:
     }
 
     void AddFrom(const TFromDesc& p, TVector<TAstNode*>& fromList) {
-        auto aliasNode = QA(std::get<1>(p));
+        auto aliasNode = QA(p.Alias);
         TVector<TAstNode*> colNamesNodes;
-        for (const auto& c : std::get<2>(p)) {
+        for (const auto& c : p.ColNames) {
             colNamesNodes.push_back(QA(c));
         }
 
         auto colNamesTuple = QVL(colNamesNodes.data(), colNamesNodes.size());
-        if (std::get<3>(p)) {
+        if (p.InjectRead) {
             auto label = "read" + ToString(ReadIndex);
-            Statements.push_back(L(A("let"), A(label), std::get<0>(p)));
+            Statements.push_back(L(A("let"), A(label), p.Source));
             Statements.push_back(L(A("let"), A("world"), L(A("Left!"), A(label))));
             fromList.push_back(QL(L(A("Right!"), A(label)), aliasNode, colNamesTuple));
             ++ReadIndex;
         } else {
-            fromList.push_back(QL(std::get<0>(p), aliasNode, colNamesTuple));
+            fromList.push_back(QL(p.Source, aliasNode, colNamesTuple));
         }
     }
 
@@ -835,6 +846,61 @@ public:
             QL()), alias, colnames, true };
     }
 
+    TFromDesc ParseRangeFunction(const RangeFunction* value) {
+        if (value->lateral) {
+            AddError("RangeFunction: unsupported lateral");
+            return {};
+        }
+
+        if (value->ordinality) {
+            AddError("RangeFunction: unsupported ordinality");
+            return {};
+        }
+
+        if (value->is_rowsfrom) {
+            AddError("RangeFunction: unsupported is_rowsfrom");
+            return {};
+        }
+
+        if (ListLength(value->coldeflist) > 0) {
+            AddError("RangeFunction: unsupported coldeflist");
+            return {};
+        }
+
+        if (ListLength(value->functions) != 1) {
+            AddError("RangeFunction: only one function is supported");
+            return {};
+        }
+
+        TString alias;
+        TVector<TString> colnames;
+        if (!ParseAlias(value->alias, alias, colnames)) {
+            return {};
+        }
+
+        auto funcNode = ListNodeNth(value->functions, 0);
+        if (NodeTag(funcNode) != T_List) {
+            AddError("RangeFunction: expected pair");
+            return {};
+        }
+
+        auto lst = CAST_NODE(List, funcNode);
+        if (ListLength(lst) != 2) {
+            AddError("RangeFunction: expected pair");
+            return {};
+        }
+        
+        TExprSettings settings;
+        settings.AllowColumns = false;
+        settings.Scope = "RANGE FUNCTION";
+        auto func = ParseExpr(ListNodeNth(lst, 0), settings);
+        if (!func) {
+            return {};
+        }
+
+        return { func, alias, colnames, false };
+    }
+
     TFromDesc ParseRangeSubselect(const RangeSubselect* value) {
         if (value->lateral) {
             AddError("RangeSubselect: unsupported lateral");