diff --git a/src/ROOTFrameWriter.cc b/src/ROOTFrameWriter.cc index e140ffc4d..3ae43f95c 100644 --- a/src/ROOTFrameWriter.cc +++ b/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& existingColls, + const std::vector& 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(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(coll)); } @@ -50,6 +78,11 @@ void ROOTFrameWriter::writeFrame(const podio::Frame& frame, const std::string& c initBranches(catInfo, collections, const_cast(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(frame.getParameters())); } diff --git a/tests/unittests/unittest.cpp b/tests/unittests/unittest.cpp index 8a313de28..ebd808748 100644 --- a/tests/unittests/unittest.cpp +++ b/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,42 @@ 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 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 otherCollsToWrite = {"clusters", "clusters2"}; + writer.writeFrame(frame, "subset2", otherCollsToWrite); + REQUIRE_NOTHROW(writer.writeFrame(frame2, "subset2", otherCollsToWrite)); + + // Make sure that restricting the second frame works. + // See https://github.com/AIDASoft/podio/issues/382 for the original issue + writer.writeFrame(frame2, "full_frame2"); + REQUIRE_NOTHROW(writer.writeFrame(frame, "full_frame2", frame2.getAvailableCollections())); +}