Allow specifying split level on a per branch basis

IOPool/Output/interface/PoolOutputModule.h

@@ -16,6 +16,7 @@
 #include <set>
 #include <string>
 #include <vector>
+#include <regex>
 
 #include "IOPool/Common/interface/RootServiceChecker.h"
 #include "FWCore/Framework/interface/Frameworkfwd.h"
@@ -107,6 +108,18 @@ namespace edm {
 
     typedef std::array<OutputItemList, NumBranchTypes> OutputItemListArray;
 
+    struct SpecialSplitLevelForBranch {
+      SpecialSplitLevelForBranch(std::string const& iBranchName, int iSplitLevel):
+      branch_(convert(iBranchName)),
+      splitLevel_(iSplitLevel < 1? 1: iSplitLevel) //minimum is 1
+      {}
+      bool match(std::string const& iBranchName) const;
+      std::regex convert(std::string const& iGlobBranchExpression )const;
+
+      std::regex branch_;
+      int splitLevel_;
+    };
+
     OutputItemListArray const& selectedOutputItemList() const {return selectedOutputItemList_;}
 
     BranchChildren const& branchChildren() const {return branchChildren_;}
@@ -157,6 +170,7 @@ namespace edm {
     RootServiceChecker rootServiceChecker_;
     AuxItemArray auxItems_;
     OutputItemListArray selectedOutputItemList_;
+    std::vector<SpecialSplitLevelForBranch> specialSplitLevelForBranches_;
     std::string const fileName_;
     std::string const logicalFileName_;
     std::string const catalog_;

IOPool/Output/src/PoolOutputModule.cc

@@ -30,6 +30,8 @@
 #include <fstream>
 #include <iomanip>
 #include <sstream>
+#include "boost/algorithm/string.hpp"
+
 
 namespace edm {
   PoolOutputModule::PoolOutputModule(ParameterSet const& pset) :
@@ -85,6 +87,14 @@ namespace edm {
       whyNotFastClonable_+= FileBlock::EventSelectionUsed;
     }
 
+    auto const& specialSplit {pset.getUntrackedParameterSetVector("overrideBranchesSplitLevel")};
+
+    specialSplitLevelForBranches_.reserve(specialSplit.size());
+    for(auto const& s: specialSplit) {
+      specialSplitLevelForBranches_.emplace_back(s.getUntrackedParameter<std::string>("branch"),
+                                                 s.getUntrackedParameter<int>("splitLevel"));
+    }
+
     // We don't use this next parameter, but we read it anyway because it is part
     // of the configuration of this module.  An external parser creates the
     // configuration by reading this source code.
@@ -155,6 +165,17 @@ namespace edm {
     return lh < rh;
   }
 
+  inline bool PoolOutputModule::SpecialSplitLevelForBranch::match( std::string const& iBranchName) const {
+    return std::regex_match(iBranchName,branch_);
+  }
+
+  std::regex PoolOutputModule::SpecialSplitLevelForBranch::convert( std::string const& iGlobBranchExpression) const {
+    std::string tmp(iGlobBranchExpression);
+    boost::replace_all(tmp, "*", ".*");
+    boost::replace_all(tmp, "?", ".");
+    return std::regex(tmp);
+  }
+
   void PoolOutputModule::fillSelectedItemList(BranchType branchType, TTree* theInputTree) {
 
     SelectedProducts const& keptVector = keptProducts()[branchType];
@@ -186,6 +207,11 @@ namespace edm {
         basketSize = theBranch->GetBasketSize();
       } else {
         splitLevel = (prod.splitLevel() == BranchDescription::invalidSplitLevel ? splitLevel_ : prod.splitLevel());
+        for(auto const& b: specialSplitLevelForBranches_) {
+          if(b.match(prod.branchName())) {
+            splitLevel =b.splitLevel_;
+          }
+        }
         basketSize = (prod.basketSize() == BranchDescription::invalidBasketSize ? basketSize_ : prod.basketSize());
       }
       outputItemList.emplace_back(&prod, kept.second, splitLevel, basketSize);
@@ -449,11 +475,18 @@ namespace edm {
                      "'DROPPED': Keep it for products produced in current process and all kept products. Drop it for dropped products produced in prior processes.\n"
                      "'PRIOR':   Keep it for products produced in current process. Drop it for products produced in prior processes.\n"
                      "'ALL':     Drop all of it.");
-    ParameterSetDescription dataSet;
-    dataSet.setAllowAnything();
-    desc.addUntracked<ParameterSetDescription>("dataset", dataSet)
-     ->setComment("PSet is only used by Data Operations and not by this module.");
-
+    {
+      ParameterSetDescription dataSet;
+      dataSet.setAllowAnything();
+      desc.addUntracked<ParameterSetDescription>("dataset", dataSet)
+      ->setComment("PSet is only used by Data Operations and not by this module.");
+    }
+    {
+      ParameterSetDescription specialSplit;
+      specialSplit.addUntracked<std::string>("branch")->setComment("Name of branch needing a special split level. The name can contain wildcards '*' and '?'");
+      specialSplit.addUntracked<int>("splitLevel")->setComment("The special split level for the branch");
+      desc.addVPSetUntracked("overrideBranchesSplitLevel",specialSplit, std::vector<ParameterSet>());
+    }
     OutputModule::fillDescription(desc);
   }