YQ-1154: support decimal and pg types in result formatter

ref:acbbaaf61064fcc18561413e7b383d96892cf289

2 files changed, 103 insertions, 9 deletions

diff --git a/ydb/core/yq/libs/result_formatter/result_formatter.cpp b/ydb/core/yq/libs/result_formatter/result_formatter.cpp
index 394832b7e50..00f84cbd49c 100644
--- a/ydb/core/yq/libs/result_formatter/result_formatter.cpp
+++ b/ydb/core/yq/libs/result_formatter/result_formatter.cpp
@@ -1,5 +1,6 @@
 #include "result_formatter.h"
 
+#include <ydb/library/mkql_proto/mkql_proto.h>
 #include <ydb/library/yql/providers/common/schema/mkql/yql_mkql_schema.h>
 #include <ydb/library/yql/providers/common/schema/expr/yql_expr_schema.h>
 #include <ydb/library/yql/providers/common/codec/yql_codec.h>
@@ -26,25 +27,36 @@ const NYql::TTypeAnnotationNode* MakePrimitiveType(NYdb::TTypeParser& parser, NY
     return ctx.MakeType<NYql::TDataExprType>(dataSlot);
 }
 
-const NYql::TTypeAnnotationNode* MakeDecimalType(NYdb::TTypeParser& parser, NYql::TExprContext& ctx)
-{
-    auto decimal = parser.GetDecimal();
-    auto dataSlot = NYql::NUdf::GetDataSlot(TStringBuilder() << "Decimal(" << (ui32)decimal.Precision << ',' << (ui32)decimal.Scale << ")");
-    return ctx.MakeType<NYql::TDataExprType>(dataSlot);
-}
-
 NKikimr::NMiniKQL::TType* MakePrimitiveType(NYdb::TTypeParser& parser, NKikimr::NMiniKQL::TTypeEnvironment& env)
 {
     auto dataSlot = NYql::NUdf::GetDataSlot(TStringBuilder() << parser.GetPrimitive());
     return TDataType::Create(GetDataTypeInfo(dataSlot).TypeId, env);
 }
 
+const NYql::TTypeAnnotationNode* MakeDecimalType(NYdb::TTypeParser& parser, NYql::TExprContext& ctx)
+{
+    auto decimal = parser.GetDecimal();
+    return ctx.MakeType<NYql::TDataExprParamsType>(NYql::EDataSlot::Decimal, ctx.AppendString(ToString(decimal.Precision)), ctx.AppendString(ToString(decimal.Scale)));
+}
+
 NKikimr::NMiniKQL::TType* MakeDecimalType(NYdb::TTypeParser& parser, NKikimr::NMiniKQL::TTypeEnvironment& env)
 {
     auto decimal = parser.GetDecimal();
     return TDataDecimalType::Create((ui8)decimal.Precision, (ui8)decimal.Scale, env);
 }
 
