From 307d7b5a14b49b04819645d559bfea3d12576ba7 Mon Sep 17 00:00:00 2001 From: Vincenzo Eduardo Padulano Date: Wed, 9 Jul 2025 18:36:05 +0200 Subject: [PATCH 1/2] [tree] Improve logic to refresh branch addresses of TChain friends Part of the logic of TChain::LoadTree deals with updating the status of friend trees. First, the current entry is loaded on all friends. In case the friend tree is a chain that needs to switch to a new tree, the branch addresses of the friend are also updated. This logic only considered the case of top-level friends of a TChain. Notably, it did not consider the case of a main TChain where one (or more) of the trees in the chain have themselves some friends, but not the main chain itself: TChain c1 --> TTree t1 TChain c2 This commit condenses the logic that updates friends in a separate method of TChain. The new method is called in TChain::LoadTree, either in the pre-existing case of top-level friends or in the new case of friends of the current tree. The new method, called RefreshFriendAddresses, needs to act: - On the main chain and on the current tree, in the case of top-level friends - On the current tree alone, in the case of friends of the current tree In the two scenarios, either we load the global entry of the chain on the friends, or the local entry of the current tree. This is done to conserve proper alignment. --- tree/tree/inc/TChain.h | 1 + tree/tree/src/TChain.cxx | 196 +++++++++++++++++----------- tree/treeplayer/test/CMakeLists.txt | 4 + tree/treeplayer/test/gh16805.cxx | 163 +++++++++++++++++++++++ 4 files changed, 290 insertions(+), 74 deletions(-) create mode 100644 tree/treeplayer/test/gh16805.cxx diff --git a/tree/tree/inc/TChain.h b/tree/tree/inc/TChain.h index 72d12df657a39..d8e6cc2da171c 100644 --- a/tree/tree/inc/TChain.h +++ b/tree/tree/inc/TChain.h @@ -49,6 +49,7 @@ class TChain : public TTree { TChain& operator=(const TChain&); // not implemented void ParseTreeFilename(const char *name, TString &filename, TString &treename, TString &query, TString &suffix) const; + Long64_t RefreshFriendAddresses(TTree &mainTree, Long64_t entry); protected: void InvalidateCurrentTree(); diff --git a/tree/tree/src/TChain.cxx b/tree/tree/src/TChain.cxx index b131210c40bce..47449de9de5ba 100644 --- a/tree/tree/src/TChain.cxx +++ b/tree/tree/src/TChain.cxx @@ -1231,6 +1231,107 @@ Int_t TChain::LoadBaskets(Long64_t /*maxmemory*/) return 0; } +//////////////////////////////////////////////////////////////////////////////// +/// Refresh branch/leaf addresses of either top-level friends or friends of +/// the current tree in the chain. +/// +/// \param[in] mainTree The main tree, either the chain itself or the current tree +/// \param[in] entry The entry loaded on the friends. If the main tree is the +/// chain, this is the global entry of the chain being read. Otherwise, it is +/// the local entry of the current tree being read. +Long64_t TChain::RefreshFriendAddresses(TTree &mainTree, Long64_t entry) +{ + auto *friendList = mainTree.GetListOfFriends(); + if (!friendList) + return 0; + + TFriendLock lock(&mainTree, kLoadTree); + + bool needUpdate = false; + + for (auto *frEl : ROOT::Detail::TRangeStaticCast(*friendList)) { + auto *frTree = frEl->GetTree(); + // Depending on whether we are traversing the friends of the chain itself + // or the friends of the current tree in the chain, the meaning of the + // entry parameter changes. In the first case, we pass the current chain + // global entry number. In the latter case, we pass the local tree entry + // number. + frTree->LoadTreeFriend(entry, &mainTree); + } + + // If the mainTree is this chain, innerTree is the current tree in the chain. + // If instead it is the current tree, innerTree points to the same. + auto *innerTreePtr = mainTree.GetTree(); + assert(innerTreePtr != nullptr); + auto &innerTree = *innerTreePtr; + + if (auto *innerFriendList = innerTree.GetListOfFriends()) { + // If the current tree has friends, check if they were mark for update + // when switching to the following tree, detect it so that we later we + // actually refresh the addresses of the friends. + for (auto *frEl : ROOT::Detail::TRangeStaticCast(*innerFriendList)) { + if (frEl->IsUpdated()) { + needUpdate = true; + frEl->ResetUpdated(); + } + } + } + + if (!needUpdate) + return 0; + + // Update the branch/leaf addresses and the list of leaves in all + // TTreeFormula of the TTreePlayer (if any). + for (auto *chainEl : ROOT::Detail::TRangeStaticCast(*fStatus)) { + // Set the branch status of all the chain elements, which may include also + // branches that are available in friends. Only set the branch status + // if it has a value provided by the user + Int_t status = chainEl->GetStatus(); + if (status != -1) + innerTree.SetBranchStatus(chainEl->GetName(), status); + + // Set the branch addresses for the newly opened file. + void *addr = chainEl->GetBaddress(); + if (!addr) + continue; + + TBranch *br = innerTree.GetBranch(chainEl->GetName()); + TBranch **pp = chainEl->GetBranchPtr(); + if (pp) { + // FIXME: What if br is zero here? + *pp = br; + } + if (!br) + continue; + + if (!chainEl->GetCheckedType()) { + Int_t res = CheckBranchAddressType(br, TClass::GetClass(chainEl->GetBaddressClassName()), + (EDataType)chainEl->GetBaddressType(), chainEl->GetBaddressIsPtr()); + if ((res & kNeedEnableDecomposedObj) && !br->GetMakeClass()) { + br->SetMakeClass(true); + } + chainEl->SetDecomposedObj(br->GetMakeClass()); + chainEl->SetCheckedType(true); + } + // FIXME: We may have to tell the branch it should + // not be an owner of the object pointed at. + br->SetAddress(addr); + if (TestBit(kAutoDelete)) { + br->SetAutoDelete(true); + } + } + if (fPlayer) { + fPlayer->UpdateFormulaLeaves(); + } + // Notify user if requested. + if (fNotify) { + if (!fNotify->Notify()) + return -6; + } + + return 0; +} + //////////////////////////////////////////////////////////////////////////////// /// Find the tree which contains entry, and set it as the current tree. /// @@ -1299,83 +1400,30 @@ Long64_t TChain::LoadTree(Long64_t entry) // next loop). fTree->LoadTree(treeReadEntry); - if (fFriends) { - // The current tree has not changed but some of its friends might. - // - TIter next(fFriends); - TFriendLock lock(this, kLoadTree); - TFriendElement* fe = nullptr; - while ((fe = (TFriendElement*) next())) { - TTree* at = fe->GetTree(); - // If the tree is a - // direct friend of the chain, it should be scanned - // used the chain entry number and NOT the tree entry - // number (treeReadEntry) hence we do: - at->LoadTreeFriend(entry, this); - } - bool needUpdate = false; - if (fTree->GetListOfFriends()) { - for(auto fetree : ROOT::Detail::TRangeStaticCast(*fTree->GetListOfFriends())) { - if (fetree->IsUpdated()) { - needUpdate = true; - fetree->ResetUpdated(); - } - } + // If we've been called while refreshing addresses of friend elements, we + // should notify them that they've been updated + if (fExternalFriends) { + for (auto external_fe : ROOT::Detail::TRangeStaticCast(*fExternalFriends)) { + external_fe->MarkUpdated(); } - if (needUpdate) { - // Update the branch/leaf addresses and - // the list of leaves in all TTreeFormula of the TTreePlayer (if any). - - // Set the branch statuses for the newly opened file. - TChainElement *frelement; - TIter fnext(fStatus); - while ((frelement = (TChainElement*) fnext())) { - Int_t status = frelement->GetStatus(); - // Only set the branch status if it has a value provided - // by the user - if (status != -1) - fTree->SetBranchStatus(frelement->GetName(), status); - } + } - // Set the branch addresses for the newly opened file. - fnext.Reset(); - while ((frelement = (TChainElement*) fnext())) { - void* addr = frelement->GetBaddress(); - if (addr) { - TBranch* br = fTree->GetBranch(frelement->GetName()); - TBranch** pp = frelement->GetBranchPtr(); - if (pp) { - // FIXME: What if br is zero here? - *pp = br; - } - if (br) { - if (!frelement->GetCheckedType()) { - Int_t res = CheckBranchAddressType(br, TClass::GetClass(frelement->GetBaddressClassName()), - (EDataType) frelement->GetBaddressType(), frelement->GetBaddressIsPtr()); - if ((res & kNeedEnableDecomposedObj) && !br->GetMakeClass()) { - br->SetMakeClass(true); - } - frelement->SetDecomposedObj(br->GetMakeClass()); - frelement->SetCheckedType(true); - } - // FIXME: We may have to tell the branch it should - // not be an owner of the object pointed at. - br->SetAddress(addr); - if (TestBit(kAutoDelete)) { - br->SetAutoDelete(true); - } - } - } - } - if (fPlayer) { - fPlayer->UpdateFormulaLeaves(); - } - // Notify user if requested. - if (fNotify) { - if(!fNotify->Notify()) return -6; - } - } + Long64_t refreshFriendAddressesRet{}; + if (fFriends) { + // If this chain has friends, we need to pass the global entry to keep + // the friends aligned to the main chain + refreshFriendAddressesRet = RefreshFriendAddresses(*this, entry); + } else if (fTree->GetListOfFriends()) { + // If instead the current tree has friends, we pass the local entry + // we're reading of the current tree to keep the alignment of the friends + // with respect to it + refreshFriendAddressesRet = RefreshFriendAddresses(*fTree, treeReadEntry); } + // At the moment RefreshFriendAddresses has one failure state if something + // went wrong when notifying the registered listeners + if (refreshFriendAddressesRet == -6) + return refreshFriendAddressesRet; + return treeReadEntry; } diff --git a/tree/treeplayer/test/CMakeLists.txt b/tree/treeplayer/test/CMakeLists.txt index c8229fbb00d60..0d84cd4c802e2 100644 --- a/tree/treeplayer/test/CMakeLists.txt +++ b/tree/treeplayer/test/CMakeLists.txt @@ -17,6 +17,8 @@ endif() ROOT_ADD_GTEST(treeplayer_gh16804 gh16804.cxx LIBRARIES TreePlayer) +ROOT_ADD_GTEST(treeplayer_gh16805 gh16805.cxx LIBRARIES TreePlayer) + ROOT_ADD_GTEST(treeplayer_leafs leafs.cxx LIBRARIES TreePlayer) add_custom_command(TARGET treeplayer_leafs POST_BUILD COMMAND ${CMAKE_COMMAND} -E copy ${CMAKE_CURRENT_SOURCE_DIR}/data.h data.h) @@ -39,3 +41,5 @@ if(imt) ROOT_ADD_GTEST(treeprocessormt_remotefiles treeprocs/treeprocessormt_remotefiles.cxx LIBRARIES TreePlayer) endif() endif() + + diff --git a/tree/treeplayer/test/gh16805.cxx b/tree/treeplayer/test/gh16805.cxx new file mode 100644 index 0000000000000..fb08fe5c3f6e8 --- /dev/null +++ b/tree/treeplayer/test/gh16805.cxx @@ -0,0 +1,163 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +struct GH16805Regression : public ::testing::Test { + + const static inline std::vector fStepZeroFileNames{"gh16805_regression_stepzero_0.root", + "gh16805_regression_stepzero_1.root"}; + const static inline std::vector> fStepZeroEntryRanges{{0, 10}, {10, 20}}; + constexpr static auto fStepZeroTreeName{"stepzero"}; + + constexpr static auto fStepOneFileName{"gh16805_regression_stepone.root"}; + constexpr static std::pair fStepOneEntryRange{100, 120}; + constexpr static auto fStepOneTreeName{"stepone"}; + + static void writeStepZero(const char *treeName, const char *fileName, int begin, int end) + { + auto file = std::make_unique(fileName, "recreate"); + auto tree = std::make_unique(treeName, treeName); + int stepZeroBr1, stepZeroBr2 = 0; + tree->Branch("stepZeroBr1", &stepZeroBr1); + tree->Branch("stepZeroBr2", &stepZeroBr2); + + for (stepZeroBr1 = begin, stepZeroBr2 = begin * 2; stepZeroBr1 < end; + ++stepZeroBr1, stepZeroBr2 = stepZeroBr1 * 2) + tree->Fill(); + file->Write(); + } + + static void writeStepOne(const char *treeName, const char *fileName, int begin, int end) + { + auto file = std::make_unique(fileName, "recreate"); + auto tree = std::make_unique(treeName, treeName); + int br1 = -1; + tree->Branch("stepOneBr1", &br1); + + for (br1 = begin; br1 < end; ++br1) + tree->Fill(); + + file->Write(); + } + + static void SetUpTestSuite() + { + for (decltype(fStepZeroFileNames.size()) i = 0; i < fStepZeroFileNames.size(); i++) { + writeStepZero(fStepZeroTreeName, fStepZeroFileNames[i].c_str(), fStepZeroEntryRanges[i].first, + fStepZeroEntryRanges[i].second); + } + writeStepOne(fStepOneTreeName, fStepOneFileName, fStepOneEntryRange.first, fStepOneEntryRange.second); + } + static void TearDownTestSuite() + { + for (const auto &f : fStepZeroFileNames) + std::remove(f.c_str()); + + std::remove(fStepOneFileName); + } +}; + +TEST_F(GH16805Regression, Test) +{ + // Merge the step0 files with a chain. + auto stepZeroChain = std::make_unique(fStepZeroTreeName); + for (const auto &f : fStepZeroFileNames) + stepZeroChain->Add(f.c_str()); + + auto stepOneChain = std::make_unique(fStepOneTreeName); + stepOneChain->Add(fStepOneFileName); + + // Load the first tree in the chain + stepOneChain->LoadTree(0); + auto stepOneFirstTree = stepOneChain->GetTree()->GetTree(); + ASSERT_NE(stepOneFirstTree, nullptr); + + stepOneFirstTree->AddFriend(stepZeroChain.get()); + + stepOneChain->LoadTree(0); // needed for the branches of the friend TChain to be available + + // Now iterate over the chain and check the contents of the branches + // inherited from the friends. + int stepZeroBr1 = -1, stepZeroBr2 = -1, stepOneBr1; + ASSERT_EQ(stepOneChain->SetBranchAddress("stepZeroBr1", &stepZeroBr1), 0); + ASSERT_EQ(stepOneChain->SetBranchAddress("stepZeroBr2", &stepZeroBr2), 0); + ASSERT_EQ(stepOneChain->SetBranchAddress("stepOneBr1", &stepOneBr1), 0); + + std::vector expectedStepZeroBr1(20); + std::vector expectedStepZeroBr2(20); + std::vector expectedStepOneBr1(20); + + std::iota(expectedStepZeroBr1.begin(), expectedStepZeroBr1.end(), 0); + std::generate(expectedStepZeroBr2.begin(), expectedStepZeroBr2.end(), [n = 0]() mutable { + auto res = n * 2; + n++; + return res; + }); + std::iota(expectedStepOneBr1.begin(), expectedStepOneBr1.end(), 100); + + for (Long64_t i = 0; i < stepOneChain->GetEntriesFast(); ++i) { + { + stepOneChain->GetEntry(i); + EXPECT_EQ(expectedStepZeroBr1[i], stepZeroBr1); + EXPECT_EQ(expectedStepZeroBr2[i], stepZeroBr2); + EXPECT_EQ(expectedStepOneBr1[i], stepOneBr1); + } + } + + // Now test with TTree::Scan + std::ostringstream strCout; + { + if (auto *treePlayer = static_cast(stepOneChain->GetPlayer())) { + struct FileRAII { + const char *fPath; + FileRAII(const char *name) : fPath(name) {} + ~FileRAII() { std::remove(fPath); } + } redirectFile{"regression_16805_redirect.txt"}; + treePlayer->SetScanRedirect(true); + treePlayer->SetScanFileName(redirectFile.fPath); + stepOneChain->Scan("stepZeroBr1:stepZeroBr2:stepOneBr1", "", "colsize=12"); + + std::ifstream redirectStream(redirectFile.fPath); + std::stringstream redirectOutput; + redirectOutput << redirectStream.rdbuf(); + + const static std::string expectedScanOut{ + R"Scan(********************************************************* +* Row * stepZeroBr1 * stepZeroBr2 * stepOneBr1 * +********************************************************* +* 0 * 0 * 0 * 100 * +* 1 * 1 * 2 * 101 * +* 2 * 2 * 4 * 102 * +* 3 * 3 * 6 * 103 * +* 4 * 4 * 8 * 104 * +* 5 * 5 * 10 * 105 * +* 6 * 6 * 12 * 106 * +* 7 * 7 * 14 * 107 * +* 8 * 8 * 16 * 108 * +* 9 * 9 * 18 * 109 * +* 10 * 10 * 20 * 110 * +* 11 * 11 * 22 * 111 * +* 12 * 12 * 24 * 112 * +* 13 * 13 * 26 * 113 * +* 14 * 14 * 28 * 114 * +* 15 * 15 * 30 * 115 * +* 16 * 16 * 32 * 116 * +* 17 * 17 * 34 * 117 * +* 18 * 18 * 36 * 118 * +* 19 * 19 * 38 * 119 * +********************************************************* +)Scan"}; + EXPECT_EQ(redirectOutput.str(), expectedScanOut); + } else + throw std::runtime_error("Could not retrieve TTreePlayer from main tree!"); + } +} From c7ff3dc712b1a2057fa889cf55228537ac6ce73d Mon Sep 17 00:00:00 2001 From: Vincenzo Eduardo Padulano Date: Mon, 14 Jul 2025 11:30:03 +0200 Subject: [PATCH 2/2] Try moving MarkUpdate loop before --- tree/tree/src/TChain.cxx | 20 ++++++-------------- 1 file changed, 6 insertions(+), 14 deletions(-) diff --git a/tree/tree/src/TChain.cxx b/tree/tree/src/TChain.cxx index 47449de9de5ba..bf3dfce13f7ac 100644 --- a/tree/tree/src/TChain.cxx +++ b/tree/tree/src/TChain.cxx @@ -1393,6 +1393,12 @@ Long64_t TChain::LoadTree(Long64_t entry) Long64_t treeReadEntry = entry - fTreeOffset[treenum]; fReadEntry = entry; + if (fExternalFriends) { + for (auto external_fe : ROOT::Detail::TRangeStaticCast(*fExternalFriends)) { + external_fe->MarkUpdated(); + } + } + // If entry belongs to the current tree return entry. if (fTree && treenum == fTreeNumber) { // First set the entry the tree on its owns friends @@ -1400,14 +1406,6 @@ Long64_t TChain::LoadTree(Long64_t entry) // next loop). fTree->LoadTree(treeReadEntry); - // If we've been called while refreshing addresses of friend elements, we - // should notify them that they've been updated - if (fExternalFriends) { - for (auto external_fe : ROOT::Detail::TRangeStaticCast(*fExternalFriends)) { - external_fe->MarkUpdated(); - } - } - Long64_t refreshFriendAddressesRet{}; if (fFriends) { // If this chain has friends, we need to pass the global entry to keep @@ -1427,12 +1425,6 @@ Long64_t TChain::LoadTree(Long64_t entry) return treeReadEntry; } - if (fExternalFriends) { - for(auto external_fe : ROOT::Detail::TRangeStaticCast(*fExternalFriends)) { - external_fe->MarkUpdated(); - } - } - // Delete the current tree and open the new tree. TTreeCache* tpf = nullptr; // Delete file unless the file owns this chain!