Fix writing an arbitrary number of collections

key4hep · Mar 10, 2024 · 92205f5 · 92205f5
1 parent 9400889
commit 92205f5
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 1 deletion.
diff --git a/k4FWCore/include/k4FWCore/FunctionalUtils.h b/k4FWCore/include/k4FWCore/FunctionalUtils.h
@@ -157,6 +157,24 @@ namespace k4FWCore {
       }
     }
 
+    template <size_t Index, typename... In, typename... Handles>
+    void deleteMapInputs(const std::tuple<Handles...>& handles, auto thisClass) {
+      if constexpr (Index < sizeof...(Handles)) {
+        if constexpr (is_map_like<std::tuple_element_t<Index, std::tuple<In...>>>::value) {
+          // In case of map types like std::map<std::string, edm4hep::MCParticleCollection&>
+          // we have to remove the reference to get the actual type
+          auto sc = thisClass->evtSvc()->unregisterObject(std::get<Index>(handles).objKey());
+          if (!sc.isSuccess()) {
+            throw GaudiException("Failed to retrieve object " + std::get<Index>(handles).objKey(), "Consumer",
+                                 StatusCode::FAILURE);
+          }
+        }
+      // Recursive call for the next index
+      deleteMapInputs<Index + 1, In...>(handles, thisClass);
+      }
+
+    }
+
   }  // namespace details
 }  // namespace k4FWCore
 

diff --git a/k4FWCore/include/k4FWCore/Transformer.h b/k4FWCore/include/k4FWCore/Transformer.h
@@ -132,6 +132,10 @@ namespace k4FWCore {
                 std::get<0>(this->m_outputs),
                 ptrOrCast(std::move(filter_evtcontext_tt<In...>::apply(*this, ctx, this->m_inputs))));
           }
+          // If any input has map type, we remove it from the store since it has been pushed
+          // but we don't want the map to be available as a map for other algorithms and
+          // all the individual collections have alreaody been pushed
+          deleteMapInputs<0, In...>(this->m_inputs, this);
           return Gaudi::Functional::FilterDecision::PASSED;
         } catch (GaudiException& e) {
           (e.code() ? this->warning() : this->error()) << e.tag() << " : " << e.message() << endmsg;
@@ -246,6 +250,10 @@ namespace k4FWCore {
           readMapInputs<0, In...>(this->m_inputs, m_inputLocations, m_inputLocationsMap, this);
           auto tmp = filter_evtcontext_tt<In...>::apply(*this, ctx, this->m_inputs);
           putMapOutputs<0>(std::move(tmp));
+          // If any input has map type, we remove it from the store since it has been pushed
+          // but we don't want the map to be available as a map for other algorithms and
+          // all the individual collections have alreaody been pushed
+          deleteMapInputs<0, In...>(this->m_inputs, this);
           return Gaudi::Functional::FilterDecision::PASSED;
         } catch (GaudiException& e) {
           (e.code() ? this->warning() : this->error()) << e.tag() << " : " << e.message() << endmsg;

diff --git a/test/k4FWCoreTest/options/CheckOutputFiles.py b/test/k4FWCoreTest/options/CheckOutputFiles.py
@@ -26,6 +26,7 @@
 
 
 def check_collections(filename, names):
+    print(f'Checking file "{filename}" for collections {names}')
     podio_reader = podio.root_io.Reader(filename)
     frames = podio_reader.get("events")
     if not len(frames) and len(names):
@@ -62,4 +63,4 @@ def check_collections(filename, names):
     ["VectorFloat", "MCParticles1", "MCParticles2", "SimTrackerHits", "TrackerHits"],
 )
 check_collections("/tmp/a/b/c/output_k4test_exampledata_producer.root", ["MCParticles"])
-check_collections("functional_transformer_runtime_empty.root", [])
+check_collections("functional_transformer_runtime_empty.root", ["MCParticles0", "MCParticles1", "MCParticles2"])