Avoid including ambiguous TTree branch names

tree/dataframe/src/RTTreeDS.cxx

@@ -110,8 +110,9 @@ void InsertBranchName(std::set<std::string> &bNamesReg, std::vector<std::string>
    foundLeaves.insert(leaf);
 }
 
-void ExploreBranch(TTree &t, std::set<std::string> &bNamesReg, std::vector<std::string> &bNames, TBranch *b,
-                   std::string prefix, std::string &friendName, bool allowDuplicates)
+void ExploreBranch(TTree &t, std::unordered_map<std::string, unsigned int> &duplicateTokens,
+                   std::set<std::string> &bNamesReg, std::vector<std::string> &bNames, TBranch *b, std::string prefix,
+                   std::string &friendName, bool allowDuplicates)
 {
    // We want to avoid situations of overlap between the prefix and the
    // sub-branch name that might happen when the branch is composite, e.g.
@@ -129,6 +130,30 @@ void ExploreBranch(TTree &t, std::set<std::string> &bNamesReg, std::vector<std::
    for (auto sb : *b->GetListOfBranches()) {
       TBranch *subBranch = static_cast<TBranch *>(sb);
       auto subBranchName = std::string(subBranch->GetName());
+      auto branchNameToken = [&subBranchName]() {
+         std::string copy{subBranchName};
+         if (!copy.empty() && !(copy == ".")) {
+            // The branch name ends with '.' which is allowed, continue the search
+            // including the last dot in the branch name
+            auto searchPos = (copy.back() == '.' ? copy.size() - 2 : copy.size());
+            if (auto pos = copy.find_last_of('.', searchPos); pos != std::string::npos)
+               return copy.substr(pos + 1);
+         }
+         return copy;
+      }();
+      // Record names of sub branches, which could reapper in different branch hierarchies of the same dataset. For
+      // example, the 'Muon' branch could have sub-branch 'pt', as well as the 'Electron' branch. Later we will
+      // disambiguate by removing the top-level 'pt' branch which TTree doesn't warn about and would end up pointing to
+      // the first sub-branch encountered in this exploration
+      if (!duplicateTokens.insert({subBranchName, 1}).second) {
+         duplicateTokens[subBranchName]++;
+      }
+      if (branchNameToken != subBranchName) {
+         if (!duplicateTokens.insert({branchNameToken, 1}).second) {
+            duplicateTokens[branchNameToken]++;
+         }
+      }
+
       auto fullName = prefix + subBranchName;
 
       if (auto subNameFirstDot = subBranchName.find_first_of('.'); subNameFirstDot != std::string::npos) {
@@ -141,7 +166,7 @@ void ExploreBranch(TTree &t, std::set<std::string> &bNamesReg, std::vector<std::
       if (!prefix.empty())
          newPrefix = fullName + ".";
 
-      ExploreBranch(t, bNamesReg, bNames, subBranch, newPrefix, friendName, allowDuplicates);
+      ExploreBranch(t, duplicateTokens, bNamesReg, bNames, subBranch, newPrefix, friendName, allowDuplicates);
 
       auto branchDirectlyFromTree = t.GetBranch(fullName.c_str());
       if (!branchDirectlyFromTree)
@@ -155,7 +180,8 @@ void ExploreBranch(TTree &t, std::set<std::string> &bNamesReg, std::vector<std::
    }
 }
 
-void GetBranchNamesImpl(TTree &t, std::set<std::string> &bNamesReg, std::vector<std::string> &bNames,
+void GetBranchNamesImpl(TTree &t, std::unordered_map<std::string, unsigned int> &duplicateTokens,
+                        std::set<std::string> &bNamesReg, std::vector<std::string> &bNames,
                         std::set<TTree *> &analysedTrees, std::string &friendName, bool allowDuplicates)
 {
    std::set<TLeaf *> foundLeaves;
@@ -192,7 +218,7 @@ void GetBranchNamesImpl(TTree &t, std::set<std::string> &bNamesReg, std::vector<
             }
          } else if (branch->IsA() == TBranchObject::Class()) {
             // TBranchObject
-            ExploreBranch(t, bNamesReg, bNames, branch, branchName + ".", friendName, allowDuplicates);
+            ExploreBranch(t, duplicateTokens, bNamesReg, bNames, branch, branchName + ".", friendName, allowDuplicates);
             InsertBranchName(bNamesReg, bNames, branchName, friendName, allowDuplicates);
          } else {
             // TBranchElement
@@ -207,9 +233,10 @@ void GetBranchNamesImpl(TTree &t, std::set<std::string> &bNamesReg, std::vector<
                dotIsImplied = true;
 
             if (dotIsImplied || branchName.back() == '.')
-               ExploreBranch(t, bNamesReg, bNames, branch, "", friendName, allowDuplicates);
+               ExploreBranch(t, duplicateTokens, bNamesReg, bNames, branch, "", friendName, allowDuplicates);
             else
-               ExploreBranch(t, bNamesReg, bNames, branch, branchName + ".", friendName, allowDuplicates);
+               ExploreBranch(t, duplicateTokens, bNamesReg, bNames, branch, branchName + ".", friendName,
+                             allowDuplicates);
 
             InsertBranchName(bNamesReg, bNames, branchName, friendName, allowDuplicates);
          }
@@ -234,19 +261,32 @@ void GetBranchNamesImpl(TTree &t, std::set<std::string> &bNamesReg, std::vector<
       else
          frName = std::string(friendTree->GetName());
 
-      GetBranchNamesImpl(*friendTree, bNamesReg, bNames, analysedTrees, frName, allowDuplicates);
+      GetBranchNamesImpl(*friendTree, duplicateTokens, bNamesReg, bNames, analysedTrees, frName, allowDuplicates);
    }
 }
 
 ///////////////////////////////////////////////////////////////////////////////
 /// Get all the branches names, including the ones of the friend trees
 std::vector<std::string> RetrieveDatasetSchema(TTree &t, bool allowDuplicates = true)
 {
+   std::unordered_map<std::string, unsigned int> duplicateTokens;
+
    std::set<std::string> bNamesSet;
    std::vector<std::string> bNames;
    std::set<TTree *> analysedTrees;
    std::string emptyFrName = "";
-   GetBranchNamesImpl(t, bNamesSet, bNames, analysedTrees, emptyFrName, allowDuplicates);
+   GetBranchNamesImpl(t, duplicateTokens, bNamesSet, bNames, analysedTrees, emptyFrName, allowDuplicates);
+
+   // Remove all sub-branches that have duplicate names between different branch hierarchies of the dataset
+   bNames.erase(std::remove_if(bNames.begin(), bNames.end(),
+                               [&duplicateTokens](const auto &name) {
+                                  if (auto it = duplicateTokens.find(name);
+                                      it != duplicateTokens.end() && it->second > 1)
+                                     return true;
+                                  return false;
+                               }),
+                bNames.end());
+
    return bNames;
 }
 } // namespace

tree/dataframe/test/CMakeLists.txt

@@ -13,7 +13,8 @@ ROOT_ADD_GTEST(dataframe_cache dataframe_cache.cxx LIBRARIES ROOTDataFrame)
 ROOT_ADD_GTEST(dataframe_callbacks dataframe_callbacks.cxx LIBRARIES ROOTDataFrame)
 ROOT_ADD_GTEST(dataframe_cloning dataframe_cloning.cxx LIBRARIES ROOTDataFrame)
 ROOT_ADD_GTEST(dataframe_histomodels dataframe_histomodels.cxx LIBRARIES ROOTDataFrame)
-ROOT_ADD_GTEST(dataframe_interface dataframe_interface.cxx LIBRARIES ROOTDataFrame)
+ROOT_GENERATE_DICTIONARY(Class19392Dict ${CMAKE_CURRENT_SOURCE_DIR}/Class19392.hxx LINKDEF ${CMAKE_CURRENT_SOURCE_DIR}/Class19392LinkDef.hxx NO_CXXMODULE DEPENDENCIES RIO)
+ROOT_ADD_GTEST(dataframe_interface dataframe_interface.cxx Class19392Dict.cxx LIBRARIES ROOTDataFrame)
 ROOT_ADD_GTEST(dataframe_nodes dataframe_nodes.cxx LIBRARIES ROOTDataFrame)
 ROOT_ADD_GTEST(dataframe_regression dataframe_regression.cxx LIBRARIES Physics ROOTDataFrame GenVector)
 ROOT_ADD_GTEST(dataframe_utils dataframe_utils.cxx LIBRARIES ROOTDataFrame)

tree/dataframe/test/Class19392.hxx

@@ -0,0 +1,24 @@
+#ifndef ROOT_DATAFRAME_TEST_SIMPLE_ELECTRON
+#define ROOT_DATAFRAME_TEST_SIMPLE_ELECTRON
+
+// Classes required for the regression test of https://github.com/root-project/root/issues/19392
+
+struct SimpleElectron {
+   float electron_pt{};
+};
+
+struct Wrapper {
+   SimpleElectron electron{};
+   float electron_pt{};
+};
+
+struct UniqueInner {
+   int uniqueInnerInt{};
+};
+
+struct UniqueWrapper {
+   UniqueInner uniqueInner{};
+   int uniqueWrapperInt{};
+};
+
+#endif

tree/dataframe/test/Class19392LinkDef.hxx

@@ -0,0 +1,8 @@
+#ifdef __CLING__
+
+#pragma link C++ class SimpleElectron + ;
+#pragma link C++ class Wrapper + ;
+#pragma link C++ class UniqueInner + ;
+#pragma link C++ class UniqueWrapper + ;
+
+#endif

tree/dataframe/test/dataframe_interface.cxx

@@ -15,6 +15,8 @@
 
 #include "gtest/gtest.h"
 
+#include "Class19392.hxx"
+
 #include <thread>
 
 using namespace ROOT;
@@ -1062,3 +1064,126 @@ TEST(RDataFrameInterface, GetNFilesFromMoreFiles)
    ROOT::RDataFrame df{"t", filenames};
    EXPECT_EQ(df.GetNFiles(), 3);
 }
+
+void expect_colnames_eq(const std::vector<std::string> &v1, const std::vector<std::string> &v2)
+{
+   ASSERT_EQ(v1.size(), v2.size()) << "Vectors 'v1' and 'v2' are of unequal length";
+   for (std::size_t i = 0ull; i < v1.size(); ++i) {
+      EXPECT_EQ(v1[i], v2[i]) << "Vectors 'v1' and 'v2' differ at index " << i;
+   }
+}
+
+// https://github.com/root-project/root/issues/19392
+TEST(RDataFrameInterface, GH19392)
+{
+   class FileRAII {
+   private:
+      std::string fPath;
+
+   public:
+      explicit FileRAII(const std::string &path) : fPath(path) {}
+      ~FileRAII() { std::remove(fPath.c_str()); }
+      auto GetPath() const { return fPath.c_str(); }
+   };
+
+   FileRAII fileraii{"dataframe_interface_gh19392.root"};
+   const auto treeName{"tree"};
+
+   {
+      auto file = std::make_unique<TFile>(fileraii.GetPath(), "RECREATE");
+      auto tree = std::make_unique<TTree>(treeName, treeName);
+
+      SimpleElectron el1;
+      el1.electron_pt = 10.f;
+
+      SimpleElectron el2;
+      el2.electron_pt = 20.f;
+
+      // The wrapper class also has the same data member named 'electron_pt'
+      // just to make the parsing exercise more complicated
+      Wrapper wrap1;
+      wrap1.electron_pt = 30.f;
+      wrap1.electron.electron_pt = 40.f;
+
+      Wrapper wrap2;
+      wrap2.electron_pt = 50.f;
+      wrap2.electron.electron_pt = 60.f;
+
+      UniqueInner ui{};
+      ui.uniqueInnerInt = 70;
+
+      UniqueWrapper uw{};
+      uw.uniqueInner.uniqueInnerInt = 80;
+      uw.uniqueWrapperInt = 90;
+
+      // Create a true top-level branch which has the same name as one of the
+      // inner sub-branches. This should not confuse the retrieval of the schema
+      SimpleElectron electron;
+      float electron_pt{100.f};
+
+      tree->Branch("electron_pt", &electron_pt);
+      tree->Branch("electron", &electron);
+      tree->Branch("el1", &el1);
+      tree->Branch("el2", &el2);
+      tree->Branch("wr1", &wrap1);
+      tree->Branch("wr2", &wrap2);
+      tree->Branch("ui", &ui);
+      tree->Branch("uw", &uw);
+      tree->Fill();
+      tree->Write();
+   }
+   ROOT::RDataFrame df(treeName, fileraii.GetPath());
+   const auto columns = df.GetColumnNames();
+   // The wrapper classes also have the same data member named 'electron_pt'
+   // just to make the parsing exercise more complicated
+   // It's thus expected to see e.g. 'wr1.electron_pt'
+   // The 'uniqueInnerInt' branch should not be exposed as top-level branch,
+   // but 'uniqueInner' and 'uniqueWrapperInt' should as they are unambiguous
+   const std::vector<std::string> expectedCols{"el1",
+                                               "el1.electron_pt",
+                                               "el2",
+                                               "el2.electron_pt",
+                                               "ui",
+                                               "ui.uniqueInnerInt",
+                                               "uniqueInner",
+                                               "uniqueInner.uniqueInnerInt",
+                                               "uniqueWrapperInt",
+                                               "uw",
+                                               "uw.uniqueInner",
+                                               "uw.uniqueInner.uniqueInnerInt",
+                                               "uw.uniqueWrapperInt",
+                                               "wr1",
+                                               "wr1.electron",
+                                               "wr1.electron.electron_pt",
+                                               "wr1.electron_pt",
+                                               "wr2",
+                                               "wr2.electron",
+                                               "wr2.electron.electron_pt",
+                                               "wr2.electron_pt"};
+   expect_colnames_eq(columns, expectedCols);
+
+   // Check all values separately, ensures that all full leaf names point to
+   // the correct values
+   auto el1_pt = df.Take<float>("el1.electron_pt");
+   auto el2_pt = df.Take<float>("el2.electron_pt");
+   auto wr1_pt = df.Take<float>("wr1.electron_pt");
+   auto wr1_el_pt = df.Take<float>("wr1.electron.electron_pt");
+   auto wr2_pt = df.Take<float>("wr2.electron_pt");
+   auto wr2_el_pt = df.Take<float>("wr2.electron.electron_pt");
+   auto ui = df.Take<int>("ui.uniqueInnerInt");
+   // Use the unambiguous sub-branch name in this case
+   auto uw_ui_ui = df.Take<int>("uniqueInner.uniqueInnerInt");
+   auto uwi = df.Take<int>("uniqueWrapperInt");
+   auto electron_pt = df.Take<float>("electron_pt");
+
+   EXPECT_FLOAT_EQ(el1_pt->at(0), 10.f);
+   EXPECT_FLOAT_EQ(el2_pt->at(0), 20.f);
+   EXPECT_FLOAT_EQ(wr1_pt->at(0), 30.f);
+   EXPECT_FLOAT_EQ(wr1_el_pt->at(0), 40.f);
+   EXPECT_FLOAT_EQ(wr2_pt->at(0), 50.f);
+   EXPECT_FLOAT_EQ(wr2_el_pt->at(0), 60.f);
+   EXPECT_EQ(ui->at(0), 70);
+   EXPECT_EQ(uw_ui_ui->at(0), 80);
+   EXPECT_EQ(uwi->at(0), 90);
+   EXPECT_FLOAT_EQ(electron_pt->at(0), 100.f);
+}