merge.C

Go to the documentation of this file.

//========================================================================================================
// merge.C
// Helper macro for NICA-Scheduler partial output processing.
// mode: 0 - save resulting TChain pointing to partial files
//       1 - merge partial files to result file and delete partial files
//       2 - merge partial files to result file and keep partial files
//========================================================================================================
 
#if !defined(__CLING__) || defined(__MAKECLING__)
// ROOT includes
#include "TChain.h"
#include "TChainElement.h"
#include "TClass.h"
#include "TDirectory.h"
#include "TFile.h"
#include "TFileMerger.h"
#include "TKey.h"
#include "TObjArray.h"
#include "TSystem.h"
#include "TTree.h"
 
#include <cstdio>
#include <iostream>
#include <sstream>
#include <string>
#include <vector>
 
using namespace std;
#endif
 
// Split a whitespace-separated file list. The first item is the final result file,
// all following items are partial output files produced by parallel_mode.
vector<TString> split_file_list(const char* pcFileList)
{
    vector<TString> files;
    string token;
    stringstream ss(pcFileList == nullptr ? "" : pcFileList);
 
    while (ss >> token) {
        TString fileName(token.c_str());
        gSystem->ExpandPathName(fileName);
        files.push_back(fileName);
    }
 
    return files;
}
 
// Find the first TTree/TNtuple object stored in a ROOT directory.
// If the tree is inside a subdirectory, return the full path: "dir/tree".
TString find_tree_name_in_directory(TDirectory* directory, const TString& prefix = "")
{
    if (directory == nullptr || directory->GetListOfKeys() == nullptr)
        return "";
 
    TIter next(directory->GetListOfKeys());
    TKey* key = nullptr;
 
    while ((key = (TKey*)next()) != nullptr) {
        TClass* cl = TClass::GetClass(key->GetClassName());
        if (cl == nullptr)
            continue;
 
        TString objectName = key->GetName();
        TString fullName = (prefix.Length() == 0) ? objectName : prefix + "/" + objectName;
 
        if (cl->InheritsFrom(TTree::Class()))
            return fullName;
 
        if (cl->InheritsFrom(TDirectory::Class())) {
            TObject* obj = key->ReadObj();
            TDirectory* subdir = dynamic_cast<TDirectory*>(obj);
            TString treeName = find_tree_name_in_directory(subdir, fullName);
            if (treeName.Length() > 0)
                return treeName;
        }
    }
 
    return "";
}
 
// Open one ROOT file and return the first TTree name stored in it.
TString find_tree_name_in_file(const TString& fileName)
{
    TFile* file = TFile::Open(fileName, "READ");
    if (file == nullptr || file->IsZombie()) {
        cout << "ERROR: cannot open ROOT file \"" << fileName << "\" to detect TTree name" << endl;
        if (file != nullptr) {
            file->Close();
            delete file;
        }
        return "";
    }
 
    TString treeName = find_tree_name_in_directory(file);
 
    file->Close();
    delete file;
 
    return treeName;
}
 
// Detect the TTree name from partial output files. It checks all partial files
// because the first partial file may be missing or corrupted after a failed subtask.
TString detect_tree_name_from_partial_files(const vector<TString>& files)
{
    // files[0] is the result file name, partial files start from files[1]
    for (size_t i = 1; i < files.size(); ++i) {
        TString treeName = find_tree_name_in_file(files[i]);
        if (treeName.Length() > 0) {
            cout << "Detected TTree name for TChain: \"" << treeName << "\" from file \"" << files[i] << "\"" << endl;
            return treeName;
        }
 
        cout << "WARNING: no TTree object was found in partial file \"" << files[i] << "\"" << endl;
    }
 
    return "";
}
 
// Check that a partial file contains the detected TTree. This prevents creating
// an apparently valid, but actually incomplete or empty, TChain.
bool file_contains_tree(const TString& fileName, const TString& treeName)
{
    TFile* file = TFile::Open(fileName, "READ");
    if (file == nullptr || file->IsZombie()) {
        cout << "ERROR: cannot open ROOT file \"" << fileName << "\"" << endl;
        if (file != nullptr) {
            file->Close();
            delete file;
        }
        return false;
    }
 
    TObject* obj = file->Get(treeName);
    bool result = (obj != nullptr && obj->InheritsFrom(TTree::Class()));
 
    file->Close();
    delete file;
 
    return result;
}
 
