Add shared utilities to retrieve collections from the store

key4hep · Jan 24, 2025 · 6462700 · 6462700
1 parent 8cbbdc6
commit 6462700
Show file tree

Hide file tree

Showing 6 changed files with 91 additions and 49 deletions.
diff --git a/k4MarlinWrapper/CMakeLists.txt b/k4MarlinWrapper/CMakeLists.txt
@@ -54,6 +54,13 @@ target_include_directories(MarlinWrapper PUBLIC
   ${LCIO_INCLUDE_DIRS}
 )
 
+# k4MarlinWrapperUtils
+gaudi_add_library(k4MarlinWrapperUtils SHARED
+  SOURCES
+    src/components/StoreUtils.cpp
+  LINK k4FWCore::k4FWCore
+)
+
 # EDM4hep2lcio
 gaudi_add_module(EDM4hep2Lcio
   SOURCES
@@ -64,6 +71,7 @@ gaudi_add_module(EDM4hep2Lcio
     ${LCIO_LIBRARIES}
     ${Marlin_LIBRARIES}
     EDM4HEP::edm4hep
+    k4MarlinWrapperUtils
 )
 
 target_include_directories(EDM4hep2Lcio PUBLIC
@@ -80,6 +88,7 @@ gaudi_add_module(Lcio2EDM4hep
     EDM4HEP::edm4hep
     k4FWCore::k4FWCore
     k4EDM4hep2LcioConv::k4EDM4hep2LcioConv
+    k4MarlinWrapperUtils
 )
 
 target_include_directories(Lcio2EDM4hep PUBLIC

diff --git a/k4MarlinWrapper/k4MarlinWrapper/converters/EDM4hep2Lcio.h b/k4MarlinWrapper/k4MarlinWrapper/converters/EDM4hep2Lcio.h
@@ -116,8 +116,6 @@ class EDM4hep2LcioTool : public AlgTool, virtual public IEDMConverter {
                   std::vector<EDM4hep2LCIOConv::ParticleIDConvData>& pidCollections,
                   std::vector<EDM4hep2LCIOConv::TrackDqdxConvData>&  dQdxCollections);
 
-  StatusCode getAvailableCollectionsFromStore();
-
   /// Get an EDM4hep collection by name, consulting either the podio based data
   /// svc or the IOSvc
   podio::CollectionBase* getEDM4hepCollection(const std::string& name) const;

diff --git a/k4MarlinWrapper/src/components/EDM4hep2Lcio.cpp b/k4MarlinWrapper/src/components/EDM4hep2Lcio.cpp
@@ -19,6 +19,7 @@
 #include "k4MarlinWrapper/converters/EDM4hep2Lcio.h"
 #include <GaudiKernel/StatusCode.h>
 #include "GlobalConvertedObjectsMap.h"
+#include "StoreUtils.h"
 
 #include "UTIL/PIDHandler.h"
 
@@ -317,48 +318,6 @@ podio::CollectionBase* EDM4hep2LcioTool::getEDM4hepCollection(const std::string&
   throw GaudiException("Collection could not be casted to the expected type", name(), StatusCode::FAILURE);
 }
 
-StatusCode EDM4hep2LcioTool::getAvailableCollectionsFromStore() {
-  SmartIF<IDataManagerSvc> mgr;
-  mgr = evtSvc();
-
-  SmartDataPtr<DataObject> root(evtSvc(), "/Event");
-  if (!root) {
-    error() << "Failed to retrieve root object /Event" << endmsg;
-    return StatusCode::FAILURE;
-  }
-
-  auto pObj = root->registry();
-  if (!pObj) {
-    error() << "Failed to retrieve the root registry object" << endmsg;
-    return StatusCode::FAILURE;
-  }
-  std::vector<IRegistry*> leaves;
-  StatusCode              sc = mgr->objectLeaves(pObj, leaves);
-  if (!sc.isSuccess()) {
-    error() << "Failed to retrieve object leaves" << endmsg;
-    return StatusCode::FAILURE;
-  }
-  for (const auto& pReg : leaves) {
-    if (pReg->name() == k4FWCore::frameLocation) {
-      continue;
-    }
-    DataObject* p;
-    sc = m_eventDataSvc->retrieveObject("/Event" + pReg->name(), p);
-    if (sc.isFailure()) {
-      error() << "Could not retrieve object " << pReg->name() << " from the EventStore" << endmsg;
-      return sc;
-    }
-    auto wrapper = dynamic_cast<AnyDataWrapper<std::unique_ptr<podio::CollectionBase>>*>(p);
-    if (!wrapper) {
-      continue;
-    }
-    // Remove the leading /
-    m_collectionNames.push_back(pReg->name().substr(1, pReg->name().size() - 1));
-    m_idToName.emplace(wrapper->getData()->getID(), pReg->name());
-  }
-  return StatusCode::SUCCESS;
-}
-
 // Select the appropriate method to convert a collection given its type
 void EDM4hep2LcioTool::convertAdd(const std::string& e4h_coll_name, const std::string& lcio_coll_name,
                                   lcio::LCEventImpl* lcio_event, CollectionPairMappings& collection_pairs,
@@ -430,11 +389,10 @@ StatusCode EDM4hep2LcioTool::convertCollections(lcio::LCEventImpl* lcio_event) {
     edmEvent          = m_podioDataSvc->getEventFrame();
     m_collectionNames = edmEvent.value().get().getAvailableCollections();
   } else if (m_collectionNames.empty()) {
-    auto sc = getAvailableCollectionsFromStore();
-    if (sc.isFailure()) {
-      warning() << "Could not retrieve available collections from the EventStore" << endmsg;
-      return sc;
-    }
+    std::optional<std::map<uint32_t, std::string>> idToNameOpt(std::move(m_idToName));
+    auto collections = getAvailableCollectionsFromStore(this, idToNameOpt);
+    m_idToName = std::move(idToNameOpt.value());
+    m_collectionNames.insert(m_collectionNames.end(), collections.begin(), collections.end());
   }
   // Start off with the pre-defined collection name mappings
   auto collsToConvert{m_collNames.value()};

diff --git a/k4MarlinWrapper/src/components/Lcio2EDM4hep.cpp b/k4MarlinWrapper/src/components/Lcio2EDM4hep.cpp
@@ -18,6 +18,7 @@
  */
 #include "k4MarlinWrapper/converters/Lcio2EDM4hep.h"
 #include "GlobalConvertedObjectsMap.h"
+#include "StoreUtils.h"
 
 #include <EVENT/LCCollection.h>
 #include <Exceptions.h>
@@ -71,6 +72,8 @@ bool Lcio2EDM4hepTool::collectionExist(const std::string& collection_name) {
   if (m_podioDataSvc) {
     collections = m_podioDataSvc->getEventFrame().getAvailableCollections();
   } else {
+    std::optional<std::map<uint32_t, std::string>> dummy = std::nullopt;
+    getAvailableCollectionsFromStore(this, dummy, true);
   }
   if (std::find(collections.begin(), collections.end(), collection_name) != collections.end()) {
     debug() << "Collection named " << collection_name << " already registered, skipping conversion." << endmsg;

diff --git a/k4MarlinWrapper/src/components/StoreUtils.cpp b/k4MarlinWrapper/src/components/StoreUtils.cpp
@@ -0,0 +1,66 @@
+#include "GaudiKernel/AlgTool.h"
+#include "GaudiKernel/AnyDataWrapper.h"
+#include "GaudiKernel/IDataManagerSvc.h"
+#include "GaudiKernel/IDataProviderSvc.h"
+#include "GaudiKernel/SmartDataPtr.h"
+
+#include "StoreUtils.h"
+#include "podio/Frame.h"
+
+#include "k4FWCore/FunctionalUtils.h"
+
+#include <vector>
+#include <string>
+
+std::vector<std::string> getAvailableCollectionsFromStore(const AlgTool*                                  thisClass,
+                                                          std::optional<std::map<uint32_t, std::string>>& idToName,
+                                                          bool returnFrameCollections
+                                                          ) {
+  std::vector<std::string> collectionNames;
+
+  SmartIF<IDataManagerSvc> mgr;
+  mgr = thisClass->evtSvc();
+
+  SmartDataPtr<DataObject> root(thisClass->evtSvc(), "/Event");
+  if (!root) {
+    thisClass->error() << "Failed to retrieve root object /Event" << endmsg;
+  }
+
+  auto pObj = root->registry();
+  if (!pObj) {
+    thisClass->error() << "Failed to retrieve the root registry object" << endmsg;
+  }
+  std::vector<IRegistry*> leaves;
+  StatusCode              sc = mgr->objectLeaves(pObj, leaves);
+  if (!sc.isSuccess()) {
+    thisClass->error() << "Failed to retrieve object leaves" << endmsg;
+  }
+  for (const auto& pReg : leaves) {
+    if (pReg->name() == k4FWCore::frameLocation) {
+      if (!returnFrameCollections)
+        continue;
+      auto wrapper = dynamic_cast<AnyDataWrapper<podio::Frame>*>(pReg->object());
+      if (!wrapper) {
+        throw std::runtime_error("Could not cast object to Frame");
+      }
+      for (const auto& name : wrapper->getData().getAvailableCollections()) {
+        collectionNames.push_back(name);
+      }
+    }
+    DataObject* p;
+    sc = thisClass->evtSvc()->retrieveObject("/Event" + pReg->name(), p);
+    if (sc.isFailure()) {
+      thisClass->error() << "Could not retrieve object " << pReg->name() << " from the EventStore" << endmsg;
+    }
+    auto wrapper = dynamic_cast<AnyDataWrapper<std::unique_ptr<podio::CollectionBase>>*>(p);
+    if (!wrapper) {
+      continue;
+    }
+    // Remove the leading /
+    collectionNames.push_back(pReg->name().substr(1, pReg->name().size() - 1));
+    if (idToName) {
+      idToName->emplace(wrapper->getData()->getID(), pReg->name());
+    }
+  }
+  return collectionNames;
+}
diff --git a/k4MarlinWrapper/src/components/StoreUtils.h b/k4MarlinWrapper/src/components/StoreUtils.h
@@ -0,0 +1,8 @@
+#include "GaudiKernel/AlgTool.h"
+
+#include <string>
+#include <vector>
+
+std::vector<std::string> getAvailableCollectionsFromStore(const AlgTool*                                  thisClass,
+                                                          std::optional<std::map<uint32_t, std::string>>& idToName,
+                                                          bool returnFrameCollections = false);