[yql] Support Pg values in evaluation

YQL-12393

ref:135b1b52d9e65c1a95493bb469352d14e23453b3

4 files changed, 125 insertions, 9 deletions

diff --git a/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp b/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp
index f210675a203..001ab87e6ee 100644
--- a/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp
+++ b/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp
@@ -1595,15 +1595,13 @@ void WriteYsonValueInTableFormatPg(TOutputBuf& buf, TPgType* type, const NUdf::T
     }
 }
 
-void WriteYsonValuePg(TYsonResultWriter& writer, const NUdf::TUnboxedValuePod& value, NKikimr::NMiniKQL::TPgType* type,
-    const TVector<ui32>* structPositions) {
+TString PgValueToString(const NUdf::TUnboxedValuePod& value, ui32 pgTypeId) {
     if (!value) {
-        writer.OnNull();
-        return;
+        return "null";
     }
 
     TString ret;
-    switch (type->GetTypeId()) {
+    switch (pgTypeId) {
     case BOOLOID:
         ret = DatumGetBool(ScalarDatumFromPod(value)) ? "true" : "false";
         break;
@@ -1636,7 +1634,7 @@ void WriteYsonValuePg(TYsonResultWriter& writer, const NUdf::TUnboxedValuePod& v
     }
     default:
         TPAllocScope call;
-        const auto& typeInfo = NPg::LookupType(type->GetTypeId());
+        const auto& typeInfo = NPg::LookupType(pgTypeId);
         auto outFuncId = typeInfo.OutFuncId;
         if (typeInfo.TypeId == typeInfo.ArrayTypeId) {
             outFuncId = NPg::LookupProc("array_out", { 0 }).ProcId;
@@ -1667,10 +1665,20 @@ void WriteYsonValuePg(TYsonResultWriter& writer, const NUdf::TUnboxedValuePod& v
         ret = str;
     }
 
-    writer.OnStringScalar(ret);
+    return ret;
 }
 
-NUdf::TUnboxedValue ReadYsonValuePg(TPgType* type, char cmd, TInputBuf& buf) {
+void WriteYsonValuePg(TYsonResultWriter& writer, const NUdf::TUnboxedValuePod& value, NKikimr::NMiniKQL::TPgType* type,
+    const TVector<ui32>* structPositions) {
+    if (!value) {
+        writer.OnNull();
+        return;
+    }
+
+    writer.OnStringScalar(PgValueToString(value, type->GetTypeId()));
+}
+
+NUdf::TUnboxedValue ReadYsonValueInTableFormatPg(TPgType* type, char cmd, TInputBuf& buf) {
     using namespace NYson::NDetail;
     if (cmd == EntitySymbol) {
         return NUdf::TUnboxedValuePod();
@@ -1762,6 +1770,77 @@ NUdf::TUnboxedValue ReadYsonValuePg(TPgType* type, char cmd, TInputBuf& buf) {
     }
 }
 
+NUdf::TUnboxedValue ReadYsonValuePg(TPgType* type, char cmd, TInputBuf& buf) {
+    using namespace NYson::NDetail;
+    if (cmd == EntitySymbol) {
+        return NUdf::TUnboxedValuePod();
+    }
+
+    CHECK_EXPECTED(cmd, StringMarker);
+    auto s = buf.ReadYtString();
+    switch (type->GetTypeId()) {
+    case BOOLOID: {
+        return ScalarDatumToPod(BoolGetDatum(FromString<bool>(s)));
+    }
+    case INT2OID: {
+        return ScalarDatumToPod(Int16GetDatum(FromString<i16>(s)));
+    }
+    case INT4OID: {
+        return ScalarDatumToPod(Int32GetDatum(FromString<i32>(s)));
+    }
+    case INT8OID: {
+        return ScalarDatumToPod(Int64GetDatum(FromString<i64>(s)));
+    }
+    case FLOAT4OID: {
+        return ScalarDatumToPod(Float4GetDatum(FromString<float>(s)));
+    }
+    case FLOAT8OID: {
+        return ScalarDatumToPod(Float8GetDatum(FromString<double>(s)));
+    }
+    case BYTEAOID:
+    case VARCHAROID:
+    case TEXTOID: {
+        auto ret = MakeVar(s);
+        return PointerDatumToPod((Datum)ret);
+    }
+    case CSTRINGOID: {
+        auto ret = MakeCString(s);
+        return PointerDatumToPod((Datum)ret);
+    }
+    default:
+        TString str{s};
+
+        TPAllocScope call;
+        const auto& typeInfo = NPg::LookupType(type->GetTypeId());
+        auto typeIOParam = MakeTypeIOParam(typeInfo);
+        auto inFuncId = typeInfo.InFuncId;
+        if (typeInfo.TypeId == typeInfo.ArrayTypeId) {
+            inFuncId = NPg::LookupProc("array_in", { 0,0,0 }).ProcId;
+        }
+
+        FmgrInfo finfo;
+        Zero(finfo);
+        Y_ENSURE(inFuncId);
+        fmgr_info(inFuncId, &finfo);
+        Y_ENSURE(!finfo.fn_retset);
+        Y_ENSURE(finfo.fn_addr);
+        Y_ENSURE(finfo.fn_nargs >= 1 && finfo.fn_nargs <= 3);
+        LOCAL_FCINFO(callInfo, 3);
+        Zero(*callInfo);
+        callInfo->flinfo = &finfo;
+        callInfo->nargs = 3;
+        callInfo->fncollation = DEFAULT_COLLATION_OID;
+        callInfo->isnull = false;
+        callInfo->args[0] = { (Datum)str.c_str(), false };
+        callInfo->args[1] = { ObjectIdGetDatum(typeIOParam), false };
+        callInfo->args[2] = { Int32GetDatum(-1), false };
+
+        auto x = finfo.fn_addr(callInfo);
+        Y_ENSURE(!callInfo->isnull);
+        return typeInfo.PassByValue ? ScalarDatumToPod(x) : PointerDatumToPod(x);
+    }
+}
+
 NKikimr::NUdf::TUnboxedValue ReadSkiffPg(NKikimr::NMiniKQL::TPgType* type, NCommon::TInputBuf& buf) {
     auto marker = buf.Read();
     if (!marker) {
diff --git a/ydb/library/yql/providers/common/codec/yql_codec.cpp b/ydb/library/yql/providers/common/codec/yql_codec.cpp
index 8f06281f45d..01e3e863881 100644
--- a/ydb/library/yql/providers/common/codec/yql_codec.cpp
+++ b/ydb/library/yql/providers/common/codec/yql_codec.cpp
@@ -1243,7 +1243,7 @@ NUdf::TUnboxedValue ReadYsonValue(TType* type,
 
     case TType::EKind::Pg: {
         auto pgType = static_cast<TPgType*>(type);
-        return ReadYsonValuePg(pgType, cmd, buf);
+        return isTableFormat ? ReadYsonValueInTableFormatPg(pgType, cmd, buf) : ReadYsonValuePg(pgType, cmd, buf);
     }
 
     default:
@@ -2419,6 +2419,27 @@ TExprNode::TPtr ValueToExprLiteral(const TTypeAnnotationNode* type, const NKikim
         });
     }
 
+    case ETypeAnnotationKind::Pg: {
+        auto pgType = type->Cast<TPgExprType>();
+        if (!value) {
+            return ctx.NewCallable(pos, "PgCast", {
+                ctx.NewCallable(pos, "Utf8", {
+                    ctx.NewAtom(pos, "null")
+                }),
+                ctx.NewCallable(pos, "PgType", {
+                    ctx.NewAtom(pos, pgType->GetName())
+                })
+            });
+        } else {
+            return ctx.NewCallable(pos, "PgConst", {
+                ctx.NewAtom(pos, PgValueToString(value, pgType->GetId())),
+                ctx.NewCallable(pos, "PgType", {
+                    ctx.NewAtom(pos, pgType->GetName())
+                })
+            });
+        }
+    }
+
     default:
         break;
     }
diff --git a/ydb/library/yql/providers/common/codec/yql_pg_codec.h b/ydb/library/yql/providers/common/codec/yql_pg_codec.h
index b20adab166e..36a09bec4fb 100644
--- a/ydb/library/yql/providers/common/codec/yql_pg_codec.h
+++ b/ydb/library/yql/providers/common/codec/yql_pg_codec.h
@@ -12,11 +12,14 @@
 namespace NYql {
 namespace NCommon {
 
+TString PgValueToString(const NUdf::TUnboxedValuePod& value, ui32 pgTypeId);
+
 void WriteYsonValuePg(TYsonResultWriter& writer, const NUdf::TUnboxedValuePod& value, NKikimr::NMiniKQL::TPgType* type,
     const TVector<ui32>* structPositions);
 
 void WriteYsonValueInTableFormatPg(TOutputBuf& buf, NKikimr::NMiniKQL::TPgType* type, const NKikimr::NUdf::TUnboxedValuePod& value);
 
+NUdf::TUnboxedValue ReadYsonValueInTableFormatPg(NKikimr::NMiniKQL::TPgType* type, char cmd, TInputBuf& buf);
 NUdf::TUnboxedValue ReadYsonValuePg(NKikimr::NMiniKQL::TPgType* type, char cmd, TInputBuf& buf);
 
 extern "C" void ReadSkiffPgValue(NKikimr::NMiniKQL::TPgType* type, NKikimr::NUdf::TUnboxedValue& value, NCommon::TInputBuf& buf);
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 1688176083e..3f97980d3e3 100644
--- a/ydb/library/yql/sql/pg_dummy/pg_sql_dummy.cpp
+++ b/ydb/library/yql/sql/pg_dummy/pg_sql_dummy.cpp
@@ -20,6 +20,12 @@ NYql::TAstParseResult PGToYql(const TString& query, const NSQLTranslation::TTran
 namespace NYql {
 namespace NCommon {
 
+TString PgValueToString(const NUdf::TUnboxedValuePod& value, ui32 pgTypeId) {
+    Y_UNUSED(value);
+    Y_UNUSED(pgTypeId);
+    throw yexception() << "PG types are not supported";
+}
+
 void WriteYsonValuePg(TYsonResultWriter& writer, const NUdf::TUnboxedValuePod& value, NKikimr::NMiniKQL::TPgType* type,
     const TVector<ui32>* structPositions) {
     Y_UNUSED(writer);
@@ -36,6 +42,13 @@ void WriteYsonValueInTableFormatPg(TOutputBuf& buf, NKikimr::NMiniKQL::TPgType*
     throw yexception() << "PG types are not supported";
 }
 
+NUdf::TUnboxedValue ReadYsonValueInTableFormatPg(NKikimr::NMiniKQL::TPgType* type, char cmd, TInputBuf& buf) {
+    Y_UNUSED(type);
+    Y_UNUSED(cmd);
+    Y_UNUSED(buf);
+    throw yexception() << "PG types are not supported";
+}
+
 NUdf::TUnboxedValue ReadYsonValuePg(NKikimr::NMiniKQL::TPgType* type, char cmd, TInputBuf& buf) {
     Y_UNUSED(type);
     Y_UNUSED(cmd);