support Tagged type in kqp

5 files changed, 46 insertions, 2 deletions

diff --git a/ydb/core/kqp/session_actor/kqp_worker_common.cpp b/ydb/core/kqp/session_actor/kqp_worker_common.cpp
index 5c64c0c50f4..27a9e93b72d 100644
--- a/ydb/core/kqp/session_actor/kqp_worker_common.cpp
+++ b/ydb/core/kqp/session_actor/kqp_worker_common.cpp
@@ -175,8 +175,10 @@ bool IsSameProtoType(const NKikimrMiniKQL::TType& actual, const NKikimrMiniKQL::
             return true;
         case NKikimrMiniKQL::ETypeKind::Pg:
             return actual.GetPg().Getoid() == expected.GetPg().Getoid();
+        case NKikimrMiniKQL::ETypeKind::Tagged:
+            return (actual.GetTagged().GetTag() == expected.GetTagged().GetTag()) &&
+                IsSameProtoType(actual.GetTagged().GetItem(), expected.GetTagged().GetItem());
         case NKikimrMiniKQL::ETypeKind::Unknown:
-        case NKikimrMiniKQL::ETypeKind::Reserved_11:
         case NKikimrMiniKQL::ETypeKind::Reserved_12:
         case NKikimrMiniKQL::ETypeKind::Reserved_13:
         case NKikimrMiniKQL::ETypeKind::Reserved_14:
diff --git a/ydb/core/kqp/ut/kqp_scan_ut.cpp b/ydb/core/kqp/ut/kqp_scan_ut.cpp
index 9ef81036b53..96f85620de3 100644
--- a/ydb/core/kqp/ut/kqp_scan_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_scan_ut.cpp
@@ -196,6 +196,17 @@ Y_UNIT_TEST_SUITE(KqpScan) {
         ])", StreamResultToYson(it));
     }
 
+    Y_UNIT_TEST(TaggedScalar) {
+        auto kikimr = DefaultKikimrRunner();
+
+        auto it = kikimr.GetTableClient().StreamExecuteScanQuery(R"(
+                SELECT AsTagged(789, "xxx") AS t;
+            )").GetValueSync();
+
+        UNIT_ASSERT_C(it.IsSuccess(), it.GetIssues().ToString());
+        CompareYson(R"([[789]])", StreamResultToYson(it));
+    }
+
     Y_UNIT_TEST(Offset) {
         auto kikimr = DefaultKikimrRunner({}, AppCfg());
         auto db = kikimr.GetTableClient();
diff --git a/ydb/core/ydb_convert/ydb_convert.cpp b/ydb/core/ydb_convert/ydb_convert.cpp
index 4676765f0e2..01c58c3109a 100644
--- a/ydb/core/ydb_convert/ydb_convert.cpp
+++ b/ydb/core/ydb_convert/ydb_convert.cpp
@@ -82,6 +82,12 @@ void ConvertMiniKQLTypeToYdbType(const NKikimrMiniKQL::TType& input, Ydb::Type&
             ConvertMiniKQLTypeToYdbType(protoOptionalType.GetItem(), *output.mutable_optional_type()->mutable_item());
             break;
         }
+        case NKikimrMiniKQL::ETypeKind::Tagged: {
+            const NKikimrMiniKQL::TTaggedType& protoTaggedType = input.GetTagged();
+            output.mutable_tagged_type()->set_tag(protoTaggedType.GetTag());
+            ConvertMiniKQLTypeToYdbType(protoTaggedType.GetItem(), *output.mutable_tagged_type()->mutable_type());
+            break;
+        }
         case NKikimrMiniKQL::ETypeKind::List: {
             const NKikimrMiniKQL::TListType& protoListType = input.GetList();
             ConvertMiniKQLTypeToYdbType(protoListType.GetItem(), *output.mutable_list_type()->mutable_item());
@@ -497,6 +503,12 @@ void ConvertMiniKQLValueToYdbValue(const NKikimrMiniKQL::TType& inputType,
             }
             break;
         }
+        case NKikimrMiniKQL::ETypeKind::Tagged: {
+            const NKikimrMiniKQL::TTaggedType& protoTaggedType = inputType.GetTagged();
+            const NKikimrMiniKQL::TType& protoItemType = protoTaggedType.GetItem();
+            ConvertMiniKQLValueToYdbValue(protoItemType, inputValue, output);
+            break;
+        }
         case NKikimrMiniKQL::ETypeKind::List: {
             const NKikimrMiniKQL::TListType& protoListType = inputType.GetList();
             const NKikimrMiniKQL::TType& protoItemType = protoListType.GetItem();
diff --git a/ydb/library/mkql_proto/mkql_proto.cpp b/ydb/library/mkql_proto/mkql_proto.cpp
index c30a8bc114a..8157b3dbde0 100644
--- a/ydb/library/mkql_proto/mkql_proto.cpp
+++ b/ydb/library/mkql_proto/mkql_proto.cpp
@@ -157,6 +157,14 @@ void ExportTypeToProtoImpl(TType* type, NKikimrMiniKQL::TType& res) {
             break;
         }
 
+        case TType::EKind::Tagged: {
+            auto taggedType = static_cast<TTaggedType *>(type);
+            res.SetKind(NKikimrMiniKQL::ETypeKind::Tagged);
+            res.MutableTagged()->SetTag(TString(taggedType->GetTag()));
+            ExportTypeToProtoImpl(taggedType->GetBaseType(), *res.MutableTagged()->MutableItem());
+            break;
+        }
+
         case TType::EKind::Optional: {
             auto optionalType = static_cast<TOptionalType *>(type);
             res.SetKind(NKikimrMiniKQL::ETypeKind::Optional);
@@ -1004,6 +1012,11 @@ TType* TProtoImporter::ImportTypeFromProto(const NKikimrMiniKQL::TType& type) {
             TOptionalType* optionalType = TOptionalType::Create(child, env);
             return optionalType;
         }
+        case NKikimrMiniKQL::ETypeKind::Tagged: {
+            const NKikimrMiniKQL::TTaggedType& protoTaggedType = type.GetTagged();
+            TType* child = ImportTypeFromProto(protoTaggedType.GetItem());
+            return TTaggedType::Create(child, protoTaggedType.GetTag(), env);
+        }
         case NKikimrMiniKQL::ETypeKind::List: {
             const NKikimrMiniKQL::TListType& protoListType = type.GetList();
             const NKikimrMiniKQL::TType& protoItemType = protoListType.GetItem();
diff --git a/ydb/library/mkql_proto/protos/minikql.proto b/ydb/library/mkql_proto/protos/minikql.proto
index b07ca18d4c1..444efdef8a2 100644
--- a/ydb/library/mkql_proto/protos/minikql.proto
+++ b/ydb/library/mkql_proto/protos/minikql.proto
@@ -16,7 +16,7 @@ enum ETypeKind {
     Variant = 8;
     Null = 9;
     Pg = 10;
-    Reserved_11 = 11;
+    Tagged = 11;
     Reserved_12 = 12;
     Reserved_13 = 13;
     Reserved_14 = 14;
@@ -69,6 +69,11 @@ message TPgType {
     required uint32 oid = 1;
 }
 
+message TTaggedType {
+    required string Tag = 1;
+    required TType Item = 2;
+}
+
 message TType {
     required ETypeKind Kind = 1;
     oneof type_type {
@@ -80,6 +85,7 @@ message TType {
         TDictType Dict = 7;
         TVariantType Variant = 8;
         TPgType Pg = 9;
+        TTaggedType Tagged = 10;
     }
 }