Make sure the contents of Frames in one category are consistent

src/ROOTFrameWriter.cc

@@ -11,6 +11,30 @@
 
 namespace podio {
 
+/**
+ * Check whether existingIds and candidateIds both contain the same collection
+ * Ids / hashes. Returns false if the two vectors differ in content. Inputs can
+ * have random order wrt each other, but the assumption is that all the ids are
+ * unique in each vector.
+ */
+bool checkConsistentColls(const std::vector<std::string>& existingColls,
+                          const std::vector<std::string>& candidateColls) {
+  if (existingColls.size() != candidateColls.size()) {
+    return false;
+  }
+
+  // Since we are guaranteed to have unique names here, we can just look for
+  // collisions brute force, which seems to be quickest approach for vector
+  // sizes we typically have here (few hundred)
+  for (const auto& id : candidateColls) {
+    if (std::find(existingColls.begin(), existingColls.end(), id) == existingColls.end()) {
+      return false;
+    }
+  }
+
+  return true;
+}
+
 ROOTFrameWriter::ROOTFrameWriter(const std::string& filename) {
   m_file = std::make_unique<TFile>(filename.c_str(), "recreate");
 }
@@ -41,6 +65,10 @@ void ROOTFrameWriter::writeFrame(const podio::Frame& frame, const std::string& c
   collections.reserve(catInfo.collsToWrite.size());
   for (const auto& name : catInfo.collsToWrite) {
     auto* coll = frame.getCollectionForWrite(name);
+    if (!coll) {
+      // Make sure all collections that we want to write are actually available
+      throw std::runtime_error("Collection '" + name + "' in category '" + category + "' is not available in Frame");
+    }
     collections.emplace_back(name, const_cast<podio::CollectionBase*>(coll));
   }
 
@@ -50,6 +78,11 @@ void ROOTFrameWriter::writeFrame(const podio::Frame& frame, const std::string& c
     initBranches(catInfo, collections, const_cast<podio::GenericParameters&>(frame.getParameters()));
 
   } else {
+    // Make sure that the category contents are consistent with the initial
+    // frame in the category
+    if (!checkConsistentColls(catInfo.collsToWrite, collsToWrite)) {
+      throw std::runtime_error("Trying to write category '" + category + "' with inconsistent collection content");
+    }
     resetBranches(catInfo.branches, collections, &const_cast<podio::GenericParameters&>(frame.getParameters()));
   }
 

tests/unittests/unittest.cpp

@@ -11,8 +11,10 @@
 
 // podio specific includes
 #include "podio/EventStore.h"
+#include "podio/Frame.h"
 #include "podio/GenericParameters.h"
 #include "podio/ROOTFrameReader.h"
+#include "podio/ROOTFrameWriter.h"
 #include "podio/ROOTLegacyReader.h"
 #include "podio/ROOTReader.h"
 #include "podio/podioVersion.h"
@@ -1134,3 +1136,37 @@ TEST_CASE("JSON", "[json]") {
 }
 
 #endif
+
+TEST_CASE("Consistency ROOTFrameWriter", "[basics]") {
+  podio::Frame frame;
+
+  frame.put(ExampleClusterCollection(), "clusters");
+  frame.put(ExampleClusterCollection(), "clusters2");
+  frame.put(ExampleHitCollection(), "hits");
+
+  podio::ROOTFrameWriter writer("unittests_frame_consistency.root");
+  writer.writeFrame(frame, "full");
+
+  // Write a frame with more collections
+  frame.put(ExampleHitCollection(), "hits2");
+  REQUIRE_THROWS_AS(writer.writeFrame(frame, "full"), std::runtime_error);
+
+  // Write a frame with less collections
+  podio::Frame frame2;
+  frame2.put(ExampleClusterCollection(), "clusters");
+  frame2.put(ExampleClusterCollection(), "clusters2");
+  REQUIRE_THROWS_AS(writer.writeFrame(frame2, "full"), std::runtime_error);
+
+  // Write only a subset of collections
+  const std::vector<std::string> collsToWrite = {"clusters", "hits"};
+  writer.writeFrame(frame, "subset", collsToWrite);
+
+  // Frame is missing a collection
+  REQUIRE_THROWS_AS(writer.writeFrame(frame2, "subset", collsToWrite), std::runtime_error);
+
+  // Don't throw if frame contents are different, but the subset that is written
+  // is consistent
+  const std::vector<std::string> otherCollsToWrite = {"clusters", "clusters2"};
+  writer.writeFrame(frame, "subset2", otherCollsToWrite);
+  REQUIRE_NOTHROW(writer.writeFrame(frame2, "subset2", otherCollsToWrite));
+}