Move joinedDim logic into middle-end for extendDataset

Author: franzpoeschel

include/openPMD/Dataset.hpp

@@ -86,5 +86,6 @@ class Dataset
     bool empty() const;
 
     std::optional<size_t> joinedDimension() const;
+    static std::optional<size_t> joinedDimension(Extent const &);
 };
 } // namespace openPMD

include/openPMD/IO/IOTask.hpp

@@ -405,7 +405,15 @@ template <>
 struct OPENPMDAPI_EXPORT Parameter<Operation::EXTEND_DATASET>
     : public AbstractParameter
 {
-    Parameter() = default;
+    Parameter(Extent e) : joinedDimension(Dataset::joinedDimension(e))
+    {
+        this->extent = std::move(e);
+    }
+
+    // default constructor, but callsites need to explicitly acknowledge that
+    // joined dimensions will not be automatically configured when using it
+    Parameter(I_dont_want_to_use_joined_dimensions_t)
+    {}
     Parameter(Parameter &&) = default;
     Parameter(Parameter const &) = default;
     Parameter &operator=(Parameter &&) = default;
@@ -418,6 +426,7 @@ struct OPENPMDAPI_EXPORT Parameter<Operation::EXTEND_DATASET>
     }
 
     Extent extent = {};
+    std::optional<size_t> joinedDimension;
 };
 
 template <>

src/Dataset.cpp

@@ -68,6 +68,11 @@ bool Dataset::empty() const
 }
 
 std::optional<size_t> Dataset::joinedDimension() const
+{
+    return joinedDimension(extent);
+}
+
+std::optional<size_t> Dataset::joinedDimension(Extent const &extent)
 {
     std::optional<size_t> res;
     for (size_t i = 0; i < extent.size(); ++i)

src/IO/ADIOS/ADIOS2IOHandler.cpp

@@ -873,7 +873,10 @@ namespace detail
     {
         template <typename T, typename... Args>
         static void call(
-            adios2::IO &IO, std::string const &variable, Extent const &newShape)
+            adios2::IO &IO,
+            std::string const &variable,
+            Extent const &newShape,
+            std::optional<size_t> const &newJoinedDim)
         {
             auto var = IO.InquireVariable<T>(variable);
             if (!var)
@@ -888,54 +891,60 @@ namespace detail
             {
                 dims.push_back(ext);
             }
-            auto joinedDim = joinedDimension(var.Shape());
+            auto oldJoinedDim = joinedDimension(var.Shape());
             auto make_runtime_error = [&](char const *message) {
                 std::stringstream s;
                 s << "[ADIOS2IOHandlerImpl::extendDataset()] " << message
                   << "\nNote: Variable '" << variable << "' has old shape ";
                 auxiliary::write_vec_to_stream(s, var.Shape());
-                if (joinedDim.has_value())
+                if (oldJoinedDim.has_value())
                 {
-                    s << " (joined dimension on index " << *joinedDim << ").";
+                    s << " (joined dimension on index " << *oldJoinedDim << ")";
                 }
                 else
                 {
                     s << " (no joined dimension)";
                 }
                 s << " and is extended to new shape ";
-                auxiliary::write_vec_to_stream(s, newShape) << ".";
+                auxiliary::write_vec_to_stream(s, newShape);
+                if (newJoinedDim.has_value())
+                {
+                    s << " (joined dimension on index " << *newJoinedDim
+                      << ").";
+                }
+                else
+                {
+                    s << " (no joined dimension).";
+                }
                 return std::runtime_error(s.str());
             };
-            if (joinedDim.has_value() ||
+            if (oldJoinedDim.has_value() ||
                 var.ShapeID() == adios2::ShapeID::JoinedArray)
             {
-                if (!joinedDim.has_value())
+                if (!oldJoinedDim.has_value())
                 {
                     throw make_runtime_error(
                         "Inconsistent state of variable: Has shape ID "
                         "JoinedArray, but its shape contains no value "
                         "adios2::JoinedDim.");
                 }
-                if (newShape.at(*joinedDim) != Dataset::JOINED_DIMENSION)
+                if (newJoinedDim != oldJoinedDim)
                 {
                     throw make_runtime_error(
                         "Variable was previously configured with a joined "
                         "dimension, so the new dataset extent must keep the "
                         "joined dimension on that index.");
                 }
-                dims[*joinedDim] = adios2::JoinedDim;
+                dims[*newJoinedDim] = adios2::JoinedDim;
             }
             else
             {
-                for (auto s : newShape)
+                if (newJoinedDim.has_value())
                 {
-                    if (s == Dataset::JOINED_DIMENSION)
-                    {
-                        throw make_runtime_error(
-                            "Variable was not previously configured with a "
-                            "joined dimension, but is now requested to change "
-                            "extent to a joined array.");
-                    }
+                    throw make_runtime_error(
+                        "Variable was not previously configured with a "
+                        "joined dimension, but is now requested to change "
+                        "extent to a joined array.");
                 }
             }
 
@@ -958,7 +967,7 @@ void ADIOS2IOHandlerImpl::extendDataset(
     auto &filedata = getFileData(file, IfFileNotOpen::ThrowError);
     Datatype dt = detail::fromADIOS2Type(filedata.m_IO.VariableType(name));
     switchAdios2VariableType<detail::DatasetExtender>(
-        dt, filedata.m_IO, name, parameters.extent);
+        dt, filedata.m_IO, name, parameters.extent, parameters.joinedDimension);
 }
 
 void ADIOS2IOHandlerImpl::openFile(

src/IO/HDF5/HDF5IOHandler.cpp

@@ -840,6 +840,12 @@ void HDF5IOHandlerImpl::extendDataset(
         throw std::runtime_error(
             "[HDF5] Extending an unwritten Dataset is not possible.");
 
+    if (parameters.joinedDimension.has_value())
+    {
+        error::throwOperationUnsupportedInBackend(
+            "HDF5", "Joined Arrays currently only supported in ADIOS2");
+    }
+
     auto res = getFile(writable);
     if (!res)
         res = getFile(writable->parent);

src/IO/JSON/JSONIOHandlerImpl.cpp

@@ -711,6 +711,13 @@ void JSONIOHandlerImpl::extendDataset(
     VERIFY_ALWAYS(
         access::write(m_handler->m_backendAccess),
         "[JSON] Cannot extend a dataset in read-only mode.")
+
+    if (parameters.joinedDimension.has_value())
+    {
+        error::throwOperationUnsupportedInBackend(
+            "JSON", "Joined Arrays currently only supported in ADIOS2");
+    }
+
     setAndGetFilePosition(writable);
     refreshFileFromParent(writable);
     auto &j = obtainJsonContents(writable);

src/RecordComponent.cpp

@@ -339,8 +339,8 @@ void RecordComponent::flush(
             }
             else
             {
-                Parameter<Operation::EXTEND_DATASET> pExtend;
-                pExtend.extent = rc.m_dataset.value().extent;
+                Parameter<Operation::EXTEND_DATASET> pExtend(
+                    rc.m_dataset.value().extent);
                 IOHandler()->enqueue(IOTask(this, std::move(pExtend)));
                 rc.m_hasBeenExtended = false;
             }

src/binding/python/Dataset.cpp

@@ -79,7 +79,8 @@ void init_Dataset(py::module &m)
                 })
 
             .def_property_readonly(
-                "joined_dimension", &Dataset::joinedDimension)
+                "joined_dimension",
+                py::overload_cast<>(&Dataset::joinedDimension, py::const_))
             .def_readonly("extent", &Dataset::extent)
             .def("extend", &Dataset::extend)
             .def_readonly("rank", &Dataset::rank)