Skip to content

Commit

Permalink
Update the reader and writer to conform to AIDASoft#405
Browse files Browse the repository at this point in the history
  • Loading branch information
jmcarcell committed Jul 10, 2023
1 parent a471949 commit 178e6a1
Show file tree
Hide file tree
Showing 2 changed files with 87 additions and 44 deletions.
50 changes: 36 additions & 14 deletions src/ROOTNTupleReader.cc
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
#include "podio/CollectionBufferFactory.h"
#include "podio/CollectionBuffers.h"
#include "podio/CollectionIDTable.h"
#include "podio/DatamodelRegistry.h"
#include "podio/GenericParameters.h"
#include "rootUtils.h"

Expand Down Expand Up @@ -136,7 +137,6 @@ std::unique_ptr<ROOTFrameData> ROOTNTupleReader::readEntry(const std::string& ca
auto collection =
std::unique_ptr<podio::CollectionBase>(static_cast<podio::CollectionBase*>(collectionClass->New()));

// const std::string bufferClassName = "std::vector<" + collection->getDataTypeName() + ">";
const auto bufferClassName = collection->getTypeName();
const auto bufferClass = m_collectionInfo[category].isSubsetCollection[i]
? nullptr
Expand All @@ -156,23 +156,27 @@ std::unique_ptr<ROOTFrameData> ROOTNTupleReader::readEntry(const std::string& ca
return nullptr;
}

if (!isSubsetColl) {
if (isSubsetColl) {
auto brName = root_utils::subsetBranch(m_collectionInfo[category].name[i]);
auto vec = new std::vector<podio::ObjectID>;
dentry->CaptureValueUnsafe(brName, vec);
tmp[{brName, 0}] = vec;
} else {
dentry->CaptureValueUnsafe(m_collectionInfo[category].name[i], collBuffers.data);
}
if (auto* refCollections = collBuffers.references) {
for (size_t j = 0; j < refCollections->size(); ++j) {

const auto relVecNames = podio::DatamodelRegistry::instance().getRelationNames(collection->getTypeName());
for (size_t j = 0; j < relVecNames.relations.size(); ++j) {
const auto relName = relVecNames.relations[j];
auto vec = new std::vector<podio::ObjectID>;
const auto brName = root_utils::refBranch(m_collectionInfo[category].name[i], j);
const auto brName = root_utils::refBranch(m_collectionInfo[category].name[i], relName);
dentry->CaptureValueUnsafe(brName, vec);
tmp[{brName, j}] = vec;
}
}

if (auto* vecMembers = collBuffers.vectorMembers) {
for (size_t j = 0; j < vecMembers->size(); ++j) {
const auto typeName = "vector<" + vecMembers->at(j).first + ">";
const auto brName = root_utils::vecBranch(m_collectionInfo[category].name[i], j);
dentry->CaptureValueUnsafe(brName, vecMembers->at(j).second);
for (size_t j = 0; j < relVecNames.vectorMembers.size(); ++j) {
const auto vecName = relVecNames.vectorMembers[j];
const auto brName = root_utils::vecBranch(m_collectionInfo[category].name[i], vecName);
dentry->CaptureValueUnsafe(brName, collBuffers.vectorMembers->at(j).second);
}
}

Expand All @@ -182,11 +186,29 @@ std::unique_ptr<ROOTFrameData> ROOTNTupleReader::readEntry(const std::string& ca
m_readers[category][0]->LoadEntry(entNum);

for (size_t i = 0; i < m_collectionInfo[category].id.size(); ++i) {
const auto collectionClass = TClass::GetClass(m_collectionInfo[category].type[i].c_str());
auto collection =
std::unique_ptr<podio::CollectionBase>(static_cast<podio::CollectionBase*>(collectionClass->New()));
const auto bufferClassName = collection->getTypeName();
const auto bufferClass = m_collectionInfo[category].isSubsetCollection[i]
? nullptr
: TClass::GetClass(std::string(bufferClassName).c_str());
auto collBuffers = buffers[m_collectionInfo[category].name[i]];
const auto relVecNames = podio::DatamodelRegistry::instance().getRelationNames(collection->getTypeName());
if (auto* refCollections = collBuffers.references) {
for (size_t j = 0; j < refCollections->size(); ++j) {
const auto brName = root_utils::refBranch(m_collectionInfo[category].name[i], j);
size_t maxj = refCollections->size();
if (bufferClass == nullptr) {
maxj = 1;
}
for (size_t j = 0; j < maxj; ++j) {
if (bufferClass != nullptr) {
auto brName = root_utils::refBranch(m_collectionInfo[category].name[i], relVecNames.relations[j]);
refCollections->at(j) = std::unique_ptr<std::vector<podio::ObjectID>>(tmp[{brName, j}]);
}
if (bufferClass == nullptr) {
auto brName = root_utils::subsetBranch(m_collectionInfo[category].name[i]);
refCollections->at(j) = std::unique_ptr<std::vector<podio::ObjectID>>(tmp[{brName, j}]);
}
}
}
}
Expand Down
81 changes: 51 additions & 30 deletions src/ROOTNTupleWriter.cc
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
#include "podio/ROOTNTupleWriter.h"
#include "podio/CollectionBase.h"
#include "podio/GenericParameters.h"
#include "podio/DatamodelRegistry.h"
#include "podio/SchemaEvolution.h"
#include "podio/podioVersion.h"
#include "rootUtils.h"
Expand Down Expand Up @@ -89,21 +90,31 @@ void ROOTNTupleWriter::writeFrame(const podio::Frame& frame, const std::string&
entry->CaptureValueUnsafe(name, (void*)collBuffers.vecPtr);
}

if (auto refColls = collBuffers.references) {
int i = 0;
for (auto& c : (*refColls)) {
const auto brName = root_utils::refBranch(name, i++);
entry->CaptureValueUnsafe(brName, c.get());
if (coll->isSubsetCollection()) {
auto& refColl = (*collBuffers.references)[0];
const auto brName = root_utils::subsetBranch(name);
entry->CaptureValueUnsafe(brName, refColl.get());
} else {

const auto relVecNames = podio::DatamodelRegistry::instance().getRelationNames(coll->getValueTypeName());
if (auto refColls = collBuffers.references) {
int i = 0;
for (auto& c : (*refColls)) {
const auto brName = root_utils::refBranch(name, relVecNames.relations[i]);
entry->CaptureValueUnsafe(brName, c.get());
++i;
}
}
}

if (auto vmInfo = collBuffers.vectorMembers) {
int i = 0;
for (auto& [type, vec] : (*vmInfo)) {
const auto typeName = "vector<" + type + ">";
const auto brName = root_utils::vecBranch(name, i++);
auto ptr = *(std::vector<int>**)vec;
entry->CaptureValueUnsafe(brName, ptr);
if (auto vmInfo = collBuffers.vectorMembers) {
int i = 0;
for (auto& [type, vec] : (*vmInfo)) {
const auto typeName = "vector<" + type + ">";
const auto brName = root_utils::vecBranch(name, relVecNames.vectorMembers[i]);
auto ptr = *(std::vector<int>**)vec;
entry->CaptureValueUnsafe(brName, ptr);
++i;
}
}
}

Expand Down Expand Up @@ -136,31 +147,41 @@ ROOTNTupleWriter::createModels(const std::vector<StoreCollection>& collections)
for (auto& [name, coll] : collections) {
const auto collBuffers = coll->getBuffers();


if (collBuffers.vecPtr) {
auto collClassName = "std::vector<" + std::string(coll->getDataTypeName()) + ">";
auto field = ROOT::Experimental::Detail::RFieldBase::Create(name, collClassName).Unwrap();
model->AddField(std::move(field));
}

if (auto refColls = collBuffers.references) {
int i = 0;
for (auto& c [[maybe_unused]] : (*refColls)) {
const auto brName = root_utils::refBranch(name, i);
auto collClassName = "vector<podio::ObjectID>";
auto field = ROOT::Experimental::Detail::RFieldBase::Create(brName, collClassName).Unwrap();
model->AddField(std::move(field));
++i;
if (coll->isSubsetCollection()) {
const auto brName = root_utils::subsetBranch(name);
auto collClassName = "vector<podio::ObjectID>";
auto field = ROOT::Experimental::Detail::RFieldBase::Create(brName, collClassName).Unwrap();
model->AddField(std::move(field));
} else {

const auto relVecNames = podio::DatamodelRegistry::instance().getRelationNames(coll->getValueTypeName());
if (auto refColls = collBuffers.references) {
int i = 0;
for (auto& c [[maybe_unused]] : (*refColls)) {
const auto brName = root_utils::refBranch(name, relVecNames.relations[i]);
auto collClassName = "vector<podio::ObjectID>";
auto field = ROOT::Experimental::Detail::RFieldBase::Create(brName, collClassName).Unwrap();
model->AddField(std::move(field));
++i;
}
}
}

if (auto vminfo = collBuffers.vectorMembers) {
int i = 0;
for (auto& [type, vec] : (*vminfo)) {
const auto typeName = "vector<" + type + ">";
const auto brName = root_utils::vecBranch(name, i);
auto field = ROOT::Experimental::Detail::RFieldBase::Create(brName, typeName).Unwrap();
model->AddField(std::move(field));
++i;
if (auto vminfo = collBuffers.vectorMembers) {
int i = 0;
for (auto& [type, vec] : (*vminfo)) {
const auto typeName = "vector<" + type + ">";
const auto brName = root_utils::vecBranch(name, relVecNames.vectorMembers[i]);
auto field = ROOT::Experimental::Detail::RFieldBase::Create(brName, typeName).Unwrap();
model->AddField(std::move(field));
++i;
}
}
}
}
Expand Down

0 comments on commit 178e6a1

Please sign in to comment.