YT-22408: Add certain attribute filter to attribute consumer

commit_hash:2a232aa75f4da6f6ff6e8fa9522c04feab7553ae

2 files changed, 18 insertions, 13 deletions

diff --git a/yt/yt/core/ytree/attribute_consumer.cpp b/yt/yt/core/ytree/attribute_consumer.cpp
index 7e28cb13f9c..5fbe929182d 100644
--- a/yt/yt/core/ytree/attribute_consumer.cpp
+++ b/yt/yt/core/ytree/attribute_consumer.cpp
@@ -9,23 +9,26 @@ using namespace NYson;
 
 ////////////////////////////////////////////////////////////////////////////////
 
-TAttributeConsumer::TAttributeConsumer(IAttributeDictionary* attributes)
-    : Attributes(attributes)
+TAttributeConsumer::TAttributeConsumer(IAttributeDictionary* attributes, std::optional<THashSet<TString>> keyWhitelist)
+    : Attributes_(attributes)
+    , KeyWhitelist_(std::move(keyWhitelist))
 { }
 
 IAttributeDictionary* TAttributeConsumer::GetAttributes() const
 {
-    return Attributes;
+    return Attributes_;
 }
 
 void TAttributeConsumer::OnMyKeyedItem(TStringBuf key)
 {
-    Writer.reset(new TBufferedBinaryYsonWriter(&Output));
-    Forward(Writer.get(), [this, key = TString(key)] {
-        Writer->Flush();
-        Writer.reset();
-        Attributes->SetYson(key, TYsonString(Output.Str()));
-        Output.clear();
+    Writer_.reset(new TBufferedBinaryYsonWriter(&Output_));
+    Forward(Writer_.get(), [this, key = TString(key)] {
+        Writer_->Flush();
+        Writer_.reset();
+        if (!KeyWhitelist_ || KeyWhitelist_->contains(key)) {
+            Attributes_->SetYson(key, TYsonString(Output_.Str()));
+        }
+        Output_.clear();
     });
 }
 
diff --git a/yt/yt/core/ytree/attribute_consumer.h b/yt/yt/core/ytree/attribute_consumer.h
index cbcaf031d29..4c1fd206df6 100644
--- a/yt/yt/core/ytree/attribute_consumer.h
+++ b/yt/yt/core/ytree/attribute_consumer.h
@@ -14,7 +14,7 @@ class TAttributeConsumer
     : public NYson::TForwardingYsonConsumer
 {
 public:
-    explicit TAttributeConsumer(IAttributeDictionary* attributes);
+    explicit TAttributeConsumer(IAttributeDictionary* attributes, std::optional<THashSet<TString>> keyWhitelist = {});
     IAttributeDictionary* GetAttributes() const;
 
 protected:
@@ -33,10 +33,12 @@ protected:
     void OnMyEndAttributes()override;
 
 private:
-    IAttributeDictionary* const Attributes;
+    IAttributeDictionary* const Attributes_;
 
-    TStringStream Output;
-    std::unique_ptr<NYson::TBufferedBinaryYsonWriter> Writer;
+    const std::optional<THashSet<TString>> KeyWhitelist_;
+
+    TStringStream Output_;
+    std::unique_ptr<NYson::TBufferedBinaryYsonWriter> Writer_;
 
     void ThrowMapExpected();