fix infering column filter for oneof columns

commit_hash:2a9c32e6eb0840b23ca468374067f0fe660eb4a3

1 files changed, 18 insertions, 1 deletions

diff --git a/yt/cpp/mapreduce/interface/protobuf_format.cpp b/yt/cpp/mapreduce/interface/protobuf_format.cpp
index 65090aa4d5b..cfd8ef9c31e 100644
--- a/yt/cpp/mapreduce/interface/protobuf_format.cpp
+++ b/yt/cpp/mapreduce/interface/protobuf_format.cpp
@@ -237,7 +237,6 @@ EWrapperOneofFlag::Enum OptionToOneofFlag(TOneofOption option)
     return std::visit(TVisitor(), option);
 }
 
-
 template <typename T, typename TOptionToFlag>
 void SetOption(TMaybe<T>& option, T newOption, TOptionToFlag optionToFlag)
 {
@@ -496,6 +495,18 @@ private:
 
 ////////////////////////////////////////////////////////////////////////////////
 
+TString GetOneofName(const ::google::protobuf::OneofDescriptor* descriptor)
+{
+    auto nameFromExtension = descriptor->options().GetExtension(variant_field_name);
+    if (nameFromExtension.empty()) {
+        return FromProto<TString>(descriptor->name());
+    } else {
+        return nameFromExtension;
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 } // namespace
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -1126,6 +1137,12 @@ TMaybe<TVector<TString>> InferColumnFilter(const ::google::protobuf::Descriptor&
             result.push_back(GetColumnName(*field));
         }
     }
+
+    for (int i = 0; i < descriptor.real_oneof_decl_count(); ++i) {
+        const auto* oneof = descriptor.oneof_decl(i);
+        result.push_back(GetOneofName(oneof));
+    }
+
     return result;
 }