[TTree] Add option to import branches not present in first tree when merging TTrees

Fixes #14558
Fixes #12510

Author: ferdymercury

io/io/test/TFileMergerTests.cxx

@@ -394,3 +394,162 @@ TEST(TFileMerger, MergeSelectiveTutorial)
       EXPECT_NE(file.Get<TNtuple>("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<TFile> 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<TFile> 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<TFile> 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<TFile> 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);
+   }
+}

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="");

tree/tree/src/TTree.cxx

@@ -6924,21 +6924,71 @@ 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<TBranch *> importedCollection;
+      for (Int_t i = 0; i < nbranches; ++i) {
+         TBranch *nbranch = static_cast<TBranch *>(newbranches->UncheckedAt(i));
+         if (nbranch->TestBit(kDoNotProcess)) {
+            continue;
+         }
+         if (!GetListOfBranches()->FindObject(nbranch->GetName())) {
+            auto addbranch = static_cast<TBranch *>(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.
 ///
 /// Trees in the list can be memory or disk-resident trees.
 /// 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)
 {
    if (!li) return nullptr;
    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;