Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
159 changes: 159 additions & 0 deletions io/io/test/TFileMergerTests.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}
1 change: 1 addition & 0 deletions tree/tree/inc/TTree.h
Original file line number Diff line number Diff line change
Expand Up @@ -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="");
Expand Down
68 changes: 65 additions & 3 deletions tree/tree/src/TTree.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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)
{
Expand Down Expand Up @@ -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).
Expand All @@ -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;
Expand All @@ -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.

Expand Down Expand Up @@ -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).
Expand All @@ -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;
Expand Down
Loading