Variable encoding in ADIOS2: Support partial datasets in ReadRandomAccess (#1746)

* Move ADIOS2 variables buffer to own struct

* Implement preparsing for variables

* Failing test

* Fixes

* Fix rankTable reading

* Fix non-ADIOS2 builds

* Skip preparsing variables if not needed

* Fixes for last commit

* Remove restrictions from docs

* Add nodiscard...

* Some testing, fixes and optimizations

* Extend test for group-based encoding

* Fix use after free

* Licensing header

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: franzpoeschel

CMakeLists.txt

@@ -441,6 +441,7 @@ set(IO_SOURCE
         src/IO/JSON/JSONFilePosition.cpp
         src/IO/ADIOS/ADIOS2IOHandler.cpp
         src/IO/ADIOS/ADIOS2PreloadAttributes.cpp
+        src/IO/ADIOS/ADIOS2PreloadVariables.cpp
         src/IO/ADIOS/ADIOS2File.cpp
         src/IO/ADIOS/ADIOS2Auxiliary.cpp
         src/IO/InvalidatableFile.cpp)

docs/source/usage/workflow.rst

@@ -77,15 +77,7 @@ The openPMD-api distinguishes between a number of different access modes:
      When using such a backend, the access mode will be coerced automatically to *linear read mode*.
      Use of Series::readIterations() is mandatory for access.
   4. *Random-access read mode* for a variable-based Series is currently experimental.
-     Support for metadata definitions that change across steps is (currently) restricted:
-
-     1. There is no support for datasets that do not exist in all Iterations. The internal Iteration layouts should be homogeneous.
-        If you need this feature, please contact the openPMD-api developers; implementing this is currently not a priority.
-        Datasets that do not exist in all steps will be skipped at read time (with an error).
-
-        This restriction affects only datasets that contain array data.
-        Datasets defined exlusively in terms of attributes (especially: constant components) may be defined and read in a subselection of Iterations.
-     2. Datasets with changing extents are supported under the condition that they do not change their dimensionality.
+     In this mode, datasets with changing extents are supported under the condition that they do not change their dimensionality.
 
 * **Read/Write mode**: Creates a new Series if not existing, otherwise opens an existing Series for reading and writing.
   New datasets and iterations will be inserted as needed.

include/openPMD/IO/ADIOS/ADIOS2File.hpp

@@ -22,6 +22,7 @@
 
 #include "openPMD/IO/ADIOS/ADIOS2Auxiliary.hpp"
 #include "openPMD/IO/ADIOS/ADIOS2PreloadAttributes.hpp"
+#include "openPMD/IO/ADIOS/ADIOS2PreloadVariables.hpp"
 #include "openPMD/IO/AbstractIOHandler.hpp"
 #include "openPMD/IO/IOTask.hpp"
 #include "openPMD/IO/InvalidatableFile.hpp"
@@ -83,7 +84,8 @@ struct DatasetReader
         adios2::IO &IO,
         adios2::Engine &engine,
         std::string const &fileName,
-        std::optional<size_t> stepSelection);
+        std::optional<size_t> stepSelection,
+        detail::AdiosVariables const &av);
 
     static constexpr char const *errorMsg = "ADIOS2: readDataset()";
 };
@@ -335,11 +337,6 @@ class ADIOS2File
     std::vector<std::string>
     availableVariablesPrefixed(std::string const &prefix);
 
-    /*
-     * See description below.
-     */
-    void invalidateVariablesMap();
-
     void markActive(Writable *);
 
     // bool isActive(std::string const & path);
@@ -420,6 +417,14 @@ class ADIOS2File
     {
         return m_attributes;
     }
+    [[nodiscard]] detail::AdiosVariables const &variables() const
+    {
+        return m_variables;
+    }
+    [[nodiscard]] detail::AdiosVariables &variables()
+    {
+        return m_variables;
+    }
 
 private:
     ADIOS2IOHandlerImpl *m_impl;
@@ -448,7 +453,7 @@ class ADIOS2File
      * IO::Available(Attributes|Variables).
      */
     AdiosAttributes m_attributes;
-    std::optional<AttributeMap_t> m_availableVariables;
+    AdiosVariables m_variables;
 
     std::set<Writable *> m_pathsMarkedAsActive;
 

include/openPMD/IO/ADIOS/ADIOS2IOHandler.hpp

@@ -24,6 +24,7 @@
 #include "openPMD/IO/ADIOS/ADIOS2Auxiliary.hpp"
 #include "openPMD/IO/ADIOS/ADIOS2FilePosition.hpp"
 #include "openPMD/IO/ADIOS/ADIOS2PreloadAttributes.hpp"
+#include "openPMD/IO/ADIOS/ADIOS2PreloadVariables.hpp"
 #include "openPMD/IO/ADIOS/macros.hpp"
 #include "openPMD/IO/AbstractIOHandler.hpp"
 #include "openPMD/IO/AbstractIOHandlerImpl.hpp"
