Add mergeChunksFromSameSourceID

Author: franzpoeschel

include/openPMD/ChunkInfo.hpp

@@ -94,6 +94,9 @@ namespace chunk_assignment
     template <typename Chunk_t>
     void mergeChunks(std::vector<Chunk_t> &);
 
+    auto mergeChunksFromSameSourceID(std::vector<WrittenChunkInfo> const &)
+        -> std::map<unsigned int, std::vector<ChunkInfo>>;
+
     struct PartialAssignment
     {
         ChunkTable notAssigned;

src/ChunkInfo.cpp

@@ -174,6 +174,21 @@ namespace chunk_assignment
         } while (stillChanging);
     }
 
+    auto mergeChunksFromSameSourceID(std::vector<WrittenChunkInfo> const &table)
+        -> std::map<unsigned int, std::vector<ChunkInfo>>
+    {
+        std::map<unsigned int, std::vector<ChunkInfo>> sortedBySourceID;
+        for (auto const &chunk : table)
+        {
+            sortedBySourceID[chunk.sourceID].emplace_back(chunk);
+        }
+        for (auto &pair : sortedBySourceID)
+        {
+            mergeChunks(pair.second);
+        }
+        return sortedBySourceID;
+    }
+
     template void mergeChunks<ChunkInfo>(std::vector<ChunkInfo> &);
     template void
     mergeChunks<WrittenChunkInfo>(std::vector<WrittenChunkInfo> &);

src/binding/python/ChunkInfo.cpp

@@ -239,7 +239,10 @@ void init_Chunk(py::module &m)
         .def("available", &host_info::methodAvailable);
 
     py::bind_vector<ChunkTable>(m, "ChunkTable")
-        .def("merge_chunks", &chunk_assignment::mergeChunks<WrittenChunkInfo>);
+        .def("merge_chunks", &chunk_assignment::mergeChunks<WrittenChunkInfo>)
+        .def(
+            "merge_chunks_from_same_sourceID",
+            &chunk_assignment::mergeChunksFromSameSourceID);
 
     using namespace chunk_assignment;