Skip to content

Commit 6cdf0d5

Browse files
vepadulanovepadulano
committed
[tree] WIP for #16805 
1 parent 10d7bd3 commit 6cdf0d5

File tree

4 files changed

+262
-74
lines changed

4 files changed

+262
-74
lines changed

‎tree/tree/inc/TChain.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ class TChain : public TTree {
5050
TChain& operator=(const TChain&); // not implemented
5151
void
5252
ParseTreeFilename(const char *name, TString &filename, TString &treename, TString &query, TString &suffix) const;
53+
Long64_t RefreshFriendAddresses(TTree &mainTree, TTree &innerTree, Long64_t entry);
5354

5455
protected:
5556
void InvalidateCurrentTree();

‎tree/tree/src/TChain.cxx

Lines changed: 97 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -1346,6 +1346,91 @@ Int_t TChain::LoadBaskets(Long64_t /*maxmemory*/)
13461346
return 0;
13471347
}
13481348

1349+
Long64_t TChain::RefreshFriendAddresses(TTree &mainTree, TTree &innerTree, Long64_t entry)
1350+
{
1351+
auto *friendList = mainTree.GetListOfFriends();
1352+
if (!friendList)
1353+
return 0;
1354+
1355+
TFriendLock lock(&mainTree, kLoadTree);
1356+
1357+
bool needUpdate = false;
1358+
1359+
for (auto *frEl : ROOT::Detail::TRangeStaticCast<TFriendElement>(*friendList)) {
1360+
auto *frTree = frEl->GetTree();
1361+
// Depending on whether we are traversing the friends of the chain itself
1362+
// or the friends of the current tree in the chain, the meaning of the
1363+
// entry parameter changes. In the first case, we pass the current chain
1364+
// global entry number. In the latter case, we pass the local tree entry
1365+
// number.
1366+
frTree->LoadTreeFriend(entry, &mainTree);
1367+
}
1368+
1369+
if (auto *innerFriendList = innerTree.GetListOfFriends()) {
1370+
for (auto *frEl : ROOT::Detail::TRangeStaticCast<TFriendElement>(*innerFriendList)) {
1371+
if (frEl->IsUpdated()) {
1372+
needUpdate = true;
1373+
frEl->ResetUpdated();
1374+
}
1375+
}
1376+
}
1377+
1378+
if (needUpdate) {
1379+
// Update the branch/leaf addresses and
1380+
// the list of leaves in all TTreeFormula of the TTreePlayer (if any).
1381+
1382+
// Set the branch statuses for the newly opened file.
1383+
for (auto *chainEl : ROOT::Detail::TRangeStaticCast<TChainElement>(*fStatus)) {
1384+
Int_t status = chainEl->GetStatus();
1385+
// Only set the branch status if it has a value provided
1386+
// by the user
1387+
if (status != -1)
1388+
innerTree.SetBranchStatus(chainEl->GetName(), status);
1389+
}
1390+
1391+
// Set the branch addresses for the newly opened file.
1392+
for (auto *chainEl : ROOT::Detail::TRangeStaticCast<TChainElement>(*fStatus)) {
1393+
void *addr = chainEl->GetBaddress();
1394+
if (addr) {
1395+
TBranch *br = innerTree.GetBranch(chainEl->GetName());
1396+
TBranch **pp = chainEl->GetBranchPtr();
1397+
if (pp) {
1398+
// FIXME: What if br is zero here?
1399+
*pp = br;
1400+
}
1401+
if (br) {
1402+
if (!chainEl->GetCheckedType()) {
1403+
Int_t res =
1404+
CheckBranchAddressType(br, TClass::GetClass(chainEl->GetBaddressClassName()),
1405+
(EDataType)chainEl->GetBaddressType(), chainEl->GetBaddressIsPtr());
1406+
if ((res & kNeedEnableDecomposedObj) && !br->GetMakeClass()) {
1407+
br->SetMakeClass(true);
1408+
}
1409+
chainEl->SetDecomposedObj(br->GetMakeClass());
1410+
chainEl->SetCheckedType(true);
1411+
}
1412+
// FIXME: We may have to tell the branch it should
1413+
// not be an owner of the object pointed at.
1414+
br->SetAddress(addr);
1415+
if (TestBit(kAutoDelete)) {
1416+
br->SetAutoDelete(true);
1417+
}
1418+
}
1419+
}
1420+
}
1421+
if (fPlayer) {
1422+
fPlayer->UpdateFormulaLeaves();
1423+
}
1424+
// Notify user if requested.
1425+
if (fNotify) {
1426+
if (!fNotify->Notify())
1427+
return -6;
1428+
}
1429+
}
1430+
1431+
return 0;
1432+
}
1433+
13491434
////////////////////////////////////////////////////////////////////////////////
13501435
/// Find the tree which contains entry, and set it as the current tree.
13511436
///
@@ -1414,83 +1499,21 @@ Long64_t TChain::LoadTree(Long64_t entry)
14141499
// next loop).
14151500
fTree->LoadTree(treeReadEntry);
14161501

