Skip to content

Commit dd72fc6

Browse files
vepadulanovepadulano
committed
[df] Allow untyped reading of TTree values
1 parent d2e7874 commit dd72fc6

File tree

5 files changed

+123
-36
lines changed

5 files changed

+123
-36
lines changed

‎tree/dataframe/inc/ROOT/RDF/RTreeColumnReader.hxx‎

Lines changed: 40 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
#include "RColumnReaderBase.hxx"
1616
#include <ROOT/RVec.hxx>
17+
#include "ROOT/RDF/Utils.hxx"
1718
#include <Rtypes.h> // Long64_t, R__CLING_PTRCHECK
1819
#include <TTreeReader.h>
1920
#include <TTreeReaderValue.h>
@@ -22,6 +23,7 @@
2223
#include <array>
2324
#include <memory>
2425
#include <string>
26+
#include <cstddef>
2527

2628
namespace ROOT {
2729
namespace Internal {
@@ -30,13 +32,13 @@ namespace RDF {
3032
/// RTreeColumnReader specialization for TTree values read via TTreeReaderValues
3133
template <typename T>
3234
class R__CLING_PTRCHECK(off) RTreeColumnReader final : public ROOT::Detail::RDF::RColumnReaderBase {
33-
std::unique_ptr<TTreeReaderValue<T>> fTreeValue;
35+
std::unique_ptr<TTreeReaderUntypedValue> fTreeValue;
3436

3537
void *GetImpl(Long64_t) final { return fTreeValue->Get(); }
3638
public:
3739
/// Construct the RTreeColumnReader. Actual initialization is performed lazily by the Init method.
3840
RTreeColumnReader(TTreeReader &r, const std::string &colName)
39-
: fTreeValue(std::make_unique<TTreeReaderValue<T>>(r, colName.c_str()))
41+
: fTreeValue(std::make_unique<TTreeReaderUntypedValue>(r, colName.c_str(), ROOT::Internal::RDF::TypeID2TypeName(typeid(T))))
4042
{
4143
}
4244
};
@@ -59,10 +61,12 @@ public:
5961
/// TTreeReaderArrays are used whenever the RDF column type is RVec<T>.
6062
template <typename T>
6163
class R__CLING_PTRCHECK(off) RTreeColumnReader<RVec<T>> final : public ROOT::Detail::RDF::RColumnReaderBase {
62-
std::unique_ptr<TTreeReaderArray<T>> fTreeArray;
64+
std::unique_ptr<TTreeReaderUntypedArray> fTreeArray;
65+
66+
using Byte_t = std::byte;
6367

6468
/// We return a reference to this RVec to clients, to guarantee a stable address and contiguous memory layout.
65-
RVec<T> fRVec;
69+
RVec<Byte_t> fRVec;
6670

6771
Long64_t fLastEntry = -1;
6872

@@ -86,11 +90,10 @@ class R__CLING_PTRCHECK(off) RTreeColumnReader<RVec<T>> final : public ROOT::Det
8690
// trigger loading of the contents of the TTreeReaderArray
8791
// the address of the first element in the reader array is not necessarily equal to
8892
// the address returned by the GetAddress method
89-
auto readerArrayAddr = &readerArray.At(0);
90-
RVec<T> rvec(readerArrayAddr, readerArraySize);
93+
RVec<Byte_t> rvec(readerArray.At(0), readerArraySize);
9194
swap(fRVec, rvec);
9295
} else {
93-
RVec<T> emptyVec{};
96+
RVec<Byte_t> emptyVec{};
9497
swap(fRVec, emptyVec);
9598
}
9699
} else {
@@ -107,10 +110,17 @@ class R__CLING_PTRCHECK(off) RTreeColumnReader<RVec<T>> final : public ROOT::Det
107110
(void)fCopyWarningPrinted;
108111
#endif
109112
if (readerArraySize > 0) {
110-
RVec<T> rvec(readerArray.begin(), readerArray.end());
111-
swap(fRVec, rvec);
113+
// Array is not contiguous, make a full copy of it.
114+
fRVec = RVec<Byte_t>();
115+
fRVec.reserve(readerArraySize * sizeof(T));
116+
for (std::size_t i{0}; i < readerArraySize; i++)
117+
{
118+
auto val = readerArray.At(i);
119+
std::copy(val, val + sizeof(T), std::back_inserter(fRVec));
120+
}
121+
fRVec.resize(readerArraySize);
112122
} else {
113-
RVec<T> emptyVec{};
123+
RVec<Byte_t> emptyVec{};
114124
swap(fRVec, emptyVec);
115125
}
116126
}
@@ -120,7 +130,7 @@ class R__CLING_PTRCHECK(off) RTreeColumnReader<RVec<T>> final : public ROOT::Det
120130

121131
public:
122132
RTreeColumnReader(TTreeReader &r, const std::string &colName)
123-
: fTreeArray(std::make_unique<TTreeReaderArray<T>>(r, colName.c_str()))
133+
: fTreeArray(std::make_unique<TTreeReaderUntypedArray>(r, colName, ROOT::Internal::RDF::TypeID2TypeName(typeid(T))))
124134
{
125135
}
126136
};
@@ -131,10 +141,12 @@ public:
131141
template <>
132142
class R__CLING_PTRCHECK(off) RTreeColumnReader<RVec<bool>> final : public ROOT::Detail::RDF::RColumnReaderBase {
133143

134-
std::unique_ptr<TTreeReaderArray<bool>> fTreeArray;
144+
using Byte_t = std::byte;
145+
146+
std::unique_ptr<TTreeReaderUntypedArray> fTreeArray;
135147

136148
/// We return a reference to this RVec to clients, to guarantee a stable address and contiguous memory layout
137-
RVec<bool> fRVec;
149+
RVec<Byte_t> fRVec;
138150

139151
// We always copy the contents of TTreeReaderArray<bool> into an RVec<bool> (never take a view into the memory
140152
// buffer) because the underlying memory buffer might be the one of a std::vector<bool>, which is not a contiguous
@@ -146,19 +158,25 @@ class R__CLING_PTRCHECK(off) RTreeColumnReader<RVec<bool>> final : public ROOT::
146158
auto &readerArray = *fTreeArray;
147159
const auto readerArraySize = readerArray.GetSize();
148160
if (readerArraySize > 0) {
149-
// always perform a copy
150-
RVec<bool> rvec(readerArray.begin(), readerArray.end());
151-
swap(fRVec, rvec);
161+
// Always perform a copy
162+
fRVec = RVec<Byte_t>();
163+
fRVec.reserve(readerArraySize * sizeof(bool));
164+
for (std::size_t i{0}; i < readerArraySize; i++)
165+
{
166+
auto val = readerArray.At(i);
167+
std::copy(val, val + sizeof(bool), std::back_inserter(fRVec));
168+
}
169+
fRVec.resize(readerArraySize);
152170
} else {
153-
RVec<bool> emptyVec{};
171+
RVec<Byte_t> emptyVec{};
154172
swap(fRVec, emptyVec);
155173
}
156174
return &fRVec;
157175
}
158176

159177
public:
160178
RTreeColumnReader(TTreeReader &r, const std::string &colName)
161-
: fTreeArray(std::make_unique<TTreeReaderArray<bool>>(r, colName.c_str()))
179+
: fTreeArray(std::make_unique<TTreeReaderUntypedArray>(r, colName.c_str(), ROOT::Internal::RDF::TypeID2TypeName(typeid(bool))))
162180
{
163181
}
164182
};
@@ -168,32 +186,18 @@ public:
168186
/// This specialization is used when the requested type for reading is std::array
169187
template <typename T, std::size_t N>
170188
class R__CLING_PTRCHECK(off) RTreeColumnReader<std::array<T, N>> final : public ROOT::Detail::RDF::RColumnReaderBase {
171-
std::unique_ptr<TTreeReaderArray<T>> fTreeArray;
172-
173-
/// We return a reference to this RVec to clients, to guarantee a stable address and contiguous memory layout
174-
RVec<T> fArray;
189+
std::unique_ptr<TTreeReaderUntypedArray> fTreeArray;
175190

176191
Long64_t fLastEntry = -1;
177192

178-
void *GetImpl(Long64_t entry) final
193+
void *GetImpl(Long64_t) final
179194
{
180-
if (entry == fLastEntry)
181-
return fArray.data();
182-
183-
// This is a non-owning view on the contents of the TTreeReaderArray
184-
RVec<T> view{&fTreeArray->At(0), fTreeArray->GetSize()};
185-
swap(fArray, view);
186-
187-
fLastEntry = entry;
188-
// The data member of this class is an RVec, to avoid an extra copy
189-
// but we need to return the array buffer as the reader expects
190-
// a std::array
191-
return fArray.data();
195+
return fTreeArray->At(0);
192196
}
193197

194198
public:
195199
RTreeColumnReader(TTreeReader &r, const std::string &colName)
196-
: fTreeArray(std::make_unique<TTreeReaderArray<T>>(r, colName.c_str()))
200+
: fTreeArray(std::make_unique<TTreeReaderUntypedArray>(r, colName.c_str(), ROOT::Internal::RDF::TypeID2TypeName(typeid(T))))
197201
{
198202
}
199203
};

