Add mergeChunksFromSameSourceID

openPMD · Jul 23, 2024 · 6063622 · 6063622
1 parent f6d0278
commit 6063622
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 1 deletion.
diff --git a/include/openPMD/ChunkInfo.hpp b/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;

diff --git a/src/ChunkInfo.cpp b/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> &);

diff --git a/src/binding/python/ChunkInfo.cpp b/src/binding/python/ChunkInfo.cpp
@@ -237,7 +237,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;