YTORM-214: Add enum yson's storage type configuration

4 files changed, 151 insertions, 15 deletions

diff --git a/yt/yt/core/yson/protobuf_interop.cpp b/yt/yt/core/yson/protobuf_interop.cpp
index c722e75de3..6cd80f3689 100644
--- a/yt/yt/core/yson/protobuf_interop.cpp
+++ b/yt/yt/core/yson/protobuf_interop.cpp
@@ -344,6 +344,7 @@ public:
         , YsonMap_(descriptor->options().GetExtension(NYT::NYson::NProto::yson_map))
         , Required_(descriptor->options().GetExtension(NYT::NYson::NProto::required))
         , Converter_(registry->FindMessageBytesFieldConverter(descriptor->containing_type(), descriptor->index()))
+        , EnumYsonStorageType_(descriptor->options().GetExtension(NYT::NYson::NProto::enum_yson_storage_type))
     {
         if (YsonMap_ && !descriptor->is_map()) {
             THROW_ERROR_EXCEPTION("Field %v is not a map and cannot be annotated with \"yson_map\" option",
@@ -454,6 +455,11 @@ public:
         return Converter_;
     }
 
+    const NYT::NYson::NProto::EEnumYsonStorageType& GetEnumYsonStorageType() const
+    {
+        return EnumYsonStorageType_;
+    }
+
 private:
     const FieldDescriptor* const Underlying_;
     const TStringBuf YsonName_;
@@ -464,6 +470,7 @@ private:
     const bool YsonMap_;
     const bool Required_;
     const std::optional<TProtobufMessageBytesFieldConverter> Converter_;
+    const NYT::NYson::NProto::EEnumYsonStorageType EnumYsonStorageType_;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -773,7 +780,7 @@ int ConvertToProtobufEnumValueUntyped(
         case NYTree::ENodeType::Int64:
         case NYTree::ENodeType::Uint64: {
             int value = NYTree::ConvertTo<int>(node);
-            THROW_ERROR_EXCEPTION_UNLESS(type->GetUnderlying()->FindValueByNumber(value),
+            THROW_ERROR_EXCEPTION_UNLESS(type->FindLiteralByValue(value),
                 "Unknown value %v of enum %Qv",
                 value,
                 type->GetUnderlying()->name());
@@ -2276,6 +2283,31 @@ private:
         int tag,
         WireFormatLite::WireType wireType)
     {
+        auto storeEnumAsInt = [this, field] (auto value) {
+            const auto* enumType = field->GetEnumType();
+            if (!enumType->FindLiteralByValue(value)) {
+                THROW_ERROR_EXCEPTION("Unknown value %v for field %v",
+                    value,
+                    YPathStack_.GetHumanReadablePath())
+                    << TErrorAttribute("ypath", YPathStack_.GetPath())
+                    << TErrorAttribute("proto_field", field->GetFullName());
+            }
+            Consumer_->OnInt64Scalar(value);
+        };
+        auto storeEnumAsString = [this, field] (auto value) {
+            auto signedValue = static_cast<int>(value);
+            const auto* enumType = field->GetEnumType();
+            auto literal = enumType->FindLiteralByValue(signedValue);
+            if (!literal) {
+                THROW_ERROR_EXCEPTION("Unknown value %v for field %v",
+                    signedValue,
+                    YPathStack_.GetHumanReadablePath())
+                    << TErrorAttribute("ypath", YPathStack_.GetPath())
+                    << TErrorAttribute("proto_field", field->GetFullName());
+            }
+            Consumer_->OnStringScalar(literal);
+        };
+
         switch (wireType) {
             case WireFormatLite::WIRETYPE_VARINT: {
                 ui64 unsignedValue;
@@ -2294,18 +2326,16 @@ private:
                         break;
 
                     case FieldDescriptor::TYPE_ENUM: {
-                        auto signedValue = static_cast<int>(unsignedValue);
-                        const auto* enumType = field->GetEnumType();
-                        auto literal = enumType->FindLiteralByValue(signedValue);
-                        if (!literal) {
-                            THROW_ERROR_EXCEPTION("Unknown value %v for field %v",
-                                signedValue,
-                                YPathStack_.GetHumanReadablePath())
-                                << TErrorAttribute("ypath", YPathStack_.GetPath())
-                                << TErrorAttribute("proto_field", field->GetFullName());
-                        }
                         ParseScalar([&] {
-                            Consumer_->OnStringScalar(literal);
+                            using NYT::NYson::NProto::EEnumYsonStorageType;
+                            switch(field->GetEnumYsonStorageType()) {
+                                case EEnumYsonStorageType::EYST_INT:
+                                    storeEnumAsInt(unsignedValue);
+                                    break;
+                                case EEnumYsonStorageType::EYST_STRING:
+                                    storeEnumAsString(unsignedValue);
+                                    break;
+                            }
                         });
                         break;
                     }
@@ -2542,7 +2572,17 @@ private:
                     }
 
                     case FieldDescriptor::TYPE_ENUM: {
-                        ParseVarintPacked<ui32>(length, field, [&] (auto value) {Consumer_->OnInt64Scalar(value);});
+                        ParseVarintPacked<ui32>(length, field, [&] (auto value) {
+                            using NYT::NYson::NProto::EEnumYsonStorageType;
+                            switch(field->GetEnumYsonStorageType()) {
+                                case EEnumYsonStorageType::EYST_INT:
+                                    storeEnumAsInt(value);
+                                    break;
+                                case EEnumYsonStorageType::EYST_STRING:
+                                    storeEnumAsString(value);
+                                    break;
+                            }
+                        });
                         break;
                     }
 
diff --git a/yt/yt/core/yson/unittests/proto/protobuf_yson_ut.proto b/yt/yt/core/yson/unittests/proto/protobuf_yson_ut.proto
index 6ccc8efe90..db781c9b12 100644
--- a/yt/yt/core/yson/unittests/proto/protobuf_yson_ut.proto
+++ b/yt/yt/core/yson/unittests/proto/protobuf_yson_ut.proto
@@ -162,6 +162,21 @@ message TPackedRepeatedMessage
     repeated fixed64 fixed64_rep = 8 [packed = true];
     repeated sfixed32 sfixed32_rep = 9 [packed = true];
     repeated sfixed64 sfixed64_rep = 10 [packed = true];
-    repeated EEnum enum_rep = 11 [packed = true];
-    
+    repeated EEnum enum_rep = 11 [packed = true, (NYT.NYson.NProto.enum_yson_storage_type) = EYST_INT];
+}
+
+message TMessageWithEnums
+{
+    enum EEnum {
+        VALUE0 = 0;
+        VALUE1 = 1;
+    };
+
+    optional EEnum enum_int = 1                   [(NYT.NYson.NProto.enum_yson_storage_type) = EYST_INT];
+    repeated EEnum enum_rep_not_packed_int = 2    [(NYT.NYson.NProto.enum_yson_storage_type) = EYST_INT, packed = false];
+    repeated EEnum enum_rep_packed_int = 3        [(NYT.NYson.NProto.enum_yson_storage_type) = EYST_INT, packed = true];
+
+    optional EEnum enum_string = 4                [(NYT.NYson.NProto.enum_yson_storage_type) = EYST_STRING];
+    repeated EEnum enum_rep_not_packed_string = 5 [(NYT.NYson.NProto.enum_yson_storage_type) = EYST_STRING, packed = false];
+    repeated EEnum enum_rep_packed_string = 6     [(NYT.NYson.NProto.enum_yson_storage_type) = EYST_STRING, packed = true];
 }
diff --git a/yt/yt/core/yson/unittests/protobuf_yson_ut.cpp b/yt/yt/core/yson/unittests/protobuf_yson_ut.cpp
index 6d020d543c..dd76e05e7a 100644
--- a/yt/yt/core/yson/unittests/protobuf_yson_ut.cpp
+++ b/yt/yt/core/yson/unittests/protobuf_yson_ut.cpp
@@ -16,6 +16,8 @@
 #include <google/protobuf/io/coded_stream.h>
 #include <google/protobuf/io/zero_copy_stream_impl_lite.h>
 
+#include <google/protobuf/util/message_differencer.h>
+
 #include <google/protobuf/wire_format.h>
 
 namespace NYT {
@@ -2615,5 +2617,77 @@ TEST(TPackedRepeatedProtobufTest, TestSerializeDeserialize)
     }
 }
 
+TEST(TEnumYsonStorageTypeTest, TestDeserializeSerialize)
+{
+    NProto::TMessageWithEnums message;
+    {
+        auto zero = NProto::TMessageWithEnums_EEnum::TMessageWithEnums_EEnum_VALUE0;
+        auto one = NProto::TMessageWithEnums_EEnum::TMessageWithEnums_EEnum_VALUE1;
+
+        message.set_enum_int(zero);
+        message.add_enum_rep_not_packed_int(zero);
+        message.add_enum_rep_not_packed_int(one);
+        message.add_enum_rep_packed_int(zero);
+        message.add_enum_rep_packed_int(one);
+
+        message.set_enum_string(one);
+        message.add_enum_rep_not_packed_string(one);
+        message.add_enum_rep_not_packed_string(zero);
+        message.add_enum_rep_packed_string(one);
+        message.add_enum_rep_packed_string(zero);
+    }
+
+    // Proto message to yson.
+    TString stringWithYson;
+    {
+        TString protobufString;
+        Y_UNUSED(message.SerializeToString(&protobufString));
+        TStringOutput ysonOutputStream(stringWithYson);
+        TYsonWriter ysonWriter(&ysonOutputStream, EYsonFormat::Pretty);
+        ArrayInputStream protobufInput(protobufString.data(), protobufString.length());
+
+        EXPECT_NO_THROW(ParseProtobuf(&ysonWriter, &protobufInput, ReflectProtobufMessageType<NProto::TMessageWithEnums>()));
+    }
+
+    // Check enum representation in yson.
+    auto resultedNode = ConvertToNode(TYsonString(stringWithYson));
+    {
+        auto expectedNode = BuildYsonNodeFluently()
+            .BeginMap()
+                .Item("enum_int").Value(0)
+                .Item("enum_rep_not_packed_int").BeginList()
+                    .Item().Value(0)
+                    .Item().Value(1)
+                .EndList()
+                .Item("enum_rep_packed_int").BeginList()
+                    .Item().Value(0)
+                    .Item().Value(1)
+                .EndList()
+                .Item("enum_string").Value("VALUE1")
+                .Item("enum_rep_not_packed_string").BeginList()
+                    .Item().Value("VALUE1")
+                    .Item().Value("VALUE0")
+                .EndList()
+                .Item("enum_rep_packed_string").BeginList()
+                    .Item().Value("VALUE1")
+                    .Item().Value("VALUE0")
+                .EndList()
+            .EndMap();
+
+        EXPECT_TRUE(AreNodesEqual(resultedNode, expectedNode));
+    }
+
+    // Yson to proto message.
+    NProto::TMessageWithEnums resultedMessage;
+    DeserializeProtobufMessage(
+        resultedMessage,
+        ReflectProtobufMessageType<NProto::TMessageWithEnums>(),
+        resultedNode,
+        TProtobufWriterOptions{});
+
+    // Check that original message is equal to its deserialized + serialized version
+    EXPECT_TRUE(google::protobuf::util::MessageDifferencer::Equals(message, resultedMessage));
+}
+
 } // namespace
 } // namespace NYT::NYson
diff --git a/yt/yt_proto/yt/core/yson/proto/protobuf_interop.proto b/yt/yt_proto/yt/core/yson/proto/protobuf_interop.proto
index efcdf8c240..42c8ffbb8b 100644
--- a/yt/yt_proto/yt/core/yson/proto/protobuf_interop.proto
+++ b/yt/yt_proto/yt/core/yson/proto/protobuf_interop.proto
@@ -5,6 +5,12 @@ option go_package = "github.com/ydb-platform/ydb/yt/go/proto/core/yson";
 
 import "google/protobuf/descriptor.proto";
 
+enum EEnumYsonStorageType
+{
+    EYST_STRING = 0;
+    EYST_INT = 1;
+}
+
 extend google.protobuf.FieldOptions
 {
     optional string field_name = 3000;
@@ -12,6 +18,7 @@ extend google.protobuf.FieldOptions
     optional bool yson_string = 3001;
     optional bool yson_map = 3002;
     optional bool required = 3004;
+    optional EEnumYsonStorageType enum_yson_storage_type = 3005;
 }
 
 extend google.protobuf.EnumValueOptions