YQL-13710 internal argument in deserializations

ref:d3708023726a63decc66682ce944f8aef1d4d776

6 files changed, 60 insertions, 11 deletions

diff --git a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
index dff33a91550..f59b3859c09 100644
--- a/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_simple1.cpp
@@ -2400,6 +2400,11 @@ TExprNode::TPtr OptimizeToFlow(const TExprNode::TPtr& node, TExprContext& ctx) {
         return ctx.ChangeChildren(*node, node->Head().ChildrenList());
     }
 
+    if (node->Head().IsCallable("WithContext")) {
+        YQL_CLOG(DEBUG, Core) << "Swap " << node->Content() << " with " << node->Head().Content();
+        return ctx.ChangeChild(node->Head(), 1, ctx.NewCallable(node->Pos(), "ToFlow", { node->Head().TailPtr() }));
+    }
+
     return node;
 }
 
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 3fa325a6b97..2d43e58fbc4 100644
--- a/ydb/library/yql/core/type_ann/type_ann_core.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_core.cpp
@@ -9668,6 +9668,17 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         return IGraphTransformer::TStatus::Ok;
     }
 
+    IGraphTransformer::TStatus PgInternal0Wrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
+        Y_UNUSED(output);
+        if (!EnsureArgsCount(*input, 0, ctx.Expr)) {
+            return IGraphTransformer::TStatus::Error;
+        }
+
+        auto internalId = NPg::LookupType("internal").TypeId;
+        input->SetTypeAnn(ctx.Expr.MakeType<TPgExprType>(internalId));
+        return IGraphTransformer::TStatus::Ok;
+    }
+
     IGraphTransformer::TStatus PgCastWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
         Y_UNUSED(output);
         if (!EnsureArgsCount(*input, 2, ctx.Expr)) {
@@ -9909,6 +9920,8 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
                         .Atom(0, deserializeFuncDesc.Name)
                         .Atom(1, ToString(aggDesc.DeserializeFuncId))
                         .Arg(2, "state")
+                        .Callable(3, "PgInternal0")
+                        .Seal()
                     .Seal()
                 .Seal()
                 .Build();
@@ -9925,17 +9938,14 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
                             .Callable(0, "Exists")
                                 .Arg(0, "state1")
                             .Seal()
-                            .Callable(1, "If")
-                                .Callable(0, "Exists")
-                                    .Arg(0, "state2")
-                                .Seal()
-                                .Callable(1, "PgResolvedCallCtx")
+                            .Callable(1, "Coalesce")
+                                .Callable(0, "PgResolvedCallCtx")
                                     .Atom(0, combineFuncDesc.Name)
                                     .Atom(1, ToString(aggDesc.CombineFuncId))
                                     .Arg(2, "state1")
                                     .Arg(3, "state2")
                                 .Seal()
-                                .Arg(2, "state1")
+                                .Arg(1, "state1")
                             .Seal()
                             .Arg(2, "state2")
                         .Seal()
@@ -13710,6 +13720,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
         Functions["PgType"] = &PgTypeWrapper;
         Functions["PgCast"] = &PgCastWrapper;
         Functions["PgAggregationTraits"] = &PgAggregationTraitsWrapper;
+        Functions["PgInternal0"] = &PgInternal0Wrapper;
         Functions["AutoDemuxList"] = &AutoDemuxListWrapper;
         Functions["AggrCountInit"] = &AggrCountInitWrapper;
         Functions["AggrCountUpdate"] = &AggrCountUpdateWrapper;
diff --git a/ydb/library/yql/minikql/mkql_program_builder.cpp b/ydb/library/yql/minikql/mkql_program_builder.cpp
index d41c29302fd..8b98460e6f5 100644
--- a/ydb/library/yql/minikql/mkql_program_builder.cpp
+++ b/ydb/library/yql/minikql/mkql_program_builder.cpp
@@ -5089,6 +5089,11 @@ TRuntimeNode TProgramBuilder::WithContext(const std::string_view& contextType, T
     return TRuntimeNode(callableBuilder.Build(), false);
 }
 
+TRuntimeNode TProgramBuilder::PgInternal0(TType* returnType) {
+    TCallableBuilder callableBuilder(Env, __func__, returnType);
+    return TRuntimeNode(callableBuilder.Build(), false);
+}
+
 bool CanExportType(TType* type, const TTypeEnvironment& env) {
     if (type->GetKind() == TType::EKind::Type) {
         return false; // Type of Type
diff --git a/ydb/library/yql/minikql/mkql_program_builder.h b/ydb/library/yql/minikql/mkql_program_builder.h
index eeba8213321..4412523dced 100644
--- a/ydb/library/yql/minikql/mkql_program_builder.h
+++ b/ydb/library/yql/minikql/mkql_program_builder.h
@@ -630,6 +630,7 @@ public:
     TRuntimeNode FromPg(TRuntimeNode input, TType* returnType);
     TRuntimeNode ToPg(TRuntimeNode input, TType* returnType);
     TRuntimeNode WithContext(const std::string_view& contextType, TRuntimeNode input);
+    TRuntimeNode PgInternal0(TType* returnType);
 
 protected:
     TRuntimeNode Invoke(const std::string_view& funcName, TType* resultType, const TArrayRef<const TRuntimeNode>& args);
diff --git a/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp b/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp
index 248dbd92695..a4870bf0225 100644
--- a/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp
+++ b/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp
@@ -199,8 +199,6 @@ public:
     }
 
     NUdf::TUnboxedValuePod DoCalculate(TComputationContext& compCtx) const {
-        SET_MEMORY_CONTEXT;
-
         if (TypeId == INT2OID) {
             return ScalarDatumToPod(Int16GetDatum(FromString<i16>(Value)));
         } else if (TypeId == INT4OID) {
@@ -211,11 +209,14 @@ public:
             return ScalarDatumToPod(Float4GetDatum(FromString<float>(Value)));
         } else if (TypeId == FLOAT8OID) {
             return ScalarDatumToPod(Float8GetDatum(FromString<double>(Value)));
-        } else if (TypeId == TEXTOID) {
+        } else if (TypeId == TEXTOID || TypeId == VARCHAROID || TypeId == BYTEAOID) {
+            SET_MEMORY_CONTEXT;
             return PointerDatumToPod(PointerGetDatum(cstring_to_text_with_len(Value.data(), Value.size())));
         } else if (TypeId == BOOLOID) {
             return ScalarDatumToPod(BoolGetDatum(!Value.empty() && Value[0] == 't'));
-        } else if (TypeId == NUMERICOID) {
+        }
+        else if (TypeId == NUMERICOID) {
+            SET_MEMORY_CONTEXT;
             LOCAL_FCINFO(callInfo, 3);
             Zero(*callInfo);
             callInfo->nargs = 3;
@@ -243,7 +244,7 @@ public:
             }
             PG_END_TRY();
         } else {
-            UdfTerminate((TStringBuilder() << "Unsupported pg type:" << NPg::LookupType(TypeId).Name).c_str());
+            UdfTerminate((TStringBuilder() << "Unsupported pg type: " << NPg::LookupType(TypeId).Name).c_str());
         }
     }
 
@@ -255,6 +256,23 @@ private:
     const TString Value;
 };
 
