Merge pull request #33565 from pieterdavid/nanoadduint32doubleandrunf…

…lattables

Allow storing uint32_t and double in NanoAOD, and flat tables in the Runs tree, for ALCANANO

DataFormats/NanoAOD/interface/FlatTable.h

@@ -41,7 +41,9 @@ namespace nanoaod {
       Float,
       Int,
       UInt8,
-      Bool
+      Bool,
+      UInt32,
+      Double
     };  // We could have other Float types with reduced mantissa, and similar
 
     FlatTable() : size_(0) {}
@@ -143,6 +145,10 @@ namespace nanoaod {
         return ColumnType::UInt8;
       else if constexpr (std::is_same<T, bool>())
         return ColumnType::Bool;
+      else if constexpr (std::is_same<T, uint32_t>())
+        return ColumnType::UInt32;
+      else if constexpr (std::is_same<T, double>())
+        return ColumnType::Double;
       else
         static_assert(dependent_false<T>::value, "unsupported type");
     }
@@ -187,6 +193,10 @@ namespace nanoaod {
         return table.uint8s_;
       else if constexpr (std::is_same<T, bool>())
         return table.uint8s_;
+      else if constexpr (std::is_same<T, uint32_t>())
+        return table.uint32s_;
+      else if constexpr (std::is_same<T, double>())
+        return table.doubles_;
       else
         static_assert(dependent_false<T>::value, "unsupported type");
     }
@@ -198,6 +208,8 @@ namespace nanoaod {
     std::vector<float> floats_;
     std::vector<int> ints_;
     std::vector<uint8_t> uint8s_;
+    std::vector<uint32_t> uint32s_;
+    std::vector<double> doubles_;
   };
 
 }  // namespace nanoaod

DataFormats/NanoAOD/src/FlatTable.cc

@@ -25,6 +25,12 @@ void nanoaod::FlatTable::addExtension(const nanoaod::FlatTable& other) {
       case ColumnType::UInt8:
         addColumn<uint8_t>(other.columnName(i), other.columnData<uint8_t>(i), other.columnDoc(i));
         break;
+      case ColumnType::UInt32:
+        addColumn<uint32_t>(other.columnName(i), other.columnData<uint32_t>(i), other.columnDoc(i));
+        break;
+      case ColumnType::Double:
+        addColumn<double>(other.columnName(i), other.columnData<double>(i), other.columnDoc(i));
+        break;
       default:
         throw cms::Exception("LogicError", "Unsupported type");
     }
@@ -43,6 +49,10 @@ double nanoaod::FlatTable::getAnyValue(unsigned int row, unsigned int column) co
       return *(beginData<bool>(column) + row);
     case ColumnType::UInt8:
       return *(beginData<uint8_t>(column) + row);
+    case ColumnType::UInt32:
+      return *(beginData<uint32_t>(column) + row);
+    case ColumnType::Double:
+      return *(beginData<double>(column) + row);
   }
   throw cms::Exception("LogicError", "Unsupported type");
 }

DataFormats/NanoAOD/src/classes_def.xml

@@ -3,7 +3,8 @@
         <version ClassVersion="3" checksum="3066258528"/>
     </class>
     <class name="std::vector<nanoaod::FlatTable::Column>" />
-    <class name="nanoaod::FlatTable" ClassVersion="3">
+    <class name="nanoaod::FlatTable" ClassVersion="4">
+        <version ClassVersion="4" checksum="656493391"/>
         <version ClassVersion="3" checksum="2443023556"/>
     </class>
     <class name="nanoaod::FlatTable::RowView" transient="true" />

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");

PhysicsTools/NanoAOD/plugins/TableOutputBranches.cc

@@ -25,6 +25,12 @@ void TableOutputBranches::defineBranchesFromFirstEvent(const nanoaod::FlatTable
       case nanoaod::FlatTable::ColumnType::Bool:
         m_uint8Branches.emplace_back(var, tab.columnDoc(i), "O");
         break;
+      case nanoaod::FlatTable::ColumnType::UInt32:
+        m_uint32Branches.emplace_back(var, tab.columnDoc(i), "i");
+        break;
+      case nanoaod::FlatTable::ColumnType::Double:
+        m_doubleBranches.emplace_back(var, tab.columnDoc(i), "D");
+        break;
       default:
         throw cms::Exception("LogicError", "Unsupported type");
     }
@@ -49,7 +55,8 @@ void TableOutputBranches::branch(TTree &tree) {
     }
   }
   std::string varsize = m_singleton ? "" : "[n" + m_baseName + "]";
-  for (std::vector<NamedBranchPtr> *branches : {&m_floatBranches, &m_intBranches, &m_uint8Branches}) {
+  for (std::vector<NamedBranchPtr> *branches :
+       {&m_floatBranches, &m_intBranches, &m_uint8Branches, &m_uint32Branches, &m_doubleBranches}) {
     for (auto &pair : *branches) {
       std::string branchName = makeBranchName(m_baseName, pair.name);
       pair.branch =
@@ -59,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();
@@ -91,4 +98,8 @@ void TableOutputBranches::fill(const edm::EventForOutput &iEvent, TTree &tree, b
     fillColumn<int>(pair, tab);
   for (auto &pair : m_uint8Branches)
     fillColumn<uint8_t>(pair, tab);
+  for (auto &pair : m_uint32Branches)
+    fillColumn<uint32_t>(pair, tab);
+  for (auto &pair : m_doubleBranches)
+    fillColumn<double>(pair, tab);
 }

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;
@@ -44,6 +44,8 @@ class TableOutputBranches {
   std::vector<NamedBranchPtr> m_floatBranches;
   std::vector<NamedBranchPtr> m_intBranches;
   std::vector<NamedBranchPtr> m_uint8Branches;
+  std::vector<NamedBranchPtr> m_uint32Branches;
+  std::vector<NamedBranchPtr> m_doubleBranches;
   bool m_branchesBooked;
 
   template <typename T>