YQL-14728 SQLv1 typemods

ref:729cd79fdfac35cf63ef902db7b92ee5ac1436bb

10 files changed, 224 insertions, 13 deletions

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 3e080509dce..b0e75ad44b3 100644
--- a/ydb/library/yql/core/type_ann/type_ann_core.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_core.cpp
@@ -11283,6 +11283,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         Functions["PgWindowTraitsTuple"] = &PgAggregationTraitsWrapper;
         Functions["PgInternal0"] = &PgInternal0Wrapper;
         Functions["PgArray"] = &PgArrayWrapper;
+        Functions["PgTypeMod"] = &PgTypeModWrapper;
 
         Functions["AutoDemuxList"] = &AutoDemuxListWrapper;
         Functions["AggrCountInit"] = &AggrCountInitWrapper;
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 ebe21f00c64..dbc889f720a 100644
--- a/ydb/library/yql/core/type_ann/type_ann_pg.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_pg.cpp
@@ -11,6 +11,9 @@
 #include <ydb/library/yql/parser/pg_catalog/catalog.h>
 
 namespace NYql {
+
+bool ParsePgIntervalModifier(const TString& str, i32& ret);
+
 namespace NTypeAnnImpl {
 
 IGraphTransformer::TStatus PgStarWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
@@ -731,11 +734,15 @@ IGraphTransformer::TStatus PgAnonWindowWrapper(const TExprNode::TPtr& input, TEx
 
 IGraphTransformer::TStatus PgConstWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
     Y_UNUSED(output);
-    if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
+    if (!EnsureMinArgsCount(*input, 2, ctx.Expr)) {
         return IGraphTransformer::TStatus::Error;
     }
 
-    if (!EnsureTypePg(input->Tail(), ctx.Expr)) {
+    if (!EnsureMaxArgsCount(*input, 3, ctx.Expr)) {
+        return IGraphTransformer::TStatus::Error;
+    }
+
+    if (!EnsureTypePg(*input->Child(1), ctx.Expr)) {
         return IGraphTransformer::TStatus::Error;
     }
 
@@ -744,7 +751,27 @@ IGraphTransformer::TStatus PgConstWrapper(const TExprNode::TPtr& input, TExprNod
         return IGraphTransformer::TStatus::Error;
     }
 
-    input->SetTypeAnn(input->Tail().GetTypeAnn()->Cast<TTypeExprType>()->GetType());
+    if (input->ChildrenSize() >= 3) {
+        auto type = input->Child(2)->GetTypeAnn();
+        ui32 typeModType;
+        bool convertToPg;
+        if (!ExtractPgType(type, typeModType, convertToPg, input->Child(2)->Pos(), ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        if (convertToPg) {
+            input->ChildRef(2) = ctx.Expr.NewCallable(input->Child(2)->Pos(), "ToPg", { input->ChildPtr(2) });
+            return IGraphTransformer::TStatus::Repeat;
+        }
+
+        if (typeModType != NPg::LookupType("int4").TypeId) {
+            ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
+                TStringBuilder() << "Expected pg int4 as typemod, but got " << NPg::LookupType(typeModType).Name));
+            return IGraphTransformer::TStatus::Error;
+        }
+    }
+
+    input->SetTypeAnn(input->Child(1)->GetTypeAnn()->Cast<TTypeExprType>()->GetType());
     return IGraphTransformer::TStatus::Ok;
 }
 
@@ -2604,5 +2631,70 @@ IGraphTransformer::TStatus PgArrayWrapper(const TExprNode::TPtr& input, TExprNod
     return IGraphTransformer::TStatus::Ok;
 }
 
+IGraphTransformer::TStatus PgTypeModWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+    if (!EnsureMinArgsCount(*input, 2, ctx.Expr)) {
+        return IGraphTransformer::TStatus::Error;
+    }
+
+    if (!EnsureTypePg(*input->Child(0), ctx.Expr)) {
+        return IGraphTransformer::TStatus::Error;
+    }
+
+    auto pgType = input->Child(0)->GetTypeAnn()->Cast<TTypeExprType>()->GetType()->Cast<TPgExprType>();
+    const auto& typeDesc = NPg::LookupType(pgType->GetId());
+    if (!typeDesc.TypeModInFuncId) {
+        ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
+            TStringBuilder() << "No modifiers for type: " << pgType->GetName()));
+        return IGraphTransformer::TStatus::Error;
+    }
+
+    TVector<TString> mods;
+    for (ui32 i = 1; i < input->ChildrenSize(); ++i) {
+        if (!EnsureAtom(*input->Child(i), ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        mods.push_back(TString(input->Child(i)->Content()));
+    }
+
+    if (pgType->GetName() == "interval") {
+        if (mods.size() != 1) {
+            ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()), "Exactly one modidifer is expected for pginterval"));
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        i32 value;
+        if (!ParsePgIntervalModifier(mods[0], value)) {
+            ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()),
+                TStringBuilder() << "Unsupported modifier for pginterval: " << mods[0]));
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        mods[0] = ToString(value);
+    }
+
+    TExprNode::TListType args;
+    for (const auto& mod : mods) {
+        args.push_back(ctx.Expr.Builder(input->Pos())
+            .Callable("PgConst")
+                .Atom(0, mod)
+                .Callable(1, "PgType")
+                    .Atom(0, "cstring")
+                .Seal()
+            .Seal()
+            .Build());
+    }
+
+    auto arr = ctx.Expr.NewCallable(input->Pos(), "PgArray", std::move(args));
+    output = ctx.Expr.Builder(input->Pos())
+        .Callable("PgCall")
+            .Atom(0, NPg::LookupProc(typeDesc.TypeModInFuncId, { 0 }).Name)
+            .Add(1, arr)
+        .Seal()
+        .Build();
+
+    return IGraphTransformer::TStatus::Repeat;
+}
+
 } // namespace NTypeAnnImpl
 }
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 3739ef8b08f..516507fcc17 100644
--- a/ydb/library/yql/core/type_ann/type_ann_pg.h
+++ b/ydb/library/yql/core/type_ann/type_ann_pg.h
@@ -31,6 +31,7 @@ IGraphTransformer::TStatus PgTypeWrapper(const TExprNode::TPtr& input, TExprNode
 IGraphTransformer::TStatus PgSetItemWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExtContext& ctx);
 IGraphTransformer::TStatus PgSelectWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExtContext& ctx);
 IGraphTransformer::TStatus PgArrayWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
