00001
00002 #include "RootOutputTree.h"
00003
00004 #include "DataFormats/Common/interface/RefCoreStreamer.h"
00005 #include "DataFormats/Provenance/interface/BranchDescription.h"
00006 #include "FWCore/Framework/interface/ConstProductRegistry.h"
00007 #include "FWCore/MessageLogger/interface/MessageLogger.h"
00008 #include "FWCore/ServiceRegistry/interface/Service.h"
00009 #include "FWCore/Utilities/interface/Algorithms.h"
00010 #include "FWCore/Utilities/interface/EDMException.h"
00011 #include "FWCore/Utilities/interface/RootHandlers.h"
00012
00013 #include "TBranch.h"
00014 #include "TBranchElement.h"
00015 #include "TCollection.h"
00016 #include "TFile.h"
00017 #include "TTreeCloner.h"
00018 #include "Rtypes.h"
00019 #include "RVersion.h"
00020
00021 #include "boost/bind.hpp"
00022 #include <limits>
00023
00024 namespace edm {
00025
00026 RootOutputTree::RootOutputTree(
00027 boost::shared_ptr<TFile> filePtr,
00028 BranchType const& branchType,
00029 int bufSize,
00030 int splitLevel,
00031 int treeMaxVirtualSize) :
00032 filePtr_(filePtr),
00033 tree_(makeTTree(filePtr.get(), BranchTypeToProductTreeName(branchType), splitLevel)),
00034 producedBranches_(),
00035 readBranches_(),
00036 auxBranches_(),
00037 unclonedReadBranches_(),
00038 unclonedReadBranchNames_(),
00039 currentlyFastCloning_(),
00040 fastCloneAuxBranches_(false) {
00041
00042 if(treeMaxVirtualSize >= 0) tree_->SetMaxVirtualSize(treeMaxVirtualSize);
00043 }
00044
00045 TTree*
00046 RootOutputTree::assignTTree(TFile* filePtr, TTree* tree) {
00047 tree->SetDirectory(filePtr);
00048
00049
00050 tree->SetAutoSave(std::numeric_limits<Long64_t>::max());
00051 return tree;
00052 }
00053
00054 TTree*
00055 RootOutputTree::makeTTree(TFile* filePtr, std::string const& name, int splitLevel) {
00056 TTree* tree = new TTree(name.c_str(), "", splitLevel);
00057 if(!tree) throw edm::Exception(errors::FatalRootError)
00058 << "Failed to create the tree: " << name << "\n";
00059 if(tree->IsZombie())
00060 throw edm::Exception(errors::FatalRootError)
00061 << "Tree: " << name << " is a zombie." << "\n";
00062
00063 return assignTTree(filePtr, tree);
00064 }
00065
00066 bool
00067 RootOutputTree::checkSplitLevelsAndBasketSizes(TTree* inputTree) const {
00068
00069 assert (inputTree != 0);
00070
00071
00072 for(std::vector<TBranch*>::const_iterator it = readBranches_.begin(), itEnd = readBranches_.end();
00073 it != itEnd; ++it) {
00074
00075 TBranch* outputBranch = *it;
00076 if(outputBranch != 0) {
00077 TBranch* inputBranch = inputTree->GetBranch(outputBranch->GetName());
00078
00079 if(inputBranch != 0) {
00080 if(inputBranch->GetSplitLevel() != outputBranch->GetSplitLevel() ||
00081 inputBranch->GetBasketSize() != outputBranch->GetBasketSize()) {
00082 return false;
00083 }
00084 }
00085 }
00086 }
00087 return true;
00088 }
00089
00090 namespace {
00091 bool checkMatchingBranches(TBranchElement* inputBranch, TBranchElement* outputBranch) {
00092 if(inputBranch->GetStreamerType() != outputBranch->GetStreamerType()) {
00093 return false;
00094 }
00095 TObjArray* inputArray = inputBranch->GetListOfBranches();
00096 TObjArray* outputArray = outputBranch->GetListOfBranches();
00097
00098 if(outputArray->GetSize() < inputArray->GetSize()) {
00099 return false;
00100 }
00101 TIter iter(outputArray);
00102 TObject* obj = 0;
00103 while((obj = iter.Next()) != 0) {
00104 TBranchElement* outBranch = dynamic_cast<TBranchElement*>(obj);
00105 if(outBranch) {
00106 TBranchElement* inBranch = dynamic_cast<TBranchElement*>(inputArray->FindObject(outBranch->GetName()));
00107 if(!inBranch) {
00108 return false;
00109 }
00110 if(!checkMatchingBranches(inBranch, outBranch)) {
00111 return false;
00112 }
00113 }
00114 }
00115 return true;
00116 }
00117 }
00118
00119 bool RootOutputTree::checkIfFastClonable(TTree* inputTree) const {
00120
00121 if(inputTree == 0) return false;
00122
00123
00124 for(std::vector<TBranch*>::const_iterator it = readBranches_.begin(), itEnd = readBranches_.end(); it != itEnd; ++it) {
00125 TBranchElement* outputBranch = dynamic_cast<TBranchElement*>(*it);
00126 if(outputBranch != 0) {
00127 TBranchElement* inputBranch = dynamic_cast<TBranchElement*>(inputTree->GetBranch(outputBranch->GetName()));
00128 if(inputBranch != 0) {
00129
00130 if(!checkMatchingBranches(inputBranch, outputBranch)) {
00131 LogInfo("FastCloning")
00132 << "Fast Cloning disabled because a data member has been added to split branch: " << inputBranch->GetName() << "\n.";
00133 }
00134 }
00135 }
00136 }
00137 return true;
00138 }
00139
00140 void
00141 RootOutputTree::fastCloneTTree(TTree* in, std::string const& option) {
00142 if(in->GetEntries() != 0) {
00143 TObjArray* branches = tree_->GetListOfBranches();
00144
00145
00146 std::map<Int_t, TBranch *> auxIndexes;
00147 if (!fastCloneAuxBranches_) {
00148 for (std::vector<TBranch *>::const_iterator it = auxBranches_.begin(), itEnd = auxBranches_.end();
00149 it != itEnd; ++it) {
00150 int auxIndex = branches->IndexOf(*it);
00151 assert (auxIndex >= 0);
00152 auxIndexes.insert(std::make_pair(auxIndex, *it));
00153 branches->RemoveAt(auxIndex);
00154 }
00155 branches->Compress();
00156 }
00157
00158 #if ROOT_VERSION_CODE >= ROOT_VERSION(5,26,0)
00159 TTreeCloner cloner(in, tree_, option.c_str(), TTreeCloner::kNoWarnings|TTreeCloner::kIgnoreMissingTopLevel);
00160 #else
00161 TTreeCloner cloner(in, tree_, option.c_str());
00162 #endif
00163
00164 if(!cloner.IsValid()) {
00165
00166 #if ROOT_VERSION_CODE >= ROOT_VERSION(5,26,0)
00167
00168 static const char* okerror = "One of the export branch";
00169 if ( strncmp(cloner.GetWarning(),okerror,strlen(okerror)) == 0 ) {
00170
00171 }
00172 else {
00173 throw edm::Exception(errors::FatalRootError)
00174 << "invalid TTreeCloner\n";
00175 }
00176 #else
00177 throw edm::Exception(errors::FatalRootError)
00178 << "invalid TTreeCloner\n";
00179 #endif
00180 }
00181 tree_->SetEntries(tree_->GetEntries() + in->GetEntries());
00182 Service<RootHandlers> rootHandler;
00183 rootHandler->enableErrorHandlerWithoutWarnings();
00184 cloner.Exec();
00185 rootHandler->enableErrorHandler();
00186 if (!fastCloneAuxBranches_) {
00187 for (std::map<Int_t, TBranch *>::const_iterator it = auxIndexes.begin(), itEnd = auxIndexes.end();
00188 it != itEnd; ++it) {
00189
00190 Int_t last = branches->GetLast();
00191 if (last >= 0) {
00192 branches->AddAtAndExpand(branches->At(last), last+1);
00193 for(Int_t ind = last-1; ind >= it->first; --ind) {
00194 branches->AddAt(branches->At(ind), ind+1);
00195 };
00196 branches->AddAt(it->second, it->first);
00197 } else {
00198 branches->Add(it->second);
00199 }
00200 }
00201 }
00202 }
00203 }
00204
00205 void
00206 RootOutputTree::writeTTree(TTree* tree) {
00207 if(tree->GetNbranches() != 0) {
00208 tree->SetEntries(-1);
00209 }
00210 setRefCoreStreamer(true);
00211 tree->AutoSave("FlushBaskets");
00212 }
00213
00214 void
00215 RootOutputTree::fillTTree(TTree* tree, std::vector<TBranch*> const& branches) {
00216 for_all(branches, boost::bind(&TBranch::Fill, _1));
00217 }
00218
00219 void
00220 RootOutputTree::writeTree() const {
00221 writeTTree(tree_);
00222 }
00223
00224 void
00225 RootOutputTree::maybeFastCloneTree(bool canFastClone, bool canFastCloneAux, TTree* tree, std::string const& option) {
00226 unclonedReadBranches_.clear();
00227 unclonedReadBranchNames_.clear();
00228 currentlyFastCloning_ = canFastClone && !readBranches_.empty();
00229 if(currentlyFastCloning_) {
00230 fastCloneAuxBranches_ = canFastCloneAux;
00231 fastCloneTTree(tree, option);
00232 for(std::vector<TBranch*>::const_iterator it = readBranches_.begin(), itEnd = readBranches_.end();
00233 it != itEnd; ++it) {
00234 if((*it)->GetEntries() != tree_->GetEntries()) {
00235 unclonedReadBranches_.push_back(*it);
00236 unclonedReadBranchNames_.insert(std::string((*it)->GetName()));
00237 }
00238 }
00239 }
00240 }
00241
00242 void
00243 RootOutputTree::fillTree() const {
00244 if (currentlyFastCloning_) {
00245 if (!fastCloneAuxBranches_)fillTTree(tree_, auxBranches_);
00246 fillTTree(tree_, producedBranches_);
00247 fillTTree(tree_, unclonedReadBranches_);
00248 } else {
00249 tree_->Fill();
00250 }
00251 }
00252
00253 void
00254 RootOutputTree::addBranch(std::string const& branchName,
00255 std::string const& className,
00256 void const*& pProd,
00257 int splitLevel,
00258 int basketSize,
00259 bool produced) {
00260 assert(splitLevel != BranchDescription::invalidSplitLevel);
00261 assert(basketSize != BranchDescription::invalidBasketSize);
00262 TBranch* branch = tree_->Branch(branchName.c_str(),
00263 className.c_str(),
00264 &pProd,
00265 basketSize,
00266 splitLevel);
00267 if(produced) {
00268 producedBranches_.push_back(branch);
00269 } else {
00270 readBranches_.push_back(branch);
00271 }
00272 }
00273
00274 void
00275 RootOutputTree::close() {
00276
00277
00278 auxBranches_.clear();
00279 producedBranches_.clear();
00280 readBranches_.clear();
00281 unclonedReadBranches_.clear();
00282 tree_ = 0;
00283 filePtr_.reset();
00284 }
00285 }