[DF] Retrieve RDatasetSpec for RDF built from JSON

Author: martamaja10

tree/dataframe/inc/ROOT/RDF/Utils.hxx

@@ -31,7 +31,9 @@
 class TTree;
 class TTreeReader;
 
-
+namespace ROOT::RDF::Experimental {
+class RDatasetSpec;
+}
 namespace ROOT {
 namespace RDF {
 using ColumnNames_t = std::vector<std::string>;
@@ -325,6 +327,16 @@ auto MakeAliasedSharedPtr(T *rawPtr)
    return std::shared_ptr<T>(fgRawPtrCtrlBlock, rawPtr);
 }
 
+
+/**
+ * \brief Function to retrieve RDatasetSpec from JSON file provided
+ * \param[in] jsonFile Path to the dataset specification JSON file.
+ *
+ * This function allows us to have access to an RDatasetSpec which needs to
+ * be created when we use the FromSpec factory function.
+ */
+ROOT::RDF::Experimental::RDatasetSpec RetrieveSpecFromJson(const std::string &jsonFile);
+
 } // end NS RDF
 } // end NS Internal
 } // end NS ROOT

tree/dataframe/src/RDFUtils.cxx

@@ -28,6 +28,8 @@
 #include "TROOT.h" // IsImplicitMTEnabled, GetThreadPoolSize
 #include "TTree.h"
 
+#include <fstream>
+#include <nlohmann/json.hpp> // nlohmann::json::parse
 #include <stdexcept>
 #include <string>
 #include <cstring>
@@ -460,6 +462,84 @@ auto RStringCache::Insert(const std::string &string) -> decltype(fStrings)::cons
 
    return fStrings.insert(string).first;
 }
+
+ROOT::RDF::Experimental::RDatasetSpec RetrieveSpecFromJson(const std::string &jsonFile)
+{
+      const nlohmann::ordered_json fullData = nlohmann::ordered_json::parse(std::ifstream(jsonFile));
+   if (!fullData.contains("samples") || fullData["samples"].empty()) {
+      throw std::runtime_error(
+         R"(The input specification does not contain any samples. Please provide the samples in the specification like:
+{
+   "samples": {
+      "sampleA": {
+         "trees": ["tree1", "tree2"],
+         "files": ["file1.root", "file2.root"],
+         "metadata": {"lumi": 1.0, }
+      },
+      "sampleB": {
+         "trees": ["tree3", "tree4"],
+         "files": ["file3.root", "file4.root"],
+         "metadata": {"lumi": 0.5, }
+      },
+      ...
+   },
+})");
+   }
+
+   ROOT::RDF::Experimental::RDatasetSpec spec;
+   for (const auto &keyValue : fullData["samples"].items()) {
+      const std::string &sampleName = keyValue.key();
+      const auto &sample = keyValue.value();
+      // TODO: if requested in https://github.com/root-project/root/issues/11624
+      // allow union-like types for trees and files, see: https://github.com/nlohmann/json/discussions/3815
+      if (!sample.contains("trees")) {
+         throw std::runtime_error("A list of tree names must be provided for sample " + sampleName + ".");
+      }
+      std::vector<std::string> trees = sample["trees"];
+      if (!sample.contains("files")) {
+         throw std::runtime_error("A list of files must be provided for sample " + sampleName + ".");
+      }
+      std::vector<std::string> files = sample["files"];
+      if (!sample.contains("metadata")) {
+         spec.AddSample( ROOT::RDF::Experimental::RSample{sampleName, trees, files});
+      } else {
+         ROOT::RDF::Experimental::RMetaData m;
+         for (const auto &metadata : sample["metadata"].items()) {
+            const auto &val = metadata.value();
+            if (val.is_string())
+               m.Add(metadata.key(), val.get<std::string>());
+            else if (val.is_number_integer())
+               m.Add(metadata.key(), val.get<int>());
+            else if (val.is_number_float())
+               m.Add(metadata.key(), val.get<double>());
+            else
+               throw std::logic_error("The metadata keys can only be of type [string|int|double].");
+         }
+         spec.AddSample( ROOT::RDF::Experimental::RSample{sampleName, trees, files, m});
+      }
+   }
+   if (fullData.contains("friends")) {
+      for (const auto &friends : fullData["friends"].items()) {
+         std::string alias = friends.key();
+         std::vector<std::string> trees = friends.value()["trees"];
+         std::vector<std::string> files = friends.value()["files"];
+         if (files.size() != trees.size() && trees.size() > 1)
+            throw std::runtime_error("Mismatch between trees and files in a friend.");
+         spec.WithGlobalFriends(trees, files, alias);
+      }
+   }
+
+   if (fullData.contains("range")) {
+      std::vector<int> range = fullData["range"];
+
+      if (range.size() == 1)
+         spec.WithGlobalRange({range[0]});
+      else if (range.size() == 2)
+         spec.WithGlobalRange({range[0], range[1]});
+   }
+   return spec;
+};
+
 } // end NS RDF
 } // end NS Internal
 } // end NS ROOT