1417-
if (fFriends) {
1418-
// The current tree has not changed but some of its friends might.
1419-
//
1420-
TIter next(fFriends);
1421-
TFriendLock lock(this, kLoadTree);
1422-
TFriendElement* fe = nullptr;
1423-
while ((fe = (TFriendElement*) next())) {
1424-
TTree* at = fe->GetTree();
1425-
// If the tree is a
1426-
// direct friend of the chain, it should be scanned
1427-
// used the chain entry number and NOT the tree entry
1428-
// number (treeReadEntry) hence we do:
1429-
at->LoadTreeFriend(entry, this);
1430-
}
1431-
bool needUpdate = false;
1432-
if (fTree->GetListOfFriends()) {
1433-
for(auto fetree : ROOT::Detail::TRangeStaticCast<TFriendElement>(*fTree->GetListOfFriends())) {
1434-
if (fetree->IsUpdated()) {
1435-
needUpdate = true;
1436-
fetree->ResetUpdated();
1437-
}
1438-
}
1502+
if (fExternalFriends) {
1503+
for (auto external_fe : ROOT::Detail::TRangeStaticCast<TFriendElement>(*fExternalFriends)) {
1504+
external_fe->MarkUpdated();
14391505
}
1440-
if (needUpdate) {
1441-
// Update the branch/leaf addresses and
1442-
// the list of leaves in all TTreeFormula of the TTreePlayer (if any).
1443-
1444-
// Set the branch statuses for the newly opened file.
1445-
TChainElement *frelement;
1446-
TIter fnext(fStatus);
1447-
while ((frelement = (TChainElement*) fnext())) {
1448-
Int_t status = frelement->GetStatus();
1449-
// Only set the branch status if it has a value provided
1450-
// by the user
1451-
if (status != -1)
1452-
fTree->SetBranchStatus(frelement->GetName(), status);
1453-
}
1506+
}
14541507

1455-
// Set the branch addresses for the newly opened file.
1456-
fnext.Reset();
1457-
while ((frelement = (TChainElement*) fnext())) {
1458-
void* addr = frelement->GetBaddress();
1459-
if (addr) {
1460-
TBranch* br = fTree->GetBranch(frelement->GetName());
1461-
TBranch** pp = frelement->GetBranchPtr();
1462-
if (pp) {
1463-
// FIXME: What if br is zero here?
1464-
*pp = br;
1465-
}
1466-
if (br) {
1467-
if (!frelement->GetCheckedType()) {
1468-
Int_t res = CheckBranchAddressType(br, TClass::GetClass(frelement->GetBaddressClassName()),
1469-
(EDataType) frelement->GetBaddressType(), frelement->GetBaddressIsPtr());
1470-
if ((res & kNeedEnableDecomposedObj) && !br->GetMakeClass()) {
1471-
br->SetMakeClass(true);
1472-
}
1473-
frelement->SetDecomposedObj(br->GetMakeClass());
1474-
frelement->SetCheckedType(true);
1475-
}
1476-
// FIXME: We may have to tell the branch it should
1477-
// not be an owner of the object pointed at.
1478-
br->SetAddress(addr);
1479-
if (TestBit(kAutoDelete)) {
1480-
br->SetAutoDelete(true);
1481-
}
1482-
}
1483-
}
1484-
}
1485-
if (fPlayer) {
1486-
fPlayer->UpdateFormulaLeaves();
1487-
}
1488-
// Notify user if requested.
1489-
if (fNotify) {
1490-
if(!fNotify->Notify()) return -6;
1491-
}
1492-
}
1508+
Long64_t refreshFriendAddressesRet{};
1509+
if (fFriends) {
1510+
refreshFriendAddressesRet = RefreshFriendAddresses(*this, *fTree, entry);
1511+
} else if (fTree->GetListOfFriends()) {
1512+
refreshFriendAddressesRet = RefreshFriendAddresses(*fTree, *fTree, treeReadEntry);
14931513
}
1514+
if (refreshFriendAddressesRet == -6)
1515+
return refreshFriendAddressesRet;
1516+
14941517
return treeReadEntry;
14951518
}
14961519

‎tree/treeplayer/test/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,3 +27,4 @@ ROOT_ADD_GTEST(ttreereader_friends ttreereader_friends.cxx LIBRARIES TreePlayer)
2727
ROOT_ADD_GTEST(ttreeindex_getlistoffriends ttreeindex_getlistoffriends.cxx LIBRARIES TreePlayer)
2828

2929
ROOT_ADD_GTEST(gh16804 gh16804.cxx LIBRARIES TreePlayer)
30+
ROOT_ADD_GTEST(gh16805 gh16805.cxx LIBRARIES TreePlayer)

‎tree/treeplayer/test/gh16805.cxx

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,163 @@
1+
#include <TChain.h>
2+
#include <TFile.h>
3+
#include <TTree.h>
4+
#include <TTreePlayer.h>
5+
#include <memory>
6+
#include <iostream>
7+
#include <vector>
8+
#include <string>
9+
#include <numeric>
10+
#include <fstream>
11+
12+
#include <gtest/gtest.h>
13+
14+
struct GH16805Regression : public ::testing::Test {
15+
16+
const static inline std::vector<std::string> fStepZeroFileNames{"gh16805_regression_stepzero_0.root",
17+
"gh16805_regression_stepzero_1.root"};
18+
const static inline std::vector<std::pair<int, int>> fStepZeroEntryRanges{{0, 10}, {10, 20}};
19+
constexpr static auto fStepZeroTreeName{"stepzero"};
20+
21+
constexpr static auto fStepOneFileName{"gh16805_regression_stepone.root"};
22+
constexpr static std::pair<int, int> fStepOneEntryRange{100, 120};
23+
constexpr static auto fStepOneTreeName{"stepone"};
24+
25+
static void writeStepZero(const char *treeName, const char *fileName, int begin, int end)
26+
{
27+
auto file = std::make_unique<TFile>(fileName, "recreate");
28+
auto tree = std::make_unique<TTree>(treeName, treeName);
29+
int stepZeroBr1, stepZeroBr2 = 0;
30+
tree->Branch("stepZeroBr1", &stepZeroBr1);
31+
tree->Branch("stepZeroBr2", &stepZeroBr2);
32+
33+
for (stepZeroBr1 = begin, stepZeroBr2 = begin * 2; stepZeroBr1 < end;
34+
++stepZeroBr1, stepZeroBr2 = stepZeroBr1 * 2)
35+
tree->Fill();
36+
file->Write();
37+
}
38+
39+
static void writeStepOne(const char *treeName, const char *fileName, int begin, int end)
40+
{
41+
auto file = std::make_unique<TFile>(fileName, "recreate");
42+
auto tree = std::make_unique<TTree>(treeName, treeName);
43+
int br1 = -1;
44+
tree->Branch("stepOneBr1", &br1);
45+
46+
for (br1 = begin; br1 < end; ++br1)
47+
tree->Fill();
48+
49+
file->Write();
50+
}
51+
52+
static void SetUpTestSuite()
53+
{
54+
for (decltype(fStepZeroFileNames.size()) i = 0; i < fStepZeroFileNames.size(); i++) {
55+
writeStepZero(fStepZeroTreeName, fStepZeroFileNames[i].c_str(), fStepZeroEntryRanges[i].first,
56+
fStepZeroEntryRanges[i].second);
57+
}
58+
writeStepOne(fStepOneTreeName, fStepOneFileName, fStepOneEntryRange.first, fStepOneEntryRange.second);
59+
}
60+
static void TearDownTestSuite()
61+
{
62+
for (const auto &f : fStepZeroFileNames)
63+
std::remove(f.c_str());
64+
65+
std::remove(fStepOneFileName);
66+
}
67+
};
68+
69+
TEST_F(GH16805Regression, Test)
70+
{
71+
// Merge the step0 files with a chain.
72+
auto stepZeroChain = std::make_unique<TChain>(fStepZeroTreeName);
73+
for (const auto &f : fStepZeroFileNames)
74+
stepZeroChain->Add(f.c_str());
75+
76+
auto stepOneChain = std::make_unique<TChain>(fStepOneTreeName);
77+
stepOneChain->Add(fStepOneFileName);
78+
79+
// Load the first tree in the chain
80+
stepOneChain->LoadTree(0);
81+
auto stepOneFirstTree = stepOneChain->GetTree()->GetTree();
82+
ASSERT_NE(stepOneFirstTree, nullptr);
83+
84+
stepOneFirstTree->AddFriend(stepZeroChain.get());
85+
86+
stepOneChain->LoadTree(0); // needed for the branches of the friend TChain to be available
87+
88+
// Now iterate over the chain and check the contents of the branches
89+
// inherited from the friends.
90+
int stepZeroBr1 = -1, stepZeroBr2 = -1, stepOneBr1;
91+
ASSERT_EQ(stepOneChain->SetBranchAddress("stepZeroBr1", &stepZeroBr1), 0);
92+
ASSERT_EQ(stepOneChain->SetBranchAddress("stepZeroBr2", &stepZeroBr2), 0);
93+
ASSERT_EQ(stepOneChain->SetBranchAddress("stepOneBr1", &stepOneBr1), 0);
94+
95+
std::vector<int> expectedStepZeroBr1(20);
96+
std::vector<int> expectedStepZeroBr2(20);
97+
std::vector<int> expectedStepOneBr1(20);
98+
99+
std::iota(expectedStepZeroBr1.begin(), expectedStepZeroBr1.end(), 0);
100+
std::generate(expectedStepZeroBr2.begin(), expectedStepZeroBr2.end(), [n = 0]() mutable {
101+
auto res = n * 2;
102+
n++;
103+
return res;
104+
});
105+
std::iota(expectedStepOneBr1.begin(), expectedStepOneBr1.end(), 100);
106+
107+
for (Long64_t i = 0; i < stepOneChain->GetEntriesFast(); ++i) {
108+
{
109+
stepOneChain->GetEntry(i);
110+
EXPECT_EQ(expectedStepZeroBr1[i], stepZeroBr1);
111+
EXPECT_EQ(expectedStepZeroBr2[i], stepZeroBr2);
112+
EXPECT_EQ(expectedStepOneBr1[i], stepOneBr1);
113+
}
114+
}
115+
116+
// Now test with TTree::Scan
117+
std::ostringstream strCout;
118+
{
119+
if (auto *treePlayer = static_cast<TTreePlayer *>(stepOneChain->GetPlayer())) {
120+
struct FileRAII {
121+
const char *fPath;
122+
FileRAII(const char *name) : fPath(name) {}
123+
~FileRAII() { std::remove(fPath); }
124+
} redirectFile{"regression_16805_redirect.txt"};
125+
treePlayer->SetScanRedirect(true);
126+
treePlayer->SetScanFileName(redirectFile.fPath);
127+
stepOneChain->Scan("stepZeroBr1:stepZeroBr2:stepOneBr1", "", "colsize=12");
128+
129+
std::ifstream redirectStream(redirectFile.fPath);
130+
std::stringstream redirectOutput;
131+
redirectOutput << redirectStream.rdbuf();
132+
133+
const static std::string expectedScanOut{
134+
R"Scan(*********************************************************
135+
* Row * stepZeroBr1 * stepZeroBr2 * stepOneBr1 *
136+
*********************************************************
137+
* 0 * 0 * 0 * 100 *
138+
* 1 * 1 * 2 * 101 *
139+
* 2 * 2 * 4 * 102 *
140+
* 3 * 3 * 6 * 103 *
141+
* 4 * 4 * 8 * 104 *
142+
* 5 * 5 * 10 * 105 *
143+
* 6 * 6 * 12 * 106 *
144+
* 7 * 7 * 14 * 107 *
145+
* 8 * 8 * 16 * 108 *
146+
* 9 * 9 * 18 * 109 *
147+
* 10 * 10 * 20 * 110 *
148+
* 11 * 11 * 22 * 111 *
149+
* 12 * 12 * 24 * 112 *
150+
* 13 * 13 * 26 * 113 *
151+
* 14 * 14 * 28 * 114 *
152+
* 15 * 15 * 30 * 115 *
153+
* 16 * 16 * 32 * 116 *
154+
* 17 * 17 * 34 * 117 *
155+
* 18 * 18 * 36 * 118 *
156+
* 19 * 19 * 38 * 119 *
157+
*********************************************************
158+
)Scan"};
159+
EXPECT_EQ(redirectOutput.str(), expectedScanOut);
160+
} else
161+
throw std::runtime_error("Could not retrieve TTreePlayer from main tree!");
162+
}
163+
}

0 commit comments

Comments
 (0)