[df] Add data source for TTree processing

A new RDataSource derived class is implemented for TTree data processing. The
class follows almost entirely the same approach for processing of other data
sources, with the exception of the MT case which still needs to be partially
handled externally by TTreeProcessorMT. Nonetheless, the class and its API can
be used throughout the RDF codebase to centralise all TTree-related usage and
processing. While implementing the class, a few missing features or limitations
of RDataSource were found which needed API extensions. For the moment, all the
API extensions of RDataSource are private and it will be later decided whether
to make them public or not.

Author: vepadulano

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,6 +108,7 @@ 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
@@ -129,6 +131,7 @@ ROOT_STANDARD_LIBRARY_PACKAGE(ROOTDataFrame
     src/RSample.cxx
     src/RTreeColumnReader.cxx
     src/RResultPtr.cxx
+    src/RTTreeDS.cxx
     src/RVariationBase.cxx
     src/RVariationReader.cxx
     src/RVariationsDescription.cxx

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

@@ -44,6 +44,7 @@ class RDataSource;
 } // ns RDF
 
 namespace Internal {
+class RSlotStack;
 namespace RDF {
 std::vector<std::string> GetBranchNames(TTree &t, bool allowDuplicates = true);
 
@@ -312,6 +313,13 @@ public:
    {
       return fSuppressErrorsForMissingBranches;
    }
+
+   /// The task run by every thread on the input entry range, for the generic RDataSource.
+   void DataSourceThreadTask(const std::pair<ULong64_t, ULong64_t> &entryRange, ROOT::Internal::RSlotStack &slotStack,
+                             std::atomic<ULong64_t> &entryCount);
+   /// The task run by every thread on an entry range (known by the input TTreeReader), for the TTree data source.
+   void
+   TTreeThreadTask(TTreeReader &treeReader, ROOT::Internal::RSlotStack &slotStack, std::atomic<ULong64_t> &entryCount);
 };
 
 /// \brief Create an RLoopManager that reads a TChain.

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

@@ -123,6 +123,8 @@ struct IsVector_t : public std::false_type {};
 template <typename T, typename A>
 struct IsVector_t<std::vector<T, A>> : public std::true_type {};
 
+std::string GetBranchOrLeafTypeName(TTree &t, const std::string &colName);
+
 const std::type_info &TypeName2TypeID(const std::string &name);
 
 std::string TypeID2TypeName(const std::type_info &id);
@@ -314,6 +316,14 @@ struct CallGuaranteedOrder {
       f(std::forward<Args>(args)...);
    }
 };
+
+template <typename T>
+auto MakeAliasedSharedPtr(T *rawPtr)
+{
+   const static std::shared_ptr<T> fgRawPtrCtrlBlock;
+   return std::shared_ptr<T>(fgRawPtrCtrlBlock, rawPtr);
+}
+
 } // end NS RDF
 } // end NS Internal
 } // end NS ROOT

tree/dataframe/inc/ROOT/RDataSource.hxx

@@ -18,13 +18,28 @@
 
 #include <algorithm> // std::transform
 #include <cassert>
+#include <optional>
+#include <set>
 #include <string>
 #include <typeinfo>
+#include <unordered_map>
+#include <variant>
 #include <vector>
+#include <functional>
+
+// Need to fwd-declare TTreeReader for CreateColumnReader
+class TTreeReader;
+namespace ROOT::Detail::RDF {
+class RLoopManager;
+}
 
 namespace ROOT {
 namespace RDF {
 class RDataSource;
+class RSampleInfo;
+namespace Experimental {
+class RSample;
+}
 }
 }
 
@@ -71,6 +86,23 @@ public:
 };
 
 } // ns TDS
