[WIP] Make it possible to read only a subset of available collections in ROOTFrameReader

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

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));
   }
 

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 <iostream>
+#include <set>
 #include <string>
+#include <vector>
+
+int test_read_frame_limited(const std::string& inputFile) {
+  auto reader = podio::ROOTFrameReader();
+  reader.openFile(inputFile);
+  const std::vector<std::string> 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<ExampleClusterCollection>("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<podio::ROOTFrameReader>(inputFile, assertBuildVersion) +
-      test_frame_aux_info<podio::ROOTFrameReader>(inputFile);
+      test_frame_aux_info<podio::ROOTFrameReader>(inputFile) + test_read_frame_limited(inputFile);
 }