[RDF] Prepare for snapshot with variations 2

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

@@ -62,7 +62,7 @@ class RDefinesWithReaders {
    // (see BookDefineJit). it is never null.
    std::shared_ptr<ROOT::Detail::RDF::RDefineBase> fDefine;
    // Column readers per variation (in the map) per slot (in the vector).
-   std::vector<std::unordered_map<std::string_view, std::unique_ptr<RDefineReader>>> fReadersPerVariation;
+   std::vector<std::unordered_map<std::string_view, std::shared_ptr<RDefineReader>>> fReadersPerVariation;
 
    // Strings that were already used to represent column names in this RDataFrame instance.
    ROOT::Internal::RDF::RStringCache &fCachedColNames;

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

@@ -882,7 +882,7 @@ public:
    ///            return an RVec of varied values, one for each variation tag, in the same order as the tags.
    /// \param[in] inputColumns the names of the columns to be passed to the callable.
    /// \param[in] nVariations number of variations returned by the expression. The corresponding tags will be `"0"`,
-   /// `"1"`, etc. 
+   /// `"1"`, etc.
    /// \param[in] variationName a generic name for this set of varied values, e.g. `"ptvariation"`.
    ///            colName is used if none is provided.
    ///
@@ -988,7 +988,7 @@ public:
    ///            return an RVec of varied values, one for each variation tag, in the same order as the tags.
    /// \param[in] inputColumns the names of the columns to be passed to the callable.
    /// \param[in] nVariations number of variations returned by the expression. The corresponding tags will be `"0"`,
-   /// `"1"`, etc. 
+   /// `"1"`, etc.
    /// \param[in] variationName a generic name for this set of varied values, e.g. `"ptvariation"`.
    ///            colName is used if none is provided.
    ///
@@ -1036,7 +1036,7 @@ public:
    /// \param[in] inputColumns the names of the columns to be passed to the callable.
    /// \param[in] inputColumns the names of the columns to be passed to the callable.
    /// \param[in] nVariations number of variations returned by the expression. The corresponding tags will be `"0"`,
-   /// `"1"`, etc. 
+   /// `"1"`, etc.
    /// \param[in] variationName a generic name for this set of varied values, e.g. `"ptvariation"`.
    ///            colName is used if none is provided.
    ///
@@ -1091,7 +1091,7 @@ public:
    /// \param[in] expression a string containing valid C++ code that evaluates to an RVec containing the varied
    ///            values for the specified column.
    /// \param[in] nVariations number of variations returned by the expression. The corresponding tags will be `"0"`,
-   /// `"1"`, etc. 
+   /// `"1"`, etc.
    /// \param[in] variationName a generic name for this set of varied values, e.g. `"ptvariation"`.
    ///            colName is used if none is provided.
    ///
@@ -1126,7 +1126,7 @@ public:
    /// \param[in] expression a string containing valid C++ code that evaluates to an RVec or RVecs containing the varied
    ///            values for the specified columns.
    /// \param[in] nVariations number of variations returned by the expression. The corresponding tags will be `"0"`,
-   /// `"1"`, etc. 
+   /// `"1"`, etc.
    /// \param[in] variationName a generic name for this set of varied values, e.g. `"ptvariation"`.
    ///
    /// This overload adds the possibility for the expression used to evaluate the varied values to be just-in-time
@@ -1163,7 +1163,7 @@ public:
    /// \param[in] expression a string containing valid C++ code that evaluates to an RVec containing the varied
    ///            values for the specified column.
    /// \param[in] nVariations number of variations returned by the expression. The corresponding tags will be `"0"`,
-   /// `"1"`, etc. 
+   /// `"1"`, etc.
    /// \param[in] variationName a generic name for this set of varied values, e.g. `"ptvariation"`.
    ///            colName is used if none is provided.
    ///
@@ -1251,6 +1251,27 @@ public:
    /// \param[in] options RSnapshotOptions struct with extra options to pass to the output TFile and TTree/RNTuple.
    /// \return a `RDataFrame` that wraps the snapshotted dataset.
    ///
+   template <typename... ColumnTypes>
+   R__DEPRECATED(
+      6, 40, "Snapshot does not need template arguments anymore, you can safely remove them from this function call.")
+   RResultPtr<RInterface<RLoopManager>> Snapshot(std::string_view treename, std::string_view filename,
+                                                 const ColumnNames_t &columnList,
+                                                 const RSnapshotOptions &options = RSnapshotOptions())
+   {
+      return Snapshot(treename, filename, columnList, options);
+   }
+
+   ////////////////////////////////////////////////////////////////////////////
+   /// \brief Save selected columns to disk, in a new TTree or RNTuple `treename` in file `filename`.
+   /// \param[in] treename The name of the output TTree or RNTuple.
+   /// \param[in] filename The name of the output TFile.
+   /// \param[in] columnList The list of names of the columns/branches/fields to be written.
+   /// \param[in] options RSnapshotOptions struct with extra options to pass to TFile and TTree/RNTuple.
+   /// \return a `RDataFrame` that wraps the snapshotted dataset.
+   ///
+   /// This function returns a `RDataFrame` built with the output TTree or RNTuple as a source.
+   /// The types of the columns are automatically inferred and do not need to be specified.
+   ///
    /// Support for writing of nested branches/fields is limited (although RDataFrame is able to read them) and dot ('.')
    /// characters in input column names will be replaced by underscores ('_') in the branches produced by Snapshot.
    /// When writing a variable size array through Snapshot, it is required that the column indicating its size is also
@@ -1306,28 +1327,6 @@ public:
    /// opts.fOutputFormat = ROOT::RDF::ESnapshotOutputFormat::kRNTuple;
    /// df.Snapshot("outputNTuple", "outputFile.root", {"x"}, opts);
    /// ~~~
-   template <typename... ColumnTypes>
-   R__DEPRECATED(
-      6, 40, "Snapshot does not need template arguments anymore, you can safely remove them from this function call.")
-   RResultPtr<RInterface<RLoopManager>> Snapshot(std::string_view treename, std::string_view filename,
-                                                 const ColumnNames_t &columnList,
-                                                 const RSnapshotOptions &options = RSnapshotOptions())
-   {
-      return Snapshot(treename, filename, columnList, options);
-   }
-
-   ////////////////////////////////////////////////////////////////////////////
-   /// \brief Save selected columns to disk, in a new TTree or RNTuple `treename` in file `filename`.
-   /// \param[in] treename The name of the output TTree or RNTuple.
-   /// \param[in] filename The name of the output TFile.
-   /// \param[in] columnList The list of names of the columns/branches/fields to be written.
-   /// \param[in] options RSnapshotOptions struct with extra options to pass to TFile and TTree/RNTuple.
-   /// \return a `RDataFrame` that wraps the snapshotted dataset.
-   ///
-   /// This function returns a `RDataFrame` built with the output TTree or RNTuple as a source.
-   /// The types of the columns are automatically inferred and do not need to be specified.
-   ///
-   /// See above for a more complete description and example usages.
    RResultPtr<RInterface<RLoopManager>> Snapshot(std::string_view treename, std::string_view filename,
                                                  const ColumnNames_t &columnList,
                                                  const RSnapshotOptions &options = RSnapshotOptions())
@@ -1464,7 +1463,7 @@ public:
    /// This function returns a `RDataFrame` built with the output TTree or RNTuple as a source.
    /// The types of the columns are automatically inferred and do not need to be specified.
    ///
-   /// See above for a more complete description and example usages.
+   /// See Snapshot(std::string_view, std::string_view, const ColumnNames_t&, const RSnapshotOptions &) for a more complete description and example usages.
    RResultPtr<RInterface<RLoopManager>> Snapshot(std::string_view treename, std::string_view filename,
                                                  std::string_view columnNameRegexp = "",
                                                  const RSnapshotOptions &options = RSnapshotOptions())
@@ -1507,7 +1506,7 @@ public:
    /// This function returns a `RDataFrame` built with the output TTree or RNTuple as a source.
    /// The types of the columns are automatically inferred and do not need to be specified.
    ///
-   /// See above for a more complete description and example usages.
+   /// See Snapshot(std::string_view, std::string_view, const ColumnNames_t&, const RSnapshotOptions &) for a more complete description and example usages.
    RResultPtr<RInterface<RLoopManager>> Snapshot(std::string_view treename, std::string_view filename,
                                                  std::initializer_list<std::string> columnList,
                                                  const RSnapshotOptions &options = RSnapshotOptions())
@@ -2362,13 +2361,13 @@ public:
    /// auto myGAE2 = myDf.GraphAsymmErrors<f, f, f, f, f, f>("xValues", "yValues", "exl", "exh", "eyl", "eyh");
    /// ~~~
    ///
-   /// `GraphAssymErrors` should also be used for the cases in which values associated only with 
-   /// one of the axes have associated errors. For example, only `ey` exist and `ex` are equal to zero. 
-   /// In such cases, user should do the following: 
+   /// `GraphAsymmErrors` should also be used for the cases in which values associated only with
+   /// one of the axes have associated errors. For example, only `ey` exist and `ex` are equal to zero.
+   /// In such cases, user should do the following:
    /// ~~~{.cpp}
    /// // Create a column of zeros in RDataFrame
-   /// auto rdf_withzeros = rdf.Define("zero", "0"); 
-   /// // or alternatively: 
+   /// auto rdf_withzeros = rdf.Define("zero", "0");
+   /// // or alternatively:
    /// auto rdf_withzeros = rdf.Define("zero", []() -> double { return 0.;});
    /// // Create the graph with y errors only
    /// auto rdf_errorsOnYOnly = rdf_withzeros.GraphAsymmErrors("xValues", "yValues", "zero", "zero", "eyl", "eyh");

tree/dataframe/src/RDefineReader.cxx

@@ -32,19 +32,28 @@ ROOT::Internal::RDF::RDefinesWithReaders::GetReader(unsigned int slot, std::stri
    if (it != defineReaders.end())
       return *it->second;
 
+   std::shared_ptr<RDefineReader> readerToReturn;
    auto *define = fDefine.get();
-   if (*nameIt != "nominal")
-      define = &define->GetVariedDefine(std::string(variationName));
-
-#if !defined(__clang__) && __GNUC__ >= 7 && __GNUC_MINOR__ >= 3
-   const auto insertion =
-      defineReaders.insert({*nameIt, std::make_unique<ROOT::Internal::RDF::RDefineReader>(slot, *define)});
-   return *insertion.first->second;
-#else
-   // gcc < 7.3 has issues with passing the non-movable std::pair temporary into the insert call
-   auto reader = std::make_unique<ROOT::Internal::RDF::RDefineReader>(slot, *define);
-   auto &ret = *reader;
-   defineReaders[*nameIt] = std::move(reader);
-   return ret;
-#endif
+   if (*nameIt == "nominal") {
+      readerToReturn = std::make_shared<RDefineReader>(slot, *define);
+   } else {
+      auto *variedDefine = &define->GetVariedDefine(std::string(variationName));
+      if (variedDefine == define) {
+         // The column in not affected by variations. We can return the same reader as for nominal
+         if (auto nominalReaderIt = defineReaders.find("nominal"); nominalReaderIt != defineReaders.end()) {
+            readerToReturn = nominalReaderIt->second;
+         } else {
+            // The nominal reader doesn't exist yet
+            readerToReturn = std::make_shared<RDefineReader>(slot, *define);
+            auto nominalNameIt = fCachedColNames.Insert("nominal");
+            defineReaders.insert({*nominalNameIt, readerToReturn});
+         }
+      } else {
+         readerToReturn = std::make_shared<RDefineReader>(slot, *variedDefine);
+      }
+   }
+
+   defineReaders.insert({*nameIt, readerToReturn});
+
+   return *readerToReturn;
 }

tree/dataframe/src/RJittedDefine.cxx

@@ -66,5 +66,10 @@ void RJittedDefine::MakeVariations(const std::vector<std::string> &variations)
 RDefineBase &RJittedDefine::GetVariedDefine(const std::string &variationName)
 {
    assert(fConcreteDefine != nullptr);
-   return fConcreteDefine->GetVariedDefine(variationName);
+
+   auto &variedDefine = fConcreteDefine->GetVariedDefine(variationName);
+   if (&variedDefine == fConcreteDefine.get())
+      return *this; // Ensures that the pointer is the same across all variations
+   else
+      return variedDefine;
 }