PodioInput: Read all collections by default, allow to limit them with…

… the `collections` property (#162) * Make PodioInput not crash on non-existant collections * Default to reading all collections and allow to limit them * Update documentation
key4hep · Dec 4, 2023 · d8e5bf4 · d8e5bf4
1 parent 740cad9
commit d8e5bf4
Show file tree

Hide file tree

Showing 6 changed files with 52 additions and 15 deletions.
diff --git a/doc/PodioInputOutput.md b/doc/PodioInputOutput.md
@@ -68,19 +68,23 @@ evtSvc = k4DataSvc("EventDataSvc")
 evtSvc.input = "/path/to/your/input-file.root"
 
 podioInput = PodioInput()
-podioInput.collections = [
-    # the complete list of collection names to read
-]
 ```
 
-**Note that currently only the collections that are inside the `collections`
-list will be read and become available for later algorithms.**
-
 It is possible to change the input file from the command line via
 ```bash
 k4run <your-options-file> --EventDataSvc.input=<input-file>
 ```
 
+By default the `PodioInput` will read all collections that are available from
+the input file. It is possible to limit the collections that should become
+available via the `collections` option
+
+```python
+podioInput.collections = [
+  # List of collection names that should be made available
+]
+```
+
 ## Writing events
 
 To write events you will need to use the `PodioOutput` algorithm in addition to

diff --git a/k4FWCore/components/PodioInput.cpp b/k4FWCore/components/PodioInput.cpp
@@ -186,15 +186,38 @@ PodioInput::PodioInput(const std::string& name, ISvcLocator* svcLoc) : Consumer(
   fillReaders();
 }
 
+StatusCode PodioInput::initialize() {
+  // If someone uses the collections property from the command line and passes
+  // an empty string we assume they want all collections (as a simple way to
+  // override whatever is in the options file)
+  if (m_collectionNames.size() == 1 && m_collectionNames[0].empty()) {
+    m_collectionNames.clear();
+  }
+
+  return StatusCode::SUCCESS;
+}
+
 void PodioInput::operator()() const {
   if (m_podioDataSvc->getEventFrame().get(edm4hep::EventHeaderName)) {
     m_readers[edm4hep::EventHeaderCollection::typeName](edm4hep::EventHeaderName);
   } else {
     info() << "No EventHeader collection found in the event. Not reading it" << endmsg;
   }
 
-  for (auto& collName : m_collectionNames) {
+  const auto& collsToRead = [&]() {
+    if (m_collectionNames.empty()) {
+      return m_podioDataSvc->getEventFrame().getAvailableCollections();
+    } else {
+      return m_collectionNames.value();
+    }
+  }();
+
+  for (const auto& collName : collsToRead) {
     debug() << "Registering collection to read " << collName << endmsg;
+    if (!m_podioDataSvc->getEventFrame().get(collName)) {
+      warning() << "Collection " << collName << " is not available from file." << endmsg;
+      continue;
+    }
     auto type = m_podioDataSvc->getCollectionType(collName);
     if (m_readers.find(type) != m_readers.end()) {
       m_readers[type](collName);

diff --git a/k4FWCore/components/PodioInput.h b/k4FWCore/components/PodioInput.h
@@ -42,11 +42,14 @@ class PodioInput final : public Gaudi::Functional::Consumer<void(), BaseClass_t>
   PodioInput(const std::string& name, ISvcLocator* svcLoc);
   void operator()() const override;
 
+  StatusCode initialize() final;
+
 private:
   template <typename T> void maybeRead(std::string_view collName) const;
   void                       fillReaders();
   // Name of collections to read. Set by option collections (this is temporary)
-  Gaudi::Property<std::vector<std::string>> m_collectionNames{this, "collections", {}, "Places of collections to read"};
+  Gaudi::Property<std::vector<std::string>> m_collectionNames{
+      this, "collections", {}, "Collections that should be read (default all)"};
   // Data service: needed to register objects and get collection IDs. Just an observing pointer.
   PodioDataSvc*                                                             m_podioDataSvc;
   mutable std::map<std::string_view, std::function<void(std::string_view)>> m_readers;

diff --git a/k4FWCore/src/PodioDataSvc.cpp b/k4FWCore/src/PodioDataSvc.cpp
@@ -153,6 +153,7 @@ const std::string_view PodioDataSvc::getCollectionType(const std::string& collNa
   const auto coll = m_eventframe.get(collName);
   if (coll == nullptr) {
     error() << "Collection " << collName << " does not exist." << endmsg;
+    return "";
   }
   return coll->getTypeName();
 }

diff --git a/test/k4FWCoreTest/CMakeLists.txt b/test/k4FWCoreTest/CMakeLists.txt
@@ -60,6 +60,7 @@ endfunction()
 add_test_with_env(CreateExampleEventData options/createExampleEventData.py)
 
 add_test_with_env(CheckExampleEventData options/checkExampleEventData.py PROPERTIES DEPENDS CreateExampleEventData)
+add_test_with_env(CheckExampleEventData_noCollections options/checkExampleEventData.py --collections= PROPERTIES DEPENDS CreateExampleEventData)
 add_test_with_env(CheckExampleEventData_toolong -n 999 options/checkExampleEventData.py PROPERTIES PASS_REGULAR_EXPRESSION
                   "Application Manager Terminated successfully with a user requested ScheduledStop" DEPENDS CreateExampleEventData)
 add_test_with_env(CheckExampleEventData_unbounded options/checkExampleEventData.py -n 999 PROPERTIES PASS_REGULAR_EXPRESSION

diff --git a/test/k4FWCoreTest/options/checkExampleEventData.py b/test/k4FWCoreTest/options/checkExampleEventData.py
@@ -26,14 +26,19 @@
 
 from Configurables import PodioInput
 
+from k4FWCore.parseArgs import parser
+
+parser.add_argument(
+    "--collections",
+    action="extend",
+    nargs="?",
+    help="The input collections to read",
+    default=["VectorFloat", "MCParticles", "SimTrackerHits", "TrackerHits", "Tracks"],
+)
+my_args = parser.parse_known_args()[0]
+
 inp = PodioInput()
-inp.collections = [
-    "VectorFloat",
-    "MCParticles",
-    "SimTrackerHits",
-    "TrackerHits",
-    "Tracks",
-]
+inp.collections = my_args.collections
 
 from Configurables import k4FWCoreTest_CheckExampleEventData