[df] Add new RDataSource for TTree-based processing

Author: vepadulano

bindings/pyroot/pythonizations/test/string_view.py

@@ -27,7 +27,7 @@ def test_rdataframe(self):
 
         df = ROOT.RDataFrame(treename, filename)
         self.assertEqual(
-            str(df), "A data frame built on top of the tree dataset.")
+            str(df), "A data frame associated to the data source \"TTree data source\"")
         os.remove(filename)
 
     def test_17497(self):

tree/dataframe/CMakeLists.txt

@@ -49,6 +49,7 @@ ROOT_STANDARD_LIBRARY_PACKAGE(ROOTDataFrame
     ROOT/RRootDS.hxx
     ROOT/RSnapshotOptions.hxx
     ROOT/RTrivialDS.hxx
+    ROOT/RTTreeDS.hxx
     ROOT/RDF/ActionHelpers.hxx
     ROOT/RDF/ColumnReaderUtils.hxx
     ROOT/RDF/GraphNode.hxx
@@ -107,7 +108,9 @@ ROOT_STANDARD_LIBRARY_PACKAGE(ROOTDataFrame
     src/RCutFlowReport.cxx
     src/RDataFrame.cxx
     src/RDatasetSpec.cxx
+    src/RDataSource.cxx
     src/RDFActionHelpers.cxx
+    src/RDFColumnReaderUtils.cxx
     src/RDFColumnRegister.cxx
     src/RDFDisplay.cxx
     src/RDFGraphUtils.cxx
@@ -127,6 +130,7 @@ ROOT_STANDARD_LIBRARY_PACKAGE(ROOTDataFrame
     src/RRangeBase.cxx
     src/RSample.cxx
     src/RResultPtr.cxx
+    src/RTTreeDS.cxx
     src/RVariationBase.cxx
     src/RVariationReader.cxx
     src/RVariationsDescription.cxx

tree/dataframe/inc/ROOT/RCsvDS.hxx

@@ -73,7 +73,6 @@ private:
    std::int64_t fDataLineNumber = 0;
    std::int64_t fLineNumber = 0;     // used to skip the last lines
    std::int64_t fMaxLineNumber = -1; // set to non-negative if fOptions.fSkipLastNLines is set
-   unsigned int fNSlots = 0U;
    std::unique_ptr<ROOT::Internal::RRawFile> fCsvFile;
    ULong64_t fEntryRangesRequested = 0ULL;
    ULong64_t fProcessedLines = 0ULL; // marks the progress of the consumption of the csv lines

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

@@ -52,6 +52,7 @@
 #include "ROOT/RNTupleDS.hxx"
 #include "ROOT/RNTupleWriter.hxx" // for SnapshotRNTupleHelper
 #endif
+#include "ROOT/RTTreeDS.hxx"
 
 #include <algorithm>
 #include <functional>
@@ -1530,12 +1531,15 @@ class R__CLING_PTRCHECK(off) SnapshotTTreeHelper : public RActionImpl<SnapshotTT
    std::vector<void *> fBranchAddresses; // Addresses of objects associated to output branches
    RBranchSet fOutputBranches;
    std::vector<bool> fIsDefine;
+   std::weak_ptr<ROOT::Detail::RDF::RLoopManager> fOutputLoopManager;
+   ROOT::RDF::RDataSource *fInputDataSource;
 
 public:
    using ColumnTypes_t = TypeList<ColTypes...>;
    SnapshotTTreeHelper(std::string_view filename, std::string_view dirname, std::string_view treename,
                        const ColumnNames_t &vbnames, const ColumnNames_t &bnames, const RSnapshotOptions &options,
-                       std::vector<bool> &&isDefine)
+                       std::vector<bool> &&isDefine, std::weak_ptr<ROOT::Detail::RDF::RLoopManager> loopManager,
+                       ROOT::RDF::RDataSource *inputDataSource)
       : fFileName(filename),
         fDirName(dirname),
         fTreeName(treename),
@@ -1544,7 +1548,9 @@ public:
         fOutputBranchNames(ReplaceDotWithUnderscore(bnames)),
         fBranches(vbnames.size(), nullptr),
         fBranchAddresses(vbnames.size(), nullptr),
-        fIsDefine(std::move(isDefine))
+        fIsDefine(std::move(isDefine)),
+        fOutputLoopManager(loopManager),
+        fInputDataSource(inputDataSource)
    {
       EnsureValidSnapshotTTreeOutput(fOptions, fTreeName, fFileName);
    }
@@ -1571,6 +1577,8 @@ public:
    {
       if (r)
          fInputTree = r->GetTree();
+      else if (auto treeDS = dynamic_cast<ROOT::Internal::RDF::RTTreeDS *>(fInputDataSource))
+         fInputTree = treeDS->GetTree();
       fBranchAddressesNeedReset = true;
    }
 
@@ -1650,6 +1658,11 @@ public:
       // must destroy the TTree first, otherwise TFile will delete it too leading to a double delete
       fOutputTree.reset();
       fOutputFile->Close();
+
+      // Now connect the data source to the loop manager so it can be used for further processing
+      auto fullTreeName = fDirName.empty() ? fTreeName : fDirName + '/' + fTreeName;
+      if (auto lmPtr = fOutputLoopManager.lock())
+         lmPtr->SetDataSource(std::make_unique<ROOT::Internal::RDF::RTTreeDS>(fullTreeName, fFileName));
    }
 
    std::string GetActionName() { return "Snapshot"; }
@@ -1673,8 +1686,15 @@ public:
    SnapshotTTreeHelper MakeNew(void *newName, std::string_view /*variation*/ = "nominal")
    {
       const std::string finalName = *reinterpret_cast<const std::string *>(newName);
-      return SnapshotTTreeHelper{
-         finalName, fDirName, fTreeName, fInputBranchNames, fOutputBranchNames, fOptions, std::vector<bool>(fIsDefine)};
+      return SnapshotTTreeHelper{finalName,
+                                 fDirName,
+                                 fTreeName,
+                                 fInputBranchNames,
+                                 fOutputBranchNames,
+                                 fOptions,
+                                 std::vector<bool>(fIsDefine),
+                                 fOutputLoopManager,
+                                 fInputDataSource};
    }
 };
 
