diff --git a/io/io/test/TFileMergerTests.cxx b/io/io/test/TFileMergerTests.cxx index 3d790f80f13a2..57e4f01af22d8 100644 --- a/io/io/test/TFileMergerTests.cxx +++ b/io/io/test/TFileMergerTests.cxx @@ -394,3 +394,162 @@ TEST(TFileMerger, MergeSelectiveTutorial) EXPECT_NE(file.Get("ntuple"), nullptr); } } + +// https://github.com/root-project/root/issues/14558 +TEST(TFileMerger, ImportBranches) +{ + { + // Case 1 - Static: ZeroBranches + 1 entry (1 branch) + 1 entry (2 branch) + TTree atree("atree", "atitle"); + int value; + atree.Branch("a", &value); + value = 11; + atree.Fill(); + TTree abtree("abtree", "abtitle"); + abtree.Branch("a", &value); + abtree.Branch("b", &value); + value = 42; + abtree.Fill(); + + TTree dummy("ztree", "zeroBranches"); + TList treelist; + treelist.Add(&dummy); + treelist.Add(&atree); + treelist.Add(&abtree); + std::unique_ptr file1(TFile::Open("b4716.root", "RECREATE")); + auto rtree = TTree::MergeTrees(&treelist, "ImportBranches"); + file1->Write(); + ASSERT_TRUE(rtree->FindBranch("a")); + EXPECT_EQ(rtree->FindBranch("a")->GetEntries(), 2); + ASSERT_TRUE(rtree->FindBranch("b")); + EXPECT_EQ(rtree->FindBranch("b")->GetEntries(), 2); + ASSERT_TRUE(atree.FindBranch("a")); + ASSERT_FALSE(atree.FindBranch("b")); + ASSERT_TRUE(abtree.FindBranch("a")); + ASSERT_TRUE(abtree.FindBranch("b")); + EXPECT_EQ(atree.FindBranch("a")->GetEntries(), 1); + EXPECT_EQ(abtree.FindBranch("a")->GetEntries(), 1); + EXPECT_EQ(abtree.FindBranch("b")->GetEntries(), 1); + EXPECT_EQ(dummy.FindBranch("a"), nullptr); + EXPECT_EQ(dummy.FindBranch("b"), nullptr); + } + { + // Case 2 - this (ZeroBranches) + 1 entry (1 branch) + 1 entry (2 branch) + TTree atree("atree", "atitle"); + int value; + atree.Branch("a", &value); + value = 11; + atree.Fill(); + TTree abtree("abtree", "abtitle"); + abtree.Branch("a", &value); + abtree.Branch("b", &value); + value = 42; + abtree.Fill(); + + TTree dummy("ztree", "zeroBranches"); + TList treelist; + treelist.Add(&atree); + treelist.Add(&abtree); + std::unique_ptr file2(TFile::Open("c4716.root", "RECREATE")); + TFileMergeInfo info2(file2.get()); + info2.fOptions += " ImportBranches"; + dummy.Merge(&treelist, &info2); + file2->Write(); + ASSERT_TRUE(dummy.FindBranch("a")); + EXPECT_EQ(dummy.FindBranch("a")->GetEntries(), 2); + ASSERT_TRUE(dummy.FindBranch("b")); + EXPECT_EQ(dummy.FindBranch("b")->GetEntries(), 2); + EXPECT_EQ(TString(atree.GetName()), + "ztree"); // As a side effect of trees with zero branches being ignored but it's name kept since it's first + // in list, the first tree with branches (name) gets ztree's name. + ASSERT_TRUE(atree.FindBranch("a")); + EXPECT_EQ(atree.FindBranch("a")->GetEntries(), 2); // As a side effect, the first with branches (atree) has been + // modified and has now 2 entries instead of 1, and ztree's name + ASSERT_TRUE(atree.FindBranch("b")); + EXPECT_EQ(atree.FindBranch("b")->GetEntries(), 2); // As a side effect, the first with branches (atree) has been + // modified and has now 2 entries instead of 1, and ztree's name + ASSERT_TRUE(abtree.FindBranch("a")); + ASSERT_TRUE(abtree.FindBranch("b")); + EXPECT_EQ(abtree.FindBranch("a")->GetEntries(), 1); + EXPECT_EQ(abtree.FindBranch("b")->GetEntries(), 1); + } + { + // Case 3 - this (0 entry / 1 branch) + 1 entry (1 branch) + 1 entry (2 branch) + TTree atree("atree", "atitle"); + int value; + atree.Branch("a", &value); + value = 11; + atree.Fill(); + TTree abtree("abtree", "abtitle"); + abtree.Branch("a", &value); + abtree.Branch("b", &value); + value = 42; + abtree.Fill(); + + TList treelist; + treelist.Add(&atree); + treelist.Add(&abtree); + TTree a0tree("a0tree", "a0title"); + a0tree.Branch("a", &value); + std::unique_ptr file3(TFile::Open("d4716.root", "RECREATE")); + TFileMergeInfo info3(file3.get()); + info3.fOptions += " ImportBranches"; + a0tree.Merge(&treelist, &info3); + file3->Write(); + ASSERT_TRUE(a0tree.FindBranch("a")); + EXPECT_EQ(a0tree.FindBranch("a")->GetEntries(), 2); + ASSERT_TRUE(a0tree.FindBranch("b")); + EXPECT_EQ(a0tree.FindBranch("b")->GetEntries(), 2); + ASSERT_TRUE(atree.FindBranch("a")); + ASSERT_FALSE(atree.FindBranch("b")); + ASSERT_TRUE(abtree.FindBranch("a")); + ASSERT_TRUE(abtree.FindBranch("b")); + EXPECT_EQ(atree.FindBranch("a")->GetEntries(), 1); + EXPECT_EQ(abtree.FindBranch("a")->GetEntries(), 1); + EXPECT_EQ(abtree.FindBranch("b")->GetEntries(), 1); + } + { + // Case 4 - this 1 entry (3 branch) + 1 entry (1 branch) + (0 entry / 1 branch) + TTree abctree("abctree", "abctitle"); + int value; + abctree.Branch("a", &value); + abctree.Branch("b", &value); + abctree.Branch("c", &value); + value = 11; + abctree.Fill(); + TTree ctree("ctree", "ctitle"); + ctree.Branch("c", &value); + value = 42; + ctree.Fill(); + TTree c0tree("c0tree", "c0title"); + c0tree.Branch("c", &value); + + std::unique_ptr file4(TFile::Open("e4716.root", "RECREATE")); + TFileMergeInfo info4(file4.get()); + info4.fOptions += " ImportBranches"; + TList treelist; + treelist.Add(&ctree); + treelist.Add(&c0tree); + ROOT::TestSupport::CheckDiagsRAII diagRAII; // ctree and c0tree don't have a/b branch so warn since they will be auto-filled + diagRAII.requiredDiag(kWarning, "TTree::CopyAddresses", "Could not find branch named 'a' in tree named 'ctree'"); + diagRAII.requiredDiag(kWarning, "TTree::CopyAddresses", "Could not find branch named 'b' in tree named 'ctree'"); + diagRAII.requiredDiag(kWarning, "TTree::CopyAddresses", "Could not find branch named 'a' in tree named 'c0tree'"); + diagRAII.requiredDiag(kWarning, "TTree::CopyAddresses", "Could not find branch named 'b' in tree named 'c0tree'"); + abctree.Merge(&treelist, &info4); + file4->Write(); + ASSERT_TRUE(abctree.FindBranch("a")); + ASSERT_TRUE(abctree.FindBranch("b")); + ASSERT_TRUE(abctree.FindBranch("c")); + EXPECT_EQ(abctree.FindBranch("a")->GetEntries(), 2); + EXPECT_EQ(abctree.FindBranch("b")->GetEntries(), 2); + EXPECT_EQ(abctree.FindBranch("c")->GetEntries(), 2); + ASSERT_FALSE(ctree.FindBranch("a")); + ASSERT_FALSE(ctree.FindBranch("b")); + ASSERT_TRUE(ctree.FindBranch("c")); + EXPECT_EQ(ctree.FindBranch("c")->GetEntries(), 1); + ASSERT_FALSE(c0tree.FindBranch("a")); + ASSERT_FALSE(c0tree.FindBranch("b")); + ASSERT_TRUE(c0tree.FindBranch("c")); + EXPECT_EQ(c0tree.FindBranch("c")->GetEntries(), 0); + } +} diff --git a/tree/tree/inc/TTree.h b/tree/tree/inc/TTree.h index 7609529df00fe..10c4b282c1a95 100644 --- a/tree/tree/inc/TTree.h +++ b/tree/tree/inc/TTree.h @@ -635,6 +635,7 @@ class TTree : public TNamed, public TAttLine, public TAttFill, public TAttMarker bool MemoryFull(Int_t nbytes); virtual Long64_t Merge(TCollection* list, Option_t* option = ""); virtual Long64_t Merge(TCollection* list, TFileMergeInfo *info); + bool ImportBranches(TTree *tree); static TTree *MergeTrees(TList* list, Option_t* option = ""); bool Notify() override; virtual void OptimizeBaskets(ULong64_t maxMemory=10000000, Float_t minComp=1.1, Option_t *option=""); diff --git a/tree/tree/src/TTree.cxx b/tree/tree/src/TTree.cxx index 3f100ace22eb5..2e7fee085af02 100644 --- a/tree/tree/src/TTree.cxx +++ b/tree/tree/src/TTree.cxx @@ -6924,6 +6924,53 @@ bool TTree::MemoryFull(Int_t nbytes) return true; } +//////////////////////////////////////////////////////////////////////////////// +/// Function merging branch structure from an outside tree into the current one +/// +/// If the same branch name already exists in "this", it's skipped, only new ones +/// are copied. +/// Entries are not copied, just branch name / type is cloned +/// Branches marked as DoNotProcess are not merged +/// If this tree had some entries already in other branches and the new tree incorporates a new branch, +/// when importing the branch, we backfill the branch value with default values until GetEntries is reached, +/// to prevent misalignments in the TTree structure between branches / entries. +/// @param tree the outside TTree whose branches will be copied into this tree +/// @return boolean true on sucess, false otherwise + +bool TTree::ImportBranches(TTree *tree) +{ + if (tree) { + TObjArray *newbranches = tree->GetListOfBranches(); + Int_t nbranches = newbranches->GetEntriesFast(); + const Long64_t nentries = GetEntries(); + std::vector importedCollection; + for (Int_t i = 0; i < nbranches; ++i) { + TBranch *nbranch = static_cast(newbranches->UncheckedAt(i)); + if (nbranch->TestBit(kDoNotProcess)) { + continue; + } + if (!GetListOfBranches()->FindObject(nbranch->GetName())) { + auto addbranch = static_cast(nbranch->Clone()); + addbranch->ResetAddress(); + addbranch->Reset(); + addbranch->SetTree(this); + fBranches.Add(addbranch); + importedCollection.push_back(addbranch); + } + } + // Backfill mechanism to realign with TTree + if (!importedCollection.empty()) { + for (Long64_t e = 0; e < nentries; ++e) { + for (auto branch : importedCollection) { + branch->BackFill(); + } + } + } + return true; + } + return false; +} + //////////////////////////////////////////////////////////////////////////////// /// Static function merging the trees in the TList into a new tree. /// @@ -6931,6 +6978,9 @@ bool TTree::MemoryFull(Int_t nbytes) /// The new tree is created in the current directory (memory if gROOT). /// Trees with no branches will be skipped, the branch structure /// will be taken from the first non-zero-branch Tree of {li} +/// Use "ImportBranches" option to incorporate branches from the (filled) trees in +/// the list that were not in the first TTree into the final result, backfilling +/// with default values to prevent misalignment TTree* TTree::MergeTrees(TList* li, Option_t* options) { @@ -6938,7 +6988,7 @@ TTree* TTree::MergeTrees(TList* li, Option_t* options) TIter next(li); TTree *newtree = nullptr; TObject *obj; - + const bool importBranches = TString(options).Contains("ImportBranches", TString::ECaseCompare::kIgnoreCase); while ((obj=next())) { if (!obj->InheritsFrom(TTree::Class())) continue; TTree *tree = (TTree*)obj; @@ -6963,6 +7013,8 @@ TTree* TTree::MergeTrees(TList* li, Option_t* options) tree->ResetBranchAddresses(); newtree->ResetBranchAddresses(); continue; + } else if (importBranches) { + newtree->ImportBranches(tree); } if (nentries == 0) continue; @@ -6980,6 +7032,9 @@ TTree* TTree::MergeTrees(TList* li, Option_t* options) /// Returns the total number of entries in the merged tree. /// Trees with no branches will be skipped, the branch structure /// will be taken from the first non-zero-branch Tree of {this+li} +/// Use "ImportBranches" option to incorporate branches from the (filled) trees in +/// the list that were not in this TTree into the final result, backfilling +/// with default values to prevent misalignment Long64_t TTree::Merge(TCollection* li, Option_t *options) { @@ -7016,6 +7071,7 @@ Long64_t TTree::Merge(TCollection* li, Option_t *options) return 0; // All trees have empty branches } if (!li) return 0; + const bool importBranches = TString(options).Contains("ImportBranches", TString::ECaseCompare::kIgnoreCase); Long64_t storeAutoSave = fAutoSave; // Disable the autosave as the TFileMerge keeps a list of key and deleting the underlying // key would invalidate its iteration (or require costly measure to not use the deleted keys). @@ -7034,7 +7090,8 @@ Long64_t TTree::Merge(TCollection* li, Option_t *options) Long64_t nentries = tree->GetEntries(); if (nentries == 0) continue; - + if (importBranches) + ImportBranches(tree); CopyEntries(tree, -1, options, true); } fAutoSave = storeAutoSave; @@ -7049,6 +7106,9 @@ Long64_t TTree::Merge(TCollection* li, Option_t *options) /// use for further merging). /// Trees with no branches will be skipped, the branch structure /// will be taken from the first non-zero-branch Tree of {this+li} +/// Use "ImportBranches" info-option to incorporate branches from the (filled) trees +/// in the list that were not in this TTree into the final result, backfilling +/// with default values to prevent misalignment /// /// Returns the total number of entries in the merged tree. @@ -7111,6 +7171,7 @@ Long64_t TTree::Merge(TCollection* li, TFileMergeInfo *info) } } if (!li) return 0; + const bool importBranches = TString(options).Contains("ImportBranches", TString::ECaseCompare::kIgnoreCase); Long64_t storeAutoSave = fAutoSave; // Disable the autosave as the TFileMerge keeps a list of key and deleting the underlying // key would invalidate its iteration (or require costly measure to not use the deleted keys). @@ -7126,7 +7187,8 @@ Long64_t TTree::Merge(TCollection* li, TFileMergeInfo *info) fAutoSave = storeAutoSave; return -1; } - + if (importBranches) + ImportBranches(tree); CopyEntries(tree, -1, options, true); } fAutoSave = storeAutoSave;