[WIP] Add a service to allow retrieving collections to which an object belongs

k4FWCore/components/CollectionFromObjectSvc.cpp

@@ -0,0 +1,81 @@
+/*
+ * Copyright (c) 2014-2024 Key4hep-Project.
+ *
+ * This file is part of Key4hep.
+ * See https://key4hep.github.io/key4hep-doc/ for further info.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "CollectionFromObjectSvc.h"
+
+#include "k4FWCore/FunctionalUtils.h"
+
+#include <podio/CollectionBase.h>
+#include <podio/ObjectID.h>
+
+#include <fmt/format.h>
+#include <fmt/ostream.h>
+
+template <>
+struct fmt::formatter<podio::ObjectID> : ostream_formatter {};
+
+StatusCode CollectionFromObjectSvc::initialize() {
+  auto sc = Service::initialize();
+  if (sc.isFailure()) {
+    error() << "Unable to initialize base class Serivce." << endmsg;
+    return sc;
+  }
+  m_dataSvc = service("EventDataSvc");
+  if (!m_dataSvc) {
+    error() << "Unable to locate the EventDataSvc" << endmsg;
+    return StatusCode::FAILURE;
+  }
+
+  return StatusCode::SUCCESS;
+}
+
+const std::optional<std::string> CollectionFromObjectSvc::getCollectionNameFor(const podio::ObjectID id) const {
+  debug() << "Trying to retrieve collection name for object " << id << endmsg;
+  const auto& idTable = k4FWCore::details::getTESCollectionIDTable(m_dataSvc, this);
+  auto name = idTable.name(id.collectionID);
+  if (!name.has_value()) {
+    error() << "Could not get a collection name for object " << id << endmsg;
+    return std::nullopt;
+  }
+  return name.value();
+}
+
+const podio::CollectionBase* CollectionFromObjectSvc::getCollectionFor(const podio::ObjectID id) const {
+  debug() << "Trying to retrieve collection for object " << id << endmsg;
+  const auto& idTable = k4FWCore::details::getTESCollectionIDTable(m_dataSvc, this);
+  auto name = idTable.name(id.collectionID);
+  if (!name.has_value()) {
+    error() << "Could not get a collection name for object " << id << endmsg;
+    return nullptr;
+  }
+
+  DataObject* p{nullptr};
+  if (m_dataSvc->retrieveObject("/Event/" + name.value(), p).isFailure()) {
+    error() << "Could not get the collection '" << name.value() << "' from the Event store" << endmsg;
+    return nullptr;
+  }
+  const auto collWrapper = dynamic_cast<AnyDataWrapper<std::unique_ptr<podio::CollectionBase>>*>(p);
+  if (!collWrapper) {
+    error() << "Could not cast data object to the necessary type for collection " << name.value() << endmsg;
+    return nullptr;
+  }
+
+  return collWrapper->getData().get();
+}
+
+DECLARE_COMPONENT(CollectionFromObjectSvc)

k4FWCore/components/CollectionFromObjectSvc.h

@@ -0,0 +1,44 @@
+/*
+ * Copyright (c) 2014-2024 Key4hep-Project.
+ *
+ * This file is part of Key4hep.
+ * See https://key4hep.github.io/key4hep-doc/ for further info.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#ifndef FWCORE_COLLECTIONFROMOBJECTSVC_H
+#define FWCORE_COLLECTIONFROMOBJECTSVC_H
+
+#include <GaudiKernel/IDataProviderSvc.h>
+#include <GaudiKernel/Service.h>
+
+#include <podio/CollectionBase.h>
+#include <podio/ObjectID.h>
+
+#include "k4FWCore/ICollectionFromObjectSvc.h"
+
+class CollectionFromObjectSvc : public extends<Service, ICollectionFromObjectSvc> {
+  using extends::extends;
+
+public:
+  StatusCode initialize() override;
+
+protected:
+  const podio::CollectionBase* getCollectionFor(const podio::ObjectID id) const override;
+  const std::optional<std::string> getCollectionNameFor(const podio::ObjectID id) const override;
+
+private:
+  SmartIF<IDataProviderSvc> m_dataSvc;
+};
+
+#endif

k4FWCore/components/Reader.cpp

@@ -25,6 +25,7 @@
 #include "GaudiKernel/StatusCode.h"
 
 #include "podio/CollectionBase.h"
+#include "podio/CollectionIDTable.h"
 #include "podio/Frame.h"
 
 #include "IIOSvc.h"
@@ -136,12 +137,19 @@ class Reader final : public CollectionPusher {
 
     auto eds = eventSvc().as<IDataProviderSvc>();
     auto frame = std::move(std::get<podio::Frame>(val));
+    auto idTable = frame.getCollectionIDTableForWrite();
 
     auto tmp = new AnyDataWrapper<podio::Frame>(std::move(frame));
     if (eds->registerObject("/Event" + k4FWCore::frameLocation, tmp).isFailure()) {
       error() << "Failed to register Frame object" << endmsg;
     }
 
+    // We also store the collection id table to make sure we can detect collisions
+    auto idTableWrapper = new AnyDataWrapper<podio::CollectionIDTable>(std::move(idTable));
+    if (eds->registerObject("/Event" + k4FWCore::idTableLocation, idTableWrapper).isFailure()) {
+      error() << "Failed to register CollectionIDTable object" << endmsg;
+    }
+
     return std::make_tuple(std::get<0>(val), std::get<1>(val));
   }
 };

k4FWCore/include/k4FWCore/FunctionalUtils.h

@@ -29,6 +29,7 @@
 #include <GaudiKernel/GaudiException.h>
 
 #include "podio/CollectionBase.h"
+#include <podio/CollectionIDTable.h>
 
 #include "k4FWCore/DataWrapper.h"
 
@@ -43,6 +44,7 @@
 namespace k4FWCore {
 
 static const std::string frameLocation = "/_Frame";
+static const std::string idTableLocation = "/_CollectionIDTable";
 
 namespace details {
 
@@ -208,6 +210,49 @@ namespace details {
     }
   }
 
+  podio::CollectionIDTable& getTESCollectionIDTable(IDataProviderSvc* eds, auto thisClass) {
+    DataObject* p;
+    auto sc = eds->retrieveObject("/Event" + k4FWCore::idTableLocation, p);
+    if (!sc.isSuccess()) {
+      // We are not reading from file so we create one on the fly
+      thisClass->debug() << "Could not retrieve CollectionIDTable for assigning a collection id" << endmsg;
+      auto idTableWrapper = new AnyDataWrapper<podio::CollectionIDTable>(podio::CollectionIDTable{});
+      if (eds->registerObject("/Event" + k4FWCore::idTableLocation, idTableWrapper).isFailure()) {
+        thisClass->error() << "Failed to place an empty CollectionIDTable into the TES" << endmsg;
+      } else {
+        return idTableWrapper->getData();
+      }
+    } else {
+      auto idTableWrapper = dynamic_cast<AnyDataWrapper<podio::CollectionIDTable>*>(p);
+      if (!idTableWrapper) {
+        thisClass->error() << "Object at /Event" << k4FWCore::idTableLocation
+                           << " is not the expected CollectionIDTable" << endmsg;
+      }
+      return idTableWrapper->getData();
+    }
+
+    throw std::runtime_error("Could neither retrieve an existing nor place an empty CollectionIDTable in the TES");
+  }
+
+  podio::CollectionIDTable& getTESCollectionIDTable(auto thisClass) {
+    auto eds = thisClass->eventSvc().template as<IDataProviderSvc>();
+    return getTESCollectionIDTable(eds, thisClass);
+  }
+
+  void putCollectionSetID(std::unique_ptr<podio::CollectionBase> coll,
+                          const DataObjectWriteHandle<std::unique_ptr<podio::CollectionBase>>& handle, auto thisClass) {
+    auto& idTable = getTESCollectionIDTable(thisClass);
+    if (idTable.present(handle.objKey())) {
+      thisClass->error() << "Collection with name " << handle.objKey() << " is already stored in the TES" << endmsg;
+    }
+    const auto id = idTable.add(handle.objKey());
+    coll->setID(id);
+    thisClass->verbose() << fmt::format("Assigning collection id {:0>8x} to collection '{}'", coll->getID(),
+                                        handle.objKey())
+                         << endmsg;
+    Gaudi::Functional::details::put(handle, std::move(coll));
+  }
+
   template <size_t Index, typename... Out, typename... Handles>
   void putVectorOutputs(std::tuple<Handles...>&& handles, const auto& m_outputs, auto thisClass) {
     if constexpr (Index < sizeof...(Handles)) {
@@ -221,12 +266,12 @@ namespace details {
           throw GaudiException(thisClass->name(), msg, StatusCode::FAILURE);
         }
         for (auto& val : std::get<Index>(handles)) {
-          Gaudi::Functional::details::put(std::get<Index>(m_outputs)[i], convertToUniquePtr(std::move(val)));
+          putCollectionSetID(convertToUniquePtr(std::move(val)), std::get<Index>(m_outputs)[i], thisClass);
           i++;
         }
       } else {
-        Gaudi::Functional::details::put(std::get<Index>(m_outputs)[0],
-                                        convertToUniquePtr(std::move(std::get<Index>(handles))));
+        putCollectionSetID(convertToUniquePtr(std::move(std::get<Index>(handles))), std::get<Index>(m_outputs)[0],
+                           thisClass);
       }
 
       // Recursive call for the next index

k4FWCore/include/k4FWCore/ICollectionFromObjectSvc.h

@@ -0,0 +1,51 @@
+/*
+ * Copyright (c) 2014-2024 Key4hep-Project.
+ *
+ * This file is part of Key4hep.
+ * See https://key4hep.github.io/key4hep-doc/ for further info.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#ifndef FWCORE_ICOLLECTIONFROMOBJECTSVC_H
+#define FWCORE_ICOLLECTIONFROMOBJECTSVC_H
+
+#include <GaudiKernel/IInterface.h>
+
+#include <podio/CollectionBase.h>
+#include <podio/ObjectID.h>
+#include <podio/utilities/TypeHelpers.h>
+
+#include <optional>
+#include <string>
+
+class ICollectionFromObjectSvc : virtual public IInterface {
+public:
+  DeclareInterfaceID(ICollectionFromObjectSvc, 1, 0);
+
+  template <podio::ObjectType O>
+  const typename O::collection_type* getCollectionFor(const O& object) const {
+    return dynamic_cast<const typename O::collection_type*>(getCollectionFor(object.id()));
+  }
+
+  // TODO: return string_view? Some form of DataHandle?
+  template <podio::ObjectType O>
+  const std::optional<std::string> getCollectionNameFor(const O& object) const {
+    return getCollectionNameFor(object.id());
+  }
+
+protected:
+  virtual const podio::CollectionBase* getCollectionFor(const podio::ObjectID) const = 0;
+  virtual const std::optional<std::string> getCollectionNameFor(const podio::ObjectID) const = 0;
+};
+
+#endif

k4FWCore/include/k4FWCore/Transformer.h

@@ -116,9 +116,9 @@ namespace details {
           std::tuple<Out> tmp = filter_evtcontext_tt<In...>::apply(*this, ctx, this->m_inputs);
           putVectorOutputs<0, Out>(std::move(tmp), m_outputs, this);
         } else {
-          Gaudi::Functional::details::put(
-              std::get<0>(this->m_outputs)[0],
-              convertToUniquePtr(std::move(filter_evtcontext_tt<In...>::apply(*this, ctx, this->m_inputs))));
+          putCollectionSetID(
+              convertToUniquePtr(std::move(filter_evtcontext_tt<In...>::apply(*this, ctx, this->m_inputs))),
+              std::get<0>(this->m_outputs)[0], this);
         }
         return Gaudi::Functional::FilterDecision::PASSED;
       } catch (GaudiException& e) {

python/k4FWCore/ApplicationMgr.py

@@ -20,7 +20,14 @@
 
 from Gaudi import Configuration
 from Configurables import ApplicationMgr as AppMgr
-from Configurables import Reader, Writer, IOSvc, Gaudi__Sequencer, EventLoopMgr
+from Configurables import (
+    Reader,
+    Writer,
+    IOSvc,
+    Gaudi__Sequencer,
+    EventLoopMgr,
+    CollectionFromObjectSvc,
+)
 from k4FWCore.utils import get_logger
 
 logger = get_logger()
@@ -118,6 +125,9 @@ def fix_properties(self):
         if "MetadataSvc" in self._mgr.allConfigurables:
             self._mgr.ExtSvc.append(self._mgr.allConfigurables["MetadataSvc"])
 
+        # Always attach the CollectionFromObjectSvc
+        self._mgr.ExtSvc.append(CollectionFromObjectSvc("CollectionFromObjectSvc"))
+
         if "IOSvc" not in self._mgr.allConfigurables:
             return
         if not isinstance(self._mgr.allConfigurables["IOSvc"], IOSvc):

test/k4FWCoreTest/CMakeLists.txt

@@ -202,6 +202,7 @@ add_test_with_env(ReadLimitedInputsAllEventsIOSvc options/ExampleIOSvcLimitInput
 
 add_test_with_env(ParticleIDMetadataFramework options/ExampleParticleIDMetadata.py)
 
+add_test_with_env(CollectionFromObjectRetrieval options/TestCollectionFromObjectRetrieval.py)
 
 # The following is done to make the tests work without installing the files in
 # the installation directory. The k4FWCore in the build directory is populated by

test/k4FWCoreTest/options/TestCollectionFromObjectRetrieval.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) 2014-2024 Key4hep-Project.
+#
+# This file is part of Key4hep.
+# See https://key4hep.github.io/key4hep-doc/ for further info.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# This is a simple test checking that the collection IDs that are assigned for
+# the Functional algorithms are the same as we get when adding to a frame.
+# Effectively, it checks whether we hash the right thing when we assign a
+# collection ID. We need this check because for Functional algorithms the
+# insertion into a Frame (where the usual assignment happens) happens just
+# before writing.
+
+from Gaudi.Configuration import VERBOSE
+from Configurables import (
+    TestCollectionFromObjectRetrieval,
+    ExampleFunctionalProducerMultiple,
+    CollectionFromObjectSvc,
+)
+
+from k4FWCore import ApplicationMgr
+from Configurables import EventDataSvc
+
+producer = ExampleFunctionalProducerMultiple(
+    "Producer", OutputCollectionParticles1=["UnconventionalCollName"], OutputLevel=VERBOSE
+)
+
+
+collid_checker = TestCollectionFromObjectRetrieval(
+    "CollectionIDChecker", InputCollection=["UnconventionalCollName"], OutputLevel=VERBOSE
+)
+
+ApplicationMgr(
+    TopAlg=[producer, collid_checker],
+    EvtSel="NONE",
+    EvtMax=1,
+    ExtSvc=[EventDataSvc(), CollectionFromObjectSvc("CollectionFromObjSvc")],
+)