plugins/rntuple/TableOutputFields.h

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