Merge branch 'branch_select' into 'master'

Allow branch selection when running mergePOOL utility

See merge request atlas/athena!23675

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();