calculate_size.C

Go to the documentation of this file.

// Calculate the disk space taken by the event tree of a BmnRoot ROOT file and by each of its branches.
//
// The macro opens the "bmndata" tree of a simulation, digi or DST file and prints the size on disk of
// the whole tree and of every top-level branch (including all its sub-branches), e.g. to estimate how
// much a new branch or a change of a data format costs in the output files.
//
// The size of a branch is the compressed size of its baskets (the event data) plus the compressed size
// of its metadata (the TBranch object with its sub-branches and leaves, which is streamed into a memory
// file with the compression settings of the input file). For an uncompressed branch the uncompressed
// size of its baskets is taken. The size of the tree is the size of its header stored in the file
// (it contains the metadata of all branches) plus the baskets of all branches, so the sizes of the
// branches add up approximately to the size of the tree.
// The event data themselves are not read, only the tree header, so even large files are processed fast.
// NB: only the "bmndata" tree is taken into account. Other objects of the file (geometry, parameter
// containers, run headers...) are not, therefore the file itself is larger than the tree.
// NB: the sizes are printed in binary units (1 KB = 1024 B)
//
// Usage: root -l -b -q 'calculate_size.C("bmndst.root")'
//        root -l -b -q '$VMCWORKDIR/macro/profiling/size/calculate_size.C("$VMCWORKDIR/macro/run/bmndst.root")'
 
R__ADD_INCLUDE_PATH($VMCWORKDIR)
#include "services/function_set.h"
 
Long64_t GetTotalSize(TBranch* b, bool ondisk, bool inclusive);
Long64_t GetBasketSize(TBranch* b, bool ondisk, bool inclusive);
 
Long64_t GetBasketSize(TObjArray* branches, bool ondisk, bool inclusive)
{
    Long64_t result = 0;
    size_t n = branches->GetEntries();
    for (size_t i = 0; i < n; i++)
        result += GetBasketSize(dynamic_cast<TBranch*>(branches->At(i)), ondisk, inclusive);
 
    return result;
}
 
Long64_t GetBasketSize(TBranch* b, bool ondisk, bool inclusive)
{
    Long64_t result = 0;
    if (b) {
        if (ondisk && b->GetZipBytes() > 0)
            result = b->GetZipBytes();
        else
            result = b->GetTotBytes();
 
        if (inclusive) {
            result += GetBasketSize(b->GetListOfBranches(), ondisk, true);
        }
        return result;
    }
    return result;
}
 
Long64_t GetTotalSize(TBranch* br, bool ondisk, bool inclusive)
{
    TMemFile f("buffer", "CREATE");
    if (br->GetTree()->GetCurrentFile())
        f.SetCompressionSettings(br->GetTree()->GetCurrentFile()->GetCompressionSettings());
 
    f.WriteObject(br, "thisbranch");
    TKey* key = f.GetKey("thisbranch");
    Long64_t size;
    if (ondisk)
        size = key->GetNbytes();
    else
        size = key->GetObjlen();
    return GetBasketSize(br, ondisk, inclusive) + size;
}
 
Long64_t GetTotalSize(TObjArray* branches, bool ondisk)
{
    Long64_t result = 0;
    size_t n = branches->GetEntries();
    for (size_t i = 0; i < n; i++) {
        result += GetTotalSize(dynamic_cast<TBranch*>(branches->At(i)), ondisk, true);
        cerr << "After " << branches->At(i)->GetName() << " " << result << endl;
    }
    return result;
}
 
Long64_t GetTotalSize(TTree* t, bool ondisk)
{
    TKey* key = 0;
    if (t->GetDirectory())
        key = t->GetDirectory()->GetKey(t->GetName());
 
    Long64_t ondiskSize = 0;
    Long64_t totalSize = 0;
    if (key) {
        ondiskSize = key->GetNbytes();
        totalSize = key->GetObjlen();
    } else {
        TMemFile f("buffer", "CREATE");
        if (t->GetCurrentFile())
            f.SetCompressionSettings(t->GetCurrentFile()->GetCompressionSettings());
 
        f.WriteTObject(t);
        key = f.GetKey(t->GetName());
        ondiskSize = key->GetNbytes();
        totalSize = key->GetObjlen();
    }
    if (t->GetBranchRef()) {
        if (ondisk)
            ondiskSize += GetBasketSize(t->GetBranchRef(), true, true);
        else
            totalSize += GetBasketSize(t->GetBranchRef(), false, true);
    }
    if (ondisk)
        return ondiskSize + GetBasketSize(t->GetListOfBranches(), /* ondisk */ true, /* inclusive */ true);
    else
        return totalSize + GetBasketSize(t->GetListOfBranches(), /* ondisk */ false, /* inclusive */ true);
}
 
