RDFGroupBy.cpp

#include <ROOT/RDataFrame.hxx>
#include <ROOT/TypeTraits.hxx>
#include <TEntryList.h>
#include <map>
#include <mutex>
#include <utility>

template <typename GroupFn, typename R = typename ROOT::TypeTraits::CallableTraits<GroupFn>::ret_type>
std::map<R, ROOT::RDF::RNode> GroupBy(ROOT::RDF::RNode df, GroupFn &&fn, const std::vector<std::string> &cols)
{
   using EListPerGroup = std::map<R, TEntryList>;
   using GroupAndEntry = std::pair<R, ULong64_t>;
   std::mutex m;
   auto agg = [&m](EListPerGroup &elistPerGroup, const  GroupAndEntry&groupAndEntry) {
      std::lock_guard<std::mutex> l(m);
      elistForGroup[groupAndEntry.first].Enter(groupAndEntry.second);
   };

   auto red = [](const EListPerGroup e1, const EListPerGroup &e2) {
      for (auto p : e1) {
         if (e2.find(p.first) != e2.end())
            e1[p.first].Add(&e2[p.first]);
      }
      for (auto p : e2) {
         if (e1.find(p.first) != e1.end())
            e1[p.first].Add(&e2[p.first]);
      }
      return e1;
   };

   EListPerGroup elists =
      df.Define("group", fn, cols)
         .Define("groupAndEntry", [](R g, ULong64_t e) { return std::make_pair(g, e); }, {"group", "rdfentry_"})
         .Aggregate(agg, red, {"groupAndEntry"})
         .GetValue();

   std::map<R, ROOT::RDF::RNode> dfs;
   for (auto p : elists) {
      ROOT::RDataFrame df(???);
      dfs[p.first] = df;
   }
      
   return dfs;
}

int main()
{
   ROOT::RDataFrame df(10);
   auto dfs = GroupBy(df, [](Ulong64_t e) { return e; }, {"rdfentry_"});
   for (auto p : dfs)
      std::cout << p.first << " " << p.second.Mean<ULong64_t>("rdfentry_").GetValue() << std::endl;
   return 0;
}