Sort all YSONStruct fields before saving, including unrecognized keys

commit_hash:4e04b10da96e7f9e6f201f51eb8f3ed249985d10

2 files changed, 47 insertions, 7 deletions

diff --git a/yt/yt/core/ytree/yson_struct.cpp b/yt/yt/core/ytree/yson_struct.cpp
index 41e6c3de2f3..e019fc22fac 100644
--- a/yt/yt/core/ytree/yson_struct.cpp
+++ b/yt/yt/core/ytree/yson_struct.cpp
@@ -91,7 +91,7 @@ void TYsonStructBase::Save(IYsonConsumer* consumer) const
     consumer->OnEndMap();
 }
 
-void TYsonStructBase::SaveAsMapFragment(NYson::IYsonConsumer* consumer) const
+void TYsonStructBase::SaveRecognizedAsMapFragment(NYson::IYsonConsumer* consumer) const
 {
     for (const auto& [name, parameter] : Meta_->GetParameterSortedList()) {
         if (!parameter->CanOmitValue(this)) {
@@ -99,14 +99,51 @@ void TYsonStructBase::SaveAsMapFragment(NYson::IYsonConsumer* consumer) const
             parameter->Save(this, consumer);
         }
     }
+}
+
+void TYsonStructBase::SaveAsMapFragment(NYson::IYsonConsumer* consumer) const
+{
+    if (!LocalUnrecognized_) {
+        // Fast path.
+        return SaveRecognizedAsMapFragment(consumer);
+    }
+
+    const auto& recognizedList = Meta_->GetParameterSortedList();
+    auto recognizedIt = recognizedList.begin();
+
+    auto unrecognizedList = LocalUnrecognized_->GetChildren();
+    SortBy(unrecognizedList, [] (const auto& item) { return item.first; });
+    auto unrecognizedIt = unrecognizedList.begin();
 
-    if (LocalUnrecognized_) {
-        auto unrecognizedList = LocalUnrecognized_->GetChildren();
-        SortBy(unrecognizedList, [] (const auto& item) { return item.first; });
-        for (const auto& [key, child] : unrecognizedList) {
-            consumer->OnKeyedItem(key);
-            Serialize(child, consumer);
+    auto saveRecognized = [&recognizedIt, this] (auto* consumer) {
+        const auto& parameter = recognizedIt->second;
+        if (!parameter->CanOmitValue(this)) {
+            consumer->OnKeyedItem(recognizedIt->first);
+            parameter->Save(this, consumer);
         }
+        ++recognizedIt;
+    };
+
+    auto saveUnrecognized = [&unrecognizedIt] (auto* consumer) {
+        consumer->OnKeyedItem(unrecognizedIt->first);
+        Serialize(unrecognizedIt->second, consumer);
+        ++unrecognizedIt;
+    };
+
+    while (recognizedIt != recognizedList.end() && unrecognizedIt != unrecognizedList.end()) {
+        if (recognizedIt->first < unrecognizedIt->first) {
+            saveRecognized(consumer);
+        } else {
+            saveUnrecognized(consumer);
+        }
+    }
+
+    while (recognizedIt != recognizedList.end()) {
+        saveRecognized(consumer);
+    }
+
+    while (unrecognizedIt != unrecognizedList.end()) {
+        saveUnrecognized(consumer);
     }
 }
 
diff --git a/yt/yt/core/ytree/yson_struct.h b/yt/yt/core/ytree/yson_struct.h
index 62837542e69..cd93566f865 100644
--- a/yt/yt/core/ytree/yson_struct.h
+++ b/yt/yt/core/ytree/yson_struct.h
@@ -90,6 +90,9 @@ public:
     // of the |Load| call.
     void SaveAsMapFragment(NYson::IYsonConsumer* consumer) const;
 
+    // Same as the above, but does not save local unrecognized parameters.
+    void SaveRecognizedAsMapFragment(NYson::IYsonConsumer* consumer) const;
+
     void Save(IOutputStream* output) const;
 
     IMapNodePtr GetLocalUnrecognized() const;