+const NYql::TTypeAnnotationNode* MakePgType(NYdb::TTypeParser& parser, NYql::TExprContext& ctx)
+{
+    auto pgType = parser.GetPg();
+    return ctx.MakeType<NYql::TPgExprType>(pgType.Oid);
+}
+
+NKikimr::NMiniKQL::TType* MakePgType(NYdb::TTypeParser& parser, NKikimr::NMiniKQL::TTypeEnvironment& env)
+{
+    auto pgType = parser.GetPg();
+    return TPgType::Create(pgType.Oid, env);
+}
+
 const NYql::TTypeAnnotationNode* MakeOptionalType(const NYql::TTypeAnnotationNode* underlying, NYql::TExprContext& ctx)
 {
     return ctx.MakeType<NYql::TOptionalExprType>(underlying);
@@ -184,6 +196,9 @@ TType MakeType(NYdb::TTypeParser& parser, TContext& env)
     case NYdb::TTypeParser::ETypeKind::Decimal: {
         return MakeDecimalType(parser, env);
     }
+    case NYdb::TTypeParser::ETypeKind::Pg: {
+        return MakePgType(parser, env);
+    }
     case NYdb::TTypeParser::ETypeKind::Optional: {
         parser.OpenOptional();
         auto underlying = MakeType<TType>(parser, env);
@@ -288,7 +303,7 @@ struct TTypePair {
 
 TTypePair FormatColumnType(
     NJson::TJsonValue& root,
-    NYdb::TType type,
+    const NYdb::TType& type,
     NKikimr::NMiniKQL::TTypeEnvironment& typeEnv,
     NYql::TExprContext& ctx)
 {
@@ -319,11 +334,29 @@ void FormatColumnValue(
     NJson::TJsonValue& root,
     const NYdb::TValue& value,
     NKikimr::NMiniKQL::TType* type,
+    const NKikimr::NMiniKQL::TTypeEnvironment&,
     const THolderFactory& holderFactory)
 {
+    if (type->GetKind() == TType::EKind::Pg) {
+        NYdb::TValueParser parser(value);
+        auto pgValue = parser.GetPg();
+        /*if (pgValue.IsNull()) {
+            root = NJson::TJsonValue();
+            return;
+        }*/
+
+        if (pgValue.IsText()) {
+            root = NJson::TJsonValue(pgValue.Content_);
+            return;
+        }
+
+        root = NJson::TJsonValue("<binary pg value>");
+        return;
+    }
+
     const Ydb::Value& rawProtoValue = NYdb::TProtoAccessor::GetProto(value);
 
-    NYql::NUdf::TUnboxedValue unboxed = ImportValueFromProto(
+    auto unboxed = ImportValueFromProto(
         type,
         rawProtoValue,
         holderFactory);
@@ -384,6 +417,7 @@ void FormatResultSet(NJson::TJsonValue& root, const NYdb::TResultSet& resultSet)
                 row[columnMeta.Name],
                 rsParser.GetValue(columnNum),
                 columnTypes[columnNum].MiniKQLType,
+                typeEnv,
                 holderFactory);
         }
     }
diff --git a/ydb/core/yq/libs/result_formatter/result_formatter_ut.cpp b/ydb/core/yq/libs/result_formatter/result_formatter_ut.cpp
index bebbf68ab36..8105aabd278 100644
--- a/ydb/core/yq/libs/result_formatter/result_formatter_ut.cpp
+++ b/ydb/core/yq/libs/result_formatter/result_formatter_ut.cpp
@@ -328,6 +328,66 @@ Y_UNIT_TEST_SUITE(ResultFormatter) {
         UNIT_ASSERT_VALUES_EQUAL(stream.Str(), expected);
     }
 
+    Y_UNIT_TEST(Decimal) {
+        Ydb::ResultSet rs;
+        {
+            auto& column = *rs.add_columns();
+            column.set_name("column0");
+            auto& type = *column.mutable_type()->mutable_decimal_type();
+            type.set_precision(21);
+            type.set_scale(7);
+        }
+        {
+            auto& value = *rs.add_rows();
+            value.add_items()->set_low_128(1);
+            value.add_items()->set_high_128(2);
+        }
+
+        NJson::TJsonValue root;
+        FormatResultSet(root, rs);
+
+        TStringStream stream;
+        NJson::WriteJson(&stream, &root);
+
+        // Cerr << stream.Str() << Endl;
+
+        TString expected = R"___({"data":[{"column0":"0.0000001"}],"columns":[{"name":"column0","type":["DataType","Decimal","21","7"]}]})___";
+
+        UNIT_ASSERT_VALUES_EQUAL(stream.Str(), expected);
+    }
+
+    Y_UNIT_TEST(Pg) {
+        Ydb::ResultSet rs;
+        {
+            auto& column = *rs.add_columns();
+            column.set_name("column0");
+            auto& pg_type = *column.mutable_type()->mutable_pg_type();
+            pg_type.set_oid(16);
+            pg_type.set_typlen(234);
+            pg_type.set_typmod(-987);
+        }
+        {
+            auto& value = *rs.add_rows();
+            auto& cell = *value.add_items();
+            cell.set_text_value("my_text_value");
+        }
+        // empty text value
+        {
+            auto& value = *rs.add_rows();
+            auto& cell = *value.add_items();
+            cell.set_text_value("");
+        }
+
+        NJson::TJsonValue root;
+        FormatResultSet(root, rs);
+
+        TStringStream stream;
+        NJson::WriteJson(&stream, &root);
+        //Cerr << "Stream >> " << stream.Str() << Endl;
+        TString expected = R"___({"data":[{"column0":"my_text_value"},{"column0":""}],"columns":[{"name":"column0","type":["PgType","bool"]}]})___";
+        UNIT_ASSERT_VALUES_EQUAL(stream.Str(), expected);
+    }
+
     Y_UNIT_TEST(VariantStruct) {
         Ydb::ResultSet rs;
         {