‎tree/treeplayer/inc/TTreeReaderArray.h‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
#include "TTreeReaderValue.h"
1616
#include "TTreeReaderUtils.h"
1717
#include <type_traits>
18+
#include <cstddef>
1819

1920
namespace ROOT {
2021
namespace Internal {
@@ -37,6 +38,9 @@ class TTreeReaderArrayBase : public TTreeReaderValueBase {
3738

3839
bool IsContiguous() const { return fImpl->IsContiguous(GetProxy()); }
3940

41+
// Returns the `sizeof` of the collection element type.
42+
UInt_t GetValueSize() const { return fImpl ? fImpl->GetValueSize(GetProxy()): 0; }
43+
4044
protected:
4145
void *UntypedAt(std::size_t idx) const { return fImpl->At(GetProxy(), idx); }
4246
void CreateProxy() override;
@@ -51,6 +55,22 @@ class TTreeReaderArrayBase : public TTreeReaderValueBase {
5155
// ClassDefOverride(TTreeReaderArrayBase, 0);//Accessor to member of an object stored in a collection
5256
};
5357

58+
class R__CLING_PTRCHECK(off) TTreeReaderUntypedArray final : public TTreeReaderArrayBase {
59+
std::string fArrayElementTypeName;
60+
61+
public:
62+
TTreeReaderUntypedArray(TTreeReader &tr, std::string_view branchName, std::string_view innerTypeName)
63+
: TTreeReaderArrayBase(&tr, branchName.data(), TDictionary::GetDictionary(innerTypeName.data())),
64+
fArrayElementTypeName(innerTypeName)
65+
{
66+
}
67+
68+
std::byte *At(std::size_t idx) const { return reinterpret_cast<std::byte *>(UntypedAt(idx)); }
69+
70+
protected:
71+
const char *GetDerivedTypeName() const final { return fArrayElementTypeName.c_str(); }
72+
};
73+
5474
} // namespace Internal
5575
} // namespace ROOT
5676

‎tree/treeplayer/inc/TTreeReaderUtils.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,7 @@ namespace Internal {
8989
virtual size_t GetSize(Detail::TBranchProxy*) = 0;
9090
virtual void* At(Detail::TBranchProxy*, size_t /*idx*/) = 0;
9191
virtual bool IsContiguous(Detail::TBranchProxy *) = 0;
92+
virtual UInt_t GetValueSize(Detail::TBranchProxy *) = 0;
9293
};
9394

9495
}

