diff --git a/Framework/Core/src/TableTreeHelpers.cxx b/Framework/Core/src/TableTreeHelpers.cxx index 187d80b0e2e06..9000a01c81c48 100644 --- a/Framework/Core/src/TableTreeHelpers.cxx +++ b/Framework/Core/src/TableTreeHelpers.cxx @@ -196,7 +196,7 @@ std::pair, std::shared_ptr> B while (readEntries < totalEntries) { auto readLast = mBranch->GetBulkRead().GetBulkEntries(readEntries, *buffer); readEntries += readLast; - status &= static_cast(mValueBuilder)->AppendValues(reinterpret_cast(buffer->GetCurrent()), (int64_t)readLast * (int64_t)mListSize); + status &= static_cast(mValueBuilder)->AppendValues(reinterpret_cast(buffer->GetCurrent()), readLast * mListSize); } if (mListSize > 1) { status &= static_cast(mListBuilder.get())->AppendValues(readEntries); @@ -214,209 +214,78 @@ std::pair, std::shared_ptr> B } } else { // other types: use serialized read to build arrays directly - if (mVLA && totalEntries > 616) { - // special case workaround - auto status = arrow::MakeBuilder(mPool, mArrowType->field(0)->type(), &mBuilder); - if (!status.ok()) { - throw runtime_error("Failed to create value builder"); - } - mListBuilder = std::make_unique(mPool, std::move(mBuilder)); - mValueBuilder = static_cast(mListBuilder.get())->value_builder(); - void* ptr = nullptr; - - switch (mType) { - case EDataType::kUChar_t: - ptr = new uint8_t[255]; - break; - case EDataType::kUShort_t: - ptr = new uint16_t[255]; - break; - case EDataType::kUInt_t: - ptr = new uint32_t[255]; - break; - case EDataType::kULong64_t: - ptr = new uint64_t[255]; - break; - case EDataType::kChar_t: - ptr = new int8_t[255]; - break; - case EDataType::kShort_t: - ptr = new int16_t[255]; - break; - case EDataType::kInt_t: - ptr = new int32_t[255]; - break; - case EDataType::kLong64_t: - ptr = new int64_t[255]; - break; - case EDataType::kFloat_t: - ptr = new float[255]; - break; - case EDataType::kDouble_t: - ptr = new double[255]; - break; - default: - throw runtime_error("Unsupported branch type"); - } + auto&& result = arrow::AllocateResizableBuffer(mBranch->GetTotBytes(), mPool); + if (!result.ok()) { + throw runtime_error("Cannot allocate values buffer"); + } + std::shared_ptr arrowValuesBuffer = std::move(result).ValueUnsafe(); + auto ptr = arrowValuesBuffer->mutable_data(); + if (ptr == nullptr) { + throw runtime_error("Invalid buffer"); + } - int sz; - auto* mSizeBranch = mBranch->GetTree()->GetBranch((std::string{mBranch->GetName()} + TableTreeHelpers::sizeBranchSuffix).c_str()); - mSizeBranch->SetAddress(&sz); - mBranch->SetAddress(ptr); - std::vector offsets; - - offsets.push_back(0); - for (auto entry = 0; entry < totalEntries; ++entry) { - mBranch->GetEntry(entry); - mSizeBranch->GetEntry(entry); - offsets.push_back(sz + offsets.back()); - arrow::Status status; - switch (mType) { - case EDataType::kUChar_t: - status = static_cast(mValueBuilder)->AppendValues(reinterpret_cast(ptr), sz); - break; - case EDataType::kUShort_t: - status = static_cast(mValueBuilder)->AppendValues(reinterpret_cast(ptr), sz); - break; - case EDataType::kUInt_t: - status = static_cast(mValueBuilder)->AppendValues(reinterpret_cast(ptr), sz); - break; - case EDataType::kULong64_t: - status = static_cast(mValueBuilder)->AppendValues(reinterpret_cast(ptr), sz); - break; - case EDataType::kChar_t: - status = static_cast(mValueBuilder)->AppendValues(reinterpret_cast(ptr), sz); - break; - case EDataType::kShort_t: - status = static_cast(mValueBuilder)->AppendValues(reinterpret_cast(ptr), sz); - break; - case EDataType::kInt_t: - status = static_cast(mValueBuilder)->AppendValues(reinterpret_cast(ptr), sz); - break; - case EDataType::kLong64_t: - status = static_cast(mValueBuilder)->AppendValues(reinterpret_cast(ptr), sz); - break; - case EDataType::kFloat_t: - status = static_cast(mValueBuilder)->AppendValues(reinterpret_cast(ptr), sz); - break; - case EDataType::kDouble_t: - status = static_cast(mValueBuilder)->AppendValues(reinterpret_cast(ptr), sz); - break; - default: - throw runtime_error("Unsupported branch type"); - } - } - status &= static_cast(mListBuilder.get())->AppendValues(offsets.data(), totalEntries); - status &= static_cast(mListBuilder.get())->Finish(&array); - - mSizeBranch->SetStatus(false); - mSizeBranch->DropBaskets("all"); - mSizeBranch->Reset(); - mSizeBranch->GetTransientBuffer(0)->Expand(0); - - switch (mType) { - case EDataType::kUChar_t: - delete[] static_cast(ptr); - break; - case EDataType::kUShort_t: - delete[] static_cast(ptr); - break; - case EDataType::kUInt_t: - delete[] static_cast(ptr); - break; - case EDataType::kULong64_t: - delete[] static_cast(ptr); - break; - case EDataType::kChar_t: - delete[] static_cast(ptr); - break; - case EDataType::kShort_t: - delete[] static_cast(ptr); - break; - case EDataType::kInt_t: - delete[] static_cast(ptr); - break; - case EDataType::kLong64_t: - delete[] static_cast(ptr); - break; - case EDataType::kFloat_t: - delete[] static_cast(ptr); - break; - case EDataType::kDouble_t: - delete[] static_cast(ptr); - break; - default: - throw runtime_error("Unsupported branch type"); - } - } else { - auto&& result = arrow::AllocateResizableBuffer(mBranch->GetTotBytes(), mPool); + auto typeSize = TDataType::GetDataType(mType)->Size(); + std::unique_ptr offsetBuffer = nullptr; + + uint32_t offset = 0; + int count = 0; + std::shared_ptr arrowOffsetBuffer; + gsl::span offsets; + int size = 0; + uint32_t totalSize = 0; + TBranch* mSizeBranch = nullptr; + if (mVLA) { + mSizeBranch = mBranch->GetTree()->GetBranch((std::string{mBranch->GetName()} + TableTreeHelpers::sizeBranchSuffix).c_str()); + offsetBuffer = std::make_unique(TBuffer::EMode::kWrite, 4 * 1024 * 1024); + result = arrow::AllocateResizableBuffer((totalEntries + 1) * (int64_t)sizeof(int), mPool); if (!result.ok()) { - throw runtime_error("Cannot allocate values buffer"); - } - std::shared_ptr arrowValuesBuffer = std::move(result).ValueUnsafe(); - auto ptr = arrowValuesBuffer->mutable_data(); - if (ptr == nullptr) { - throw runtime_error("Invalid buffer"); - } - - auto typeSize = TDataType::GetDataType(mType)->Size(); - std::unique_ptr offsetBuffer; - - uint32_t offset = 0; - uint32_t lastOffset; - int count = 0; - std::shared_ptr arrowOffsetBuffer; - gsl::span offsets; - int size = 0; - uint32_t totalSize = 0; - if (mVLA) { - offsetBuffer = std::make_unique(TBuffer::EMode::kWrite, 4 * 1024 * 1024); - result = arrow::AllocateResizableBuffer((int64_t)(sizeof(int) * (totalEntries + 1)), mPool); - if (!result.ok()) { - throw runtime_error("Cannot allocate offset buffer"); - } - arrowOffsetBuffer = std::move(result).ValueUnsafe(); - unsigned char* ptrOffset = arrowOffsetBuffer->mutable_data(); - auto* tPtrOffset = reinterpret_cast(ptrOffset); - offsets = gsl::span{tPtrOffset, tPtrOffset + totalEntries + 1}; + throw runtime_error("Cannot allocate offset buffer"); } + arrowOffsetBuffer = std::move(result).ValueUnsafe(); + unsigned char* ptrOffset = arrowOffsetBuffer->mutable_data(); + auto* tPtrOffset = reinterpret_cast(ptrOffset); + offsets = gsl::span{tPtrOffset, tPtrOffset + totalEntries + 1}; + // read sizes first while (readEntries < totalEntries) { - auto readLast = mBranch->GetBulkRead().GetEntriesSerialized(readEntries, *buffer, offsetBuffer.get()); + auto readLast = mSizeBranch->GetBulkRead().GetEntriesSerialized(readEntries, *offsetBuffer); readEntries += readLast; - - if (mVLA) { - lastOffset = offset; - for (auto i = 0; i < readLast; ++i) { - offsets[count++] = (int)offset; - offset += swap32_(reinterpret_cast(offsetBuffer->GetCurrent())[i]); - } - size = (int)(offset - lastOffset); - } else { - size = readLast * mListSize; + for (auto i = 0; i < readLast; ++i) { + offsets[count++] = (int)offset; + offset += swap32_(reinterpret_cast(offsetBuffer->GetCurrent())[i]); } - swapCopy(ptr, buffer->GetCurrent(), size, typeSize); - ptr += (ptrdiff_t)(size * typeSize); } + offsets[count] = (int)offset; + totalSize = offset; + readEntries = 0; + } + + while (readEntries < totalEntries) { + auto readLast = mBranch->GetBulkRead().GetEntriesSerialized(readEntries, *buffer); if (mVLA) { - offsets[count] = (int)offset; - totalSize = offset; + size = offsets[readEntries + readLast] - offsets[readEntries]; } else { - totalSize = readEntries * mListSize; - } - std::shared_ptr varray; - switch (mListSize) { - case -1: - varray = std::make_shared(mArrowType->field(0)->type(), totalSize, arrowValuesBuffer); - array = std::make_shared(mArrowType, readEntries, arrowOffsetBuffer, varray); - break; - case 1: - array = std::make_shared(mArrowType, readEntries, arrowValuesBuffer); - break; - default: - varray = std::make_shared(mArrowType->field(0)->type(), totalSize, arrowValuesBuffer); - array = std::make_shared(mArrowType, readEntries, varray); + size = readLast * mListSize; } + readEntries += readLast; + swapCopy(ptr, buffer->GetCurrent(), size, typeSize); + ptr += (ptrdiff_t)(size * typeSize); + } + if (!mVLA) { + totalSize = readEntries * mListSize; + } + std::shared_ptr varray; + switch (mListSize) { + case -1: + varray = std::make_shared(mArrowType->field(0)->type(), totalSize, arrowValuesBuffer); + array = std::make_shared(mArrowType, readEntries, arrowOffsetBuffer, varray); + break; + case 1: + array = std::make_shared(mArrowType, readEntries, arrowValuesBuffer); + break; + default: + varray = std::make_shared(mArrowType->field(0)->type(), totalSize, arrowValuesBuffer); + array = std::make_shared(mArrowType, readEntries, varray); } } diff --git a/Framework/Core/test/test_TreeToTable.cxx b/Framework/Core/test/test_TreeToTable.cxx index 816f8f3b3c967..bbb03a1d32713 100644 --- a/Framework/Core/test/test_TreeToTable.cxx +++ b/Framework/Core/test/test_TreeToTable.cxx @@ -183,13 +183,13 @@ BOOST_AUTO_TEST_CASE(VariableLists) int empty[] = {3, 7, 10}; auto count = 0; - for (auto i = 1; i < 11; ++i) { + for (auto i = 1; i < 1000; ++i) { iv.clear(); fv.clear(); dv.clear(); ui.clear(); if (i != empty[count]) { - for (auto j = 0; j < i; ++j) { + for (auto j = 0; j < i % 10 + 1; ++j) { iv.push_back(j + 2); fv.push_back((j + 2) * 0.2134f); dv.push_back((j + 4) * 0.192873819237); @@ -218,22 +218,22 @@ BOOST_AUTO_TEST_CASE(VariableLists) int i = 1; count = 0; for (auto& row : v) { - auto iv = row.ivec(); - auto fv = row.fvec(); - auto dv = row.dvec(); - auto uv = row.uivec(); + auto ivr = row.ivec(); + auto fvr = row.fvec(); + auto dvr = row.dvec(); + auto uvr = row.uivec(); if (i != empty[count]) { - for (auto j = 0; j < i; ++j) { - BOOST_CHECK_EQUAL(iv[j], j + 2); - BOOST_CHECK_EQUAL(fv[j], (j + 2) * 0.2134f); - BOOST_CHECK_EQUAL(dv[j], (j + 4) * 0.192873819237); - BOOST_CHECK_EQUAL(uv[j], j); + for (auto j = 0; j < i % 10 + 1; ++j) { + BOOST_CHECK_EQUAL(ivr[j], j + 2); + BOOST_CHECK_EQUAL(fvr[j], (j + 2) * 0.2134f); + BOOST_CHECK_EQUAL(dvr[j], (j + 4) * 0.192873819237); + BOOST_CHECK_EQUAL(uvr[j], j); } } else { - BOOST_CHECK_EQUAL(iv.size(), 0); - BOOST_CHECK_EQUAL(fv.size(), 0); - BOOST_CHECK_EQUAL(dv.size(), 0); - BOOST_CHECK_EQUAL(uv.size(), 0); + BOOST_CHECK_EQUAL(ivr.size(), 0); + BOOST_CHECK_EQUAL(fvr.size(), 0); + BOOST_CHECK_EQUAL(dvr.size(), 0); + BOOST_CHECK_EQUAL(uvr.size(), 0); count++; } ++i;