Skip to content

Commit 78067b5

Browse files
aalkinaalkin
authored
DPL Analysis: fix VLA reading by reading the size branch separetely (#7918)
1 parent 0ad36ef commit 78067b5

2 files changed

Lines changed: 77 additions & 208 deletions

File tree

‎Framework/Core/src/TableTreeHelpers.cxx‎

Lines changed: 62 additions & 193 deletions
Original file line numberDiff line numberDiff line change
@@ -196,7 +196,7 @@ std::pair<std::shared_ptr<arrow::ChunkedArray>, std::shared_ptr<arrow::Field>> B
196196
while (readEntries < totalEntries) {
197197
auto readLast = mBranch->GetBulkRead().GetBulkEntries(readEntries, *buffer);
198198
readEntries += readLast;
199-
status &= static_cast<arrow::BooleanBuilder*>(mValueBuilder)->AppendValues(reinterpret_cast<uint8_t const*>(buffer->GetCurrent()), (int64_t)readLast * (int64_t)mListSize);
199+
status &= static_cast<arrow::BooleanBuilder*>(mValueBuilder)->AppendValues(reinterpret_cast<uint8_t const*>(buffer->GetCurrent()), readLast * mListSize);
200200
}
201201
if (mListSize > 1) {
202202
status &= static_cast<arrow::FixedSizeListBuilder*>(mListBuilder.get())->AppendValues(readEntries);
@@ -214,209 +214,78 @@ std::pair<std::shared_ptr<arrow::ChunkedArray>, std::shared_ptr<arrow::Field>> B
214214
}
215215
} else {
216216
// other types: use serialized read to build arrays directly
217-
if (mVLA && totalEntries > 616) {
218-
// special case workaround
219-
auto status = arrow::MakeBuilder(mPool, mArrowType->field(0)->type(), &mBuilder);
220-
if (!status.ok()) {
221-
throw runtime_error("Failed to create value builder");
222-
}
223-
mListBuilder = std::make_unique<arrow::ListBuilder>(mPool, std::move(mBuilder));
224-
mValueBuilder = static_cast<arrow::ListBuilder*>(mListBuilder.get())->value_builder();
225-
void* ptr = nullptr;
226-
227-
switch (mType) {
228-
case EDataType::kUChar_t:
229-
ptr = new uint8_t[255];
230-
break;
231-
case EDataType::kUShort_t:
232-
ptr = new uint16_t[255];
233-
break;
234-
case EDataType::kUInt_t:
235-
ptr = new uint32_t[255];
236-
break;
237-
case EDataType::kULong64_t:
238-
ptr = new uint64_t[255];
239-
break;
240-
case EDataType::kChar_t:
241-
ptr = new int8_t[255];
242-
break;
243-
case EDataType::kShort_t:
244-
ptr = new int16_t[255];
245-
break;
246-
case EDataType::kInt_t:
247-
ptr = new int32_t[255];
248-
break;
249-
case EDataType::kLong64_t:
250-
ptr = new int64_t[255];
251-
break;
252-
case EDataType::kFloat_t:
253-
ptr = new float[255];
254-
break;
255-
case EDataType::kDouble_t:
256-
ptr = new double[255];
257-
break;
258-
default:
259-
throw runtime_error("Unsupported branch type");
260-
}
217+
auto&& result = arrow::AllocateResizableBuffer(mBranch->GetTotBytes(), mPool);
218+
if (!result.ok()) {
219+
throw runtime_error("Cannot allocate values buffer");
220+
}
221+
std::shared_ptr<arrow::Buffer> arrowValuesBuffer = std::move(result).ValueUnsafe();
222+
auto ptr = arrowValuesBuffer->mutable_data();
223+
if (ptr == nullptr) {
224+
throw runtime_error("Invalid buffer");
225+
}
261226

262-
int sz;
263-
auto* mSizeBranch = mBranch->GetTree()->GetBranch((std::string{mBranch->GetName()} + TableTreeHelpers::sizeBranchSuffix).c_str());
264-
mSizeBranch->SetAddress(&sz);
265-
mBranch->SetAddress(ptr);
266-
std::vector<int> offsets;
267-
268-
offsets.push_back(0);
269-
for (auto entry = 0; entry < totalEntries; ++entry) {
270-
mBranch->GetEntry(entry);
271-
mSizeBranch->GetEntry(entry);
272-
offsets.push_back(sz + offsets.back());
273-
arrow::Status status;
274-
switch (mType) {
275-
case EDataType::kUChar_t:
276-
status = static_cast<arrow::UInt8Builder*>(mValueBuilder)->AppendValues(reinterpret_cast<uint8_t const*>(ptr), sz);
277-
break;
278-
case EDataType::kUShort_t:
279-
status = static_cast<arrow::UInt16Builder*>(mValueBuilder)->AppendValues(reinterpret_cast<uint16_t const*>(ptr), sz);
280-
break;
281-
case EDataType::kUInt_t:
282-
status = static_cast<arrow::UInt32Builder*>(mValueBuilder)->AppendValues(reinterpret_cast<uint32_t const*>(ptr), sz);
283-
break;
284-
case EDataType::kULong64_t:
285-
status = static_cast<arrow::UInt64Builder*>(mValueBuilder)->AppendValues(reinterpret_cast<uint64_t const*>(ptr), sz);
286-
break;
287-
case EDataType::kChar_t:
288-
status = static_cast<arrow::Int8Builder*>(mValueBuilder)->AppendValues(reinterpret_cast<int8_t const*>(ptr), sz);
289-
break;
290-
case EDataType::kShort_t:
291-
status = static_cast<arrow::Int16Builder*>(mValueBuilder)->AppendValues(reinterpret_cast<int16_t const*>(ptr), sz);
292-
break;
293-
case EDataType::kInt_t:
294-
status = static_cast<arrow::Int32Builder*>(mValueBuilder)->AppendValues(reinterpret_cast<int32_t const*>(ptr), sz);
295-
break;
296-
case EDataType::kLong64_t:
297-
status = static_cast<arrow::Int64Builder*>(mValueBuilder)->AppendValues(reinterpret_cast<int64_t const*>(ptr), sz);
298-
break;
299-
case EDataType::kFloat_t:
300-
status = static_cast<arrow::FloatBuilder*>(mValueBuilder)->AppendValues(reinterpret_cast<float const*>(ptr), sz);
301-
break;
302-
case EDataType::kDouble_t:
303-
status = static_cast<arrow::DoubleBuilder*>(mValueBuilder)->AppendValues(reinterpret_cast<double const*>(ptr), sz);
304-
break;
305-
default:
306-
throw runtime_error("Unsupported branch type");
307-
}
308-
}
309-
status &= static_cast<arrow::ListBuilder*>(mListBuilder.get())->AppendValues(offsets.data(), totalEntries);
310-
status &= static_cast<arrow::ListBuilder*>(mListBuilder.get())->Finish(&array);
311-
312-
mSizeBranch->SetStatus(false);
313-
mSizeBranch->DropBaskets("all");
314-
mSizeBranch->Reset();
315-
mSizeBranch->GetTransientBuffer(0)->Expand(0);
316-
317-
switch (mType) {
318-
case EDataType::kUChar_t:
319-
delete[] static_cast<uint8_t*>(ptr);
320-
break;
321-
case EDataType::kUShort_t:
322-
delete[] static_cast<uint16_t*>(ptr);
323-
break;
324-
case EDataType::kUInt_t:
325-
delete[] static_cast<uint32_t*>(ptr);
326-
break;
327-
case EDataType::kULong64_t:
328-
delete[] static_cast<uint64_t*>(ptr);
329-
break;
330-
case EDataType::kChar_t:
331-
delete[] static_cast<int8_t*>(ptr);
332-
break;
333-
case EDataType::kShort_t:
334-
delete[] static_cast<int16_t*>(ptr);
335-
break;
336-
case EDataType::kInt_t:
337-
delete[] static_cast<int32_t*>(ptr);
338-
break;
339-
case EDataType::kLong64_t:
340-
delete[] static_cast<int64_t*>(ptr);
341-
break;
342-
case EDataType::kFloat_t:
343-
delete[] static_cast<float*>(ptr);
344-
break;
345-
case EDataType::kDouble_t:
346-
delete[] static_cast<double*>(ptr);
347-
break;
348-
default:
349-
throw runtime_error("Unsupported branch type");
350-
}
351-
} else {
352-
auto&& result = arrow::AllocateResizableBuffer(mBranch->GetTotBytes(), mPool);
227+
auto typeSize = TDataType::GetDataType(mType)->Size();
228+
std::unique_ptr<TBufferFile> offsetBuffer = nullptr;
229+
230+
uint32_t offset = 0;
231+
int count = 0;
232+
std::shared_ptr<arrow::Buffer> arrowOffsetBuffer;
233+
gsl::span<int> offsets;
234+
int size = 0;
235+
uint32_t totalSize = 0;
236+
TBranch* mSizeBranch = nullptr;
237+
if (mVLA) {
238+
mSizeBranch = mBranch->GetTree()->GetBranch((std::string{mBranch->GetName()} + TableTreeHelpers::sizeBranchSuffix).c_str());
239+
offsetBuffer = std::make_unique<TBufferFile>(TBuffer::EMode::kWrite, 4 * 1024 * 1024);
240+
result = arrow::AllocateResizableBuffer((totalEntries + 1) * (int64_t)sizeof(int), mPool);
353241
if (!result.ok()) {
354-
throw runtime_error("Cannot allocate values buffer");
355-
}
356-
std::shared_ptr<arrow::Buffer> arrowValuesBuffer = std::move(result).ValueUnsafe();
357-
auto ptr = arrowValuesBuffer->mutable_data();
358-
if (ptr == nullptr) {
359-
throw runtime_error("Invalid buffer");
360-
}
361-
362-
auto typeSize = TDataType::GetDataType(mType)->Size();
363-
std::unique_ptr<TBufferFile> offsetBuffer;
364-
365-
uint32_t offset = 0;
366-
uint32_t lastOffset;
367-
int count = 0;
368-
std::shared_ptr<arrow::Buffer> arrowOffsetBuffer;
369-
gsl::span<int> offsets;
370-
int size = 0;
371-
uint32_t totalSize = 0;
372-
if (mVLA) {
373-
offsetBuffer = std::make_unique<TBufferFile>(TBuffer::EMode::kWrite, 4 * 1024 * 1024);
374-
result = arrow::AllocateResizableBuffer((int64_t)(sizeof(int) * (totalEntries + 1)), mPool);
375-
if (!result.ok()) {
376-
throw runtime_error("Cannot allocate offset buffer");
377-
}
378-
arrowOffsetBuffer = std::move(result).ValueUnsafe();
379-
unsigned char* ptrOffset = arrowOffsetBuffer->mutable_data();
380-
auto* tPtrOffset = reinterpret_cast<int*>(ptrOffset);
381-
offsets = gsl::span<int>{tPtrOffset, tPtrOffset + totalEntries + 1};
242+
throw runtime_error("Cannot allocate offset buffer");
382243
}
244+
arrowOffsetBuffer = std::move(result).ValueUnsafe();
245+
unsigned char* ptrOffset = arrowOffsetBuffer->mutable_data();
246+
auto* tPtrOffset = reinterpret_cast<int*>(ptrOffset);
247+
offsets = gsl::span<int>{tPtrOffset, tPtrOffset + totalEntries + 1};
383248

249+
// read sizes first
384250
while (readEntries < totalEntries) {
385-
auto readLast = mBranch->GetBulkRead().GetEntriesSerialized(readEntries, *buffer, offsetBuffer.get());
251+
auto readLast = mSizeBranch->GetBulkRead().GetEntriesSerialized(readEntries, *offsetBuffer);
386252
readEntries += readLast;
387-
388-
if (mVLA) {
389-
lastOffset = offset;
390-
for (auto i = 0; i < readLast; ++i) {
391-
offsets[count++] = (int)offset;
392-
offset += swap32_(reinterpret_cast<uint32_t*>(offsetBuffer->GetCurrent())[i]);
393-
}
394-
size = (int)(offset - lastOffset);
395-
} else {
396-
size = readLast * mListSize;
253+
for (auto i = 0; i < readLast; ++i) {
254+
offsets[count++] = (int)offset;
255+
offset += swap32_(reinterpret_cast<uint32_t*>(offsetBuffer->GetCurrent())[i]);
397256
}
398-
swapCopy(ptr, buffer->GetCurrent(), size, typeSize);
399-
ptr += (ptrdiff_t)(size * typeSize);
400257
}
258+
offsets[count] = (int)offset;
259+
totalSize = offset;
260+
readEntries = 0;
261+
}
262+
263+
while (readEntries < totalEntries) {
264+
auto readLast = mBranch->GetBulkRead().GetEntriesSerialized(readEntries, *buffer);
401265
if (mVLA) {
402-
offsets[count] = (int)offset;
403-
totalSize = offset;
266+
size = offsets[readEntries + readLast] - offsets[readEntries];
404267
} else {
405-
totalSize = readEntries * mListSize;
406-
}
407-
std::shared_ptr<arrow::PrimitiveArray> varray;
408-
switch (mListSize) {
409-
case -1:
410-
varray = std::make_shared<arrow::PrimitiveArray>(mArrowType->field(0)->type(), totalSize, arrowValuesBuffer);
411-
array = std::make_shared<arrow::ListArray>(mArrowType, readEntries, arrowOffsetBuffer, varray);
412-
break;
413-
case 1:
414-
array = std::make_shared<arrow::PrimitiveArray>(mArrowType, readEntries, arrowValuesBuffer);
415-
break;
416-
default:
417-
varray = std::make_shared<arrow::PrimitiveArray>(mArrowType->field(0)->type(), totalSize, arrowValuesBuffer);
418-
array = std::make_shared<arrow::FixedSizeListArray>(mArrowType, readEntries, varray);
268+
size = readLast * mListSize;
419269
}
270+
readEntries += readLast;
271+
swapCopy(ptr, buffer->GetCurrent(), size, typeSize);
272+
ptr += (ptrdiff_t)(size * typeSize);
273+
}
274+
if (!mVLA) {
275+
totalSize = readEntries * mListSize;
276+
}
277+
std::shared_ptr<arrow::PrimitiveArray> varray;
278+
switch (mListSize) {
279+
case -1:
280+
varray = std::make_shared<arrow::PrimitiveArray>(mArrowType->field(0)->type(), totalSize, arrowValuesBuffer);
281+
array = std::make_shared<arrow::ListArray>(mArrowType, readEntries, arrowOffsetBuffer, varray);
282+
break;
283+
case 1:
284+
array = std::make_shared<arrow::PrimitiveArray>(mArrowType, readEntries, arrowValuesBuffer);
285+
break;
286+
default:
287+
varray = std::make_shared<arrow::PrimitiveArray>(mArrowType->field(0)->type(), totalSize, arrowValuesBuffer);
288+
array = std::make_shared<arrow::FixedSizeListArray>(mArrowType, readEntries, varray);
420289
}
421290
}
422291

‎Framework/Core/test/test_TreeToTable.cxx‎

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -183,13 +183,13 @@ BOOST_AUTO_TEST_CASE(VariableLists)
183183

184184
int empty[] = {3, 7, 10};
185185
auto count = 0;
186-
for (auto i = 1; i < 11; ++i) {
186+
for (auto i = 1; i < 1000; ++i) {
187187
iv.clear();
188188
fv.clear();
189189
dv.clear();
190190
ui.clear();
191191
if (i != empty[count]) {
192-
for (auto j = 0; j < i; ++j) {
192+
for (auto j = 0; j < i % 10 + 1; ++j) {
193193
iv.push_back(j + 2);
194194
fv.push_back((j + 2) * 0.2134f);
195195
dv.push_back((j + 4) * 0.192873819237);
@@ -218,22 +218,22 @@ BOOST_AUTO_TEST_CASE(VariableLists)
218218
int i = 1;
219219
count = 0;
220220
for (auto& row : v) {
221-
auto iv = row.ivec();
222-
auto fv = row.fvec();
223-
auto dv = row.dvec();
224-
auto uv = row.uivec();
221+
auto ivr = row.ivec();
222+
auto fvr = row.fvec();
223+
auto dvr = row.dvec();
224+
auto uvr = row.uivec();
225225
if (i != empty[count]) {
226-
for (auto j = 0; j < i; ++j) {
227-
BOOST_CHECK_EQUAL(iv[j], j + 2);
228-
BOOST_CHECK_EQUAL(fv[j], (j + 2) * 0.2134f);
229-
BOOST_CHECK_EQUAL(dv[j], (j + 4) * 0.192873819237);
230-
BOOST_CHECK_EQUAL(uv[j], j);
226+
for (auto j = 0; j < i % 10 + 1; ++j) {
227+
BOOST_CHECK_EQUAL(ivr[j], j + 2);
228+
BOOST_CHECK_EQUAL(fvr[j], (j + 2) * 0.2134f);
229+
BOOST_CHECK_EQUAL(dvr[j], (j + 4) * 0.192873819237);
230+
BOOST_CHECK_EQUAL(uvr[j], j);
231231
}
232232
} else {
233-
BOOST_CHECK_EQUAL(iv.size(), 0);
234-
BOOST_CHECK_EQUAL(fv.size(), 0);
235-
BOOST_CHECK_EQUAL(dv.size(), 0);
236-
BOOST_CHECK_EQUAL(uv.size(), 0);
233+
BOOST_CHECK_EQUAL(ivr.size(), 0);
234+
BOOST_CHECK_EQUAL(fvr.size(), 0);
235+
BOOST_CHECK_EQUAL(dvr.size(), 0);
236+
BOOST_CHECK_EQUAL(uvr.size(), 0);
237237
count++;
238238
}
239239
++i;

0 commit comments

Comments
 (0)