‎tree/treeplayer/inc/TTreeReaderValue.h‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -171,6 +171,30 @@ class R__CLING_PTRCHECK(off) TTreeReaderOpaqueValue final : public ROOT::Interna
171171
const char *GetDerivedTypeName() const { return ""; }
172172
};
173173

174+
class R__CLING_PTRCHECK(off) TTreeReaderUntypedValue final : public TTreeReaderValueBase {
175+
std::string fElementTypeName;
176+
177+
public:
178+
TTreeReaderUntypedValue(TTreeReader &tr, std::string_view branchName, std::string_view typeName)
179+
: TTreeReaderValueBase(&tr, branchName.data(), TDictionary::GetDictionary(typeName.data())),
180+
fElementTypeName(typeName)
181+
{
182+
}
183+
184+
void *Get()
185+
{
186+
if (!fProxy) {
187+
ErrorAboutMissingProxyIfNeeded();
188+
return nullptr;
189+
}
190+
void *address = GetAddress(); // Needed to figure out if it's a pointer
191+
return fProxy->IsaPointer() ? *(void **)address : (void *)address;
192+
}
193+
194+
protected:
195+
const char *GetDerivedTypeName() const final { return fElementTypeName.c_str(); }
196+
};
197+
174198
} // namespace Internal
175199
} // namespace ROOT
176200