+class TPgInternal0 : public TMutableComputationNode<TPgInternal0> {
+    typedef TMutableComputationNode<TPgInternal0> TBaseComputation;
+public:
+    TPgInternal0(TComputationMutables& mutables)
+        : TBaseComputation(mutables)
+    {
+    }
+
+    NUdf::TUnboxedValuePod DoCalculate(TComputationContext& compCtx) const {
+        return ScalarDatumToPod(PointerGetDatum(nullptr));
+    }
+
+private:
+    void RegisterDependencies() const final {
+    }
+};
+
 class TFunctionCallInfo {
 public:
     TFunctionCallInfo(ui32 numArgs, const FmgrInfo* finfo)
@@ -777,6 +795,10 @@ TComputationNodeFactory GetPgFactory() {
                 return new TPgConst(ctx.Mutables, typeId, value);
             }
 
+            if (name == "PgInternal0") {
+                return new TPgInternal0(ctx.Mutables);
+            }
+
             if (name == "PgResolvedCall") {
                 const auto useContextData = AS_VALUE(TDataLiteral, callable.GetInput(0));
                 const auto nameData = AS_VALUE(TDataLiteral, callable.GetInput(1));
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 285c7ed2321..e085f8e97e0 100644
--- a/ydb/library/yql/providers/common/mkql/yql_provider_mkql.cpp
+++ b/ydb/library/yql/providers/common/mkql/yql_provider_mkql.cpp
@@ -2242,6 +2242,11 @@ TMkqlCommonCallableCompiler::TShared::TShared() {
         return ctx.ProgramBuilder.PgConst(type, node.Head().Content());
     });
 
+    AddCallable("PgInternal0", [](const TExprNode& node, TMkqlBuildContext& ctx) {
+        auto returnType = BuildType(node, *node.GetTypeAnn(), ctx.ProgramBuilder);
+        return ctx.ProgramBuilder.PgInternal0(returnType);
+    });
+
     AddCallable({"PgResolvedCall","PgResolvedCallCtx" }, [](const TExprNode& node, TMkqlBuildContext& ctx) {
         auto name = node.Head().Content();
         auto id = FromString<ui32>(node.Child(1)->Content());