Long64_t sizeOnDisk(TTree* t)
{
    // Return the size on disk on this TTree
    return GetTotalSize(t, true);
}
 
Long64_t sizeOnDisk(TBranch* branch, bool inclusive)
{
    // Return the size on disk on this branch.
    // If 'inclusive' is true, include also the size of all its sub-branches
    return GetTotalSize(branch, true, inclusive);
}
 
void printBranchSummary(TBranch* br)
{
    cout << "The branch \"" << br->GetName() << "\" takes " << byte_size_to_human(sizeOnDisk(br, true)) << " on disk"
         << endl;
    size_t n = br->GetListOfBranches()->GetEntries();
    for (size_t i = 0; i < n; i++) {
        TBranch* subbr = dynamic_cast<TBranch*>(br->GetListOfBranches()->At(i));
        cout << "  Its sub-branch \"" << subbr->GetName() << "\" takes " << byte_size_to_human(sizeOnDisk(subbr, true))
             << " on disk" << endl;
    }
}
 
void printTreeSummary(TTree* t)
{
    cout << "The TTree \"" << t->GetName() << "\" takes " << byte_size_to_human(sizeOnDisk(t)) << " on disk" << endl;
    size_t n = t->GetListOfBranches()->GetEntries();
    for (size_t i = 0; i < n; i++) {
        TBranch* br = dynamic_cast<TBranch*>(t->GetListOfBranches()->At(i));
        cout << "  Its branch \"" << br->GetName() << "\" takes " << byte_size_to_human(sizeOnDisk(br, true))
             << " on disk" << endl;
    }
}
 
// reserved approach
void cloned_branch_sizez(TTree* t)
{
    TStopwatch timer;
    timer.Start();
 
    size_t n = t->GetListOfBranches()->GetEntries();
    for (size_t i = 0; i < n; i++) {
        TBranch* br = dynamic_cast<TBranch*>(t->GetListOfBranches()->At(i));
        t->SetBranchStatus("*", 0);
        t->SetBranchStatus(br->GetName(), 1);
 
        TFile* branch_file = new TFile("cloned_branch_size.root", "recreate");
        TTree* branch_tree = t->CloneTree();
        branch_tree->Write();
        branch_file->Close();
        delete branch_file;
 
        // print file size and remove it
        // ... //
    }
 
    timer.Stop();
    Double_t rtime = timer.RealTime(), ctime = timer.CpuTime();
    printf("RealTime=%f seconds, CpuTime=%f seconds\n", rtime, ctime);
}
 
void calculate_size(TString root_file_path)
{
    gSystem->ExpandPathName(root_file_path);
 
    // check that the file exists and can be opened
    if (gSystem->AccessPathName(root_file_path.Data())) {
        cerr << "Error: file does not exist or is not accessible: " << root_file_path << endl;
        return;
    }
 
    TChain* bmnChain = new TChain("bmndata");
    bmnChain->Add(root_file_path);
 
    // force opening the file and reading the tree header;
    // if the "bmndata" tree is missing, the chain stays empty
    if (bmnChain->GetEntries() <= 0) {
        cerr << "Error: TTree \"bmndata\" not found or empty in file: " << root_file_path << endl;
        delete bmnChain;
        return;
    }
 
    // measure the tree stored in the file, not the chain: the chain is not written to the file,
    // so the tree header with the metadata of all branches would be missed
    bmnChain->LoadTree(0);
    printTreeSummary(bmnChain->GetTree());
 
    delete bmnChain;
}