tree/dataframe/src/RDataFrame.cxx

@@ -1953,79 +1953,7 @@ namespace Experimental {
 ///~~~
 ROOT::RDataFrame FromSpec(const std::string &jsonFile)
 {
-   const nlohmann::ordered_json fullData = nlohmann::ordered_json::parse(std::ifstream(jsonFile));
-   if (!fullData.contains("samples") || fullData["samples"].empty()) {
-      throw std::runtime_error(
-         R"(The input specification does not contain any samples. Please provide the samples in the specification like:
-{
-   "samples": {
-      "sampleA": {
-         "trees": ["tree1", "tree2"],
-         "files": ["file1.root", "file2.root"],
-         "metadata": {"lumi": 1.0, }
-      },
-      "sampleB": {
-         "trees": ["tree3", "tree4"],
-         "files": ["file3.root", "file4.root"],
-         "metadata": {"lumi": 0.5, }
-      },
-      ...
-   },
-})");
-   }
-
-   RDatasetSpec spec;
-   for (const auto &keyValue : fullData["samples"].items()) {
-      const std::string &sampleName = keyValue.key();
-      const auto &sample = keyValue.value();
-      // TODO: if requested in https://github.com/root-project/root/issues/11624
-      // allow union-like types for trees and files, see: https://github.com/nlohmann/json/discussions/3815
-      if (!sample.contains("trees")) {
-         throw std::runtime_error("A list of tree names must be provided for sample " + sampleName + ".");
-      }
-      std::vector<std::string> trees = sample["trees"];
-      if (!sample.contains("files")) {
-         throw std::runtime_error("A list of files must be provided for sample " + sampleName + ".");
-      }
-      std::vector<std::string> files = sample["files"];
-      if (!sample.contains("metadata")) {
-         spec.AddSample(RSample{sampleName, trees, files});
-      } else {
-         RMetaData m;
-         for (const auto &metadata : sample["metadata"].items()) {
-            const auto &val = metadata.value();
-            if (val.is_string())
-               m.Add(metadata.key(), val.get<std::string>());
-            else if (val.is_number_integer())
-               m.Add(metadata.key(), val.get<int>());
-            else if (val.is_number_float())
-               m.Add(metadata.key(), val.get<double>());
-            else
-               throw std::logic_error("The metadata keys can only be of type [string|int|double].");
-         }
-         spec.AddSample(RSample{sampleName, trees, files, m});
-      }
-   }
-   if (fullData.contains("friends")) {
-      for (const auto &friends : fullData["friends"].items()) {
-         std::string alias = friends.key();
-         std::vector<std::string> trees = friends.value()["trees"];
-         std::vector<std::string> files = friends.value()["files"];
-         if (files.size() != trees.size() && trees.size() > 1)
-            throw std::runtime_error("Mismatch between trees and files in a friend.");
-         spec.WithGlobalFriends(trees, files, alias);
-      }
-   }
-
-   if (fullData.contains("range")) {
-      std::vector<int> range = fullData["range"];
-
-      if (range.size() == 1)
-         spec.WithGlobalRange({range[0]});
-      else if (range.size() == 2)
-         spec.WithGlobalRange({range[0], range[1]});
-   }
-   return ROOT::RDataFrame(spec);
+   return ROOT::RDataFrame(ROOT::Internal::RDF::RetrieveSpecFromJson(jsonFile));
 }
 
 } // namespace Experimental