Make it possible to restrict collections to read

AIDASoft · Oct 11, 2023 · 1ab6723 · 1ab6723
1 parent 9def6a3
commit 1ab6723
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 8 deletions.
diff --git 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<podio::ROOTFrameData> readNextEntry(const std::string& name);
+  std::unique_ptr<podio::ROOTFrameData> readNextEntry(const std::string& name,
+                                                      const std::vector<std::string>& 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<podio::ROOTFrameData> readEntry(const std::string& name, const unsigned entry);
+  std::unique_ptr<podio::ROOTFrameData> readEntry(const std::string& name, const unsigned entry,
+                                                  const std::vector<std::string>& 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<podio::ROOTFrameData> readEntry(ROOTFrameReader::CategoryInfo& catInfo);
+  std::unique_ptr<podio::ROOTFrameData> readEntry(ROOTFrameReader::CategoryInfo& catInfo,
+                                                  const std::vector<std::string>& collsToRead);
 
   /**
    * Get / read the buffers at index iColl in the passed category information

diff --git a/src/ROOTFrameReader.cc b/src/ROOTFrameReader.cc
@@ -14,6 +14,7 @@
 #include "TTree.h"
 #include "TTreeCache.h"
 
+#include <algorithm>
 #include <stdexcept>
 #include <unordered_map>
 
@@ -46,18 +47,21 @@ GenericParameters ROOTFrameReader::readEntryParameters(ROOTFrameReader::Category
   return params;
 }
 
-std::unique_ptr<ROOTFrameData> ROOTFrameReader::readNextEntry(const std::string& name) {
+std::unique_ptr<ROOTFrameData> ROOTFrameReader::readNextEntry(const std::string& name,
+                                                              const std::vector<std::string>& collsToRead) {
   auto& catInfo = getCategoryInfo(name);
-  return readEntry(catInfo);
+  return readEntry(catInfo, collsToRead);
 }
 
-std::unique_ptr<ROOTFrameData> ROOTFrameReader::readEntry(const std::string& name, const unsigned entNum) {
+std::unique_ptr<ROOTFrameData> ROOTFrameReader::readEntry(const std::string& name, const unsigned entNum,
+                                                          const std::vector<std::string>& collsToRead) {
   auto& catInfo = getCategoryInfo(name);
   catInfo.entry = entNum;
-  return readEntry(catInfo);
+  return readEntry(catInfo, collsToRead);
 }
 
-std::unique_ptr<ROOTFrameData> ROOTFrameReader::readEntry(ROOTFrameReader::CategoryInfo& catInfo) {
+std::unique_ptr<ROOTFrameData> ROOTFrameReader::readEntry(ROOTFrameReader::CategoryInfo& catInfo,
+                                                          const std::vector<std::string>& collsToRead) {
   if (!catInfo.chain) {
     return nullptr;
   }
@@ -77,6 +81,10 @@ std::unique_ptr<ROOTFrameData> 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));
   }