Fix: remap RootTreeWriter branches only once

The automatic branch substitution mechanism used to store MC labels should not be allowed to remap branches with >0 entries, otherwise only the last entry will be stored
AliceO2Group · Jan 15, 2021 · 867968b · 867968b
1 parent 65b4b1e
commit 867968b
Show file tree

Hide file tree

Showing 6 changed files with 23 additions and 12 deletions.
diff --git a/Detectors/ITSMFT/common/workflow/src/DigitWriterSpec.cxx b/Detectors/ITSMFT/common/workflow/src/DigitWriterSpec.cxx
@@ -70,7 +70,8 @@ DataProcessorSpec getDigitWriterSpec(bool mctruth, o2::header::DataOrigin detOri
     LOG(INFO) << "WRITING " << labels.getNElements() << " LABELS ";
 
     o2::dataformats::IOMCTruthContainerView outputcontainer;
-    auto br = framework::RootTreeWriter::remapBranch(branch, &outputcontainer);
+    auto ptr = &outputcontainer;
+    auto br = framework::RootTreeWriter::remapBranch(branch, &ptr);
     outputcontainer.adopt(labelbuffer);
     br->Fill();
     br->ResetAddress();

diff --git a/Detectors/TPC/workflow/src/RecoWorkflow.cxx b/Detectors/TPC/workflow/src/RecoWorkflow.cxx
@@ -323,7 +323,8 @@ framework::WorkflowSpec getWorkflow(CompletionPolicyData* policyData, std::vecto
   auto fillLabels = [](TBranch& branch, std::vector<char> const& labelbuffer, DataRef const& /*ref*/) {
     o2::dataformats::ConstMCTruthContainerView<o2::MCCompLabel> labels(labelbuffer);
     o2::dataformats::IOMCTruthContainerView outputcontainer;
-    auto br = framework::RootTreeWriter::remapBranch(branch, &outputcontainer);
+    auto ptr = &outputcontainer;
+    auto br = framework::RootTreeWriter::remapBranch(branch, &ptr);
     outputcontainer.adopt(labelbuffer);
     br->Fill();
     br->ResetAddress();

diff --git a/Detectors/TRD/workflow/src/TRDDigitWriterSpec.cxx b/Detectors/TRD/workflow/src/TRDDigitWriterSpec.cxx
@@ -57,7 +57,8 @@ o2::framework::DataProcessorSpec getTRDDigitWriterSpec(bool mctruth)
     // make the actual output object by adopting/casting the buffer
     // into a split format
     o2::dataformats::IOMCTruthContainerView outputcontainer(labeldata);
-    auto br = framework::RootTreeWriter::remapBranch(branch, &outputcontainer);
+    auto ptr = &outputcontainer;
+    auto br = framework::RootTreeWriter::remapBranch(branch, &ptr);
     br->Fill();
     br->ResetAddress();
   };

diff --git a/Framework/Utils/include/DPLUtils/RootTreeWriter.h b/Framework/Utils/include/DPLUtils/RootTreeWriter.h
@@ -347,14 +347,20 @@ class RootTreeWriter
   /// The function needs to be used with care. The user should ensure that "branch" is no longer used
   /// after a call to this function.
   template <typename T>
-  static TBranch* remapBranch(TBranch& branch, T* newdata)
+  static TBranch* remapBranch(TBranch& branchRef, T** newdata)
   {
-    auto name = branch.GetName();
-    auto branchleaves = branch.GetListOfLeaves();
-    auto tree = branch.GetTree();
-    branch.DropBaskets("all");
-    branch.DeleteBaskets("all");
-    tree->GetListOfBranches()->Remove(&branch);
+    auto tree = branchRef.GetTree();
+    auto name = branchRef.GetName();
+    auto branch = tree->GetBranch(name); // the input branch might actually no belong to the tree but to TreeWriter cache
+    assert(branch);
+    if (branch->GetEntries()) { // if it has entries, then it was already remapped/filled at prevous event
+      branch->SetAddress(newdata);
+      return branch;
+    }
+    auto branchleaves = branch->GetListOfLeaves();
+    branch->DropBaskets("all");
+    branch->DeleteBaskets("all");
+    tree->GetListOfBranches()->Remove(branch);
     for (auto entry : *branchleaves) {
       tree->GetListOfLeaves()->Remove(entry);
     }

diff --git a/Steer/DigitizerWorkflow/src/FDDDigitWriterSpec.h b/Steer/DigitizerWorkflow/src/FDDDigitWriterSpec.h
@@ -56,7 +56,8 @@ o2::framework::DataProcessorSpec getFDDDigitWriterSpec(bool mctruth = true)
     // make the actual output object by adopting/casting the buffer
     // into a split format
     o2::dataformats::IOMCTruthContainerView outputcontainer(labeldata);
-    auto br = framework::RootTreeWriter::remapBranch(branch, &outputcontainer);
+    auto ptr = &outputcontainer;
+    auto br = framework::RootTreeWriter::remapBranch(branch, &ptr);
     br->Fill();
     br->ResetAddress();
   };

diff --git a/Steer/DigitizerWorkflow/src/TPCDigitRootWriterSpec.cxx b/Steer/DigitizerWorkflow/src/TPCDigitRootWriterSpec.cxx
@@ -203,7 +203,8 @@ DataProcessorSpec getTPCDigitRootWriterSpec(std::vector<int> const& laneConfigur
     // first of all redefine the output format (special to labels)
     auto tree = branch.GetTree();
     auto sector = extractSector(ref);
-    auto br = framework::RootTreeWriter::remapBranch(branch, &outputcontainer);
+    auto ptr = &outputcontainer;
+    auto br = framework::RootTreeWriter::remapBranch(branch, &ptr);
 
     auto const* dh = DataRefUtils::getHeader<DataHeader*>(ref);
     LOG(INFO) << "HAVE LABEL DATA FOR SECTOR " << sector << " ON CHANNEL " << dh->subSpecification;