Skip to content

Commit 570e576

Browse files
committed
[ntuple] add streaming vector tutorial
1 parent 373837c commit 570e576

File tree

1 file changed

+161
-0
lines changed

1 file changed

+161
-0
lines changed
Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
1+
/// \file
2+
/// \ingroup tutorial_ntuple
3+
///
4+
/// Example of a streaming vector: a special purpose container that reads large vectors piece-wise.
5+
///
6+
/// \macro_code
7+
///
8+
/// \date November 2024
9+
/// \author Peter van Gemmeren, the ROOT Team
10+
11+
// NOTE: The RNTuple classes are experimental at this point.
12+
// Functionality and interface are still subject to changes.
13+
14+
#include <ROOT/RNTupleModel.hxx>
15+
#include <ROOT/RNTupleReader.hxx>
16+
#include <ROOT/RNTupleUtil.hxx>
17+
#include <ROOT/RNTupleView.hxx>
18+
#include <ROOT/RNTupleWriter.hxx>
19+
20+
#include <TRandom3.h>
21+
22+
#include <cstdint>
23+
#include <iostream>
24+
#include <vector>
25+
#include <utility>
26+
27+
using namespace ROOT::Experimental;
28+
29+
constexpr char const *kFileName = "ntpl015_streaming_vector.root";
30+
constexpr char const *kNTupleName = "ntpl";
31+
constexpr char const *kFieldName = "LargeVector";
32+
constexpr unsigned int kNEvents = 10;
33+
constexpr unsigned int kVectorSize = 1000000;
34+
35+
void CreateRNTuple()
36+
{
37+
auto model = RNTupleModel::Create();
38+
auto ptrLargeVector = model->MakeField<std::vector<std::uint32_t>>(kFieldName);
39+
auto writer = RNTupleWriter::Recreate(std::move(model), kNTupleName, kFileName);
40+
41+
auto prng = std::make_unique<TRandom3>();
42+
prng->SetSeed();
43+
44+
for (NTupleSize_t i = 0; i < kNEvents; i++) {
45+
ptrLargeVector->clear();
46+
for (std::size_t j = 0; j < kVectorSize; j++)
47+
ptrLargeVector->emplace_back(prng->Integer(-1));
48+
writer->Fill();
49+
}
50+
std::cout << "RNTuple written" << std::endl;
51+
}
52+
53+
/*
54+
* ==================================================================================================
55+
*/
56+
57+
void ReadRNTupleSimple()
58+
{
59+
auto reader = RNTupleReader::Open(kNTupleName, kFileName);
60+
61+
const auto nEntries = reader->GetNEntries();
62+
std::cout << "Simple reading, found " << nEntries << " entries" << std::endl;
63+
64+
auto ptrLargeVector = reader->GetModel().GetDefaultEntry().GetPtr<std::vector<std::uint32_t>>(kFieldName);
65+
for (NTupleSize_t i = 0; i < nEntries; i++) {
66+
reader->LoadEntry(i);
67+
68+
const auto vectorSize = ptrLargeVector->size();
69+
uint64_t sum = 0;
70+
for (auto val : *ptrLargeVector)
71+
sum += val;
72+
73+
std::cout << "Size and sum of vector: " << vectorSize << " " << sum << std::endl;
74+
}
75+
std::cout << "RNTuple simple read" << std::endl;
76+
}
77+
78+
/*
79+
* ==================================================================================================
80+
*/
81+
82+
template <class T>
83+
class RStreamingVector {
84+
RNTupleCollectionView fVectorView;
85+
RNTupleView<T> fItemView;
86+
RNTupleClusterRange fRange;
87+
NTupleSize_t fEntry{0};
88+
NTupleSize_t fSize{0};
89+
90+
public:
91+
class iterator {
92+
RNTupleClusterRange::RIterator fIndex;
93+
RNTupleView<T> &fView;
94+
95+
public:
96+
iterator(RNTupleClusterRange::RIterator index, RNTupleView<T> &view) : fIndex(index), fView(view) {}
97+
~iterator() = default;
98+
99+
iterator operator++(int) /* postfix */
100+
{
101+
auto r = *this;
102+
++(*this);
103+
return r;
104+
}
105+
iterator &operator++() /* prefix */
106+
{
107+
++fIndex;
108+
return *this;
109+
}
110+
const T &operator*() { return fView.operator()(*fIndex); }
111+
const T *operator->() { return &fView.operator()(*fIndex); }
112+
bool operator==(const iterator &rh) const { return fIndex == rh.fIndex; }
113+
bool operator!=(const iterator &rh) const { return fIndex != rh.fIndex; }
114+
};
115+
116+
RStreamingVector(RNTupleCollectionView &&vectorView)
117+
: fVectorView(std::move(vectorView)), fItemView(fVectorView.GetView<T>("_0"))
118+
{
119+
}
120+
~RStreamingVector() = default;
121+
122+
NTupleSize_t size() const { return fSize; }
123+
124+
iterator begin() { return iterator(fRange.begin(), fItemView); }
125+
iterator end() { return iterator(fRange.end(), fItemView); }
126+
127+
void LoadEntry(NTupleSize_t entry)
128+
{
129+
fEntry = entry;
130+
fRange = fVectorView.GetCollectionRange(fEntry);
131+
fSize = fVectorView.operator()(fEntry);
132+
}
133+
};
134+
135+
void ReadRNTupleStreamingVector()
136+
{
137+
auto reader = RNTupleReader::Open(kNTupleName, kFileName);
138+
139+
const auto nEntries = reader->GetNEntries();
140+
std::cout << "Streamed reading, found " << nEntries << " entries" << std::endl;
141+
142+
RStreamingVector<std::uint32_t> streamingVector(reader->GetCollectionView(kFieldName));
143+
for (NTupleSize_t i = 0; i < nEntries; i++) {
144+
streamingVector.LoadEntry(i);
145+
146+
const auto vectorSize = streamingVector.size();
147+
uint64_t sum = 0;
148+
for (auto val : streamingVector)
149+
sum += val;
150+
151+
std::cout << "Size and sum of vector: " << vectorSize << " " << sum << std::endl;
152+
}
153+
std::cout << "RNTuple streaming read" << std::endl;
154+
}
155+
156+
void ntpl015_streaming_vector()
157+
{
158+
CreateRNTuple();
159+
ReadRNTupleSimple();
160+
ReadRNTupleStreamingVector();
161+
}

0 commit comments

Comments
 (0)