From 1ab67232feccb730345fc1e8d6ed46be54309460 Mon Sep 17 00:00:00 2001 From: tmadlener Date: Wed, 11 Oct 2023 16:15:35 +0200 Subject: [PATCH 1/2] Make it possible to restrict collections to read --- include/podio/ROOTFrameReader.h | 9 ++++++--- src/ROOTFrameReader.cc | 18 +++++++++++++----- 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/include/podio/ROOTFrameReader.h b/include/podio/ROOTFrameReader.h index c152c3772..93381fd5c 100644 --- a/include/podio/ROOTFrameReader.h +++ b/include/podio/ROOTFrameReader.h @@ -80,14 +80,16 @@ class ROOTFrameReader { * given name. In case there are no more entries left for this name or in * case there is no data for this name, this returns a nullptr. */ - std::unique_ptr readNextEntry(const std::string& name); + std::unique_ptr readNextEntry(const std::string& name, + const std::vector& collsToRead = {}); /** * Read the specified data entry from which a Frame can be constructed for * the given name. In case the entry does not exist for this name or in case * there is no data for this name, this returns a nullptr. */ - std::unique_ptr readEntry(const std::string& name, const unsigned entry); + std::unique_ptr readEntry(const std::string& name, const unsigned entry, + const std::vector& collsToRead = {}); /// Returns number of entries for the given name unsigned getEntries(const std::string& name) const; @@ -154,7 +156,8 @@ class ROOTFrameReader { * counter aferwards. In case the requested entry is larger than the * available number of entries, return a nullptr. */ - std::unique_ptr readEntry(ROOTFrameReader::CategoryInfo& catInfo); + std::unique_ptr readEntry(ROOTFrameReader::CategoryInfo& catInfo, + const std::vector& collsToRead); /** * Get / read the buffers at index iColl in the passed category information diff --git a/src/ROOTFrameReader.cc b/src/ROOTFrameReader.cc index cb1c24aea..da60df922 100644 --- a/src/ROOTFrameReader.cc +++ b/src/ROOTFrameReader.cc @@ -14,6 +14,7 @@ #include "TTree.h" #include "TTreeCache.h" +#include #include #include @@ -46,18 +47,21 @@ GenericParameters ROOTFrameReader::readEntryParameters(ROOTFrameReader::Category return params; } -std::unique_ptr ROOTFrameReader::readNextEntry(const std::string& name) { +std::unique_ptr ROOTFrameReader::readNextEntry(const std::string& name, + const std::vector& collsToRead) { auto& catInfo = getCategoryInfo(name); - return readEntry(catInfo); + return readEntry(catInfo, collsToRead); } -std::unique_ptr ROOTFrameReader::readEntry(const std::string& name, const unsigned entNum) { +std::unique_ptr ROOTFrameReader::readEntry(const std::string& name, const unsigned entNum, + const std::vector& collsToRead) { auto& catInfo = getCategoryInfo(name); catInfo.entry = entNum; - return readEntry(catInfo); + return readEntry(catInfo, collsToRead); } -std::unique_ptr ROOTFrameReader::readEntry(ROOTFrameReader::CategoryInfo& catInfo) { +std::unique_ptr ROOTFrameReader::readEntry(ROOTFrameReader::CategoryInfo& catInfo, + const std::vector& collsToRead) { if (!catInfo.chain) { return nullptr; } @@ -77,6 +81,10 @@ std::unique_ptr ROOTFrameReader::readEntry(ROOTFrameReader::Categ ROOTFrameData::BufferMap buffers; for (size_t i = 0; i < catInfo.storedClasses.size(); ++i) { + if (!collsToRead.empty() && + std::find(collsToRead.begin(), collsToRead.end(), catInfo.storedClasses[i].first) == collsToRead.end()) { + continue; + } buffers.emplace(catInfo.storedClasses[i].first, getCollectionBuffers(catInfo, i, reloadBranches, localEntry)); } From 71eeece346118d9bdf02ec8b08801ac52e0499b9 Mon Sep 17 00:00:00 2001 From: tmadlener Date: Wed, 11 Oct 2023 16:52:48 +0200 Subject: [PATCH 2/2] Add test case for checking whether reading limited set works --- tests/root_io/read_frame_root.cpp | 43 ++++++++++++++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/tests/root_io/read_frame_root.cpp b/tests/root_io/read_frame_root.cpp index b687fa83e..935200736 100644 --- a/tests/root_io/read_frame_root.cpp +++ b/tests/root_io/read_frame_root.cpp @@ -1,10 +1,51 @@ #include "read_frame.h" #include "read_frame_auxiliary.h" +#include "podio/Frame.h" #include "podio/ROOTFrameReader.h" #include +#include #include +#include + +int test_read_frame_limited(const std::string& inputFile) { + auto reader = podio::ROOTFrameReader(); + reader.openFile(inputFile); + const std::vector collsToRead = {"mcparticles", "clusters"}; + + const auto event = podio::Frame(reader.readNextEntry("events", collsToRead)); + + const auto& availColls = event.getAvailableCollections(); + + const bool validColls = + std::set(availColls.begin(), availColls.end()) != std::set(collsToRead.begin(), collsToRead.end()); + + if (validColls) { + std::cerr << "The available collections are not as expected" << std::endl; + return 1; + } + + if (!event.get("mcparticles")) { + std::cerr << "Collection 'mcparticles' should be available" << std::endl; + return 1; + } + + if (event.get("hits")) { + std::cerr << "Collection 'hits' is available, but should not be" << std::endl; + return 1; + } + + const auto& clusters = event.get("clusters"); + const auto clu0 = clusters[0]; + const auto hits = clu0.Hits(); + if (hits.size() != 1 || hits[0].isAvailable()) { + std::cerr << "Hit in clusters are available but shouldn't be" << std::endl; + return 1; + } + + return 0; +} int main(int argc, char* argv[]) { std::string inputFile = "example_frame.root"; @@ -15,5 +56,5 @@ int main(int argc, char* argv[]) { } return read_frames(inputFile, assertBuildVersion) + - test_frame_aux_info(inputFile); + test_frame_aux_info(inputFile) + test_read_frame_limited(inputFile); }