@@ -1699,12 +1719,17 @@ class R__CLING_PTRCHECK(off) SnapshotTTreeHelperMT : public RActionImpl<Snapshot
    std::vector<std::vector<void *>> fBranchAddresses;
    std::vector<RBranchSet> fOutputBranches;
    std::vector<bool> fIsDefine;
+   std::weak_ptr<ROOT::Detail::RDF::RLoopManager> fOutputLoopManager;
+   ROOT::RDF::RDataSource *fInputDataSource;
 
 public:
    using ColumnTypes_t = TypeList<ColTypes...>;
+
    SnapshotTTreeHelperMT(const unsigned int nSlots, std::string_view filename, std::string_view dirname,
                          std::string_view treename, const ColumnNames_t &vbnames, const ColumnNames_t &bnames,
-                         const RSnapshotOptions &options, std::vector<bool> &&isDefine)
+                         const RSnapshotOptions &options, std::vector<bool> &&isDefine,
+                         std::weak_ptr<ROOT::Detail::RDF::RLoopManager> loopManager,
+                         ROOT::RDF::RDataSource *inputDataSource)
       : fNSlots(nSlots),
         fOutputFiles(fNSlots),
         fOutputTrees(fNSlots),
@@ -1719,7 +1744,9 @@ public:
         fBranches(fNSlots, std::vector<TBranch *>(vbnames.size(), nullptr)),
         fBranchAddresses(fNSlots, std::vector<void *>(vbnames.size(), nullptr)),
         fOutputBranches(fNSlots),
-        fIsDefine(std::move(isDefine))
+        fIsDefine(std::move(isDefine)),
+        fOutputLoopManager(loopManager),
+        fInputDataSource(inputDataSource)
    {
       EnsureValidSnapshotTTreeOutput(fOptions, fTreeName, fFileName);
    }
@@ -1766,7 +1793,9 @@ public:
       if (r) {
          // not an empty-source RDF
          fInputTrees[slot] = r->GetTree();
-      }
+      } else if (auto treeDS = dynamic_cast<ROOT::Internal::RDF::RTTreeDS *>(fInputDataSource))
+         fInputTrees[slot] = treeDS->GetTree();
+
       fBranchAddressesNeedReset[slot] = 1; // reset first event flag for this slot
    }
 
@@ -1855,6 +1884,11 @@ public:
       // flush all buffers to disk by destroying the TBufferMerger
       fOutputFiles.clear();
       fMerger.reset();
