From eab6cf954e96a2a9b17ec2bc0f27d587b5a2e954 Mon Sep 17 00:00:00 2001 From: Vincenzo Eduardo Padulano Date: Thu, 17 Jul 2025 16:01:09 +0200 Subject: [PATCH] [df] Avoid including ambiguous TTree branch names When the input data source is a TTree, GetColumnNames gathers the list of all the available TTree branches. In case there are two branches in the tree (e.g. `el1` and `el2`), each of them has a sub-branch with the same name (e.g. `electron_pt`), TTree allows calling `GetBranch("electron_pt")` and returns the pointer to the sub-branch of the first main branch (i.e. `el1.electron_pt`). This behaviour can lead to ambiguities, thus avoid exposing the ambiguous column name via RDF. A test is added to exemplify this case. --- tree/dataframe/src/RTTreeDS.cxx | 58 +++++++-- tree/dataframe/test/CMakeLists.txt | 3 +- tree/dataframe/test/Class19392.hxx | 24 ++++ tree/dataframe/test/Class19392LinkDef.hxx | 8 ++ tree/dataframe/test/dataframe_interface.cxx | 125 ++++++++++++++++++++ 5 files changed, 208 insertions(+), 10 deletions(-) create mode 100644 tree/dataframe/test/Class19392.hxx create mode 100644 tree/dataframe/test/Class19392LinkDef.hxx diff --git a/tree/dataframe/src/RTTreeDS.cxx b/tree/dataframe/src/RTTreeDS.cxx index 7d2ecd8aef91f..9988fb75c44a8 100644 --- a/tree/dataframe/src/RTTreeDS.cxx +++ b/tree/dataframe/src/RTTreeDS.cxx @@ -110,8 +110,9 @@ void InsertBranchName(std::set &bNamesReg, std::vector foundLeaves.insert(leaf); } -void ExploreBranch(TTree &t, std::set &bNamesReg, std::vector &bNames, TBranch *b, - std::string prefix, std::string &friendName, bool allowDuplicates) +void ExploreBranch(TTree &t, std::unordered_map &duplicateTokens, + std::set &bNamesReg, std::vector &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 &bNamesReg, std::vectorGetListOfBranches()) { TBranch *subBranch = static_cast(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 &bNamesReg, std::vector &bNamesReg, std::vector &bNamesReg, std::vector &bNames, +void GetBranchNamesImpl(TTree &t, std::unordered_map &duplicateTokens, + std::set &bNamesReg, std::vector &bNames, std::set &analysedTrees, std::string &friendName, bool allowDuplicates) { std::set foundLeaves; @@ -192,7 +218,7 @@ void GetBranchNamesImpl(TTree &t, std::set &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 &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,7 +261,7 @@ void GetBranchNamesImpl(TTree &t, std::set &bNamesReg, std::vector< else frName = std::string(friendTree->GetName()); - GetBranchNamesImpl(*friendTree, bNamesReg, bNames, analysedTrees, frName, allowDuplicates); + GetBranchNamesImpl(*friendTree, duplicateTokens, bNamesReg, bNames, analysedTrees, frName, allowDuplicates); } } @@ -242,11 +269,24 @@ void GetBranchNamesImpl(TTree &t, std::set &bNamesReg, std::vector< /// Get all the branches names, including the ones of the friend trees std::vector RetrieveDatasetSchema(TTree &t, bool allowDuplicates = true) { + std::unordered_map duplicateTokens; + std::set bNamesSet; std::vector bNames; std::set 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 diff --git a/tree/dataframe/test/CMakeLists.txt b/tree/dataframe/test/CMakeLists.txt index fb2e05122eab0..ca2d080f7505a 100644 --- a/tree/dataframe/test/CMakeLists.txt +++ b/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) diff --git a/tree/dataframe/test/Class19392.hxx b/tree/dataframe/test/Class19392.hxx new file mode 100644 index 0000000000000..3412b43800235 --- /dev/null +++ b/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 diff --git a/tree/dataframe/test/Class19392LinkDef.hxx b/tree/dataframe/test/Class19392LinkDef.hxx new file mode 100644 index 0000000000000..fcff2fd6ec226 --- /dev/null +++ b/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 diff --git a/tree/dataframe/test/dataframe_interface.cxx b/tree/dataframe/test/dataframe_interface.cxx index 8b55073d1ebbd..0ea02ba7f536d 100644 --- a/tree/dataframe/test/dataframe_interface.cxx +++ b/tree/dataframe/test/dataframe_interface.cxx @@ -15,6 +15,8 @@ #include "gtest/gtest.h" +#include "Class19392.hxx" + #include 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 &v1, const std::vector &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(fileraii.GetPath(), "RECREATE"); + auto tree = std::make_unique(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 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("el1.electron_pt"); + auto el2_pt = df.Take("el2.electron_pt"); + auto wr1_pt = df.Take("wr1.electron_pt"); + auto wr1_el_pt = df.Take("wr1.electron.electron_pt"); + auto wr2_pt = df.Take("wr2.electron_pt"); + auto wr2_el_pt = df.Take("wr2.electron.electron_pt"); + auto ui = df.Take("ui.uniqueInnerInt"); + // Use the unambiguous sub-branch name in this case + auto uw_ui_ui = df.Take("uniqueInner.uniqueInnerInt"); + auto uwi = df.Take("uniqueWrapperInt"); + auto electron_pt = df.Take("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); +}