@@ -426,6 +427,74 @@ class ADIOS2IOHandlerImpl
 
     void dropFileData(InvalidatableFile const &file);
 
+    template <typename T>
+    static void setStepSelectionForVariable(
+        adios2::Variable<T> var,
+        std::string const &varName,
+        size_t step_selection,
+        size_t file_steps,
+        detail::AdiosVariables const &av)
+    {
+        auto var_steps = var.Steps();
+        if (var_steps == 1 && step_selection == 0)
+        {
+            // variable has no steps
+            return;
+        }
+        if (file_steps != var_steps)
+        {
+            if (!av.m_preparsed.has_value())
+            {
+                throw error::ReadError(
+                    error::AffectedObject::Dataset,
+                    error::Reason::UnexpectedContent,
+                    "ADIOS2",
+                    "The opened file contains different data per step, but "
+                    "variable data was not preparsed. ERROR: Variable " +
+                        varName + "' has " + std::to_string(var_steps) +
+                        " step(s), but the file has " +
+                        std::to_string(file_steps) + " step(s).");
+            }
+            auto preparsed = av.m_preparsed->m_partialVariables.find(varName);
+            if (preparsed == av.m_preparsed->m_partialVariables.end())
+            {
+                throw error::ReadError(
+                    error::AffectedObject::Dataset,
+                    error::Reason::UnexpectedContent,
+                    "ADIOS2",
+                    "The opened file contains different data per step, but "
+                    "variable data contains no preparsing info on '" +
+                        varName + "'. Has " + std::to_string(var_steps) +
+                        " step(s), but the file has " +
+                        std::to_string(file_steps) + " step(s).");
+            }
+            auto step_index = std::find(
+                preparsed->second.begin(),
+                preparsed->second.end(),
+                step_selection);
+            if (step_index == preparsed->second.end())
+            {
+                throw error::ReadError(
+                    error::AffectedObject::Dataset,
+                    error::Reason::UnexpectedContent,
+                    "ADIOS2",
+                    "Tried selecting global step " +
+                        std::to_string(step_selection) + " for variable '" +
+                        varName +
+                        "', but variable is not defined for that step (only "
+                        "for steps " +
+                        auxiliary::vec_as_string(preparsed->second) +
+                        "). Has " + std::to_string(var_steps) +
+                        " step(s), but the file has " +
+                        std::to_string(file_steps) + " step(s).");
+            }
+            // We need to replace the (global) step selection with the
+            // (local) step index
+            step_selection = step_index - preparsed->second.begin();
+        }
+        var.SetStepSelection({step_selection, 1});
+    }
+
     /*
      * Prepare a variable that already exists for an IO
      * operation, including:
@@ -439,8 +508,10 @@ class ADIOS2IOHandlerImpl
         Offset const &offset,
         Extent const &extent,
         adios2::IO &IO,
+        adios2::Engine &engine,
         std::string const &varName,
-        std::optional<size_t> stepSelection)
+        std::optional<size_t> stepSelection,
+        detail::AdiosVariables const &av)
     {
         {
             auto requiredType = adios2::GetType<T>();
@@ -469,7 +540,9 @@ class ADIOS2IOHandlerImpl
         }
         if (stepSelection.has_value())
         {
-            var.SetStepSelection({*stepSelection, 1});
+            auto file_steps = engine.Steps();
+            setStepSelectionForVariable(
+                var, varName, *stepSelection, file_steps, av);
         }
         // TODO leave this check to ADIOS?
         adios2::Dims shape = var.Shape();
@@ -615,7 +688,8 @@ namespace detail
             Parameter<Operation::OPEN_DATASET> &parameters,
             std::optional<size_t> stepSelection,
             std::vector<ADIOS2IOHandlerImpl::ParameterizedOperator> const
-                &operators);
+                &operators,
+            detail::AdiosVariables const &);
 
         static constexpr char const *errorMsg = "ADIOS2: openDataset()";
     };

include/openPMD/IO/ADIOS/ADIOS2PreloadVariables.hpp

@@ -0,0 +1,67 @@
+/* Copyright 2025 Franz Poeschel
+ *
+ * This file is part of openPMD-api.
+ *
+ * openPMD-api is free software: you can redistribute it and/or modify
+ * it under the terms of of either the GNU General Public License or
+ * the GNU Lesser General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * openPMD-api is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License and the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * and the GNU Lesser General Public License along with openPMD-api.
+ * If not, see <http://www.gnu.org/licenses/>.
+ */
+#pragma once
+
+#include "openPMD/auxiliary/Variant.hpp"
+#include "openPMD/config.hpp"
+#include <deque>
+#include <optional>
+#include <variant>
+#if openPMD_HAVE_ADIOS2
+
+#include <adios2.h>
+#include <map>
+
+namespace openPMD::detail
+{
+struct AdiosVariables
+{
+    // Buffered map for current step
+    using AttributeMap_t = std::map<std::string, adios2::Params>;
+    std::optional<AttributeMap_t> m_availableVariables;
+    // For which step were the above variables buffered?
+    size_t currentStep;
+    // Optimization: If variable definitions do not vary across steps, no need
+    // to recompute them
+    bool variables_are_static = false;
+
+    // Preparsed step data
+    struct RandomAccessPreparsed_t
+    {
+        // Variable only defined in these steps
+        std::map<std::string, std::vector<size_t>> m_partialVariables;
+
+        AttributeMap_t m_allVariables;
+    };
+    std::optional<RandomAccessPreparsed_t> m_preparsed;
+
+    /*
+     * If use_step_selection is false, but preparsed step data is available,
+     * this means that Advance(stepSelection = null) was executed previously.
+     * So, we can return m_preparsed->m_allVariables.
+     */
+    auto
+    availableVariables(size_t step, bool use_step_selection, adios2::IO &IO)
+        -> AttributeMap_t const &;
+};
+} // namespace openPMD::detail
+
+#endif // openPMD_HAVE_ADIOS2

