Benchmark which shows that TTreeReader is as fast as SetBranchAddress when reading

readComparison.cpp

#include "TTreeReader.h"
#include "TTreeReaderArray.h"
#include "TTree.h"
#include "TLeafF.h"
#include "TStopwatch.h"
#include "TFile.h"
#include "ROOT/TSeq.hxx"
#include <ROOT/RDataFrame.hxx>
#include <ROOT/RVec.hxx>
#include <iostream>

class TimerRAAI
{
    public:
    TimerRAAI(const char* msg =""):fMsg(msg) {fSw.Start();}
    ~TimerRAAI() {fSw.Stop();std::cout << fMsg << "\n"; fSw.Print();}
    private:
    TStopwatch fSw;
    const std::string fMsg;

};

const auto filename = "VBF_HToMuMu_nano2016.root";
const auto treename = "Events";

void readTTreeReader()
{
   TFile f(filename);
   TTreeReader r(treename, &f);
   TTreeReaderArray<float> b0(r, "GenJet_phi");
   TTreeReaderArray<float> b1(r, "GenJet_eta");
   TTreeReaderArray<float> b2(r, "GenJet_pt");
   TTreeReaderArray<float> b3(r, "Jet_phi");
   TTreeReaderArray<float> b4(r, "Jet_eta");
   TTreeReaderArray<float> b5(r, "Jet_pt");
   TTreeReaderArray<float> b6(r, "FatJet_phi");
   TTreeReaderArray<float> b7(r, "FatJet_eta");
   TTreeReaderArray<float> b8(r, "FatJet_pt");
   double tot(0);
   TimerRAAI t("TTreeReader");
   while(r.Next()) {
      for (auto i : ROOT::TSeqI(b0.GetSize())) {
          tot += b0[i] + b1[i] + b2[i];
      }
      for (auto i : ROOT::TSeqI(b3.GetSize())) {
          tot += b3[i] + b4[i] + b5[i];
      }
      for (auto i : ROOT::TSeqI(b6.GetSize())) {
          tot += b6[i] + b7[i] + b8[i];
      }
   }
   std::cout << "TTreeReader dummy sum " << tot << std::endl;
}

std::vector<TBranch*> getBranches(TTree* t)
{
   std::vector<TBranch*> brs;
   for (auto name : {"GenJet_phi","GenJet_eta","GenJet_pt","Jet_phi","Jet_eta","Jet_pt","FatJet_phi","FatJet_eta","FatJet_pt"}) {
       brs.emplace_back(t->GetBranch(name));
   }
   return brs;
} 

void readLeaf()
{
   TFile f(filename);
   auto t = (TTree*) f.Get(treename);
   const auto nEntries = t->GetEntries();
   auto b0 = (TLeafF*)t->GetLeaf("GenJet_phi");
   auto b1 = (TLeafF*)t->GetLeaf("GenJet_eta");
   auto b2 = (TLeafF*)t->GetLeaf("GenJet_pt");
   auto b3 = (TLeafF*)t->GetLeaf("Jet_phi");
   auto b4 = (TLeafF*)t->GetLeaf("Jet_eta");
   auto b5 = (TLeafF*)t->GetLeaf("Jet_pt");
   auto b6 = (TLeafF*)t->GetLeaf("FatJet_phi");
   auto b7 = (TLeafF*)t->GetLeaf("FatJet_eta");
   auto b8 = (TLeafF*)t->GetLeaf("FatJet_pt");

   auto brs = getBranches(t);

   auto at = [](TLeafF *l, int i) {return (float) l->GetValue(i);};

   double tot(0);
   TimerRAAI tm("TLeaf");
   for(auto e : ROOT::TSeqU(nEntries)) {
      for (auto br : brs) br->GetEntry(e);
      for (auto i : ROOT::TSeqI(b0->GetLen())) {
          tot += at(b0,i) + at(b1,i) + at(b2,i);
      }
      for (auto i : ROOT::TSeqI(b3->GetLen())) {
          tot += at(b3,i)+ at(b4,i)+ at(b5,i);
      }
      for (auto i : ROOT::TSeqI(b6->GetLen())) {
          tot += at(b6,i)+ at(b7,i)+ at(b8,i);
      }
   }
   std::cout << "TLeaf dummy sum " << tot << std::endl;
}