// pcFileList - file list with space delimiter. The first file is the result file;
//              all following files are partial result files.
// mode: 0 - save only resulting TChain to the first file name with the data trees pointing to partial files
//       1 - merge to the first file name and delete partial files
//       2 - merge to the first file name and not delete partial files
void merge(const char* pcFileList, int mode = 1)
{
    vector<TString> files = split_file_list(pcFileList);
 
    if ((mode < 0) || (mode > 2)) {
        cout << "Second parameter ('mode') can be 0-2 only. The mode is assigned to a default value (1)." << endl;
        mode = 1;
    }
 
    if (files.empty()) {
        cout << "ERROR: empty file list, nothing has been merged" << endl;
        return;
    }
 
    if (files.size() == 1) {
        cout << "Only one file name in the input string, nothing has been merged" << endl;
        return;
    }
 
    TString sResultFile = files[0];
 
    // MERGE FILES: mode 1 = merge and delete partials, mode 2 = merge and keep partials
    if (mode > 0) {
        Bool_t res;
        TString sFirstFile = files[1];
 
        bool isError = false;
        TFileMerger mergerFull,   // full merge with the first partial file to keep metadata objects
            mergerPart;           // partial merge for other files
 
        // The first partial file is merged fully later, because it usually contains
        // metadata objects such as FairBaseParSet. Other partial files are merged first.
        mergerFull.AddFile(sFirstFile.Data());
        for (size_t i = 2; i < files.size(); ++i)
            mergerPart.AddFile(files[i].Data());
 
        // Merge only the data tree from the other partial files (kOnlyListed requires the object
        // names to be set explicitly); metadata objects are taken from the first partial file.
        // If no TTree was found, fall back to full merging of all objects.
        Int_t mergeType = TFileMerger::kAll | TFileMerger::kIncremental;
        TString treeName = detect_tree_name_from_partial_files(files);
        if (treeName.Length() > 0) {
            // kOnlyListed matches plain key names at any directory level, so strip a possible
            // subdirectory path from the detected tree name
            Ssiz_t slash = treeName.Last('/');
            if (slash != kNPOS)
                treeName.Remove(0, slash + 1);
            mergerPart.AddObjectNames(treeName.Data());
            mergeType |= TFileMerger::kOnlyListed;
        } else
            cout << "WARNING: no TTree was found in partial output files, merging all their objects" << endl;
 
        mergerPart.OutputFile(sResultFile);
        res = mergerPart.PartialMerge(mergeType);
 
        if (res == kFALSE) {
            cout << "ERROR: partial merging to the result file failed" << endl;
            mergerPart.Reset();
            mergerFull.Reset();
            return;
        }
 
        mergerPart.Reset();
 
        mergerFull.AddFile(sResultFile.Data());
        mergerFull.OutputFile(sResultFile);
        res = mergerFull.Merge();
 
        if (res == kFALSE) {
            cout << "ERROR: full merging to the result file failed" << endl;
            mergerFull.Reset();
            return;
        }
 
        mergerFull.Reset();
 
        // Delete partial files only after both merge steps have succeeded.
        if (mode == 1) {
            for (size_t i = 1; i < files.size(); ++i) {
                if (remove(files[i].Data()) != 0) {
                    isError = true;
                    cout << "ERROR: deleting temporary file " << files[i] << " failed" << endl;
                }
            }
 
            if (isError)
                cout << "There were errors while temporary files removing" << endl;
            else
                cout << endl << "Temporary files were successfully removed" << endl;
        } else
            cout << "Partial result files were merged and preserved" << endl;
    }
 
    // SAVE ONLY TCHAIN: mode 0
    else {
        TString treeName = detect_tree_name_from_partial_files(files);
        if (treeName.Length() == 0) {
            cout << "ERROR: no TTree was found in partial output files. "
                 << "TChain cannot be created. Check that merge=\"chain\" is used only for ROOT files containing TTree "
                    "objects."
                 << endl;
            return;
        }
 
        // Verify that all partial files contain the same tree before writing the TChain.
        for (size_t i = 1; i < files.size(); ++i) {
            if (!file_contains_tree(files[i], treeName)) {
                cout << "ERROR: partial file \"" << files[i] << "\" does not contain TTree \"" << treeName
                     << "\". TChain was not created." << endl;
                return;
            }
        }
 
        TChain chainUnion(treeName.Data());
 
        for (size_t i = 1; i < files.size(); ++i)
            chainUnion.Add(files[i].Data());
 
        TFile fChain(sResultFile, "RECREATE");
        if (fChain.IsZombie()) {
            cout << "ERROR: result TChain file \"" << sResultFile << "\" was not created" << endl;
            return;
        }
 
        chainUnion.Write();
        fChain.Close();
 
        Long64_t events = chainUnion.GetEntries();
        cout << "The Chain with " << events << " event(s) was written to file \"" << sResultFile
             << "\" using TTree name \"" << treeName << "\" and pointing to following files:" << endl;
 
        TObjArray* fileElements = chainUnion.GetListOfFiles();
        TIter next(fileElements);
        TChainElement* chEl = nullptr;
 
        while ((chEl = (TChainElement*)next()) != nullptr) {
            const char* pc = chEl->GetTitle();
            cout << pc << endl;
        }
    }
}