Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2023-10-25 09:58:02

0001 #ifndef PhysicsTools_NanoAOD_TableOutputFields_h
0002 #define PhysicsTools_NanoAOD_TableOutputFields_h
0003 
0004 #include "RNTupleFieldPtr.h"
0005 
0006 #include "FWCore/Framework/interface/EventForOutput.h"
0007 #include "DataFormats/NanoAOD/interface/FlatTable.h"
0008 #include "FWCore/Utilities/interface/EDGetToken.h"
0009 
0010 #include <algorithm>
0011 
0012 #include <ROOT/RNTuple.hxx>
0013 #include <ROOT/RNTupleModel.hxx>
0014 using ROOT::Experimental::RCollectionNTupleWriter;
0015 using ROOT::Experimental::RNTupleModel;
0016 using ROOT::Experimental::RNTupleWriter;
0017 
0018 template <typename T>
0019 class FlatTableField {
0020 public:
0021   FlatTableField() = default;
0022   FlatTableField(const nanoaod::FlatTable& table, std::size_t i, RNTupleModel& model) {
0023     m_flatTableName = table.columnName(i);
0024     // case 1: field has a name (table may or may not have a name)
0025     if (!table.columnName(i).empty()) {
0026       m_field = RNTupleFieldPtr<T>(table.columnName(i), table.columnDoc(i), model);
0027       return;
0028     }
0029     // case 2: field doesn't have a name: use the table name as the RNTuple field name
0030     if (table.name().empty()) {
0031       throw cms::Exception("LogicError", "Empty FlatTable name and field name");
0032     }
0033     m_field = RNTupleFieldPtr<T>(table.name(), table.doc(), model);
0034   }
0035   // For collection fields and singleton fields
0036   void fill(const nanoaod::FlatTable& table, std::size_t i) {
0037     int col_idx = table.columnIndex(m_flatTableName);
0038     if (col_idx == -1) {
0039       throw cms::Exception("LogicError", "Missing column in input for " + table.name() + "_" + m_flatTableName);
0040     }
0041     m_field.fill(table.columnData<T>(col_idx)[i]);
0042   }
0043   // For vector fields without a collection, we have to buffer the results
0044   // internally and then fill the RNTuple field
0045   void fillVectored(const nanoaod::FlatTable& table) {
0046     int col_idx = table.columnIndex(m_flatTableName);
0047     if (col_idx == -1) {
0048       throw cms::Exception("LogicError", "Missing column in input for " + table.name() + "_" + m_flatTableName);
0049     }
0050     std::vector<typename T::value_type> buf(table.size());
0051     for (std::size_t i = 0; i < table.size(); i++) {
0052       buf[i] = table.columnData<typename T::value_type>(col_idx)[i];
0053     }
0054     m_field.fill(buf);
0055   }
0056   const std::string& getFlatTableName() const { return m_flatTableName; }
0057 
0058 private:
0059   RNTupleFieldPtr<T> m_field;
0060   std::string m_flatTableName;
0061 };
0062 
0063 class TableOutputFields {
0064 public:
0065   TableOutputFields() = default;
0066   explicit TableOutputFields(const edm::EDGetToken& token) : m_token(token) {}
0067   void print() const;
0068   void createFields(const edm::EventForOutput& event, RNTupleModel& model);
0069   void fillEntry(const nanoaod::FlatTable& table, std::size_t i);
0070   const edm::EDGetToken& getToken() const;
0071 
0072 private:
0073   edm::EDGetToken m_token;
0074   std::vector<FlatTableField<float>> m_floatFields;
0075   std::vector<FlatTableField<int>> m_intFields;
0076   std::vector<FlatTableField<std::uint8_t>> m_uint8Fields;
0077   std::vector<FlatTableField<bool>> m_boolFields;
0078 };
0079 
0080 class TableOutputVectorFields {
0081 public:
0082   TableOutputVectorFields() = default;
0083   explicit TableOutputVectorFields(const edm::EDGetToken& token) : m_token(token) {}
0084   void createFields(const edm::EventForOutput& event, RNTupleModel& model);
0085   void fill(const edm::EventForOutput& event);
0086 
0087 private:
0088   edm::EDGetToken m_token;
0089   std::vector<FlatTableField<std::vector<float>>> m_vfloatFields;
0090   std::vector<FlatTableField<std::vector<int>>> m_vintFields;
0091   std::vector<FlatTableField<std::vector<std::uint8_t>>> m_vuint8Fields;
0092   std::vector<FlatTableField<std::vector<bool>>> m_vboolFields;
0093 };
0094 
0095 class TableCollection {
0096 public:
0097   TableCollection() = default;
0098   // Invariants:
0099   // * table has a non-empty base name
0100   // * table has at least one column
0101   void add(const edm::EDGetToken& table_token, const nanoaod::FlatTable& table);
0102   // Invariants:
0103   // * m_main not null
0104   // * m_collectionName not empty
0105   void createFields(const edm::EventForOutput& event, RNTupleModel& eventModel);
0106   void fill(const edm::EventForOutput& event);
0107   void print() const;
0108   bool hasMainTable();
0109   const std::string& getCollectionName() const;
0110 
0111 private:
0112   std::string m_collectionName;
0113   std::shared_ptr<RCollectionNTupleWriter> m_collection;
0114   TableOutputFields m_main;
0115   std::vector<TableOutputFields> m_extensions;
0116 };
0117 
0118 class TableCollectionSet {
0119 public:
0120   void add(const edm::EDGetToken& table_token, const nanoaod::FlatTable& table);
0121   void createFields(const edm::EventForOutput& event, RNTupleModel& eventModel);
0122   void fill(const edm::EventForOutput& event);
0123   void print() const;
0124 
0125 private:
0126   // Returns true if the FlatTable has an anonymous column. Throws a cms::Exception
0127   // if there is more than one anonymous column.
0128   static bool hasAnonymousColumn(const nanoaod::FlatTable& table);
0129   std::vector<TableCollection> m_collections;
0130   std::vector<TableOutputFields> m_singletonFields;
0131   std::vector<TableOutputVectorFields> m_vectorFields;
0132 };
0133 
0134 #endif