From eac6cac6dd18fe9e65e714ca12ab72f947c23089 Mon Sep 17 00:00:00 2001
From: Enric Tejedor Saavedra <enric.tejedor.saavedra@cern.ch>
Date: Tue, 28 Feb 2017 15:18:38 +0100
Subject: [PATCH] Add new constructor in TTreeProcessorMT that receives a
TEntryList.
This is intended to support the corresponding constructor in TDataFrame.
---
tree/treeplayer/inc/ROOT/TTreeProcessorMT.hxx | 72 +++++++++++++++++--
tree/treeplayer/src/TTreeProcessorMT.cxx | 9 ++-
2 files changed, 75 insertions(+), 6 deletions(-)
diff --git a/tree/treeplayer/inc/ROOT/TTreeProcessorMT.hxx b/tree/treeplayer/inc/ROOT/TTreeProcessorMT.hxx
index 70780436cab..4c053cf984b 100644
--- a/tree/treeplayer/inc/ROOT/TTreeProcessorMT.hxx
+++ b/tree/treeplayer/inc/ROOT/TTreeProcessorMT.hxx
@@ -18,6 +18,7 @@
#include "TChain.h"
#include "TTreeReader.h"
#include "TError.h"
+#include "TEntryList.h"
#include "ROOT/TThreadedObject.hxx"
#include <string.h>
@@ -50,7 +51,9 @@ namespace ROOT {
std::string fTreeName; ///< Name of the tree
std::unique_ptr<TFile> fCurrentFile; ///<! Current file object of this view.
std::unique_ptr<TTree> fCurrentTree; ///<! Current tree object of this view.
- int fCurrentIdx; ///<! Index of the current file.
+ unsigned int fCurrentIdx; ///<! Index of the current file.
+ std::vector<TEntryList> fEntryLists; ///< Entry numbers to be processed per tree/file
+ TEntryList fCurrentEntryList; ///< Entry numbers for the current range being processed
////////////////////////////////////////////////////////////////////////////////
/// Initialize the file and the tree for this view, first looking for a tree in
@@ -134,6 +137,42 @@ namespace ROOT {
}
}
+ //////////////////////////////////////////////////////////////////////////
+ /// Constructor based on a TTree and a TEntryList.
+ /// \param[in] tree Tree or chain of files containing the tree to process.
+ /// \param[in] entries List of entry numbers to process.
+ TTreeView(TTree& tree, TEntryList& entries) : TTreeView(tree)
+ {
+ static const TClassRef clRefTChain("TChain");
+ if (clRefTChain == tree.IsA()) {
+ // We need to convert the global entry numbers to per-tree entry numbers.
+ // This will allow us to build a TEntryList for a given entry range of a tree of the chain.
+ size_t nTrees = fFileNames.size();
+ fEntryLists.resize(nTrees);
+
+ TChain *chain = dynamic_cast<TChain*>(&tree);
+ Long64_t currListEntry = entries.GetEntry(0);
+ Long64_t currTreeOffset = 0;
+
+ for (unsigned int treeNum = 0; treeNum < nTrees && currListEntry >= 0; ++treeNum) {
+ chain->LoadTree(currTreeOffset);
+ TTree *currTree = chain->GetTree();
+ Long64_t currTreeEntries = currTree->GetEntries();
+ Long64_t nextTreeOffset = currTreeOffset + currTreeEntries;
+
+ while (currListEntry >= 0 && currListEntry < nextTreeOffset) {
+ fEntryLists[treeNum].Enter(currListEntry - currTreeOffset);
+ currListEntry = entries.Next();
+ }
+
+ currTreeOffset = nextTreeOffset;
+ }
+ }
+ else {
+ fEntryLists.emplace_back(entries);
+ }
+ }
+
//////////////////////////////////////////////////////////////////////////
/// Copy constructor.
/// \param[in] view Object to copy.
@@ -141,6 +180,10 @@ namespace ROOT {
{
for (auto& fn : view.fFileNames)
fFileNames.emplace_back(fn);
+
+ for (auto& el : view.fEntryLists)
+ fEntryLists.emplace_back(el);
+
Init();
}
@@ -154,9 +197,29 @@ namespace ROOT {
//////////////////////////////////////////////////////////////////////////
/// Get a TTreeReader for the current tree of this view.
- std::unique_ptr<TTreeReader> GetTreeReader() const
+ std::unique_ptr<TTreeReader> GetTreeReader(Long64_t start, Long64_t end)
{
- return std::unique_ptr<TTreeReader>(new TTreeReader(fCurrentTree.get()));
+ TTreeReader *reader;
+ if (fEntryLists.size() > 0 && fEntryLists[fCurrentIdx].GetN() > 0) {
+ // TEntryList and SetEntriesRange do not work together (the former has precedence).
+ // We need to construct a TEntryList that contains only those entry numbers
+ // in our desired range.
+ fCurrentEntryList.Reset();
+ Long64_t entry = fEntryLists[fCurrentIdx].GetEntry(0);
+ do {
+ if (entry >= start && entry < end) fCurrentEntryList.Enter(entry);
+ } while ((entry = fEntryLists[fCurrentIdx].Next()) >= 0);
+
+ reader = new TTreeReader(fCurrentTree.get(), &fCurrentEntryList);
+
+ }
+ else {
+ // If no TEntryList is involved we can safely set the range in the reader
+ reader = new TTreeReader(fCurrentTree.get());
+ reader->SetEntriesRange(start, end);
+ }
+
+ return std::unique_ptr<TTreeReader>(reader);
}
//////////////////////////////////////////////////////////////////////////
@@ -175,7 +238,7 @@ namespace ROOT {
//////////////////////////////////////////////////////////////////////////
/// Set the current file and tree of this view.
- void SetCurrent(int i)
+ void SetCurrent(unsigned int i)
{
if (i != fCurrentIdx) {
fCurrentIdx = i;
@@ -197,6 +260,7 @@ namespace ROOT {
TTreeProcessorMT(std::string_view filename, std::string_view treename = "");
TTreeProcessorMT(const std::vector<std::string_view>& filenames, std::string_view treename = "");
TTreeProcessorMT(TTree& tree);
+ TTreeProcessorMT(TTree& tree, TEntryList& entries);
void Process(std::function<void(TTreeReader&)> func);
diff --git a/tree/treeplayer/src/TTreeProcessorMT.cxx b/tree/treeplayer/src/TTreeProcessorMT.cxx
index ef7dbf494bf..a848b73ac3b 100644
--- a/tree/treeplayer/src/TTreeProcessorMT.cxx
+++ b/tree/treeplayer/src/TTreeProcessorMT.cxx
@@ -52,6 +52,12 @@ TTreeProcessorMT::TTreeProcessorMT(const std::vector<std::string_view>& filename
/// \param[in] tree Tree or chain of files containing the tree to process.
TTreeProcessorMT::TTreeProcessorMT(TTree& tree) : treeView(tree) {}
+////////////////////////////////////////////////////////////////////////
+/// Constructor based on a TTree and a TEntryList.
+/// \param[in] tree Tree or chain of files containing the tree to process.
+/// \param[in] entries List of entry numbers to process.
+TTreeProcessorMT::TTreeProcessorMT(TTree& tree, TEntryList& entries) : treeView(tree, entries) {}
+
//////////////////////////////////////////////////////////////////////////////
/// Process the entries of a TTree in parallel. The user-provided function
/// receives a TTreeReader which can be used to iterate on a subrange of
@@ -89,8 +95,7 @@ void TTreeProcessorMT::Process(std::function<void(TTreeReader&)> func)
g.run([this, &func, start, end, i]() {
treeView->SetCurrent(i);
- auto tr = treeView->GetTreeReader();
- tr->SetEntriesRange(start, end);
+ auto tr = treeView->GetTreeReader(start, end);
func(*tr);
});
}
--
GitLab