Pickle API: Cache unpickled Series to avoid repeated file access (#1654)

* Cache the Series per thread * Fixes for Windows
openPMD · Oct 30, 2024 · 2a36ef8 · 2a36ef8
1 parent de67eff
commit 2a36ef8
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 3 deletions.
diff --git a/include/openPMD/binding/python/Pickle.hpp b/include/openPMD/binding/python/Pickle.hpp
@@ -22,6 +22,7 @@
 
 #include "openPMD/IO/Access.hpp"
 #include "openPMD/Series.hpp"
+#include "openPMD/auxiliary/StringManip.hpp"
 #include "openPMD/backend/Attributable.hpp"
 
 #include "Common.hpp"
@@ -67,9 +68,40 @@ add_pickle(pybind11::class_<T_Args...> &cl, T_SeriesAccessor &&seriesAccessor)
             std::vector<std::string> const group =
                 t[1].cast<std::vector<std::string> >();
 
-            openPMD::Series series(
-                filename, Access::READ_ONLY, "defer_iteration_parsing = true");
-            return seriesAccessor(std::move(series), group);
+            /*
+             * Cache the Series per thread.
+             */
+            thread_local std::optional<openPMD::Series> series;
+            bool re_initialize = [&]() {
+                try
+                {
+                    return !series.has_value() || !series->operator bool() ||
+                        auxiliary::replace_all(
+                            series->myPath().filePath(), "\\", "/") !=
+                        auxiliary::replace_all(filename, "\\", "/");
+                }
+                /*
+                 * Better safe than sorry, if anything goes wrong because the
+                 * Series is in a weird state, just reinitialize it.
+                 */
+                catch (...)
+                {
+                    return true;
+                }
+            }();
+            if (re_initialize)
+            {
+                /*
+                 * Do NOT close the old Series, it might still be active in
+                 * terms of handed-out handles.
+                 */
+                series = std::make_optional<Series>(
+                    filename,
+                    Access::READ_ONLY,
+                    "defer_iteration_parsing = true");
+            }
+
+            return seriesAccessor(*series, group);
         }));
 }
 } // namespace openPMD
diff --git a/test/SerialIOTest.cpp b/test/SerialIOTest.cpp
@@ -177,6 +177,9 @@ void write_and_read_many_iterations(
 
     {
         Series write(filename, Access::CREATE);
+        REQUIRE(
+            auxiliary::replace_all(write.myPath().filePath(), "\\", "/") ==
+            auxiliary::replace_all(filename, "\\", "/"));
         for (unsigned int i = 0; i < nIterations; ++i)
         {
             // std::cout << "Putting iteration " << i << std::endl;
@@ -1851,6 +1854,12 @@ inline void fileBased_write_test(const std::string &backend)
             "../samples/subdir/serial_fileBased_write%03T." + backend,
             Access::CREATE,
             jsonCfg);
+        REQUIRE(
+            auxiliary::replace_all(o.myPath().filePath(), "\\", "/") ==
+            auxiliary::replace_all(
+                "../samples/subdir/serial_fileBased_write%03T." + backend,
+                "\\",
+                "/"));
 
         ParticleSpecies &e_1 = o.iterations[1].particles["e"];