+
+      // Now connect the data source to the loop manager so it can be used for further processing
+      auto fullTreeName = fDirName.empty() ? fTreeName : fDirName + '/' + fTreeName;
+      if (auto lmPtr = fOutputLoopManager.lock())
+         lmPtr->SetDataSource(std::make_unique<ROOT::Internal::RDF::RTTreeDS>(fullTreeName, fFileName));
    }
 
    std::string GetActionName() { return "Snapshot"; }
@@ -1878,8 +1912,16 @@ public:
    SnapshotTTreeHelperMT MakeNew(void *newName, std::string_view /*variation*/ = "nominal")
    {
       const std::string finalName = *reinterpret_cast<const std::string *>(newName);
-      return SnapshotTTreeHelperMT{fNSlots,           finalName,          fDirName, fTreeName,
-                                   fInputBranchNames, fOutputBranchNames, fOptions, std::vector<bool>(fIsDefine)};
+      return SnapshotTTreeHelperMT{fNSlots,
+                                   finalName,
+                                   fDirName,
+                                   fTreeName,
+                                   fInputBranchNames,
+                                   fOutputBranchNames,
+                                   fOptions,
+                                   std::vector<bool>(fIsDefine),
+                                   fOutputLoopManager,
+                                   fInputDataSource};
    }
 };
 
@@ -1906,7 +1948,7 @@ class R__CLING_PTRCHECK(off) SnapshotRNTupleHelper : public RActionImpl<Snapshot
    std::unique_ptr<TFile> fOutputFile{nullptr};
 
    RSnapshotOptions fOptions;
-   ROOT::Detail::RDF::RLoopManager *fLoopManager;
+   std::weak_ptr<ROOT::Detail::RDF::RLoopManager> fOutputLoopManager;
    ColumnNames_t fInputFieldNames; // This contains the resolved aliases
    ColumnNames_t fOutputFieldNames;
    std::unique_ptr<ROOT::Experimental::RNTupleWriter> fWriter{nullptr};
@@ -1919,11 +1961,11 @@ public:
    using ColumnTypes_t = TypeList<ColTypes...>;
    SnapshotRNTupleHelper(std::string_view filename, std::string_view ntuplename, const ColumnNames_t &vfnames,
                          const ColumnNames_t &fnames, const RSnapshotOptions &options,
-                         ROOT::Detail::RDF::RLoopManager *lm, std::vector<bool> &&isDefine)
+                         std::weak_ptr<ROOT::Detail::RDF::RLoopManager> lm, std::vector<bool> &&isDefine)
       : fFileName(filename),
         fNTupleName(ntuplename),
         fOptions(options),
-        fLoopManager(lm),
+        fOutputLoopManager(lm),
         fInputFieldNames(vfnames),
         fOutputFieldNames(ReplaceDotWithUnderscore(fnames)),
         fIsDefine(std::move(isDefine))
@@ -1937,7 +1979,7 @@ public:
    SnapshotRNTupleHelper &operator=(SnapshotRNTupleHelper &&) = default;
    ~SnapshotRNTupleHelper()
    {
-      if (!fNTupleName.empty() && !fLoopManager->GetDataSource() && fOptions.fLazy)
+      if (auto lm = fOutputLoopManager.lock(); lm && !fNTupleName.empty() && !lm->GetDataSource() && fOptions.fLazy)
          Warning("Snapshot", "A lazy Snapshot action was booked but never triggered.");
    }
 
@@ -1993,7 +2035,8 @@ public:
    {
       fWriter.reset();
       // We can now set the data source of the loop manager for the RDataFrame that is returned by the Snapshot call.
-      fLoopManager->SetDataSource(std::make_unique<ROOT::Experimental::RNTupleDS>(fNTupleName, fFileName));
+      if (auto lm = fOutputLoopManager.lock())
+         lm->SetDataSource(std::make_unique<ROOT::Experimental::RNTupleDS>(fNTupleName, fFileName));
    }
 
    std::string GetActionName() { return "Snapshot"; }
@@ -2022,7 +2065,7 @@ public:
                                    fInputFieldNames,
                                    fOutputFieldNames,
                                    fOptions,
-                                   fLoopManager,
+                                   fOutputLoopManager,
                                    std::vector<bool>(fIsDefine)};
    }
 };

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

