YT-19714: Parallel load values for few instances of EntityMap (chunks and chunk lists)

3 files changed, 23 insertions, 0 deletions

diff --git a/yt/yt/core/misc/serialize-inl.h b/yt/yt/core/misc/serialize-inl.h
index a3c985f4cc..db38f68fbd 100644
--- a/yt/yt/core/misc/serialize-inl.h
+++ b/yt/yt/core/misc/serialize-inl.h
@@ -376,6 +376,7 @@ TEntitySerializationKey TEntityStreamSaveContext::RegisterRefCountedEntity(const
 template <class T>
 inline TEntitySerializationKey TEntityStreamLoadContext::RegisterRawEntity(T* entity)
 {
+    YT_VERIFY(!ParentContext_);
     auto key = TEntitySerializationKey(std::ssize(RawPtrs_));
     RawPtrs_.push_back(entity);
     return key;
@@ -384,6 +385,7 @@ inline TEntitySerializationKey TEntityStreamLoadContext::RegisterRawEntity(T* en
 template <class T>
 TEntitySerializationKey TEntityStreamLoadContext::RegisterRefCountedEntity(const TIntrusivePtr<T>& entity)
 {
+    YT_VERIFY(!ParentContext_);
     auto* ptr = entity.Get();
     RefCountedPtrs_.push_back(entity);
     return RegisterRawEntity(ptr);
@@ -392,6 +394,10 @@ TEntitySerializationKey TEntityStreamLoadContext::RegisterRefCountedEntity(const
 template <class T>
 T* TEntityStreamLoadContext::GetRawEntity(TEntitySerializationKey key) const
 {
+    if (ParentContext_) {
+        return ParentContext_->GetRawEntity<T>(key);
+    }
+
     YT_ASSERT(key.Index >= 0);
     YT_ASSERT(key.Index < std::ssize(RawPtrs_));
     return static_cast<T*>(RawPtrs_[key.Index]);
diff --git a/yt/yt/core/misc/serialize.cpp b/yt/yt/core/misc/serialize.cpp
index 15b975c194..b5d8720a35 100644
--- a/yt/yt/core/misc/serialize.cpp
+++ b/yt/yt/core/misc/serialize.cpp
@@ -177,5 +177,16 @@ TEntityStreamSaveContext::TEntityStreamSaveContext(
 
 ////////////////////////////////////////////////////////////////////////////////
 
+TEntityStreamLoadContext::TEntityStreamLoadContext(
+    IZeroCopyInput* input,
+    const TEntityStreamLoadContext* parentContext)
+    : TStreamLoadContext(input)
+    , ParentContext_(parentContext)
+{
+    SetVersion(ParentContext_->GetVersion());
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 } // namespace NYT
 
diff --git a/yt/yt/core/misc/serialize.h b/yt/yt/core/misc/serialize.h
index 09a9bb6e9a..165a4442d2 100644
--- a/yt/yt/core/misc/serialize.h
+++ b/yt/yt/core/misc/serialize.h
@@ -263,6 +263,10 @@ class TEntityStreamLoadContext
 public:
     using TStreamLoadContext::TStreamLoadContext;
 
+    TEntityStreamLoadContext(
+        IZeroCopyInput* input,
+        const TEntityStreamLoadContext* parentContext);
+
     template <class T>
     TEntitySerializationKey RegisterRawEntity(T* entity);
     template <class T>
@@ -274,6 +278,8 @@ public:
     TIntrusivePtr<T> GetRefCountedEntity(TEntitySerializationKey key) const;
 
 private:
+    const TEntityStreamLoadContext* const ParentContext_ = nullptr;
+
     std::vector<void*> RawPtrs_;
     std::vector<TIntrusivePtr<TRefCounted>> RefCountedPtrs_;
 };