Skip to content

Commit d7617f5

Browse files
committed
[ntuple] optimize bulk read for RVec of array of PoD
1 parent d3e97b7 commit d7617f5

File tree

3 files changed

+43
-14
lines changed

3 files changed

+43
-14
lines changed

tree/ntuple/inc/ROOT/RField/RFieldSequenceContainer.hxx

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,11 @@ protected:
143143
ROOT::Internal::RColumnIndex fNWritten;
144144
std::size_t fValueSize;
145145

146+
// For bulk read optimzation
147+
std::size_t fBulkNRepetition = 1;
148+
/// May be a direct PoD subfield or a sub-subfield of a fixed-size array of PoD
149+
RFieldBase *fBulkSubfield = nullptr;
150+
146151
std::unique_ptr<RFieldBase> CloneImpl(std::string_view newName) const final;
147152
const RColumnRepresentations &GetColumnRepresentations() const final;
148153
void GenerateColumns() final;

tree/ntuple/src/RFieldSequenceContainer.cxx

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -237,6 +237,19 @@ ROOT::RRVecField::RRVecField(std::string_view fieldName, std::unique_ptr<RFieldB
237237
fItemDeleter = GetDeleterOf(*itemField);
238238
Attach(std::move(itemField));
239239
fValueSize = EvalRVecValueSize(fSubfields[0]->GetAlignment(), fSubfields[0]->GetValueSize(), GetAlignment());
240+
241+
// Determine if we can optimimize bulk reading
242+
if (fSubfields[0]->IsSimple()) {
243+
fBulkSubfield = fSubfields[0].get();
244+
} else {
245+
if (auto f = dynamic_cast<RArrayField *>(fSubfields[0].get())) {
246+
auto grandChildFields = fSubfields[0]->GetMutableSubfields();
247+
if (grandChildFields[0]->IsSimple()) {
248+
fBulkSubfield = grandChildFields[0];
249+
fBulkNRepetition = f->GetLength();
250+
}
251+
}
252+
}
240253
}
241254

242255
std::unique_ptr<ROOT::RFieldBase> ROOT::RRVecField::CloneImpl(std::string_view newName) const
@@ -351,14 +364,14 @@ void ROOT::RRVecField::ReadGlobalImpl(ROOT::NTupleSize_t globalIndex, void *to)
351364

352365
std::size_t ROOT::RRVecField::ReadBulkImpl(const RBulkSpec &bulkSpec)
353366
{
354-
if (!fSubfields[0]->IsSimple())
367+
if (!fBulkSubfield)
355368
return RFieldBase::ReadBulkImpl(bulkSpec);
356369

357370
if (bulkSpec.fAuxData->empty()) {
358371
/// Initialize auxiliary memory: the first sizeof(size_t) bytes store the value size of the item field.
359372
/// The following bytes store the item values, consecutively.
360373
bulkSpec.fAuxData->resize(sizeof(std::size_t));
361-
*reinterpret_cast<std::size_t *>(bulkSpec.fAuxData->data()) = fSubfields[0]->GetValueSize();
374+
*reinterpret_cast<std::size_t *>(bulkSpec.fAuxData->data()) = fBulkNRepetition * fBulkSubfield->GetValueSize();
362375
}
363376
const auto itemValueSize = *reinterpret_cast<std::size_t *>(bulkSpec.fAuxData->data());
364377
unsigned char *itemValueArray = bulkSpec.fAuxData->data() + sizeof(std::size_t);
@@ -410,7 +423,8 @@ std::size_t ROOT::RRVecField::ReadBulkImpl(const RBulkSpec &bulkSpec)
410423
}
411424
}
412425

413-
GetPrincipalColumnOf(*fSubfields[0])->ReadV(firstItemIndex, nItems, itemValueArray - delta);
426+
GetPrincipalColumnOf(*fBulkSubfield)
427+
->ReadV(firstItemIndex * fBulkNRepetition, nItems * fBulkNRepetition, itemValueArray - delta);
414428
return RBulkSpec::kAllSet;
415429
}
416430

tree/ntuple/test/ntuple_bulk.cxx