@@ -17,13 +17,11 @@
 #include "RDefineReader.hxx"
 #include "RDSColumnReader.hxx"
 #include "RLoopManager.hxx"
-#include "RTreeColumnReader.hxx"
 #include "RVariationBase.hxx"
 #include "RVariationReader.hxx"
 
 #include <ROOT/RDataSource.hxx>
 #include <ROOT/TypeTraits.hxx>
-#include <TTreeReader.h>
 
 #include <array>
 #include <cassert>
@@ -33,32 +31,21 @@
 #include <typeinfo> // for typeid
 #include <vector>
 
+class TTreeReader;
+
 namespace ROOT {
 namespace Internal {
 namespace RDF {
 
 using namespace ROOT::TypeTraits;
 namespace RDFDetail = ROOT::Detail::RDF;
 
-template <typename T>
 RDFDetail::RColumnReaderBase *GetColumnReader(unsigned int slot, RColumnReaderBase *defineOrVariationReader,
-                                              RLoopManager &lm, TTreeReader *r, const std::string &colName)
-{
-   if (defineOrVariationReader != nullptr)
-      return defineOrVariationReader;
-
-   // Check if we already inserted a reader for this column in the dataset column readers (RDataSource or Tree/TChain
-   // readers)
-   auto *datasetColReader = lm.GetDatasetColumnReader(slot, colName, typeid(T));
-   if (datasetColReader != nullptr)
-      return datasetColReader;
+                                              RLoopManager &lm, const std::string &colName, const std::type_info &ti);
 
-   assert(r != nullptr && "We could not find a reader for this column, this should never happen at this point.");
-
-   // Make a RTreeColumnReader for this column and insert it in RLoopManager's map
-   auto treeColReader = std::make_unique<RTreeColumnReader<T>>(*r, colName);
-   return lm.AddTreeColumnReader(slot, colName, std::move(treeColReader), typeid(T));
-}
+RDFDetail::RColumnReaderBase *GetColumnReaderForTTreeIMT(unsigned int slot, RColumnReaderBase *defineOrVariationReader,
+                                                         RLoopManager &lm, TTreeReader &treeReader,
+                                                         const std::string &colName, const std::type_info &ti);
 
 /// This type aggregates some of the arguments passed to GetColumnReaders.
 /// We need to pass a single RColumnReadersInfo object rather than each argument separately because with too many
@@ -74,7 +61,7 @@ struct RColumnReadersInfo {
 /// Create a group of column readers, one per type in the parameter pack.
 template <typename... ColTypes>
 std::array<RDFDetail::RColumnReaderBase *, sizeof...(ColTypes)>
-GetColumnReaders(unsigned int slot, TTreeReader *r, TypeList<ColTypes...>, const RColumnReadersInfo &colInfo,
+GetColumnReaders(unsigned int slot, TTreeReader *treeReader, TypeList<ColTypes...>, const RColumnReadersInfo &colInfo,
                  const std::string &variationName = "nominal")
 {
    // see RColumnReadersInfo for why we pass these arguments like this rather than directly as function arguments
@@ -83,10 +70,18 @@ GetColumnReaders(unsigned int slot, TTreeReader *r, TypeList<ColTypes...>, const
    auto &colRegister = colInfo.fColRegister;
 
    int i = -1;
-   std::array<RDFDetail::RColumnReaderBase *, sizeof...(ColTypes)> ret{
-      (++i, GetColumnReader<ColTypes>(slot, colRegister.GetReader(slot, colNames[i], variationName, typeid(ColTypes)),
-                                      lm, r, colNames[i]))...};
-   return ret;
+   if (!treeReader) {
+      std::array<RDFDetail::RColumnReaderBase *, sizeof...(ColTypes)> ret{
+         (++i, GetColumnReader(slot, colRegister.GetReader(slot, colNames[i], variationName, typeid(ColTypes)), lm,
+                               colNames[i], typeid(ColTypes)))...};
+      return ret;
+   } else {
+      std::array<RDFDetail::RColumnReaderBase *, sizeof...(ColTypes)> ret{
+         (++i,
+          GetColumnReaderForTTreeIMT(slot, colRegister.GetReader(slot, colNames[i], variationName, typeid(ColTypes)),
+                                     lm, *treeReader, colNames[i], typeid(ColTypes)))...};
+      return ret;
+   }
 }
 
 // Shortcut overload for the case of no columns