+IGraphTransformer::TStatus PgTypeModWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
 
 } // namespace NTypeAnnImpl
 } // namespace NYql
diff --git a/ydb/library/yql/minikql/mkql_program_builder.cpp b/ydb/library/yql/minikql/mkql_program_builder.cpp
index f66a975fdb7..feee860636b 100644
--- a/ydb/library/yql/minikql/mkql_program_builder.cpp
+++ b/ydb/library/yql/minikql/mkql_program_builder.cpp
@@ -5034,7 +5034,7 @@ TRuntimeNode TProgramBuilder::Replicate(TRuntimeNode item, TRuntimeNode count, c
     return TRuntimeNode(callableBuilder.Build(), false);
 }
 
-TRuntimeNode TProgramBuilder::PgConst(TPgType* pgType, const std::string_view& value) {
+TRuntimeNode TProgramBuilder::PgConst(TPgType* pgType, const std::string_view& value, TRuntimeNode typeMod) {
     if constexpr (RuntimeVersion < 30U) {
         THROW yexception() << "Runtime version (" << RuntimeVersion << ") too old for " << __func__;
     }
@@ -5042,6 +5042,10 @@ TRuntimeNode TProgramBuilder::PgConst(TPgType* pgType, const std::string_view& v
     TCallableBuilder callableBuilder(Env, __func__, pgType);
     callableBuilder.Add(NewDataLiteral(pgType->GetTypeId()));
     callableBuilder.Add(NewDataLiteral<NUdf::EDataSlot::String>(value));
+    if (typeMod) {
+        callableBuilder.Add(typeMod);
+    }
+
     return TRuntimeNode(callableBuilder.Build(), false);
 }
 
diff --git a/ydb/library/yql/minikql/mkql_program_builder.h b/ydb/library/yql/minikql/mkql_program_builder.h
index bf21ca94969..80f0c328dee 100644
--- a/ydb/library/yql/minikql/mkql_program_builder.h
+++ b/ydb/library/yql/minikql/mkql_program_builder.h
@@ -623,7 +623,7 @@ public:
     typedef TRuntimeNode (TProgramBuilder::*ProcessFunctionMethod)(TRuntimeNode, const TUnaryLambda&);
     typedef TRuntimeNode (TProgramBuilder::*NarrowFunctionMethod)(TRuntimeNode, const TNarrowLambda&);
 
-    TRuntimeNode PgConst(TPgType* pgType, const std::string_view& value);
+    TRuntimeNode PgConst(TPgType* pgType, const std::string_view& value, TRuntimeNode typeMod = {});
     TRuntimeNode PgResolvedCall(bool useContext, const std::string_view& name, ui32 id,
         const TArrayRef<const TRuntimeNode>& args, TType* returnType);
     TRuntimeNode PgCast(TRuntimeNode input, TType* returnType, TRuntimeNode typeMod = {});
diff --git a/ydb/library/yql/parser/pg_catalog/catalog.cpp b/ydb/library/yql/parser/pg_catalog/catalog.cpp
index 9b2842e9813..a7688935b4c 100644
--- a/ydb/library/yql/parser/pg_catalog/catalog.cpp
+++ b/ydb/library/yql/parser/pg_catalog/catalog.cpp
@@ -57,7 +57,7 @@ TString ArgTypesList(const TVector<ui32>& ids) {
             str << ',';
         }
 
-        str << LookupType(ids[i]).Name;
+        str << ids[i] ? LookupType(ids[i]).Name : "NULL";
     }
 
     str << ')';
diff --git a/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp b/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp
index b1482f0ec80..f210675a203 100644
--- a/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp
+++ b/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp
@@ -29,6 +29,7 @@ extern "C" {
 #include "utils/array.h"
 #include "utils/arrayaccess.h"
 #include "utils/lsyscache.h"
+#include "utils/datetime.h"
 #include "nodes/execnodes.h"
 #include "executor/executor.h"
 #include "lib/stringinfo.h"
@@ -241,10 +242,11 @@ inline ui32 MakeTypeIOParam(const NPg::TTypeDesc& desc) {
 class TPgConst : public TMutableComputationNode<TPgConst> {
     typedef TMutableComputationNode<TPgConst> TBaseComputation;
 public:
-    TPgConst(TComputationMutables& mutables, ui32 typeId, const std::string_view& value)
+    TPgConst(TComputationMutables& mutables, ui32 typeId, const std::string_view& value, IComputationNode* typeMod)
         : TBaseComputation(mutables)
         , TypeId(typeId)
         , Value(value)
+        , TypeMod(typeMod)
         , TypeDesc(NPg::LookupType(TypeId))
     {
         Zero(FInfo);
@@ -262,6 +264,11 @@ public:
     }
 
     NUdf::TUnboxedValuePod DoCalculate(TComputationContext& compCtx) const {
+        i32 typeMod = -1;
+        if (TypeMod) {
+            typeMod = DatumGetInt32(ScalarDatumFromPod(TypeMod->GetValue(compCtx)));
+        }
+
         LOCAL_FCINFO(callInfo, 3);
         Zero(*callInfo);
         FmgrInfo copyFmgrInfo = FInfo;
@@ -271,7 +278,7 @@ public:
         callInfo->isnull = false;
         callInfo->args[0] = { (Datum)Value.c_str(), false };
         callInfo->args[1] = { ObjectIdGetDatum(TypeIOParam), false };
-        callInfo->args[2] = { Int32GetDatum(-1), false };
+        callInfo->args[2] = { Int32GetDatum(typeMod), false };
 
         TPAllocScope call;
         PG_TRY();
@@ -294,10 +301,14 @@ public:
 
 private:
     void RegisterDependencies() const final {
+        if (TypeMod) {
+            DependsOn(TypeMod);
+        }
     }
 
     const ui32 TypeId;
     const TString Value;
+    IComputationNode* const TypeMod;
     const NPg::TTypeDesc TypeDesc;
     FmgrInfo FInfo;
     ui32 TypeIOParam;
@@ -1360,7 +1371,12 @@ TComputationNodeFactory GetPgFactory() {
                 const auto valueData = AS_VALUE(TDataLiteral, callable.GetInput(1));
                 ui32 typeId = typeIdData->AsValue().Get<ui32>();
                 auto value = valueData->AsValue().AsStringRef();
-                return new TPgConst(ctx.Mutables, typeId, value);
+                IComputationNode* typeMod = nullptr;
+                if (callable.GetInputsCount() >= 3) {
+                    typeMod = LocateNode(ctx.NodeLocator, callable, 2);
+                }
+
+                return new TPgConst(ctx.Mutables, typeId, value, typeMod);
             }
 
             if (name == "PgInternal0") {
@@ -2020,6 +2036,41 @@ TMaybe<NUdf::EDataSlot> ConvertFromPgType(ui32 typeId) {
     return Nothing();
 }
 
+bool ParsePgIntervalModifier(const TString& str, i32& ret) {
+    auto ustr = to_upper(str);

+    if (ustr == "YEAR") {

+        ret = INTERVAL_MASK(YEAR);

+    } else if (ustr == "MONTH") {

+        ret = INTERVAL_MASK(YEAR);

+    } else if (ustr == "DAY") {

+        ret = INTERVAL_MASK(DAY);

+    } else if (ustr == "HOUR") {

+        ret = INTERVAL_MASK(HOUR);

+    } else if (ustr == "MINUTE") {

+        ret = INTERVAL_MASK(MINUTE);

+    } else if (ustr == "SECOND") {

+        ret = INTERVAL_MASK(SECOND);

+    } else if (ustr == "YEAR TO MONTH") {

+        ret = INTERVAL_MASK(YEAR) | INTERVAL_MASK(MONTH);

+    } else if (ustr == "DAY TO HOUR") {

+        ret = INTERVAL_MASK(DAY) | INTERVAL_MASK(HOUR);

+    } else if (ustr == "DAY TO MINUTE") {

+        ret = INTERVAL_MASK(DAY) | INTERVAL_MASK(HOUR) | INTERVAL_MASK(MINUTE);

+    } else if (ustr == "DAY TO SECOND") {

+        ret = INTERVAL_MASK(DAY) | INTERVAL_MASK(HOUR) | INTERVAL_MASK(MINUTE) | INTERVAL_MASK(SECOND);

+    } else if (ustr == "HOUR TO MINUTE") {

+        ret = INTERVAL_MASK(HOUR) | INTERVAL_MASK(MINUTE);

+    } else if (ustr == "HOUR TO SECOND") {

+        ret = INTERVAL_MASK(HOUR) | INTERVAL_MASK(MINUTE) | INTERVAL_MASK(SECOND);

+    } else if (ustr == "MINUTE TO SECOND") {

+        ret = INTERVAL_MASK(MINUTE) | INTERVAL_MASK(SECOND);

+    } else {
+        return false;

+    }
+
+    return true;
+}
+
 } // NYql
 
 namespace NKikimr {
@@ -2677,7 +2728,6 @@ void get_type_io_data(Oid typid,
     }
 }
 
-
 }
 
 
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 2c710cc99bb..08bdb52f2fc 100644
--- a/ydb/library/yql/providers/common/mkql/yql_provider_mkql.cpp
+++ b/ydb/library/yql/providers/common/mkql/yql_provider_mkql.cpp
@@ -2240,7 +2240,12 @@ TMkqlCommonCallableCompiler::TShared::TShared() {
 
     AddCallable("PgConst", [](const TExprNode& node, TMkqlBuildContext& ctx) {
         auto type = AS_TYPE(TPgType, BuildType(node, *node.GetTypeAnn(), ctx.ProgramBuilder));
-        return ctx.ProgramBuilder.PgConst(type, node.Head().Content());
+        TRuntimeNode typeMod;
+        if (node.ChildrenSize() >= 3) {
+            typeMod = MkqlBuildExpr(*node.Child(2), ctx);
+        }
+
+        return ctx.ProgramBuilder.PgConst(type, node.Head().Content(), typeMod);
     });
 
     AddCallable("PgInternal0", [](const TExprNode& node, TMkqlBuildContext& ctx) {
diff --git a/ydb/library/yql/sql/pg_dummy/pg_sql_dummy.cpp b/ydb/library/yql/sql/pg_dummy/pg_sql_dummy.cpp
index 840c2b5910c..1688176083e 100644
--- a/ydb/library/yql/sql/pg_dummy/pg_sql_dummy.cpp
+++ b/ydb/library/yql/sql/pg_dummy/pg_sql_dummy.cpp
@@ -160,4 +160,10 @@ TMaybe<NKikimr::NUdf::EDataSlot> ConvertFromPgType(ui32 typeId) {
     return Nothing();
 }
 
+bool ParsePgIntervalModifier(const TString& str, i32& ret) {
+    Y_UNUSED(str);
+    Y_UNUSED(ret);
+    return false;
+}
+
 } // NYql
diff --git a/ydb/library/yql/sql/v1/builtin.cpp b/ydb/library/yql/sql/v1/builtin.cpp
index ff0278535b1..d26811c75d7 100644
--- a/ydb/library/yql/sql/v1/builtin.cpp
+++ b/ydb/library/yql/sql/v1/builtin.cpp
@@ -607,7 +607,7 @@ public:
 class TYqlPgConst : public TCallNode {
 public:
     TYqlPgConst(TPosition pos, const TVector<TNodePtr>& args)
-        : TCallNode(pos, "PgConst", 2, 2, args)
+        : TCallNode(pos, "PgConst", 2, -1, args)
     {
     }
 
@@ -627,6 +627,22 @@ public:
             Args[0] = value;
         }
 
+        if (Args.size() > 2) {
+            TVector<TNodePtr> typeModArgs;
+            typeModArgs.push_back(Args[1]);
+            for (ui32 i = 2; i < Args.size(); ++i) {
+                if (!Args[i]->IsLiteral()) {
+                    ctx.Error(Args[i]->GetPos()) << "Expecting literal";
+                    return false;
+                }
+
+                typeModArgs.push_back(BuildQuotedAtom(Args[i]->GetPos(), Args[i]->GetLiteralValue()));
+            }
+
+            Args.erase(Args.begin() + 2, Args.end());
+            Args.push_back(new TCallNodeImpl(Pos, "PgTypeMod", typeModArgs));
+        }
+
         return TCallNode::DoInit(ctx, src);
     }
 
@@ -635,6 +651,42 @@ public:
     }
 };
 
+class TYqlPgCast : public TCallNode {
+public:
+    TYqlPgCast(TPosition pos, const TVector<TNodePtr>& args)
+        : TCallNode(pos, "PgCast", 2, -1, args)
+    {
+    }
+
+    bool DoInit(TContext& ctx, ISource* src) final {
+        if (!ValidateArguments(ctx)) {
+            return false;
+        }
+
+        if (Args.size() > 2) {
+            TVector<TNodePtr> typeModArgs;
+            typeModArgs.push_back(Args[1]);
+            for (ui32 i = 2; i < Args.size(); ++i) {
+                if (!Args[i]->IsLiteral()) {
+                    ctx.Error(Args[i]->GetPos()) << "Expecting literal";
+                    return false;
+                }
+
+                typeModArgs.push_back(BuildQuotedAtom(Args[i]->GetPos(), Args[i]->GetLiteralValue()));
+            }
+
+            Args.erase(Args.begin() + 2, Args.end());
+            Args.push_back(new TCallNodeImpl(Pos, "PgTypeMod", typeModArgs));
+        }
+
+        return TCallNode::DoInit(ctx, src);
+    }
+
+    TNodePtr DoClone() const final {
+        return new TYqlPgCast(Pos, CloneContainer(Args));
+    }
+};
+
 class TYqlPgOp : public TCallNode {
 public:
     TYqlPgOp(TPosition pos, const TVector<TNodePtr>& args)
@@ -2727,7 +2779,7 @@ struct TBuiltinFuncData {
             {"pgconst", BuildSimpleBuiltinFactoryCallback<TYqlPgConst>() },
             {"pgop", BuildSimpleBuiltinFactoryCallback<TYqlPgOp>() },
             {"pgcall", BuildSimpleBuiltinFactoryCallback<TYqlPgCall>() },
-            {"pgcast", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("PgCast", 2, 3) },
+            {"pgcast", BuildSimpleBuiltinFactoryCallback<TYqlPgCast>() },
             {"frompg", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("FromPg", 1, 1) },
             {"topg", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("ToPg", 1, 1) },
             {"pgor", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("PgOr", 2, 2) },