Fix VarShape() for Joined Arrays (#4246)

* Fix VarShape() for Joined Arrays * Reset join state
ornladios · Jul 18, 2024 · 0e0489f · 0e0489f
1 parent c424740
commit 0e0489f
Showing 1 changed file with 56 additions and 15 deletions.
diff --git a/source/adios2/toolkit/format/bp5/BP5Deserializer.cpp b/source/adios2/toolkit/format/bp5/BP5Deserializer.cpp
@@ -659,13 +659,16 @@ void BP5Deserializer::SetupForStep(size_t Step, size_t WriterCount)
             m_Engine->m_IO.RemoveDerivedVariable(RecPair.second->VarName);
 #endif
             RecPair.second->Variable = NULL;
-            if (RecPair.second->OrigShapeID == ShapeID::JoinedArray)
-            {
-                auto VarRec = RecPair.second;
-                VarRec->JoinedDimen = SIZE_MAX;
-                VarRec->LastJoinedOffset = NULL;
-                VarRec->LastJoinedShape = NULL;
-            }
+        }
+    }
+    for (auto RecPair : VarByKey)
+    {
+        if (RecPair.second->OrigShapeID == ShapeID::JoinedArray)
+        {
+            auto VarRec = RecPair.second;
+            VarRec->JoinedDimen = SIZE_MAX;
+            VarRec->LastJoinedOffset = NULL;
+            VarRec->LastJoinedShape = NULL;
         }
     }
     m_CurrentWriterCohortSize = WriterCount;
@@ -2623,19 +2626,57 @@ bool BP5Deserializer::VarShape(const VariableBase &Var, const size_t RelStep, Di
             AbsStep = VarRec->AbsStepFromRel[RelStep];
         }
     }
-    for (size_t WriterRank = 0; WriterRank < WriterCohortSize(AbsStep); WriterRank++)
+    if (VarRec->OrigShapeID == ShapeID::GlobalArray)
     {
-        MetaArrayRec *writer_meta_base =
-            (MetaArrayRec *)GetMetadataBase(VarRec, AbsStep, WriterRank);
-        if (writer_meta_base && writer_meta_base->Shape)
+        // Take the first shape in any writer that wrote it, they should all be the same
+        for (size_t WriterRank = 0; WriterRank < WriterCohortSize(AbsStep); WriterRank++)
         {
-            Shape.resize(writer_meta_base->Dims);
-            for (size_t i = 0; i < writer_meta_base->Dims; i++)
+            MetaArrayRec *writer_meta_base =
+                (MetaArrayRec *)GetMetadataBase(VarRec, AbsStep, WriterRank);
+            if (writer_meta_base && writer_meta_base->Shape)
             {
-                Shape[i] = writer_meta_base->Shape[i];
+                Shape.resize(writer_meta_base->Dims);
+                for (size_t i = 0; i < writer_meta_base->Dims; i++)
+                {
+                    Shape[i] = writer_meta_base->Shape[i];
+                }
+                return true;
+            }
+        }
+    }
+    else
+    {
+        // Joined array case, first Shape gives us structure, then add joined dimensions
+        bool first = true;
+        for (size_t WriterRank = 0; WriterRank < WriterCohortSize(AbsStep); WriterRank++)
+        {
+            MetaArrayRec *writer_meta_base =
+                (MetaArrayRec *)GetMetadataBase(VarRec, AbsStep, WriterRank);
+            if (writer_meta_base && writer_meta_base->Shape)
+            {
+                for (size_t Block = 0; Block < writer_meta_base->BlockCount; Block++)
+                {
+                    if (first)
+                    {
+                        Shape.resize(writer_meta_base->Dims);
+                        for (size_t i = 0; i < writer_meta_base->Dims; i++)
+                        {
+                            Shape[i] = writer_meta_base->Shape[i];
+                        }
+                        Shape[VarRec->JoinedDimen] =
+                            writer_meta_base->Count[VarRec->JoinedDimen + VarRec->DimCount * Block];
+                        first = false;
+                    }
+                    else
+                    {
+                        Shape[VarRec->JoinedDimen] +=
+                            writer_meta_base->Count[VarRec->JoinedDimen + VarRec->DimCount * Block];
+                    }
+                }
             }
-            return true;
         }
+        if (!first)
+            return true;
     }
     return false;
 }