|
| 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 | + |
| 25 | +#include <vector> |
| 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 | + std::unique_ptr<RNTupleModel> 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 | + std::unique_ptr<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(RNTupleReader &reader, const std::string &fieldName) |
| 117 | + : fVectorView(reader.GetCollectionView(fieldName)), 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 = std::make_unique<RNTupleClusterRange>(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, 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