Add support for the new RNTuple format (#395)

* Add a RNTuple writer * Cleanup and add a reader * Add compilation instructions for RNTuple * Add tests * Fix the reader and writer so that they pass most of the tests * Commit missing changes in the header * Add support for Generic Parameters * Add an ugly workaround to the unique_ptr issue * Read also vector members and remove some comments * Do a bit of cleanup * Do more cleanup, also compiler warnings * Add names in rootUtils.h, fix a few compiler warnings * Add a few minor changes * Add missing changes in the headers * Change map -> unordered_map and use append in CMakeLists.txt * Simplify writing and reading of generic parameters * Only create the ID table once * Add CollectionInfo structs * Add a ROOT version check * Add missing endif() * Add Name at the end of some names * Add missing Name at the end * Cast to rvalue * Cache entries and reserve * Add comment and remove old comments * Remove a few couts * Remove intermediate variables and use std::move * Run clang-format * Use clang-format on tests too * Enable RNTuple I/O in Key4hep CI * Check if dev3 workflows come with recent enough ROOT * Change MakeField to the new signature * Update the RNTuple reader and writer to use the buffer factory * Run clang-format * Update the RNTuple writer to use a bare model * Add friends for Generic Parameters * Update changes after the changes in the collectionID and string_view * Run clang-format * Update the reader and writer to conform to #405 * Reorganize and clean up code in the reader * Run clang-format * Simplify how the references are filled --------- Co-authored-by: jmcarcell <[email protected]> Co-authored-by: tmadlener <[email protected]>
AIDASoft · Jul 11, 2023 · 1e8a8f9 · 1e8a8f9
1 parent d8294d2
commit 1e8a8f9
Show file tree

Hide file tree

Showing 16 changed files with 762 additions and 6 deletions.
diff --git a/.github/workflows/key4hep.yml b/.github/workflows/key4hep.yml
@@ -30,7 +30,8 @@ jobs:
             -DCMAKE_INSTALL_PREFIX=../install \
             -DCMAKE_CXX_STANDARD=17 \
             -DCMAKE_CXX_FLAGS=" -fdiagnostics-color=always -Werror -Wno-error=deprecated-declarations " \
-            -DUSE_EXTERNAL_CATCH2=AUTO \
+            -DUSE_EXTERNAL_CATCH2=ON \
+            -DENABLE_RNTUPLE=ON \
             -G Ninja ..
           echo "::endgroup::"
           echo "::group::Build"

diff --git a/.github/workflows/ubuntu.yml b/.github/workflows/ubuntu.yml
@@ -30,9 +30,10 @@ jobs:
             -DCMAKE_INSTALL_PREFIX=../install \
             -DCMAKE_CXX_STANDARD=17 \
             -DCMAKE_CXX_FLAGS=" -fdiagnostics-color=always -Werror -Wno-error=deprecated-declarations " \
-            -DUSE_EXTERNAL_CATCH2=OFF \
-            -DPODIO_SET_RPATH=ON \
-            -G Ninja ..
+                 -DUSE_EXTERNAL_CATCH2=OFF \
+                 -DPODIO_SET_RPATH=ON \
+                 -DENABLE_RNTUPLE=ON \
+                 -G Ninja ..
           echo "::endgroup::"
           echo "::group::Build"
           ninja -k0

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -68,12 +68,20 @@ ADD_CLANG_TIDY()
 option(CREATE_DOC        "Whether or not to create doxygen doc target." OFF)
 option(ENABLE_SIO        "Build SIO I/O support" OFF)
 option(PODIO_RELAX_PYVER "Do not require exact python version match with ROOT" OFF)
+option(ENABLE_RNTUPLE    "Build with support for the new ROOT NTtuple format" OFF)
 
 
 #--- Declare ROOT dependency ---------------------------------------------------
 list(APPEND CMAKE_PREFIX_PATH $ENV{ROOTSYS})
 set(CMAKE_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/cmake)
-find_package(ROOT REQUIRED COMPONENTS RIO Tree)
+if(NOT ENABLE_RNTUPLE)
+  find_package(ROOT REQUIRED COMPONENTS RIO Tree)
+else()
+  find_package(ROOT REQUIRED COMPONENTS RIO Tree ROOTNTuple)
+  if(${ROOT_VERSION} VERSION_LESS 6.28.02)
+    message(FATAL_ERROR "You are trying to build podio with support for the new ROOT NTuple format, but your ROOT version is too old. Please update ROOT to at least version 6.28.02")
+  endif()
+endif()
 
 # Check that root is compiled with a modern enough c++ standard
 get_target_property(ROOT_COMPILE_FEATURES ROOT::Core INTERFACE_COMPILE_FEATURES)

diff --git a/include/podio/CollectionBuffers.h b/include/podio/CollectionBuffers.h
@@ -27,6 +27,7 @@ using VectorMembersInfo = std::vector<std::pair<std::string, void*>>;
  */
 struct CollectionWriteBuffers {
   void* data{nullptr};
+  void* vecPtr{nullptr};
   CollRefCollection* references{nullptr};
   VectorMembersInfo* vectorMembers{nullptr};
 

diff --git a/include/podio/GenericParameters.h b/include/podio/GenericParameters.h
@@ -18,6 +18,11 @@ class write_device;
 using version_type = uint32_t; // from sio/definitions
 } // namespace sio
 
+namespace podio {
+class ROOTNTupleReader;
+class ROOTNTupleWriter;
+} // namespace podio
+
 #define DEPR_NON_TEMPLATE                                                                                              \
   [[deprecated("Non-templated access will be removed. Switch to templated access functionality")]]
 
@@ -145,6 +150,8 @@ class GenericParameters {
 
   friend void writeGenericParameters(sio::write_device& device, const GenericParameters& parameters);
   friend void readGenericParameters(sio::read_device& device, GenericParameters& parameters, sio::version_type version);
+  friend ROOTNTupleReader;
+  friend ROOTNTupleWriter;
 
   /// Get a reference to the internal map for a given type
   template <typename T>
@@ -187,6 +194,7 @@ class GenericParameters {
     }
   }
 
+private:
   /// Get the mutex that guards the map for the given type
   template <typename T>
   std::mutex& getMutex() const {

diff --git a/include/podio/ROOTNTupleReader.h b/include/podio/ROOTNTupleReader.h
@@ -0,0 +1,104 @@
+#ifndef PODIO_ROOTNTUPLEREADER_H
+#define PODIO_ROOTNTUPLEREADER_H
+
+#include "podio/CollectionBranches.h"
+#include "podio/ICollectionProvider.h"
+#include "podio/ROOTFrameData.h"
+#include "podio/SchemaEvolution.h"
+#include "podio/podioVersion.h"
+#include "podio/utilities/DatamodelRegistryIOHelpers.h"
+
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include <ROOT/RNTuple.hxx>
+#include <ROOT/RNTupleModel.hxx>
+
+namespace podio {
+
+/**
+This class has the function to read available data from disk
+and to prepare collections and buffers.
+**/
+class ROOTNTupleReader {
+
+public:
+  ROOTNTupleReader() = default;
+  ~ROOTNTupleReader() = default;
+
+  ROOTNTupleReader(const ROOTNTupleReader&) = delete;
+  ROOTNTupleReader& operator=(const ROOTNTupleReader&) = delete;
+
+  void openFile(const std::string& filename);
+  void openFiles(const std::vector<std::string>& filename);
+
+  /**
+   * Read the next data entry from which a Frame can be constructed for the
+   * given name. In case there are no more entries left for this name or in
+   * case there is no data for this name, this returns a nullptr.
+   */
+  std::unique_ptr<podio::ROOTFrameData> readNextEntry(const std::string& name);
+
+  /**
+   * Read the specified data entry from which a Frame can be constructed for
+   * the given name. In case the entry does not exist for this name or in case
+   * there is no data for this name, this returns a nullptr.
+   */
+  std::unique_ptr<podio::ROOTFrameData> readEntry(const std::string& name, const unsigned entry);
+
+  /// Returns number of entries for the given name
+  unsigned getEntries(const std::string& name);
+
+  /// Get the build version of podio that has been used to write the current file
+  podio::version::Version currentFileVersion() const {
+    return m_fileVersion;
+  }
+
+  void closeFile();
+
+private:
+  /**
+   * Initialize the given category by filling the maps with metadata information
+   * that will be used later
+   */
+  bool initCategory(const std::string& category);
+
+  /**
+   * Read and reconstruct the generic parameters of the Frame
+   */
+  GenericParameters readEventMetaData(const std::string& name, unsigned entNum);
+
+  template <typename T>
+  void readParams(const std::string& name, unsigned entNum, GenericParameters& params);
+
+  std::unique_ptr<ROOT::Experimental::RNTupleReader> m_metadata{};
+
+  podio::version::Version m_fileVersion{};
+  DatamodelDefinitionHolder m_datamodelHolder{};
+
+  std::unordered_map<std::string, std::vector<std::unique_ptr<ROOT::Experimental::RNTupleReader>>> m_readers{};
+  std::unordered_map<std::string, std::unique_ptr<ROOT::Experimental::RNTupleReader>> m_metadata_readers{};
+  std::vector<std::string> m_filenames{};
+
+  std::unordered_map<std::string, int> m_entries{};
+  std::unordered_map<std::string, unsigned> m_totalEntries{};
+
+  struct CollectionInfo {
+    std::vector<unsigned int> id{};
+    std::vector<std::string> name{};
+    std::vector<std::string> type{};
+    std::vector<short> isSubsetCollection{};
+    std::vector<SchemaVersionT> schemaVersion{};
+  };
+
+  std::unordered_map<std::string, CollectionInfo> m_collectionInfo{};
+
+  std::vector<std::string> m_availableCategories{};
+
+  std::shared_ptr<podio::CollectionIDTable> m_table{};
+};
+
+} // namespace podio
+
+#endif
diff --git a/include/podio/ROOTNTupleWriter.h b/include/podio/ROOTNTupleWriter.h
@@ -0,0 +1,74 @@
+#ifndef PODIO_ROOTNTUPLEWRITER_H
+#define PODIO_ROOTNTUPLEWRITER_H
+
+#include "podio/CollectionBase.h"
+#include "podio/Frame.h"
+#include "podio/GenericParameters.h"
+#include "podio/SchemaEvolution.h"
+#include "podio/utilities/DatamodelRegistryIOHelpers.h"
+
+#include "TFile.h"
+#include <ROOT/RNTuple.hxx>
+#include <ROOT/RNTupleModel.hxx>
+
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+namespace podio {
+
+class ROOTNTupleWriter {
+public:
+  ROOTNTupleWriter(const std::string& filename);
+  ~ROOTNTupleWriter();
+
+  ROOTNTupleWriter(const ROOTNTupleWriter&) = delete;
+  ROOTNTupleWriter& operator=(const ROOTNTupleWriter&) = delete;
+
+  template <typename T>
+  void fillParams(GenericParameters& params, ROOT::Experimental::REntry* entry);
+
+  void writeFrame(const podio::Frame& frame, const std::string& category);
+  void writeFrame(const podio::Frame& frame, const std::string& category, const std::vector<std::string>& collsToWrite);
+  void finish();
+
+private:
+  using StoreCollection = std::pair<const std::string&, podio::CollectionBase*>;
+  std::unique_ptr<ROOT::Experimental::RNTupleModel> createModels(const std::vector<StoreCollection>& collections);
+
+  std::unique_ptr<ROOT::Experimental::RNTupleModel> m_metadata{};
+  std::unordered_map<std::string, std::unique_ptr<ROOT::Experimental::RNTupleWriter>> m_writers{};
+  std::unique_ptr<ROOT::Experimental::RNTupleWriter> m_metadataWriter{};
+
+  std::unique_ptr<TFile> m_file{};
+
+  DatamodelDefinitionCollector m_datamodelCollector{};
+
+  struct CollectionInfo {
+    std::vector<unsigned int> id{};
+    std::vector<std::string> name{};
+    std::vector<std::string> type{};
+    std::vector<short> isSubsetCollection{};
+    std::vector<SchemaVersionT> schemaVersion{};
+  };
+
+  std::unordered_map<std::string, CollectionInfo> m_collectionInfo{};
+
+  std::set<std::string> m_categories{};
+
+  bool m_finished{false};
+
+  std::vector<std::string> m_intkeys{}, m_floatkeys{}, m_doublekeys{}, m_stringkeys{};
+
+  std::vector<std::vector<int>> m_intvalues{};
+  std::vector<std::vector<float>> m_floatvalues{};
+  std::vector<std::vector<double>> m_doublevalues{};
+  std::vector<std::vector<std::string>> m_stringvalues{};
+
+  template <typename T>
+  std::pair<std::vector<std::string>&, std::vector<std::vector<T>>&> getKeyValueVectors();
+};
+
+} // namespace podio
+
+#endif // PODIO_ROOTNTUPLEWRITER_H
diff --git a/include/podio/UserDataCollection.h b/include/podio/UserDataCollection.h
@@ -123,7 +123,7 @@ class UserDataCollection : public CollectionBase {
   /// Get the collection buffers for this collection
   podio::CollectionWriteBuffers getBuffers() override {
     _vecPtr = &_vec; // Set the pointer to the correct internal vector
-    return {&_vecPtr, &m_refCollections, &m_vecmem_info};
+    return {&_vecPtr, _vecPtr, &m_refCollections, &m_vecmem_info};
   }
 
   /// check for validity of the container after read

diff --git a/python/templates/CollectionData.cc.jinja2 b/python/templates/CollectionData.cc.jinja2
@@ -92,6 +92,7 @@ podio::CollectionWriteBuffers {{ class_type }}::getCollectionBuffers(bool isSubs
 
   return {
     isSubsetColl ? nullptr : (void*)&m_data,
+    isSubsetColl ? nullptr : (void*)m_data.get(),
     &m_refCollections, // only need to store the ObjectIDs of the referenced objects
     &m_vecmem_info
   };

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -34,6 +34,9 @@ FUNCTION(PODIO_ADD_LIB_AND_DICT libname headers sources selection )
     $<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}/include>
     $<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>)
   target_link_libraries(${dictname} PUBLIC podio::${libname} podio::podio ROOT::Core ROOT::Tree)
+  if(ENABLE_RNTUPLE)
+    target_link_libraries(${dictname} PUBLIC ROOT::ROOTNTuple)
+  endif()
   PODIO_GENERATE_DICTIONARY(${dictname} ${headers} SELECTION ${selection}
     OPTIONS --library ${CMAKE_SHARED_LIBRARY_PREFIX}${dictname}${CMAKE_SHARED_LIBRARY_SUFFIX}
     )
@@ -83,15 +86,30 @@ SET(root_sources
   ROOTFrameReader.cc
   ROOTLegacyReader.cc
 )
+if(ENABLE_RNTUPLE)
+  list(APPEND root_sources
+      ROOTNTupleReader.cc
+      ROOTNTupleWriter.cc
+     )
+endif()
 
 SET(root_headers
   ${CMAKE_SOURCE_DIR}/include/podio/ROOTFrameReader.h
   ${CMAKE_SOURCE_DIR}/include/podio/ROOTLegacyReader.h
   ${CMAKE_SOURCE_DIR}/include/podio/ROOTFrameWriter.h
   )
+if(ENABLE_RNTUPLE)
+  list(APPEND root_headers
+      ${CMAKE_SOURCE_DIR}/include/podio/ROOTNTupleReader.h
+      ${CMAKE_SOURCE_DIR}/include/podio/ROOTNTupleWriter.h
+     )
+endif()
 
 PODIO_ADD_LIB_AND_DICT(podioRootIO "${root_headers}" "${root_sources}" root_selection.xml)
 target_link_libraries(podioRootIO PUBLIC podio::podio ROOT::Core ROOT::RIO ROOT::Tree)
+if(ENABLE_RNTUPLE)
+  target_link_libraries(podioRootIO PUBLIC ROOT::ROOTNTuple)
+endif()
 
 
 # --- Python EventStore for enabling (legacy) python bindings