diff --git a/source/CMakeLists.txt b/source/CMakeLists.txt index 95f97a7..0af15cc 100644 --- a/source/CMakeLists.txt +++ b/source/CMakeLists.txt @@ -76,7 +76,7 @@ file(GLOB headers ${PROJECT_SOURCE_DIR}/source/include/*.h) #---Add library----------------------------------------------------------------- set(ROOT_DEPS ROOT::Core ROOT::RIO ROOT::Tree ROOT::Physics ROOT::Geom ROOT::EG) add_library(${library_name} ${sources} ${root_dict} ${headers}) -target_link_libraries(${library_name} ${ROOT_DEPS}) +target_link_libraries(${library_name} ${ROOT_DEPS} ROOT::ROOTNTuple) set_target_properties(${library_name} PROPERTIES INTERFACE_LINK_LIBRARIES "${ROOT_DEPS}") target_include_directories(${library_name} INTERFACE $) diff --git a/source/include/TMCRootManager.h b/source/include/TMCRootManager.h index cb69920..e6bd5d0 100644 --- a/source/include/TMCRootManager.h +++ b/source/include/TMCRootManager.h @@ -18,9 +18,22 @@ #include "TMCtls.h" #include +#include "TFile.h" +#include "TTree.h" +#include +#include +#include +#include +#include +#include +#include + class TParticle; -class TFile; -class TTree; + +using ROOT::RNTupleFillStatus; +using REntry = ROOT::REntry; +using RNTupleModel = ROOT::RNTupleModel; +using RNTupleWriter = ROOT::RNTupleWriter; /// \brief The Root IO manager for VMC examples for both sequential and /// multi-threaded applications. @@ -35,6 +48,10 @@ class TMCRootManager { kRead, // Read mode kWrite // Write mode }; + enum StorageMode { + kTTree, // TTree storage + kRNTuple // RNTuple storage + }; public: // static access method @@ -45,9 +62,12 @@ class TMCRootManager { static Bool_t GetDebug(); TMCRootManager(const char *projectName, FileMode fileMode = kWrite, Int_t threadRank = -1); + TMCRootManager(const char *projectName, StorageMode storageMode, FileMode fileMode = kWrite, Int_t threadRank = -1); virtual ~TMCRootManager(); // methods + template + void Register(const char *name, T *&obj); void Register(const char *name, const char *className, void *objAddress); void Register(const char *name, const char *className, const void *objAddress); void Fill(); @@ -56,6 +76,8 @@ class TMCRootManager { void WriteAndClose(); void ReadEvent(Int_t i); + void CreateRNTuple(); + private: // not implemented TMCRootManager(const TMCRootManager &rhs); @@ -76,10 +98,19 @@ class TMCRootManager { void OpenFile(const char *projectName, FileMode fileMode, Int_t threadRank); // data members - Int_t fId; // This manager ID - TFile *fFile; // Root output file - TTree *fTree; // Root output tree - Bool_t fIsClosed; // Info whether its file was closed + Int_t fId; // This manager ID + TFile *fFile{nullptr}; // Root output file + TTree *fTree{nullptr}; // Root output tree + + std::string fStorageName{}; + std::vector> fNameAddress; + std::unique_ptr fEntry; + std::unique_ptr fModel; + std::unique_ptr fWriter; + + StorageMode fStorageMode{kTTree}; + + Bool_t fIsClosed{false}; // Info whether its file was closed }; // inline functions @@ -94,4 +125,24 @@ inline Bool_t TMCRootManager::GetDebug() return fgDebug; } +template +void TMCRootManager::Register(const char *brname, T *&obj) +{ + if (fStorageMode == kTTree) { + fFile->cd(); + if (!fTree->GetBranch(brname)) + fTree->Branch(brname, &obj, 32000, 99); + else + fTree->GetBranch(brname)->SetAddress(&obj); + } + if (fStorageMode == kRNTuple) { + if (fModel) { + fModel->MakeField(brname); + } + std::string oString; + oString = brname; + fNameAddress.push_back(std::make_pair(oString, obj)); + } +} + #endif // ROOT_TMCRootManager diff --git a/source/src/TMCRootManager.cxx b/source/src/TMCRootManager.cxx index 8879fea..4130c5a 100644 --- a/source/src/TMCRootManager.cxx +++ b/source/src/TMCRootManager.cxx @@ -15,17 +15,14 @@ #include "TMCRootManager.h" #include "Riostream.h" #include "TError.h" -#include "TFile.h" #include "TMCAutoLock.h" #include "TThread.h" -#include "TTree.h" #include #include #include #include - namespace { // Define mutexes per operation which modify shared data TMCMutex createMutex = TMCMUTEX_INITIALIZER; @@ -34,21 +31,21 @@ TMCMutex deleteMutex = TMCMUTEX_INITIALIZER; // A global counter to assign numbers sequentially std::atomic global_thread_counter{0}; -int get_clean_thread_id() { - // This variable is unique to each thread. - // It initializes ONLY the first time this function is called on that thread. - thread_local int my_id = global_thread_counter++; - return my_id; +int get_clean_thread_id() +{ + // This variable is unique to each thread. + // It initializes ONLY the first time this function is called on that thread. + thread_local int my_id = global_thread_counter++; + return my_id; } -void threadWorker() { - // No arguments passed, but the thread can still get its 0, 1, 2 ID - std::cout << "Thread " << std::this_thread::get_id() - << " assigned itself Custom ID: " << get_clean_thread_id() << "\n"; +void threadWorker() +{ + // No arguments passed, but the thread can still get its 0, 1, 2 ID + std::cout << "Thread " << std::this_thread::get_id() << " assigned itself Custom ID: " << get_clean_thread_id() + << "\n"; } - - } // namespace // @@ -73,7 +70,14 @@ TMCRootManager *TMCRootManager::Instance() //_____________________________________________________________________________ TMCRootManager::TMCRootManager(const char *projectName, TMCRootManager::FileMode fileMode, Int_t threadRank) - : fFile(0), fTree(0), fIsClosed(false) + : TMCRootManager(projectName, TMCRootManager::kTTree, fileMode, threadRank) +{ +} + +//_____________________________________________________________________________ +TMCRootManager::TMCRootManager(const char *projectName, TMCRootManager::StorageMode storageMode, + TMCRootManager::FileMode fileMode, Int_t threadRank) + : fStorageMode(storageMode) { /// Standard constructor /// \param projectName The project name (passed as the Root tree name) @@ -155,7 +159,8 @@ void TMCRootManager::OpenFile(const char *projectName, FileMode fileMode, Int_t switch (fileMode) { case TMCRootManager::kRead: fFile = new TFile(fileName); - fTree = (TTree *)fFile->Get(projectName); + if (fStorageMode == kTTree) + fTree = (TTree *)fFile->Get(projectName); break; case TMCRootManager::kWrite: @@ -167,7 +172,12 @@ void TMCRootManager::OpenFile(const char *projectName, FileMode fileMode, Int_t if (fgDebug) printf("Going to create TTree \n"); - fTree = new TTree(projectName, treeTitle); + if (fStorageMode == kTTree) + fTree = new TTree(projectName, treeTitle); + if (fStorageMode == kRNTuple) { + fStorageName = projectName; + fModel = RNTupleModel::Create(); + } if (fgDebug) printf("Done: TTree %p \n", fTree); ; @@ -206,21 +216,54 @@ void TMCRootManager::Register(const char *name, const char *className, const voi } //_____________________________________________________________________________ -void TMCRootManager::Fill() +void TMCRootManager::CreateRNTuple() { - /// Fill the Root tree. + if (fStorageMode == kRNTuple) { + fWriter = RNTupleWriter::Append(std::move(fModel), fStorageName.c_str(), *fFile); + fEntry = fWriter->GetModel().CreateBareEntry(); + for (auto nameAddress : fNameAddress) { + fEntry->BindRawPtr(nameAddress.first, nameAddress.second); + } + } +} - fFile->cd(); - fTree->Fill(); +//_____________________________________________________________________________ +void TMCRootManager::Fill() +{ + if (fStorageMode == kTTree) { + /// Fill the Root tree. + fFile->cd(); + fTree->Fill(); + } else if (fStorageMode == kRNTuple) { + /// Fill the RNTuple. + RNTupleFillStatus status; + fWriter->FillNoFlush(*fEntry, status); + if (status.ShouldFlushCluster()) { + // If we are asked to flush, first try to do as much work as possible outside of the critical section: + // FlushColumns() will flush column data and trigger compression, but not actually write to storage. + // (A framework may of course also decide to flush more often.) + fWriter->FlushColumns(); + { + // FlushCluster() will flush data to the underlying TFile, so it requires synchronization. + fWriter->FlushCluster(); + } + } + } } //_____________________________________________________________________________ void TMCRootManager::WriteAll() { - /// Write the Root tree in the file. - fFile->cd(); - fFile->Write(); + if (fStorageMode == kTTree) { + /// Write the Root tree in the file. + fFile->cd(); + fFile->Write(); + } else if (fStorageMode == kRNTuple) { + /// Write the RNTuple in the file. + fModel.reset(); + fWriter.reset(); + } } //_____________________________________________________________________________