Logo ROOT  
Reference Guide
 
Loading...
Searching...
No Matches
RNTupleDS.hxx
Go to the documentation of this file.
1/// \file RNTupleDS.hxx
2/// \ingroup NTuple ROOT7
3/// \author Jakob Blomer <jblomer@cern.ch>
4/// \author Enrico Guiraud <enrico.guiraud@cern.ch>
5/// \date 2018-10-04
6/// \warning This is part of the ROOT 7 prototype! It will change without notice. It might trigger earthquakes. Feedback
7/// is welcome!
8
9/*************************************************************************
10 * Copyright (C) 1995-2020, Rene Brun and Fons Rademakers. *
11 * All rights reserved. *
12 * *
13 * For the licensing terms see $ROOTSYS/LICENSE. *
14 * For the list of contributors see $ROOTSYS/README/CREDITS. *
15 *************************************************************************/
16
17#ifndef ROOT_RNTupleDS
18#define ROOT_RNTupleDS
19
20#include <ROOT/RDataSource.hxx>
21#include <ROOT/RNTupleUtil.hxx>
23#include <string_view>
24
25#include <condition_variable>
26#include <cstdint>
27#include <memory>
28#include <mutex>
29#include <string>
30#include <thread>
31#include <vector>
32#include <unordered_map>
33
34namespace ROOT {
35class RFieldBase;
36class RDataFrame;
37class RNTuple;
38} // namespace ROOT
39namespace ROOT::Internal::RDF {
40class RNTupleColumnReader;
41}
42namespace ROOT::Internal {
43class RPageSource;
44}
45
46namespace ROOT::RDF {
49
50 /// The PrepareNextRanges() method populates the fNextRanges list with REntryRangeDS records.
51 /// The GetEntryRanges() swaps fNextRanges and fCurrentRanges and uses the list of
52 /// REntryRangeDS records to return the list of ranges ready to use by the RDF loop manager.
54 std::unique_ptr<ROOT::Internal::RPageSource> fSource;
55 ULong64_t fFirstEntry = 0; ///< First entry index in fSource
56 /// End entry index in fSource, e.g. the number of entries in the range is fLastEntry - fFirstEntry
58 std::string_view fFileName; ///< Storage location of the current RNTuple
59 };
60
61 /// A clone of the first pages source's descriptor.
63
64 /// The data source may be constructed with an ntuple name and a list of files
65 std::string fNTupleName;
66 std::vector<std::string> fFileNames;
67 /// The staging area is relevant for chains of files, i.e. when fFileNames is not empty. In this case,
68 /// files are opened in the background in batches of size `fNSlots` and kept in the staging area.
69 /// The first file (chains or no chains) is always opened on construction in order to process the schema.
70 /// For all subsequent files, the corresponding page sources in the staging area only executed `LoadStructure()`,
71 /// i.e. they should have a compressed buffer of the meta-data available.
72 /// Concretely:
73 /// 1. We open the first file on construction to read the schema and then move the corresponding page source
74 /// in the staging area.
75 /// 2. On `Initialize()`, we start the I/O background thread, which in turn opens the first batch of files.
76 /// 3. At the beginning of `GetEntryRanges()`, we
77 /// a) wait for the I/O thread to finish,
78 /// b) call `PrepareNextRanges()` in the main thread to move the page sources from the staging area
79 /// into `fNextRanges`; this will also call `Attach()` on the page sources (i.e., deserialize the meta-data),
80 /// and
81 /// c) trigger staging of the next batch of files in the I/O background thread.
82 /// 4. On `Finalize()`, the I/O background thread is stopped.
83 std::vector<std::unique_ptr<ROOT::Internal::RPageSource>> fStagingArea;
84 std::size_t fNextFileIndex = 0; ///< Index into fFileNames to the next file to process
85
86 /// We prepare a prototype field for every column. If a column reader is actually requested
87 /// in GetColumnReaders(), we move a clone of the field into a new column reader for RDataFrame.
88 /// Only the clone connects to the backing page store and acquires I/O resources.
89 /// The field IDs are set in the context of the first source and used as keys in fFieldId2QualifiedName.
90 std::vector<std::unique_ptr<ROOT::RFieldBase>> fProtoFields;
91 /// Columns may be requested with types other than with which they were initially added as proto fields. For example,
92 /// a column with a `ROOT::RVec<float>` proto field may instead be requested as a `std::vector<float>`. In case this
93 /// happens, we create an alternative proto field and store it here, with the original index in `fProtoFields` as
94 /// key. A single column can have more than one alternative proto fields.
95 std::unordered_map<std::size_t, std::vector<std::unique_ptr<ROOT::RFieldBase>>> fAlternativeProtoFields;
96 /// Connects the IDs of active proto fields and their subfields to their fully qualified name (a.b.c.d).
97 /// This enables the column reader to rewire the field IDs when the file changes (chain),
98 /// using the fully qualified name as a search key in the descriptor of the other page sources.
99 std::unordered_map<ROOT::DescriptorId_t, std::string> fFieldId2QualifiedName;
100 std::vector<std::string> fColumnNames;
101 std::vector<std::string> fColumnTypes;
102 /// List of column readers returned by GetColumnReaders() organized by slot. Used to reconnect readers
103 /// to new page sources when the files in the chain change.
104 std::vector<std::vector<ROOT::Internal::RDF::RNTupleColumnReader *>> fActiveColumnReaders;
105
106 ULong64_t fSeenEntries = 0; ///< The number of entries so far returned by GetEntryRanges()
107 std::vector<REntryRangeDS> fCurrentRanges; ///< Basis for the ranges returned by the last GetEntryRanges() call
108 std::vector<REntryRangeDS> fNextRanges; ///< Basis for the ranges populated by the PrepareNextRanges() call
109 /// Maps the first entries from the ranges of the last GetEntryRanges() call to their corresponding index in
110 /// the fCurrentRanges vectors. This is necessary because the returned ranges get distributed arbitrarily
111 /// onto slots. In the InitSlot method, the column readers use this map to find the correct range to connect to.
112 std::unordered_map<ULong64_t, std::size_t> fFirstEntry2RangeIdx;
113 /// One element per slot, corresponding to the current range index for that slot, as filled by InitSlot
114 std::vector<std::size_t> fSlotsToRangeIdxs;
115
116 /// The background thread that runs StageNextSources()
117 std::thread fThreadStaging;
118 /// Protects the shared state between the main thread and the I/O thread
119 std::mutex fMutexStaging;
120 /// Signal for the state information of fIsReadyForStaging and fHasNextSources
121 std::condition_variable fCvStaging;
122 /// Is true when the staging thread should start working
123 bool fIsReadyForStaging = false;
124 /// Is true when the staging thread has populated the next batch of files to fStagingArea
125 bool fHasNextSources = false;
126 /// Is true when the I/O thread should quit
128
129 /// \brief Holds useful information about fields added to the RNTupleDS
130 struct RFieldInfo {
132 std::size_t fNRepetitions;
133 // Enable `std::vector::emplace_back` for this type
138 };
139
140 /// Provides the RDF column "colName" given the field identified by fieldID. For records and collections,
141 /// AddField recurses into the sub fields. The fieldInfos argument is a list of objects holding info
142 /// about the fields of the outer collection(s) (w.r.t. fieldId). For instance, if fieldId refers to an
143 /// `std::vector<Jet>`, with
144 /// ~~~{.cpp}
145 /// struct Jet {
146 /// float pt;
147 /// float eta;
148 /// };
149 /// ~~~
150 /// AddField will recurse into `Jet.pt` and `Jet.eta` and provide the two inner fields as `ROOT::VecOps::RVec<float>`
151 /// each.
152 ///
153 /// In case the field is a collection of type `ROOT::VecOps::RVec`, `std::vector` or `std::array`, its corresponding
154 /// column is added as a `ROOT::VecOps::RVec`. Otherwise, the collection field's on-disk type is used. Note, however,
155 /// that inner record members of such collections will still be added as `ROOT::VecOps::RVec` (e.g., `std::set<Jet>
156 /// will be added as a `std::set`, but `Jet.[pt|eta] will be added as `ROOT::VecOps::RVec<float>).
157 void AddField(const ROOT::RNTupleDescriptor &desc, std::string_view colName, ROOT::DescriptorId_t fieldId,
158 std::vector<RFieldInfo> fieldInfos, bool convertToRVec = true);
159
160 /// The main function of the fThreadStaging background thread
161 void ExecStaging();
162 /// Starting from `fNextFileIndex`, opens the next `fNSlots` files. Calls `LoadStructure()` on the opened files.
163 /// The very first file is already available from the constructor.
164 void StageNextSources();
165 /// Populates fNextRanges with the next set of entry ranges. Moves files from the staging area as necessary
166 /// and aligns ranges with cluster boundaries for scheduling the tail of files.
167 /// Upon return, the fNextRanges list is ordered. It has usually fNSlots elements; fewer if there
168 /// is not enough work to give at least one cluster to every slot.
169 void PrepareNextRanges();
170
171 explicit RNTupleDS(std::unique_ptr<ROOT::Internal::RPageSource> pageSource);
172
173public:
174 RNTupleDS(std::string_view ntupleName, std::string_view fileName);
175 RNTupleDS(std::string_view ntupleName, const std::vector<std::string> &fileNames);
176 // Rule of five
177 RNTupleDS(const RNTupleDS &) = delete;
178 RNTupleDS &operator=(const RNTupleDS &) = delete;
179 RNTupleDS(RNTupleDS &&) = delete;
182
183 void SetNSlots(unsigned int nSlots) final;
184 std::size_t GetNFiles() const final { return fFileNames.empty() ? 1 : fFileNames.size(); }
185 const std::vector<std::string> &GetColumnNames() const final { return fColumnNames; }
186 bool HasColumn(std::string_view colName) const final;
187 std::string GetTypeName(std::string_view colName) const final;
188 std::vector<std::pair<ULong64_t, ULong64_t>> GetEntryRanges() final;
189 std::string GetLabel() final { return "RNTupleDS"; }
190
191 void Initialize() final;
192 void InitSlot(unsigned int slot, ULong64_t firstEntry) final;
193 void FinalizeSlot(unsigned int slot) final;
194 void Finalize() final;
195
196 std::unique_ptr<ROOT::Detail::RDF::RColumnReaderBase>
197 GetColumnReaders(unsigned int /*slot*/, std::string_view /*name*/, const std::type_info &) final;
198
199 ROOT::RDF::RSampleInfo
200 CreateSampleInfo(unsigned int,
201 const std::unordered_map<std::string, ROOT::RDF::Experimental::RSample *> &) const final;
202
203 // Old API, unused
204 bool SetEntry(unsigned int, ULong64_t) final { return true; }
205
206protected:
207 Record_t GetColumnReadersImpl(std::string_view name, const std::type_info &) final;
208};
209} // namespace ROOT::RDF
210
211namespace ROOT::RDF {
212RDataFrame FromRNTuple(std::string_view ntupleName, std::string_view fileName);
213RDataFrame FromRNTuple(std::string_view ntupleName, const std::vector<std::string> &fileNames);
214} // namespace ROOT::RDF
215
216#endif
unsigned long long ULong64_t
Definition RtypesCore.h:70
char name[80]
Definition TGX11.cxx:110
Pure virtual base class for all column reader types.
Every RDF column is represented by exactly one RNTuple field.
RDataSource defines an API that RDataFrame can use to read arbitrary data formats.
friend ROOT::RDF::RSampleInfo ROOT::Internal::RDF::CreateSampleInfo(const ROOT::RDF::RDataSource &, unsigned int, const std::unordered_map< std::string, ROOT::RDF::Experimental::RSample * > &)
std::vector< void * > Record_t
The RDataSource implementation for RNTuple.
Definition RNTupleDS.hxx:47
bool fHasNextSources
Is true when the staging thread has populated the next batch of files to fStagingArea.
const std::vector< std::string > & GetColumnNames() const final
Returns a reference to the collection of the dataset's column names.
void AddField(const ROOT::RNTupleDescriptor &desc, std::string_view colName, ROOT::DescriptorId_t fieldId, std::vector< RFieldInfo > fieldInfos, bool convertToRVec=true)
Provides the RDF column "colName" given the field identified by fieldID.
std::vector< std::pair< ULong64_t, ULong64_t > > GetEntryRanges() final
Return ranges of entries to distribute to tasks.
std::size_t GetNFiles() const final
Returns the number of files from which the dataset is constructed.
std::vector< REntryRangeDS > fNextRanges
Basis for the ranges populated by the PrepareNextRanges() call.
std::unordered_map< ULong64_t, std::size_t > fFirstEntry2RangeIdx
Maps the first entries from the ranges of the last GetEntryRanges() call to their corresponding index...
void ExecStaging()
The main function of the fThreadStaging background thread.
bool fStagingThreadShouldTerminate
Is true when the I/O thread should quit.
RNTupleDS(RNTupleDS &&)=delete
std::vector< std::vector< ROOT::Internal::RDF::RNTupleColumnReader * > > fActiveColumnReaders
List of column readers returned by GetColumnReaders() organized by slot.
std::vector< std::unique_ptr< ROOT::Internal::RPageSource > > fStagingArea
The staging area is relevant for chains of files, i.e.
Definition RNTupleDS.hxx:83
std::unique_ptr< ROOT::Detail::RDF::RColumnReaderBase > GetColumnReaders(unsigned int, std::string_view, const std::type_info &) final
If the other GetColumnReaders overload returns an empty vector, this overload will be called instead.
std::vector< std::size_t > fSlotsToRangeIdxs
One element per slot, corresponding to the current range index for that slot, as filled by InitSlot.
std::vector< std::unique_ptr< ROOT::RFieldBase > > fProtoFields
We prepare a prototype field for every column.
Definition RNTupleDS.hxx:90
void SetNSlots(unsigned int nSlots) final
Inform RDataSource of the number of processing slots (i.e.
RNTupleDS & operator=(const RNTupleDS &)=delete
std::vector< std::string > fFileNames
Definition RNTupleDS.hxx:66
bool fIsReadyForStaging
Is true when the staging thread should start working.
void InitSlot(unsigned int slot, ULong64_t firstEntry) final
Convenience method called at the start of the data processing associated to a slot.
ROOT::RNTupleDescriptor fPrincipalDescriptor
A clone of the first pages source's descriptor.
Definition RNTupleDS.hxx:62
std::vector< REntryRangeDS > fCurrentRanges
Basis for the ranges returned by the last GetEntryRanges() call.
RNTupleDS(std::unique_ptr< ROOT::Internal::RPageSource > pageSource)
std::string GetTypeName(std::string_view colName) const final
Type of a column as a string, e.g.
std::size_t fNextFileIndex
Index into fFileNames to the next file to process.
Definition RNTupleDS.hxx:84
std::unordered_map< ROOT::DescriptorId_t, std::string > fFieldId2QualifiedName
Connects the IDs of active proto fields and their subfields to their fully qualified name (a....
Definition RNTupleDS.hxx:99
std::string fNTupleName
The data source may be constructed with an ntuple name and a list of files.
Definition RNTupleDS.hxx:65
void PrepareNextRanges()
Populates fNextRanges with the next set of entry ranges.
void StageNextSources()
Starting from fNextFileIndex, opens the next fNSlots files.
std::condition_variable fCvStaging
Signal for the state information of fIsReadyForStaging and fHasNextSources.
RNTupleDS(const RNTupleDS &)=delete
void Finalize() final
Convenience method called after concluding an event-loop.
std::string GetLabel() final
Return a string representation of the datasource type.
RNTupleDS & operator=(RNTupleDS &&)=delete
std::thread fThreadStaging
The background thread that runs StageNextSources()
std::mutex fMutexStaging
Protects the shared state between the main thread and the I/O thread.
std::unordered_map< std::size_t, std::vector< std::unique_ptr< ROOT::RFieldBase > > > fAlternativeProtoFields
Columns may be requested with types other than with which they were initially added as proto fields.
Definition RNTupleDS.hxx:95
bool SetEntry(unsigned int, ULong64_t) final
Advance the "cursors" returned by GetColumnReaders to the selected entry for a particular slot.
std::vector< std::string > fColumnTypes
void Initialize() final
Convenience method called before starting an event-loop.
std::vector< std::string > fColumnNames
ULong64_t fSeenEntries
The number of entries so far returned by GetEntryRanges()
bool HasColumn(std::string_view colName) const final
Checks if the dataset has a certain column.
Record_t GetColumnReadersImpl(std::string_view name, const std::type_info &) final
type-erased vector of pointers to pointers to column values - one per slot
void FinalizeSlot(unsigned int slot) final
Convenience method called at the end of the data processing associated to a slot.
This type represents a sample identifier, to be used in conjunction with RDataFrame features such as ...
The on-storage metadata of an RNTuple.
tbb::task_arena is an alias of tbb::interface7::task_arena, which doesn't allow to forward declare tb...
std::uint64_t DescriptorId_t
Distriniguishes elements of the same type within a descriptor, e.g. different fields.
The PrepareNextRanges() method populates the fNextRanges list with REntryRangeDS records.
Definition RNTupleDS.hxx:53
std::string_view fFileName
Storage location of the current RNTuple.
Definition RNTupleDS.hxx:58
ULong64_t fLastEntry
End entry index in fSource, e.g. the number of entries in the range is fLastEntry - fFirstEntry.
Definition RNTupleDS.hxx:57
std::unique_ptr< ROOT::Internal::RPageSource > fSource
Definition RNTupleDS.hxx:54
ULong64_t fFirstEntry
First entry index in fSource.
Definition RNTupleDS.hxx:55
Holds useful information about fields added to the RNTupleDS.
RFieldInfo(ROOT::DescriptorId_t fieldId, std::size_t nRepetitions)
ROOT::DescriptorId_t fFieldId