NanoAOD/plugins/TableOutputBranches.h

0001 #ifndef PhysicsTools_NanoAOD_TableOutputBranches_h
0002 #define PhysicsTools_NanoAOD_TableOutputBranches_h
0003
0004 #include <string>
0005 #include <vector>
0006 #include <TTree.h>
0007 #include "FWCore/Framework/interface/OccurrenceForOutput.h"
0008 #include "DataFormats/NanoAOD/interface/FlatTable.h"
0009 #include "DataFormats/Provenance/interface/ProductDescription.h"
0010 #include "FWCore/Utilities/interface/EDGetToken.h"
0011
0012 class TableOutputBranches {
0013 public:
0014   TableOutputBranches(const edm::ProductDescription *desc, const edm::EDGetToken &token)
0015       : m_token(token), m_extension(DontKnowYetIfMainOrExtension), m_branchesBooked(false) {
0016     if (desc->className() != "nanoaod::FlatTable")
0017       throw cms::Exception("Configuration", "NanoAODOutputModule can only write out nanoaod::FlatTable objects");
0018   }
0019
0020   void defineBranchesFromFirstEvent(const nanoaod::FlatTable &tab);
0021   void branch(TTree &tree);
0022
0023   /// Fill the current table, if extensions == table.extension().
0024   /// This parameter is used so that the fill is called first for non-extensions and then for extensions
0025   void fill(const edm::OccurrenceForOutput &iWhatever, TTree &tree, bool extensions);
0026
0027 private:
0028   edm::EDGetToken m_token;
0029   std::string m_baseName;
0030   bool m_singleton = false;
0031   enum { IsMain = 0, IsExtension = 1, DontKnowYetIfMainOrExtension = 2 } m_extension;
0032   std::string m_doc;
0033   typedef Int_t CounterType;
0034   CounterType m_counter;
0035   struct NamedBranchPtr {
0036     std::string name, title, rootTypeCode;
0037     TBranch *branch;
0038     NamedBranchPtr(const std::string &aname,
0039                    const std::string &atitle,
0040                    const std::string &rootType,
0041                    TBranch *branchptr = nullptr)
0042         : name(aname), title(atitle), rootTypeCode(rootType), branch(branchptr) {}
0043   };
0044   TBranch *m_counterBranch = nullptr;
0045   std::vector<NamedBranchPtr> m_uint8Branches;
0046   std::vector<NamedBranchPtr> m_int16Branches;
0047   std::vector<NamedBranchPtr> m_uint16Branches;
0048   std::vector<NamedBranchPtr> m_int32Branches;
0049   std::vector<NamedBranchPtr> m_uint32Branches;
0050   std::vector<NamedBranchPtr> m_floatBranches;
0051   std::vector<NamedBranchPtr> m_doubleBranches;
0052   bool m_branchesBooked;
0053
0054   template <typename T>
0055   void fillColumn(NamedBranchPtr &pair, const nanoaod::FlatTable &tab) {
0056     int idx = tab.columnIndex(pair.name);
0057     if (idx == -1)
0058       throw cms::Exception("LogicError", "Missing column in input for " + m_baseName + "_" + pair.name);
0059     pair.branch->SetAddress(
0060         tab.size() == 0 ? static_cast<T *>(nullptr)
0061                         : const_cast<T *>(&tab.columnData<T>(idx).front()));  // SetAddress should take a const * !
0062   }
0063 };
0064
0065 #endif