src/IO/ADIOS/ADIOS2File.cpp

@@ -61,10 +61,17 @@ void DatasetReader::call(
     adios2::IO &IO,
     adios2::Engine &engine,
     std::string const &fileName,
-    std::optional<size_t> stepSelection)
+    std::optional<size_t> stepSelection,
+    detail::AdiosVariables const &av)
 {
     adios2::Variable<T> var = impl->verifyDataset<T>(
-        bp.param.offset, bp.param.extent, IO, bp.name, stepSelection);
+        bp.param.offset,
+        bp.param.extent,
+        IO,
+        engine,
+        bp.name,
+        stepSelection,
+        av);
     if (!var)
     {
         throw std::runtime_error(
@@ -91,15 +98,18 @@ void WriteDataset::call(ADIOS2File &ba, detail::BufferedPut &bp)
             if constexpr (std::is_same_v<ptr_type, std::shared_ptr<void const>>)
             {
                 auto ptr = static_cast<T const *>(arg.get());
+                auto &engine = ba.getEngine();
 
                 adios2::Variable<T> var = ba.m_impl->verifyDataset<T>(
                     bp.param.offset,
                     bp.param.extent,
                     ba.m_IO,
+                    engine,
                     bp.name,
-                    std::nullopt);
+                    std::nullopt,
+                    ba.variables());
 
-                ba.getEngine().Put(var, ptr);
+                engine.Put(var, ptr);
             }
             else if constexpr (std::is_same_v<
                                    ptr_type,
@@ -146,7 +156,8 @@ void BufferedGet::run(ADIOS2File &ba)
         ba.m_IO,
         ba.getEngine(),
         ba.m_file,
-        this->stepSelection);
+        this->stepSelection,
+        ba.variables());
 }
 
 void BufferedPut::run(ADIOS2File &ba)
@@ -160,13 +171,16 @@ struct RunUniquePtrPut
     static void call(BufferedUniquePtrPut &bufferedPut, ADIOS2File &ba)
     {
         auto ptr = static_cast<T const *>(bufferedPut.data.get());
+        auto &engine = ba.getEngine();
         adios2::Variable<T> var = ba.m_impl->verifyDataset<T>(
             bufferedPut.offset,
             bufferedPut.extent,
             ba.m_IO,
+            engine,
             bufferedPut.name,
-            std::nullopt);
-        ba.getEngine().Put(var, ptr);
+            std::nullopt,
+            ba.variables());
+        engine.Put(var, ptr);
     }
 
     static constexpr char const *errorMsg = "RunUniquePtrPut";
@@ -1336,7 +1350,6 @@ Be aware of the performance implications described above.)");
         case adios2::StepStatus::OtherError:
             throw std::runtime_error("[ADIOS2] Unexpected step status.");
         }
-        invalidateVariablesMap();
         m_pathsMarkedAsActive.clear();
         return res;
     }
@@ -1387,22 +1400,10 @@ ADIOS2File::availableVariablesPrefixed(std::string const &prefix)
         prefix, ADIOS2File::availableVariables());
 }
 
-void ADIOS2File::invalidateVariablesMap()
-{
-    m_availableVariables = std::optional<AttributeMap_t>();
-}
-
 ADIOS2File::AttributeMap_t const &ADIOS2File::availableVariables()
 {
-    if (m_availableVariables)
-    {
-        return m_availableVariables.value();
-    }
-    else
-    {
-        m_availableVariables = std::make_optional(m_IO.AvailableVariables());
-        return m_availableVariables.value();
-    }
+    return m_variables.availableVariables(
+        currentStep(), useStepSelection, m_IO);
 }
 
 void ADIOS2File::markActive(Writable *writable)