From 6aa1f8901db0610786f178e0a9fee02e3cb2b506 Mon Sep 17 00:00:00 2001
From: Jack Cranshaw <cranshaw@anl.gov>
Date: Sun, 26 May 2019 11:10:01 +0000
Subject: [PATCH] Allow branch selection when running mergePOOL utility
---
.../AthenaPOOL/RootFileTools/src/merge.cxx | 65 +++++++++++++++++--
1 file changed, 59 insertions(+), 6 deletions(-)
diff --git a/Database/AthenaPOOL/RootFileTools/src/merge.cxx b/Database/AthenaPOOL/RootFileTools/src/merge.cxx
index 896b268fd36..3e81a53a0eb 100644
--- a/Database/AthenaPOOL/RootFileTools/src/merge.cxx
+++ b/Database/AthenaPOOL/RootFileTools/src/merge.cxx
@@ -15,6 +15,7 @@
#include "TTreeCloner.h"
#include "Compression.h"
#include <iostream>
+#include <fstream>
#include <iomanip>
// C++ includes
@@ -24,6 +25,7 @@
#include <set>
#include <stdexcept>
#include <sstream>
+#include <algorithm>
using namespace std;
@@ -211,7 +213,9 @@ public:
/// Close output file
DbStatus close();
/// Merge new input to existing output
- DbStatus merge(const std::string& fid, const std::set<std::string>& exclTrees);
+ DbStatus merge(const std::string& fid,
+ const std::set<std::string>& exclTrees,
+ const std::set<std::string>& branchList);
/// Dump collected database sections
void dumpSections();
/// Save new sections to the output file
@@ -456,7 +460,9 @@ void DbDatabaseMerger::dumpSections() {
/// Merge new input to existing output
-DbStatus DbDatabaseMerger::merge(const string& fid, const std::set<std::string>& exclTrees) {
+DbStatus DbDatabaseMerger::merge(const string& fid,
+ const std::set<std::string>& exclTrees,
+ const std::set<std::string>& branchList) {
if ( m_output ) {
TFile* source = TFile::Open(fid.c_str());
if ( source && !source->IsZombie() ) {
@@ -527,15 +533,40 @@ DbStatus DbDatabaseMerger::merge(const string& fid, const std::set<std::string>&
s0.length = 0;
s0.offset = 0;
+ if (std::string(src_tree->GetName())=="CollectionTree") {
+ // If no branches listed, then take all branches
+ if (!branchList.empty()) {
+ TObjArray* blist = src_tree->GetListOfBranches();
+ for (Int_t i=0; i<blist->GetEntriesFast(); i++) {
+ if (s_dbg) cout << "Processing leaf "
+ << blist->At(i)->ClassName() << endl;
+ // search list from file
+ if (std::find( branchList.begin(),
+ branchList.end(),
+ blist->At(i)->GetName()
+ ) != branchList.end()) {
+ src_tree->GetBranch(blist->At(i)->GetName())->SetStatus(1);
+ } else {
+ // if not in list, turn off that branch
+ src_tree->GetBranch(blist->At(i)->GetName())->SetStatus(0);
+ }
+ }
+ }
+ }
+
TTree *out_tree = (TTree*)m_output->Get(key->GetName());
if ( out_tree == 0 ) {
out_tree = src_tree->CloneTree(-1,"fast");
- if ( s_dbg ) cout << "+++ Created new Tree " << out_tree->GetName() << endl;
+ if ( s_dbg ) cout << "+++ Created new Tree "
+ << out_tree->GetName() << " with "
+ << out_tree->GetListOfBranches()->GetEntriesFast()
+ << " branches" << endl;
} else {
m_output->GetObject(key->GetName(),out_tree);
- s.start = (int) out_tree->GetEntries();
if (name == "##Params") {
if ( s_dbg ) cout << "+++ Slow merge for " << name << endl;
+ Long64_t out_entries = out_tree->GetEntries();
+ s.start = static_cast<int>(out_entries);
out_tree->CopyAddresses(src_tree);
for (Long64_t i=0; i<src_entries; i++) {
src_tree->GetEntry(i);
@@ -599,9 +630,10 @@ DbStatus DbDatabaseMerger::merge(const string& fid, const std::set<std::string>&
static int usage() {
cout << "POOL merge facility for ROOT tree based files.\n"
" Usage: \n"
- "mergePool -o <output-file> -i <input-file 1> [ -i <input-file 2> ...] [-e <exclude-tree 1>] -C LZ4/LZMA/ZLIB\n\n"
+ "mergePool -o <output-file> -i <input-file 1> [ -i <input-file 2> ...] [-e <exclude-tree 1>] -C LZ4/LZMA/ZLIB -B branchfile\n\n"
"input- and output files may consist of any legal file name.\n"
" -debug Switch debug flag on.\n"
+ " branchfile text file with single branchname per line"
<< endl;
return 2;
}
@@ -610,8 +642,10 @@ int main(int argc, char** argv) {
bool dbg = false;
string output;
string compalg("Input");
+ string branchfile("");
vector<string> input;
set<string> exclTrees;
+ set<string> branchList;
for(int i=1; i < argc; ++i) {
if ( *argv[i] == '-' ) {
switch(::toupper(*(argv[i]+1))) {
@@ -645,12 +679,19 @@ int main(int argc, char** argv) {
}
++i;
break;
+
+ case 'B':
+ if ( i+1 < argc ) branchfile = std::string(argv[i+1]);
+ ++i;
+ break;
default:
return usage();
}
}
}
+ // Check that arguments make sense and are complete
+ // Input/Output
if ( input.empty() ) {
cout << "No input files supplied" << endl;
return usage();
@@ -658,6 +699,18 @@ int main(int argc, char** argv) {
cout << "No output file supplied" << endl;
return usage();
}
+ // Branch selections
+ if (!branchfile.empty()) {
+ try {
+ std::ifstream file(branchfile);
+ std::string branchname;
+ while (std::getline(file,branchname)) {
+ branchList.insert(branchname);
+ }
+ } catch (...) {
+ cout << "Problem reading list of branches from " << branchfile << endl;
+ }
+ }
s_dbg = dbg;
gROOT->SetBatch(kTRUE);
DbDatabaseMerger m;
@@ -671,7 +724,7 @@ int main(int argc, char** argv) {
}
DbStatus ret = m.exists(output.c_str(), s_dbg) ? m.attach(output) : m.create(output);
if ( ret == SUCCESS ) {
- ret = m.merge(in, exclTrees);
+ ret = m.merge(in, exclTrees, branchList);
if ( ret == SUCCESS ) {
m.dumpSections();
if ( fixup ) m.createFID();
--
GitLab