void readBranchAddress()
{
   TFile f(filename);
   auto t = (TTree*) f.Get(treename);
   float b0[100];
   float b1[100];
   float b2[100];
   float b3[100];
   float b4[100];
   float b5[100];
   float b6[100];
   float b7[100];
   float b8[100];
   unsigned int nJet, nGenJet, nFatJet;
   const auto nEntries = t->GetEntries();
   t->SetBranchAddress("GenJet_phi", b0);
   t->SetBranchAddress("GenJet_eta", b1);
   t->SetBranchAddress("GenJet_pt", b2);
   t->SetBranchAddress("nGenJet", &nGenJet);
   t->SetBranchAddress("Jet_phi", b3);
   t->SetBranchAddress("Jet_eta", b4);
   t->SetBranchAddress("Jet_pt", b5);
   t->SetBranchAddress("nJet", &nJet);
   t->SetBranchAddress("FatJet_phi", b6);
   t->SetBranchAddress("FatJet_eta", b7);
   t->SetBranchAddress("FatJet_pt", b8);
   t->SetBranchAddress("nFatJet", &nFatJet);

   auto brs = getBranches(t);
   auto at = [](float *a, int i) {return a[i];};
   double tot(0);
   TimerRAAI tm("SetBranchAddress");
   for(auto e : ROOT::TSeqU(nEntries)) {
      for (auto br : brs) br->GetEntry(e);
      for (auto i : ROOT::TSeqI(nGenJet)) {
          tot += at(b0,i) + at(b1,i) + at(b2,i);
      }
      for (auto i : ROOT::TSeqI(nJet)) {
          tot += at(b3,i)+ at(b4,i)+ at(b5,i);
      }
      for (auto i : ROOT::TSeqI(nFatJet)) {
          tot += at(b6,i)+ at(b7,i)+ at(b8,i);
      }
   }
   std::cout << "SetBranchAddress dummy sum " << tot << std::endl;
}

void readDataFrame()
{
    ROOT::RDataFrame r(treename, filename);
    double tot(0);
    auto at = [](const ROOT::RVec<float> &a, int i) {return a[i];};
    TimerRAAI tm("DataFrame");
    r.Foreach([&](const ROOT::RVec<float> &b0, 
    const ROOT::RVec<float> &b1, 
    const ROOT::RVec<float> &b2, 
    const ROOT::RVec<float> &b3, 
    const ROOT::RVec<float> &b4, 
    const ROOT::RVec<float> &b5, 
    const ROOT::RVec<float> &b6, 
    const ROOT::RVec<float> &b7, 
    const ROOT::RVec<float> &b8){
      for (auto i : ROOT::TSeqI(b1.size())) {
          tot += at(b0,i) + at(b1,i) + at(b2,i);
      }
      for (auto i : ROOT::TSeqI(b3.size())) {
          tot += at(b3,i)+ at(b4,i)+ at(b5,i);
      }
      for (auto i : ROOT::TSeqI(b6.size())) {
          tot += at(b6,i)+ at(b7,i)+ at(b8,i);
      }
    }, {"GenJet_phi","GenJet_eta","GenJet_pt","Jet_phi","Jet_eta","Jet_pt","FatJet_phi","FatJet_eta","FatJet_pt"});

    std::cout << "DataFrame dummy sum " << tot << std::endl;
}

void bench()
{
    readLeaf();
    readBranchAddress();
    readTTreeReader();
    readDataFrame();
}

int main()
{
    bench();
    std::cout << "---------- Start Real Profiling ------ \n\n";
    for (auto i : ROOT::TSeqI(5)) {
       std::cout << "Run " << i << std::endl;
       bench();
    }

    return 0;
}