Allow writing nanoaod::FlatTable to the Runs tree

cms-sw · Apr 28, 2021 · d7319f8 · d7319f8
1 parent 310ca91
commit d7319f8
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 4 deletions.
diff --git a/PhysicsTools/NanoAOD/plugins/NanoAODOutputModule.cc b/PhysicsTools/NanoAOD/plugins/NanoAODOutputModule.cc
@@ -126,6 +126,7 @@ class NanoAODOutputModule : public edm::one::OutputModule<> {
 
   std::vector<SummaryTableOutputBranches> m_runTables;
   std::vector<SummaryTableOutputBranches> m_lumiTables;
+  std::vector<TableOutputBranches> m_runFlatTables;
 
   std::vector<std::pair<std::string, edm::EDGetToken>> m_nanoMetadata;
 };
@@ -243,6 +244,11 @@ void NanoAODOutputModule::writeRun(edm::RunForOutput const& iRun) {
   for (auto& t : m_runTables)
     t.fill(iRun, *m_runTree);
 
+  for (unsigned int extensions = 0; extensions <= 1; ++extensions) {
+    for (auto& t : m_runFlatTables)
+      t.fill(iRun, *m_runTree, extensions);
+  }
+
   edm::Handle<nanoaod::UniqueString> hstring;
   for (const auto& p : m_nanoMetadata) {
     iRun.getByToken(p.second, hstring);
@@ -293,6 +299,7 @@ void NanoAODOutputModule::openFile(edm::FileBlock const&) {
   m_evstrings.clear();
   m_runTables.clear();
   m_lumiTables.clear();
+  m_runFlatTables.clear();
   const auto& keeps = keptProducts();
   for (const auto& keep : keeps[edm::InEvent]) {
     if (keep.first->className() == "nanoaod::FlatTable")
@@ -322,6 +329,8 @@ void NanoAODOutputModule::openFile(edm::FileBlock const&) {
       m_runTables.push_back(SummaryTableOutputBranches(keep.first, keep.second));
     else if (keep.first->className() == "nanoaod::UniqueString" && keep.first->moduleLabel() == "nanoMetadata")
       m_nanoMetadata.emplace_back(keep.first->productInstanceName(), keep.second);
+    else if (keep.first->className() == "nanoaod::FlatTable")
+      m_runFlatTables.emplace_back(keep.first, keep.second);
     else
       throw cms::Exception("Configuration",
                            "NanoAODOutputModule cannot handle class " + keep.first->className() + " in Run branch");

diff --git a/PhysicsTools/NanoAOD/plugins/TableOutputBranches.cc b/PhysicsTools/NanoAOD/plugins/TableOutputBranches.cc
@@ -66,14 +66,14 @@ void TableOutputBranches::branch(TTree &tree) {
   }
 }
 
-void TableOutputBranches::fill(const edm::EventForOutput &iEvent, TTree &tree, bool extensions) {
+void TableOutputBranches::fill(const edm::OccurrenceForOutput &iWhatever, TTree &tree, bool extensions) {
   if (m_extension != DontKnowYetIfMainOrExtension) {
     if (extensions != m_extension)
       return;  // do nothing, wait to be called with the proper flag
   }
 
   edm::Handle<nanoaod::FlatTable> handle;
-  iEvent.getByToken(m_token, handle);
+  iWhatever.getByToken(m_token, handle);
   const nanoaod::FlatTable &tab = *handle;
   m_counter = tab.size();
   m_singleton = tab.singleton();

diff --git a/PhysicsTools/NanoAOD/plugins/TableOutputBranches.h b/PhysicsTools/NanoAOD/plugins/TableOutputBranches.h
@@ -4,7 +4,7 @@
 #include <string>
 #include <vector>
 #include <TTree.h>
-#include "FWCore/Framework/interface/EventForOutput.h"
+#include "FWCore/Framework/interface/OccurrenceForOutput.h"
 #include "DataFormats/NanoAOD/interface/FlatTable.h"
 #include "DataFormats/Provenance/interface/BranchDescription.h"
 #include "FWCore/Utilities/interface/EDGetToken.h"
@@ -22,7 +22,7 @@ class TableOutputBranches {
 
   /// Fill the current table, if extensions == table.extension().
   /// This parameter is used so that the fill is called first for non-extensions and then for extensions
-  void fill(const edm::EventForOutput &iEvent, TTree &tree, bool extensions);
+  void fill(const edm::OccurrenceForOutput &iWhatever, TTree &tree, bool extensions);
 
 private:
   edm::EDGetToken m_token;