Lines changed: 21 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -151,21 +151,25 @@ TEST(RNTupleBulk, RVec)
151151
FileRaii fileGuard("test_ntuple_bulk_rvec.root");
152152
{
153153
auto model = RNTupleModel::Create();
154-
auto fldVecI = model->MakeField<ROOT::RVecI>("vint");
155-
auto fldVecS = model->MakeField<ROOT::RVec<CustomStruct>>("vs");
156-
auto fldVecVI = model->MakeField<ROOT::RVec<ROOT::RVecI>>("vvint");
154+
auto ptrVecI = model->MakeField<ROOT::RVecI>("vint");
155+
auto ptrVecS = model->MakeField<ROOT::RVec<CustomStruct>>("vs");
156+
auto ptrVecVI = model->MakeField<ROOT::RVec<ROOT::RVecI>>("vvint");
157+
auto ptrVecArrI = model->MakeField<ROOT::RVec<std::array<int, 2>>>("varrint");
157158
auto writer = RNTupleWriter::Recreate(std::move(model), "ntpl", fileGuard.GetPath());
158159
for (int i = 0; i < 10; ++i) {
159-
fldVecI->resize(i);
160-
fldVecS->resize(i);
161-
fldVecVI->resize(i);
160+
ptrVecI->resize(i);
161+
ptrVecS->resize(i);
162+
ptrVecVI->resize(i);
163+
ptrVecArrI->resize(i);
162164
for (int j = 0; j < i; ++j) {
163-
fldVecI->at(j) = j;
164-
fldVecS->at(j).a = j;
165-
fldVecVI->at(j).resize(j);
165+
ptrVecI->at(j) = j;
166+
ptrVecS->at(j).a = j;
167+
ptrVecVI->at(j).resize(j);
166168
for (int k = 0; k < j; ++k) {
167-
fldVecVI->at(j).at(k) = k;
169+
ptrVecVI->at(j).at(k) = k;
168170
}
171+
ptrVecArrI->at(j).at(0) = 1000 * i + 2 * j;
172+
ptrVecArrI->at(j).at(1) = 1000 * i + 2 * j + 1;
169173
}
170174
writer->Fill();
171175
}
@@ -177,6 +181,7 @@ TEST(RNTupleBulk, RVec)
177181
RFieldBase::RBulkValues bulkI = model.CreateBulk("vint");
178182
RFieldBase::RBulkValues bulkS = model.CreateBulk("vs");
179183
RFieldBase::RBulkValues bulkVI = model.CreateBulk("vvint");
184+
RFieldBase::RBulkValues bulkVArrI = model.CreateBulk("varrint");
180185

181186
auto mask = std::make_unique<bool[]>(10);
182187
std::fill(mask.get(), mask.get() + 10, true);
@@ -185,12 +190,17 @@ TEST(RNTupleBulk, RVec)
185190
auto iArr = static_cast<ROOT::RVecI *>(bulkI.ReadBulk(RNTupleLocalIndex(0, 0), mask.get(), 10));
186191
auto sArr = static_cast<ROOT::RVec<CustomStruct> *>(bulkS.ReadBulk(RNTupleLocalIndex(0, 0), mask.get(), 10));
187192
auto viArr = static_cast<ROOT::RVec<ROOT::RVecI> *>(bulkVI.ReadBulk(RNTupleLocalIndex(0, 0), mask.get(), 10));
193+
auto arriArr =
194+
static_cast<ROOT::RVec<std::array<int, 2>> *>(bulkVArrI.ReadBulk(RNTupleLocalIndex(0, 0), mask.get(), 10));
188195
for (int i = 0; i < 10; ++i) {
189196
EXPECT_EQ(i, iArr[i].size());
197+
EXPECT_EQ(i, arriArr[i].size());
190198
EXPECT_EQ(i == 1 ? 0 : i, sArr[i].size());
191199
EXPECT_EQ(i == 1 ? 0 : i, viArr[i].size());
192-
for (std::size_t j = 0; j < iArr[i].size(); ++j) {
200+
for (int j = 0; j < i; ++j) {
193201
EXPECT_EQ(j, iArr[i].at(j));
202+
EXPECT_EQ(1000 * i + 2 * j, arriArr[i].at(j).at(0));
203+
EXPECT_EQ(1000 * i + 2 * j + 1, arriArr[i].at(j).at(1));
194204
}
195205
// RVec<PoD> should have all the vector elements of the bulk stored consecutively in memory
196206
if (i > 1) {

0 commit comments

Comments
 (0)