‎tree/treeplayer/src/TTreeReaderArray.cxx‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232

3333
#include <memory>
3434
#include <optional>
35+
#include <iostream>
3536

3637
// pin vtable
3738
ROOT::Internal::TVirtualCollectionReader::~TVirtualCollectionReader() {}
@@ -72,6 +73,8 @@ class TClonesReader : public TVirtualCollectionReader {
7273
}
7374

7475
bool IsContiguous(ROOT::Detail::TBranchProxy *) override { return false; }
76+
77+
UInt_t GetValueSize(ROOT::Detail::TBranchProxy *) override { return sizeof(TObject *); }
7578
};
7679

7780
bool IsCPContiguous(const TVirtualCollectionProxy &cp)
@@ -86,6 +89,11 @@ bool IsCPContiguous(const TVirtualCollectionProxy &cp)
8689
}
8790
}
8891

92+
UInt_t GetCPValueSize(const TVirtualCollectionProxy &cp)
93+
{
94+
return cp.Sizeof();
95+
}
96+
8997
// Reader interface for STL
9098
class TSTLReader final : public TVirtualCollectionReader {
9199
public:
@@ -131,6 +139,12 @@ class TSTLReader final : public TVirtualCollectionReader {
131139
auto cp = GetCP(proxy);
132140
return IsCPContiguous(*cp);
133141
}
142+
143+
UInt_t GetValueSize(ROOT::Detail::TBranchProxy *proxy) override
144+
{
145+
auto cp = GetCP(proxy);
146+
return GetCPValueSize(*cp);
147+
}
134148
};
135149

136150
class TCollectionLessSTLReader final : public TVirtualCollectionReader {
@@ -190,6 +204,12 @@ class TCollectionLessSTLReader final : public TVirtualCollectionReader {
190204
auto cp = GetCP(proxy);
191205
return IsCPContiguous(*cp);
192206
}
207+
208+
UInt_t GetValueSize(ROOT::Detail::TBranchProxy *proxy) override
209+
{
210+
auto cp = GetCP(proxy);
211+
return GetCPValueSize(*cp);
212+
}
193213
};
194214

195215
// Reader interface for leaf list
@@ -243,6 +263,12 @@ class TObjectArrayReader : public TVirtualCollectionReader {
243263
void SetBasicTypeSize(Int_t size) { fBasicTypeSize = size; }
244264

245265
bool IsContiguous(ROOT::Detail::TBranchProxy *) override { return true; }
266+
267+
UInt_t GetValueSize(ROOT::Detail::TBranchProxy *proxy) override
268+
{
269+
auto cp = GetCP(proxy);
270+
return GetCPValueSize(*cp);
271+
}
246272
};
247273

248274
template <class BASE>
@@ -387,6 +413,12 @@ class TBasicTypeArrayReader final : public TVirtualCollectionReader {
387413
}
388414

389415
bool IsContiguous(ROOT::Detail::TBranchProxy *) override { return false; }
416+
417+
UInt_t GetValueSize(ROOT::Detail::TBranchProxy *proxy) override
418+
{
419+
auto cp = GetCP(proxy);
420+
return GetCPValueSize(*cp);
421+
}
390422
};
391423

392424
class TBasicTypeClonesReader final : public TClonesReader {
@@ -434,6 +466,12 @@ class TLeafReader : public TVirtualCollectionReader {
434466

435467
bool IsContiguous(ROOT::Detail::TBranchProxy *) override { return true; }
436468

469+
UInt_t GetValueSize(ROOT::Detail::TBranchProxy *) override
470+
{
471+
auto *leaf = fValueReader->GetLeaf();
472+
return leaf ? leaf->GetLenType(): 0;
473+
}
474+
437475
protected:
438476
void ProxyRead() { fValueReader->ProxyRead(); }
439477
};

0 commit comments

Comments
 (0)