From cedefaa1cba32be23934de479ccdc3e0ba463a4e Mon Sep 17 00:00:00 2001 From: tmadlener Date: Tue, 25 Apr 2023 11:13:29 +0200 Subject: [PATCH] Make it possible to read index based files as well - Keep branch names around for resetting them (since invalidated branches can no longer provide that information) --- include/podio/CollectionBranches.h | 3 + src/ROOTFrameReader.cc | 89 ++++++++++++++++++++++-------- src/rootUtils.h | 17 ++++++ 3 files changed, 87 insertions(+), 22 deletions(-) diff --git a/include/podio/CollectionBranches.h b/include/podio/CollectionBranches.h index ddbce7ec2..9b20016aa 100644 --- a/include/podio/CollectionBranches.h +++ b/include/podio/CollectionBranches.h @@ -3,6 +3,7 @@ #include "TBranch.h" +#include #include namespace podio::root_utils { @@ -15,6 +16,8 @@ struct CollectionBranches { TBranch* data{nullptr}; std::vector refs{}; std::vector vecs{}; + std::vector refNames{}; ///< The names of the relation branches + std::vector vecNames{}; ///< The names of the vector member branches }; } // namespace podio::root_utils diff --git a/src/ROOTFrameReader.cc b/src/ROOTFrameReader.cc index 08819f2f9..bb7d2b3a3 100644 --- a/src/ROOTFrameReader.cc +++ b/src/ROOTFrameReader.cc @@ -23,6 +23,10 @@ std::tuple, std::vector& collInfo); +std::tuple, std::vector>> +createCollectionBranchesIndexBased(TChain* chain, const podio::CollectionIDTable& idTable, + const std::vector& collInfo); + GenericParameters ROOTFrameReader::readEntryParameters(ROOTFrameReader::CategoryInfo& catInfo, bool reloadBranches, unsigned int localEntry) { // Parameter branch is always the last one @@ -95,23 +99,7 @@ podio::CollectionReadBuffers ROOTFrameReader::getCollectionBuffers(ROOTFrameRead auto collBuffers = maybeBuffers.value_or(podio::CollectionReadBuffers{}); if (reloadBranches) { - branches.data = root_utils::getBranch(catInfo.chain.get(), name.c_str()); - - // reference collections - if (auto* refCollections = collBuffers.references) { - for (size_t i = 0; i < refCollections->size(); ++i) { - const auto brName = root_utils::refBranch(name, i); - branches.refs[i] = root_utils::getBranch(catInfo.chain.get(), branches.refs[i]->GetName()); - } - } - - // vector members - if (auto* vecMembers = collBuffers.vectorMembers) { - for (size_t i = 0; i < vecMembers->size(); ++i) { - const auto brName = root_utils::vecBranch(name, i); - branches.vecs[i] = root_utils::getBranch(catInfo.chain.get(), branches.vecs[i]->GetName()); - } - } + root_utils::resetBranches(catInfo.chain.get(), branches, name); } // set the addresses and read the data @@ -152,8 +140,15 @@ void ROOTFrameReader::initCategory(CategoryInfo& catInfo, const std::string& cat collInfoBranch->SetAddress(&collInfo); collInfoBranch->GetEntry(0); - std::tie(catInfo.branches, catInfo.storedClasses) = - createCollectionBranches(catInfo.chain.get(), *catInfo.table, *collInfo); + // For backwards compatibility make it possible to read the index based files + // from older versions + if (m_fileVersion <= podio::version::Version{0, 16, 5}) { + std::tie(catInfo.branches, catInfo.storedClasses) = + createCollectionBranchesIndexBased(catInfo.chain.get(), *catInfo.table, *collInfo); + } else { + std::tie(catInfo.branches, catInfo.storedClasses) = + createCollectionBranches(catInfo.chain.get(), *catInfo.table, *collInfo); + } delete collInfo; @@ -242,6 +237,53 @@ std::vector ROOTFrameReader::getAvailableCategories() const { return cats; } +std::tuple, std::vector>> +createCollectionBranchesIndexBased(TChain* chain, const podio::CollectionIDTable& idTable, + const std::vector& collInfo) { + + size_t collectionIndex{0}; + std::vector collBranches; + collBranches.reserve(collInfo.size() + 1); + std::vector> storedClasses; + storedClasses.reserve(collInfo.size()); + + for (const auto& [collID, collType, isSubsetColl, collSchemaVersion] : collInfo) { + // We only write collections that are in the collectionIDTable, so no need + // to check here + const auto name = idTable.name(collID); + + const auto collectionClass = TClass::GetClass(collType.c_str()); + // Need the collection here to setup all the branches. Have to manage the + // temporary collection ourselves + auto collection = + std::unique_ptr(static_cast(collectionClass->New())); + root_utils::CollectionBranches branches{}; + if (!isSubsetColl) { + // This branch is guaranteed to exist since only collections that are + // also written to file are in the info metadata that we work with here + branches.data = root_utils::getBranch(chain, name.c_str()); + } + + const auto buffers = collection->getBuffers(); + for (size_t i = 0; i < buffers.references->size(); ++i) { + auto brName = root_utils::refBranch(name, i); + branches.refs.push_back(root_utils::getBranch(chain, brName.c_str())); + branches.refNames.emplace_back(std::move(brName)); + } + + for (size_t i = 0; i < buffers.vectorMembers->size(); ++i) { + auto brName = root_utils::vecBranch(name, i); + branches.vecs.push_back(root_utils::getBranch(chain, brName.c_str())); + branches.vecNames.emplace_back(std::move(brName)); + } + + storedClasses.emplace_back(name, std::make_tuple(collType, isSubsetColl, collSchemaVersion, collectionIndex++)); + collBranches.emplace_back(std::move(branches)); + } + + return {collBranches, storedClasses}; +} + std::tuple, std::vector>> createCollectionBranches(TChain* chain, const podio::CollectionIDTable& idTable, const std::vector& collInfo) { @@ -260,8 +302,9 @@ createCollectionBranches(TChain* chain, const podio::CollectionIDTable& idTable, root_utils::CollectionBranches branches{}; if (isSubsetColl) { // Only one branch will exist and we can trivially get its name - const auto brName = root_utils::subsetBranch(name); + auto brName = root_utils::subsetBranch(name); branches.refs.push_back(root_utils::getBranch(chain, brName.c_str())); + branches.refNames.emplace_back(std::move(brName)); } else { // This branch is guaranteed to exist since only collections that are // also written to file are in the info metadata that we work with here @@ -269,12 +312,14 @@ createCollectionBranches(TChain* chain, const podio::CollectionIDTable& idTable, const auto relVecNames = podio::DatamodelRegistry::instance().getRelationNames(collType); for (const auto& relName : relVecNames.relations) { - const auto brName = root_utils::refBranch(name, relName); + auto brName = root_utils::refBranch(name, relName); branches.refs.push_back(root_utils::getBranch(chain, brName.c_str())); + branches.refNames.emplace_back(std::move(brName)); } for (const auto& vecName : relVecNames.vectorMembers) { - const auto brName = root_utils::refBranch(name, vecName); + auto brName = root_utils::refBranch(name, vecName); branches.vecs.push_back(root_utils::getBranch(chain, brName.c_str())); + branches.vecNames.emplace_back(std::move(brName)); } } diff --git a/src/rootUtils.h b/src/rootUtils.h index c666eb0f0..692cfddce 100644 --- a/src/rootUtils.h +++ b/src/rootUtils.h @@ -93,6 +93,23 @@ inline std::string subsetBranch(const std::string& name) { return name + "_objIdx"; } +/** + * Reset all the branches that by getting them from the TTree again + */ +inline void resetBranches(TTree* chain, CollectionBranches& branches, const std::string& name) { + if (branches.data) { + branches.data = getBranch(chain, name); + } + + for (size_t i = 0; i < branches.refs.size(); ++i) { + branches.refs[i] = getBranch(chain, branches.refNames[i]); + } + + for (size_t i = 0; i < branches.vecs.size(); ++i) { + branches.vecs[i] = getBranch(chain, branches.vecNames[i]); + } +} + template inline void setCollectionAddresses(const BufferT& collBuffers, const CollectionBranches& branches) {