From 82eea7aa5a74be0d55187e2fb3e3afff8b38abde Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Franz=20P=C3=B6schel?= <franz.poeschel@gmail.com>
Date: Tue, 14 Jan 2025 13:47:49 +0100
Subject: [PATCH] Reset steps before reading the rankTable

---
 include/openPMD/IO/ADIOS/ADIOS2File.hpp |  2 +-
 include/openPMD/IO/IOTask.hpp           |  2 +-
 src/IO/ADIOS/ADIOS2File.cpp             | 14 +++++++++++---
 src/IO/AbstractIOHandlerImpl.cpp        |  6 +++++-
 src/Series.cpp                          |  8 ++++++++
 test/SerialIOTest.cpp                   |  9 ++++++++-
 6 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/include/openPMD/IO/ADIOS/ADIOS2File.hpp b/include/openPMD/IO/ADIOS/ADIOS2File.hpp
index 2f53ba6686..91c5773ba8 100644
--- a/include/openPMD/IO/ADIOS/ADIOS2File.hpp
+++ b/include/openPMD/IO/ADIOS/ADIOS2File.hpp
@@ -413,7 +413,7 @@ class ADIOS2File
     StreamStatus streamStatus = StreamStatus::OutsideOfStep;
 
     size_t currentStep();
-    void setStepSelection(size_t);
+    void setStepSelection(std::optional<size_t>);
     [[nodiscard]] std::optional<size_t> stepSelection() const;
 
 private:
diff --git a/include/openPMD/IO/IOTask.hpp b/include/openPMD/IO/IOTask.hpp
index d65f0dda72..5ac4c9f6c5 100644
--- a/include/openPMD/IO/IOTask.hpp
+++ b/include/openPMD/IO/IOTask.hpp
@@ -672,7 +672,7 @@ struct OPENPMDAPI_EXPORT Parameter<Operation::ADVANCE>
 
     struct StepSelection
     {
-        size_t step;
+        std::optional<size_t> step;
     };
 
     //! input parameter
diff --git a/src/IO/ADIOS/ADIOS2File.cpp b/src/IO/ADIOS/ADIOS2File.cpp
index 22823fc130..50b5024cc3 100644
--- a/src/IO/ADIOS/ADIOS2File.cpp
+++ b/src/IO/ADIOS/ADIOS2File.cpp
@@ -353,7 +353,7 @@ size_t ADIOS2File::currentStep()
     }
 }
 
-void ADIOS2File::setStepSelection(size_t step)
+void ADIOS2File::setStepSelection(std::optional<size_t> step)
 {
     if (streamStatus != StreamStatus::ReadWithoutStream)
     {
@@ -361,8 +361,16 @@ void ADIOS2File::setStepSelection(size_t step)
             "ADIOS2 backend: Cannot only use random-access step selections "
             "when reading without streaming mode.");
     }
-    m_currentStep = step;
-    useStepSelection = true;
+    if (!step.has_value())
+    {
+        m_currentStep = 0;
+        useStepSelection = false;
+    }
+    else
+    {
+        m_currentStep = *step;
+        useStepSelection = true;
+    }
 }
 
 std::optional<size_t> ADIOS2File::stepSelection() const
diff --git a/src/IO/AbstractIOHandlerImpl.cpp b/src/IO/AbstractIOHandlerImpl.cpp
index d362c91b2e..95fb7dfa09 100644
--- a/src/IO/AbstractIOHandlerImpl.cpp
+++ b/src/IO/AbstractIOHandlerImpl.cpp
@@ -421,7 +421,11 @@ std::future<void> AbstractIOHandlerImpl::flush()
                                 [](Parameter<Operation::ADVANCE>::StepSelection
                                        step) {
                                     std::stringstream s;
-                                    s << "RANDOMACCESS '" << step.step << "'";
+                                    s << "RANDOMACCESS '"
+                                      << (step.step.has_value()
+                                              ? std::to_string(*step.step)
+                                              : std::string("RESET"))
+                                      << "'";
                                     return s.str();
                                 }},
                             parameter.mode);
diff --git a/src/Series.cpp b/src/Series.cpp
index b2aa65e121..9a9d9a2972 100644
--- a/src/Series.cpp
+++ b/src/Series.cpp
@@ -256,6 +256,14 @@ chunk_assignment::RankMeta Series::rankTable([[maybe_unused]] bool collective)
         rankTable.m_bufferedRead = chunk_assignment::RankMeta{};
         return {};
     }
+    if (iterationEncoding() == IterationEncoding::variableBased &&
+        IOHandler()->m_backendAccess == Access::READ_RANDOM_ACCESS)
+    {
+        Parameter<Operation::ADVANCE> advance;
+        advance.mode =
+            Parameter<Operation::ADVANCE>::StepSelection{std::nullopt};
+        IOHandler()->enqueue(IOTask(this, std::move(advance)));
+    }
     Parameter<Operation::OPEN_DATASET> openDataset;
     openDataset.name = "rankTable";
     IOHandler()->enqueue(IOTask(&rankTable.m_attributable, openDataset));
diff --git a/test/SerialIOTest.cpp b/test/SerialIOTest.cpp
index 512c91f45d..1fd779a794 100644
--- a/test/SerialIOTest.cpp
+++ b/test/SerialIOTest.cpp
@@ -5246,7 +5246,8 @@ void serial_iterator(std::string const &file)
 {
     constexpr Extent::value_type extent = 1000;
     {
-        Series writeSeries(file, Access::CREATE);
+        Series writeSeries(
+            file, Access::CREATE, "rank_table = \"posix_hostname\"");
         auto iterations = writeSeries.writeIterations();
         for (size_t i = 0; i < 10; ++i)
         {
@@ -5277,6 +5278,12 @@ void serial_iterator(std::string const &file)
         }
         last_iteration_index = iteration.iterationIndex;
     }
+    if (readSeries.iterationEncoding() == IterationEncoding::variableBased)
+        for (auto const &[rank, host] : readSeries.rankTable(true))
+        {
+            std::cout << "POST Rank '" << rank << "' written from host '"
+                      << host << "'\n";
+        }
     REQUIRE(last_iteration_index == 9);
     REQUIRE(numberOfIterations == 10);
 }