[DF][RDatasetSpec] Initial version of friend trees handling

Author: ikabadzhov

tree/dataframe/inc/ROOT/RDF/RDatasetSpec.hxx

@@ -38,6 +38,31 @@ struct RDatasetSpec {
       }
    };
 
+   struct RFriendInfo {
+      std::vector<std::string> fTreeNames{};
+      std::vector<std::string> fFileNameGlobs{};
+      RFriendInfo() {}
+
+      RFriendInfo(const std::string &treeName, const std::string &fileName)
+         : fTreeNames(std::vector<std::string>{treeName}), fFileNameGlobs(std::vector<std::string>{fileName})
+      {
+      }
+
+      RFriendInfo(const std::string &treeName, const std::vector<std::string> &fileNames)
+         : fTreeNames(std::vector<std::string>{treeName}), fFileNameGlobs(fileNames)
+      {
+      }
+
+      RFriendInfo(const std::vector<std::string> &treeNames, const std::vector<std::string> &fileNames)
+         : fTreeNames(
+              fileNames.size() != treeNames.size() && treeNames.size() != 1
+                 ? throw std::logic_error("RFriendInfo exepcts either N trees and N files, or 1 tree and N files.")
+                 : treeNames),
+           fFileNameGlobs(fileNames)
+      {
+      }
+   };
+
    /**
     * A list of names of trees.
     * This list should go in lockstep with fFileNameGlobs, only in case this dataset is a TChain where each file
@@ -49,31 +74,35 @@ struct RDatasetSpec {
     * A list of file names.
     * They can contain the globbing characters supported by TChain. See TChain::Add for more information.
     */
-
    std::vector<std::string> fFileNameGlobs{};
 
    ULong64_t fStartEntry{}; ///< The entry where the dataset processing should start (inclusive).
    ULong64_t fEndEntry{};   ///< The entry where the dataset processing should end (exclusive).
 
-   RDatasetSpec(const std::string &treeName, const std::string &fileName, REntryRange entryRange = {})
+   std::vector<RFriendInfo> fFriendInfos{}; ///< List of friends
+
+   RDatasetSpec(const std::string &treeName, const std::string &fileName, REntryRange entryRange = {},
+                const std::vector<RFriendInfo> &friendInfos = {})
       : fTreeNames(std::vector<std::string>{treeName}), fFileNameGlobs(std::vector<std::string>{fileName}),
-        fStartEntry(entryRange.fStartEntry), fEndEntry(entryRange.fEndEntry)
+        fStartEntry(entryRange.fStartEntry), fEndEntry(entryRange.fEndEntry), fFriendInfos(friendInfos)
    {
    }
 
-   RDatasetSpec(const std::string &treeName, const std::vector<std::string> &fileNames, REntryRange entryRange = {})
+   RDatasetSpec(const std::string &treeName, const std::vector<std::string> &fileNames, REntryRange entryRange = {},
+                const std::vector<RFriendInfo> &friendInfos = {})
       : fTreeNames(std::vector<std::string>{treeName}), fFileNameGlobs(fileNames), fStartEntry(entryRange.fStartEntry),
-        fEndEntry(entryRange.fEndEntry)
+        fEndEntry(entryRange.fEndEntry), fFriendInfos(friendInfos)
    {
    }
 
    RDatasetSpec(const std::vector<std::string> &treeNames, const std::vector<std::string> &fileNames,
-                REntryRange entryRange = {})
+                REntryRange entryRange = {}, const std::vector<RFriendInfo> &friendInfos = {})
       : fTreeNames(
            fileNames.size() != treeNames.size() && treeNames.size() != 1
-              ? throw std::runtime_error("RDatasetSpec exepcts either N trees and N files, or 1 tree and N files.")
+              ? throw std::logic_error("RDatasetSpec exepcts either N trees and N files, or 1 tree and N files.")
               : treeNames),
-        fFileNameGlobs(fileNames), fStartEntry(entryRange.fStartEntry), fEndEntry(entryRange.fEndEntry)
+        fFileNameGlobs(fileNames), fStartEntry(entryRange.fStartEntry), fEndEntry(entryRange.fEndEntry),
+        fFriendInfos(friendInfos)
    {
    }
 };

tree/dataframe/src/RLoopManager.cxx

@@ -370,6 +370,17 @@ RLoopManager::RLoopManager(const ROOT::RDF::RDatasetSpec &spec)
       chain->Add(fullpath.c_str());
    }
    SetTree(chain);
+   for (auto i = 0u; i < spec.fFriendInfos.size(); ++i) {
+      auto temp_friend_chain_name = "f" + std::to_string(i);
+      auto temp_friend_chain = std::make_shared<TChain>(temp_friend_chain_name.c_str());
+      for (auto j = 0u; j < spec.fFriendInfos[i].fFileNameGlobs[j].size(); ++j) {
+         const auto fullpath = spec.fFriendInfos[i].fFileNameGlobs[j] + "?#" +
+                               spec.fFriendInfos[i].fTreeNames[spec.fFriendInfos[i].fTreeNames.size() == 1 ? 0 : j];
+         temp_friend_chain->Add(fullpath.c_str());
+      }
+      SetTree(temp_friend_chain);
+      chain->AddFriend(temp_friend_chain_name.c_str());
+   }
 }
 
 struct RSlotRAII {

tree/dataframe/test/dataframe_datasetspec.cxx

@@ -52,11 +52,11 @@ TEST(RDFDatasetSpec, SingleFileSingleColConstructor)
    EXPECT_THROW(
       try {
          RDatasetSpec({"tree"s, "anothertree"s}, {"file.root"s}, {2, 4});
-      } catch (const std::runtime_error &err) {
+      } catch (const std::logic_error &err) {
          EXPECT_EQ(std::string(err.what()), "RDatasetSpec exepcts either N trees and N files, or 1 tree and N files.");
          throw;
       },
-      std::runtime_error);
+      std::logic_error);
 
    // specify range [2, 2) (3 is a valid index) => range is disregarded
    const auto dfRDS7 = RDataFrame(RDatasetSpec("tree", "file.root", {2, 2})).Display<int>({"x"})->AsString();
@@ -281,3 +281,22 @@ TEST(RDFDatasetSpec, MultipleFiles)
 
    gSystem->Exec("rm file0.root file1.root file2.root");
 }
+
+// TODO: test the friends
+/*
+TEST(RDFDatasetSpec, FriendTrees)
+{
+   auto dfWriter0 = RDataFrame(3)
+                       .Define("x", [](ULong64_t e) { return int(e); }, {"rdfentry_"})
+                       .Define("y", [](ULong64_t e) { return int(e) + 1; }, {"rdfentry_"});
+   dfWriter0.Snapshot<int, int>("treeA", "file0.root", {"x", "y"});
+   dfWriter0.Snapshot<int, int>("treeA", "file1.root", {"x", "y"});
+   dfWriter0.Snapshot<int, int>("treeB", "file2.root", {"x", "y"}); // different tree's name
+   auto dfWriter1 = RDataFrame(2)
+                       .Define("x", [](ULong64_t e) { return int(e) + 2; }, {"rdfentry_"})
+                       .Define("y", [](ULong64_t e) { return int(e) + 3; }, {"rdfentry_"});
+   dfWriter0.Snapshot<int, int>("treeA", "file3.root", {"x", "y"});
+   dfWriter0.Snapshot<int, int>("treeB", "file4.root", {"x", "y"});
+   dfWriter0.Snapshot<int, int>("treeC", "file5.root", {"x", "y"});
+}
+*/