+
+namespace RDF {
+std::string GetTypeNameWithOpts(const ROOT::RDF::RDataSource &ds, std::string_view colName, bool vector2RVec);
+const std::vector<std::string> &GetTopLevelFieldNames(const ROOT::RDF::RDataSource &ds);
+const std::vector<std::string> &GetColumnNamesNoDuplicates(const ROOT::RDF::RDataSource &ds);
+void CallInitializeWithOpts(ROOT::RDF::RDataSource &ds, const std::set<std::string> &suppressErrorsForMissingColumns);
+std::string DescribeDataset(ROOT::RDF::RDataSource &ds);
+ROOT::RDF::RSampleInfo
+CreateSampleInfo(const ROOT::RDF::RDataSource &ds,
+                 const std::unordered_map<std::string, ROOT::RDF::Experimental::RSample *> &sampleMap);
+void RunFinalChecks(const ROOT::RDF::RDataSource &ds, bool nodesLeftNotRun);
+void ProcessMT(ROOT::RDF::RDataSource &ds, ROOT::Detail::RDF::RLoopManager &lm);
+std::unique_ptr<ROOT::Detail::RDF::RColumnReaderBase>
+CreateColumnReader(ROOT::RDF::RDataSource &ds, unsigned int slot, std::string_view col, const std::type_info &tid,
+                   TTreeReader *treeReader);
+} // namespace RDF
+
 } // ns Internal
 
 namespace RDF {
@@ -117,6 +149,57 @@ protected:
 
    unsigned int fNSlots{};
 
+   std::optional<std::pair<ULong64_t, ULong64_t>> fGlobalEntryRange{};
+
+   friend std::string ROOT::Internal::RDF::GetTypeNameWithOpts(const RDataSource &, std::string_view, bool);
+   virtual std::string GetTypeNameWithOpts(std::string_view colName, bool) const { return GetTypeName(colName); }
+
+   friend const std::vector<std::string> &ROOT::Internal::RDF::GetTopLevelFieldNames(const ROOT::RDF::RDataSource &);
+   virtual const std::vector<std::string> &GetTopLevelFieldNames() const { return GetColumnNames(); }
+
+   friend const std::vector<std::string> &
+   ROOT::Internal::RDF::GetColumnNamesNoDuplicates(const ROOT::RDF::RDataSource &);
+   virtual const std::vector<std::string> &GetColumnNamesNoDuplicates() const { return GetColumnNames(); }
+
+   friend void ROOT::Internal::RDF::CallInitializeWithOpts(ROOT::RDF::RDataSource &, const std::set<std::string> &);
+   virtual void InitializeWithOpts(const std::set<std::string> &) { Initialize(); }
+
+   friend std::string ROOT::Internal::RDF::DescribeDataset(ROOT::RDF::RDataSource &);
+   virtual std::string DescribeDataset() { return "Dataframe from datasource " + GetLabel(); }
+
+   friend ROOT::RDF::RSampleInfo
+   ROOT::Internal::RDF::CreateSampleInfo(const ROOT::RDF::RDataSource &,
+                                         const std::unordered_map<std::string, ROOT::RDF::Experimental::RSample *> &);
+   virtual ROOT::RDF::RSampleInfo
+   CreateSampleInfo(const std::unordered_map<std::string, ROOT::RDF::Experimental::RSample *> &) const;
+
+   friend void ROOT::Internal::RDF::RunFinalChecks(const ROOT::RDF::RDataSource &, bool);
+   virtual void RunFinalChecks(bool) const {}
+
+   friend void ROOT::Internal::RDF::ProcessMT(RDataSource &, ROOT::Detail::RDF::RLoopManager &);
+   virtual void ProcessMT(ROOT::Detail::RDF::RLoopManager &);
+
+   friend std::unique_ptr<ROOT::Detail::RDF::RColumnReaderBase>
+   ROOT::Internal::RDF::CreateColumnReader(ROOT::RDF::RDataSource &, unsigned int, std::string_view,
+                                           const std::type_info &, TTreeReader *);
+   /**
+    * \brief Creates a column reader for the requested column
+    *
+    * In the general case, this is just a redirect to the right GetColumnReaders overload. The signature notably also
+    * has a TTreeReader * parameter. This is currently necessary to still allow the TTree-based MT scheduling via
+    * TTreeProcessorMT. We use the TTreeProcessorMT::Process method to launch the same kernel across all threads. In
+    * each thread task, TTreeProcessorMT creates a thread-local instance of a TTreeReader which is going to read the
+    * range of events assigned to that task. That TTreeReader instance is what is passed to this method whenever a
+    * column reader needs to be created in a thread task. In the future this method might be removed by either allowing
+    * to request a handle to the thread-local TTreeReader instance programmatically from the TTreeProcessorMT, or
+    * refactoring the TTreeProcessorMT scheduling into RTTreeDS altogether.
+    */
+   virtual std::unique_ptr<ROOT::Detail::RDF::RColumnReaderBase>
+   CreateColumnReader(unsigned int slot, std::string_view col, const std::type_info &tid, TTreeReader *)
+   {
+      return GetColumnReaders(slot, col, tid);
+   }
+
 public:
    RDataSource() = default;
    // Rule of five
@@ -242,6 +325,13 @@ public:
    /// Concrete datasources can override the default implementation.
    virtual std::string GetLabel() { return "Custom Datasource"; }
 
+   /// \brief Restrict processing to a [begin, end) range of entries.
+   /// \param entryRange The range of entries to process.
+   virtual void SetGlobalEntryRange(std::pair<ULong64_t, ULong64_t> entryRange)
+   {
+      fGlobalEntryRange = std::move(entryRange);
+   };
+
 protected:
    /// type-erased vector of pointers to pointers to column values - one per slot
    virtual Record_t GetColumnReadersImpl(std::string_view name, const std::type_info &) = 0;

tree/dataframe/inc/ROOT/RTTreeDS.hxx

@@ -0,0 +1,154 @@
+/**
+ \file ROOT/RTTreeDS.hxx
+ \ingroup dataframe
+ \author Vincenzo Eduardo Padulano
+ \date 2024-12
+*/
+
+/*************************************************************************
+ * Copyright (C) 1995-2024, Rene Brun and Fons Rademakers.               *
+ * All rights reserved.                                                  *
+ *                                                                       *
+ * For the licensing terms see $ROOTSYS/LICENSE.                         *
+ * For the list of contributors see $ROOTSYS/README/CREDITS.             *
+ *************************************************************************/
+
+#ifndef ROOT_INTERNAL_RDF_RTTREEDS
+#define ROOT_INTERNAL_RDF_RTTREEDS
+
+#include "ROOT/RDataSource.hxx"
+
+#include <memory>
+#include <string>
+#include <vector>
+#include <stdexcept>
+#include <string_view>
+
+// Begin forward decls
+
+namespace ROOT {
+class RDataFrame;
+}
+
+namespace ROOT::Detail::RDF {
+class RLoopManager;
+}
+
+namespace ROOT::RDF {
+class RSampleInfo;
+}
+
+namespace ROOT::RDF::Experimental {
+class RSample;
+}
+
+namespace ROOT::TreeUtils {
+struct RFriendInfo;
+}
+
+class TChain;
+class TDirectory;
+class TTree;
+class TTreeReader;
+
+// End forward decls
+
+namespace ROOT::Internal::RDF {
+
+class RTTreeDS final : public ROOT::RDF::RDataSource {
+   std::vector<std::string> fBranchNamesWithDuplicates{};
+   std::vector<std::string> fBranchNamesWithoutDuplicates{};
+   std::vector<std::string> fTopLevelBranchNames{};
+
+   std::shared_ptr<TTree> fTree;
+
+   std::unique_ptr<TTreeReader> fTreeReader;
+
+   std::vector<std::unique_ptr<TChain>> fFriends;
+
+   ROOT::RDF::RSampleInfo
+   CreateSampleInfo(const std::unordered_map<std::string, ROOT::RDF::Experimental::RSample *> &sampleMap) const final;
+
+   void RunFinalChecks(bool nodesLeftNotRun) const final;
+
+   void Setup(std::shared_ptr<TTree> &&tree, const ROOT::TreeUtils::RFriendInfo *friendInfo = nullptr);
+
+   std::vector<std::pair<ULong64_t, ULong64_t>> GetTTreeEntryRange(TTree &tree);
+   std::vector<std::pair<ULong64_t, ULong64_t>> GetTChainEntryRange(TChain &chain);
+
+public:
+   RTTreeDS(std::shared_ptr<TTree> tree);
+   RTTreeDS(std::shared_ptr<TTree> tree, const ROOT::TreeUtils::RFriendInfo &friendInfo);
+   RTTreeDS(std::string_view treeName, TDirectory *dirPtr);
+   RTTreeDS(std::string_view treeName, std::string_view fileNameGlob);
+   RTTreeDS(std::string_view treeName, const std::vector<std::string> &fileNameGlobs);
+
+   // Rule of five
+   RTTreeDS(const RTTreeDS &) = delete;
+   RTTreeDS &operator=(const RTTreeDS &) = delete;
+   RTTreeDS(RTTreeDS &&) = delete;
+   RTTreeDS &operator=(RTTreeDS &&) = delete;
+   ~RTTreeDS() final; // Define destructor where data member types are defined
+
+   void Initialize() final;
+
+   void Finalize() final;
+
+   std::vector<std::pair<ULong64_t, ULong64_t>> GetEntryRanges() final;
+
+   const std::vector<std::string> &GetColumnNames() const final { return fBranchNamesWithDuplicates; }
+
+   bool HasColumn(std::string_view colName) const final
+   {
+      return std::find(fBranchNamesWithDuplicates.begin(), fBranchNamesWithDuplicates.end(), colName) !=
+             fBranchNamesWithDuplicates.end();
+   }
+
+   std::string GetTypeName(std::string_view colName) const final;
+
+   std::string GetTypeNameWithOpts(std::string_view colName, bool vector2RVec) const final;
+
+   bool SetEntry(unsigned int, ULong64_t entry) final;
+
+   Record_t GetColumnReadersImpl(std::string_view /* name */, const std::type_info & /* ti */) final
+   {
+      // This datasource uses the newer GetColumnReaders() API
+      return {};
+   }
+
+   std::unique_ptr<ROOT::Detail::RDF::RColumnReaderBase>
+   GetColumnReaders(unsigned int, std::string_view, const std::type_info &) final
+   {
+      // This data source creates column readers via CreateColumnReader
+      throw std::runtime_error("GetColumnReaders should not be called on this data source, something wrong happened!");
+   }
+
+   std::unique_ptr<ROOT::Detail::RDF::RColumnReaderBase> CreateColumnReader(unsigned int slot, std::string_view col,
+                                                                            const std::type_info &tid,
+                                                                            TTreeReader *treeReader) final;
+
+   std::string GetLabel() final { return "TTreeDS"; }
+
+   TTree *GetTree();
+
+   const std::vector<std::string> &GetTopLevelFieldNames() const final { return fTopLevelBranchNames; }
+
+   const std::vector<std::string> &GetColumnNamesNoDuplicates() const final { return fBranchNamesWithoutDuplicates; }
+
+   void InitializeWithOpts(const std::set<std::string> &suppressErrorsForMissingBranches) final;
+
+   std::string DescribeDataset() final;
+
+   std::string AsString() final { return "TTree data source"; }
+
+   std::size_t GetNFiles() const final;
+
+   void ProcessMT(ROOT::Detail::RDF::RLoopManager &lm) final;
+};
+
+ROOT::RDataFrame FromTTree(std::string_view treeName, std::string_view fileNameGlob);
+ROOT::RDataFrame FromTTree(std::string_view treeName, const std::vector<std::string> &fileNameGlobs);
+
+} // namespace ROOT::Internal::RDF
+
+#endif