Add utility function to load parameters in RDataFrame

AIDASoft · Jun 12, 2024 · b4e581f · b4e581f
1 parent b38c1d6
commit b4e581f
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 4 deletions.
diff --git a/include/podio/GenericParameters.h b/include/podio/GenericParameters.h
@@ -105,8 +105,8 @@ class GenericParameters {
   }
 
   /// Load multiple key value pairs simultaneously
-  template <typename T>
-  void loadFrom(std::vector<std::string> keys, std::vector<std::vector<T>> values);
+  template <typename T, template <typename...> typename VecLike>
+  void loadFrom(VecLike<std::string> keys, VecLike<std::vector<T>> values);
 
   /// Get the number of elements stored under the given key for a type
   template <typename T, typename = EnableIfValidGenericDataType<T>>
@@ -277,8 +277,8 @@ std::vector<std::vector<T>> GenericParameters::getValues() const {
   return values;
 }
 
-template <typename T>
-void GenericParameters::loadFrom(std::vector<std::string> keys, std::vector<std::vector<T>> values) {
+template <typename T, template <typename...> typename VecLike>
+void GenericParameters::loadFrom(VecLike<std::string> keys, VecLike<std::vector<T>> values) {
   auto& map = getMap<T>();
   auto& mtx = getMutex<T>();
 

diff --git a/include/podio/utilities/RootHelpers.h b/include/podio/utilities/RootHelpers.h
@@ -1,6 +1,9 @@
 #ifndef PODIO_UTILITIES_ROOTHELPERS_H
 #define PODIO_UTILITIES_ROOTHELPERS_H
 
+#include "podio/GenericParameters.h"
+
+#include "ROOT/RVec.hxx"
 #include "TBranch.h"
 
 #include <string>
@@ -77,6 +80,12 @@ namespace root_utils {
     std::vector<std::vector<T>>* m_valuesPtr{nullptr};
   };
 
+  GenericParameters
+  loadParamsFrom(ROOT::VecOps::RVec<std::string> intKeys, ROOT::VecOps::RVec<std::vector<int>> intValues,
+                 ROOT::VecOps::RVec<std::string> floatKeys, ROOT::VecOps::RVec<std::vector<float>> floatValues,
+                 ROOT::VecOps::RVec<std::string> doubleKeys, ROOT::VecOps::RVec<std::vector<double>> doubleValues,
+                 ROOT::VecOps::RVec<std::string> stringKeys, ROOT::VecOps::RVec<std::vector<std::string>> stringValues);
+
 } // namespace root_utils
 } // namespace podio
 

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -81,6 +81,7 @@ SET(root_sources
   ROOTReader.cc
   ROOTLegacyReader.cc
   ROOTFrameData.cc
+  RootHelpers.cc
 )
 if(ENABLE_RNTUPLE)
   list(APPEND root_sources

diff --git a/tests/root_io/CMakeLists.txt b/tests/root_io/CMakeLists.txt
@@ -62,3 +62,7 @@ if (ENABLE_RNTUPLE)
 
   set_property(TEST read_python_frame_rntuple PROPERTY DEPENDS write_python_frame_rntuple)
 endif()
+
+add_test(NAME param_reading_rdataframe COMMAND python3 ${CMAKE_CURRENT_SOURCE_DIR}/param_reading_rdataframe.py example_frame.root)
+PODIO_SET_TEST_ENV(param_reading_rdataframe)
+set_property(TEST param_reading_rdataframe PROPERTY DEPENDS write_frame_root)