Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2024-04-06 12:23:40

0001 #include "TableOutputFields.h"
0002 
0003 namespace {
0004   void printTable(const nanoaod::FlatTable& table) {
0005     std::cout << "FlatTable {\n";
0006     std::cout << "  name: " << (table.name().empty() ? "// anon" : table.name()) << "\n";
0007     std::cout << "  singleton: " << (table.singleton() ? "true" : "false") << "\n";
0008     std::cout << "  size: " << table.size() << "\n";
0009     std::cout << "  extension: " << (table.extension() ? "true" : "false") << "\n";
0010     std::cout << "  fields: {\n";
0011     for (std::size_t i = 0; i < table.nColumns(); i++) {
0012       std::cout << "    " << (table.columnName(i).empty() ? "// anon" : table.columnName(i)) << ": ";
0013       switch (table.columnType(i)) {
0014         case nanoaod::FlatTable::ColumnType::Float:
0015           std::cout << "f32,";
0016           break;
0017         case nanoaod::FlatTable::ColumnType::Int32:
0018           std::cout << "i32,";
0019           break;
0020         case nanoaod::FlatTable::ColumnType::UInt8:
0021           std::cout << "u8,";
0022           break;
0023         case nanoaod::FlatTable::ColumnType::Bool:
0024           std::cout << "bool,";
0025           break;
0026         default:
0027           throw cms::Exception("LogicError", "Unsupported type");
0028       }
0029       std::cout << "\n";
0030     }
0031     std::cout << "  }\n}\n";
0032   }
0033 }  // anonymous namespace
0034 
0035 void TableOutputFields::print() const {
0036   for (const auto& field : m_floatFields) {
0037     std::cout << "  " << field.getFlatTableName() << ": f32,\n";
0038   }
0039   for (const auto& field : m_intFields) {
0040     std::cout << "  " << field.getFlatTableName() << ": i32,\n";
0041   }
0042   for (const auto& field : m_uint8Fields) {
0043     std::cout << "  " << field.getFlatTableName() << ": u8,\n";
0044   }
0045   for (const auto& field : m_boolFields) {
0046     std::cout << "  " << field.getFlatTableName() << ": bool,\n";
0047   }
0048 }
0049 
0050 void TableOutputFields::createFields(const edm::EventForOutput& event, RNTupleModel& model) {
0051   edm::Handle<nanoaod::FlatTable> handle;
0052   event.getByToken(m_token, handle);
0053   const nanoaod::FlatTable& table = *handle;
0054   for (std::size_t i = 0; i < table.nColumns(); i++) {
0055     switch (table.columnType(i)) {
0056       case nanoaod::FlatTable::ColumnType::Float:
0057         m_floatFields.emplace_back(FlatTableField<float>(table, i, model));
0058         break;
0059       case nanoaod::FlatTable::ColumnType::Int32:
0060         m_intFields.emplace_back(FlatTableField<int>(table, i, model));
0061         break;
0062       case nanoaod::FlatTable::ColumnType::UInt8:
0063         m_uint8Fields.emplace_back(FlatTableField<std::uint8_t>(table, i, model));
0064         break;
0065       case nanoaod::FlatTable::ColumnType::Bool:
0066         m_boolFields.emplace_back(FlatTableField<bool>(table, i, model));
0067         break;
0068       default:
0069         throw cms::Exception("LogicError", "Unsupported type");
0070     }
0071   }
0072 }
0073 
0074 void TableOutputFields::fillEntry(const nanoaod::FlatTable& table, std::size_t i) {
0075   for (auto& field : m_floatFields) {
0076     field.fill(table, i);
0077   }
0078   for (auto& field : m_intFields) {
0079     field.fill(table, i);
0080   }
0081   for (auto& field : m_uint8Fields) {
0082     field.fill(table, i);
0083   }
0084   for (auto& field : m_boolFields) {
0085     field.fill(table, i);
0086   }
0087 }
0088 
0089 const edm::EDGetToken& TableOutputFields::getToken() const { return m_token; }
0090 
0091 ///////////////////////////////////////////////////////////////////////////////
0092 
0093 void TableOutputVectorFields::createFields(const edm::EventForOutput& event, RNTupleModel& model) {
0094   edm::Handle<nanoaod::FlatTable> handle;
0095   event.getByToken(m_token, handle);
0096   const nanoaod::FlatTable& table = *handle;
0097   for (std::size_t i = 0; i < table.nColumns(); i++) {
0098     switch (table.columnType(i)) {
0099       case nanoaod::FlatTable::ColumnType::Float:
0100         m_vfloatFields.emplace_back(FlatTableField<std::vector<float>>(table, i, model));
0101         break;
0102       case nanoaod::FlatTable::ColumnType::Int32:
0103         m_vintFields.emplace_back(FlatTableField<std::vector<int>>(table, i, model));
0104         break;
0105       case nanoaod::FlatTable::ColumnType::UInt8:
0106         m_vuint8Fields.emplace_back(FlatTableField<std::vector<std::uint8_t>>(table, i, model));
0107         break;
0108       case nanoaod::FlatTable::ColumnType::Bool:
0109         m_vboolFields.emplace_back(FlatTableField<std::vector<bool>>(table, i, model));
0110         break;
0111       default:
0112         throw cms::Exception("LogicError", "Unsupported type");
0113     }
0114   }
0115 }
0116 void TableOutputVectorFields::fill(const edm::EventForOutput& event) {
0117   edm::Handle<nanoaod::FlatTable> handle;
0118   event.getByToken(m_token, handle);
0119   const auto& table = *handle;
0120   for (auto& field : m_vfloatFields) {
0121     field.fillVectored(table);
0122   }
0123   for (auto& field : m_vintFields) {
0124     field.fillVectored(table);
0125   }
0126   for (auto& field : m_vuint8Fields) {
0127     field.fillVectored(table);
0128   }
0129   for (auto& field : m_vboolFields) {
0130     field.fillVectored(table);
0131   }
0132 }
0133 
0134 ///////////////////////////////////////////////////////////////////////////////
0135 
0136 void TableCollection::add(const edm::EDGetToken& table_token, const nanoaod::FlatTable& table) {
0137   if (m_collectionName.empty()) {
0138     m_collectionName = table.name();
0139   }
0140   if (table.extension()) {
0141     m_extensions.emplace_back(TableOutputFields(table_token));
0142     return;
0143   }
0144   if (hasMainTable()) {
0145     throw cms::Exception("LogicError", "Trying to save multiple main tables for " + m_collectionName + "\n");
0146   }
0147   m_main = TableOutputFields(table_token);
0148 }
0149 
0150 void TableCollection::createFields(const edm::EventForOutput& event, RNTupleModel& eventModel) {
0151   auto collectionModel = RNTupleModel::Create();
0152   m_main.createFields(event, *collectionModel);
0153   for (auto& extension : m_extensions) {
0154     extension.createFields(event, *collectionModel);
0155   }
0156   edm::Handle<nanoaod::FlatTable> handle;
0157   event.getByToken(m_main.getToken(), handle);
0158   const nanoaod::FlatTable& table = *handle;
0159   collectionModel->SetDescription(table.doc());
0160   m_collection = eventModel.MakeCollection(m_collectionName, std::move(collectionModel));
0161 }
0162 
0163 void TableCollection::fill(const edm::EventForOutput& event) {
0164   edm::Handle<nanoaod::FlatTable> handle;
0165   event.getByToken(m_main.getToken(), handle);
0166   const auto& main_table = *handle;
0167   auto table_size = main_table.size();
0168   for (std::size_t i = 0; i < table_size; i++) {
0169     m_main.fillEntry(main_table, i);
0170     for (auto& ext : m_extensions) {
0171       edm::Handle<nanoaod::FlatTable> handle;
0172       event.getByToken(ext.getToken(), handle);
0173       const auto& ext_table = *handle;
0174       if (ext_table.size() != table_size) {
0175         throw cms::Exception("LogicError",
0176                              "Mismatch in number of entries between extension and main table for " + m_collectionName);
0177       }
0178       ext.fillEntry(ext_table, i);
0179     }
0180     m_collection->Fill();
0181   }
0182 }
0183 
0184 void TableCollection::print() const {
0185   std::cout << "Collection: " << m_collectionName << " {\n";
0186   m_main.print();
0187   for (const auto& ext : m_extensions) {
0188     ext.print();
0189   }
0190   std::cout << "}\n";
0191 }
0192 
0193 bool TableCollection::hasMainTable() { return !m_main.getToken().isUninitialized(); }
0194 
0195 const std::string& TableCollection::getCollectionName() const { return m_collectionName; }
0196 
0197 ///////////////////////////////////////////////////////////////////////////////
0198 
0199 void TableCollectionSet::add(const edm::EDGetToken& table_token, const nanoaod::FlatTable& table) {
0200   // skip empty tables -- requirement of RNTuple to define schema before filling
0201   if (table.nColumns() == 0) {
0202     std::cout << "Warning: skipping empty table: \n";
0203     printTable(table);
0204     return;
0205   }
0206   // Can handle either anonymous table or anonymous column but not both
0207   // - anonymous table: use column names directly as top-level fields
0208   // - anonymous column: use the table name as the field name
0209   if (table.name().empty() && hasAnonymousColumn(table)) {
0210     throw cms::Exception("LogicError", "Anonymous FlatTable and anonymous field");
0211   }
0212   // case 1: create a top-level RNTuple field for each table column
0213   if (table.name().empty() || hasAnonymousColumn(table)) {
0214     if (table.singleton()) {
0215       m_singletonFields.emplace_back(TableOutputFields(table_token));
0216     } else {
0217       m_vectorFields.emplace_back(TableOutputVectorFields(table_token));
0218     }
0219     return;
0220   }
0221   // case 2: Named singleton and vector tables are both written as RNTuple collections.
0222   auto collection = std::find_if(m_collections.begin(), m_collections.end(), [&](const TableCollection& c) {
0223     return c.getCollectionName() == table.name();
0224   });
0225   if (collection == m_collections.end()) {
0226     m_collections.emplace_back(TableCollection());
0227     m_collections.back().add(table_token, table);
0228     return;
0229   }
0230   collection->add(table_token, table);
0231 }
0232 
0233 void TableCollectionSet::print() const {
0234   for (const auto& collection : m_collections) {
0235     collection.print();
0236     std::cout << "\n";
0237   }
0238 }
0239 
0240 void TableCollectionSet::createFields(const edm::EventForOutput& event, RNTupleModel& eventModel) {
0241   for (auto& collection : m_collections) {
0242     if (!collection.hasMainTable()) {
0243       throw cms::Exception("LogicError",
0244                            "Trying to save an extension table for " + collection.getCollectionName() +
0245                                " without the corresponding main table\n");
0246     }
0247     collection.createFields(event, eventModel);
0248   }
0249   for (auto& table : m_singletonFields) {
0250     table.createFields(event, eventModel);
0251   }
0252   for (auto& table : m_vectorFields) {
0253     table.createFields(event, eventModel);
0254   }
0255 }
0256 
0257 void TableCollectionSet::fill(const edm::EventForOutput& event) {
0258   for (auto& collection : m_collections) {
0259     collection.fill(event);
0260   }
0261   for (auto& fields : m_singletonFields) {
0262     edm::Handle<nanoaod::FlatTable> handle;
0263     event.getByToken(fields.getToken(), handle);
0264     const auto& table = *handle;
0265     fields.fillEntry(table, 0);
0266   }
0267   for (auto& fields : m_vectorFields) {
0268     fields.fill(event);
0269   }
0270 }
0271 
0272 bool TableCollectionSet::hasAnonymousColumn(const nanoaod::FlatTable& table) {
0273   int num_anon = 0;
0274   for (std::size_t i = 0; i < table.nColumns(); i++) {
0275     if (table.columnName(i).empty()) {
0276       num_anon++;
0277     }
0278   }
0279   if (num_anon > 1) {
0280     throw cms::Exception("LogicError",
0281                          "FlatTable `" + table.name() + "` has " + std::to_string(num_anon) + "anonymous fields");
0282   }
0283   return num_anon;
0284 }