[RDF] Add internal utility to retrieve cluster entry ranges

Author: siliataider

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

@@ -130,7 +130,6 @@ protected:
       }
    }
 
-   RDFDetail::RLoopManager *GetLoopManager() const { return fLoopManager.get(); }
    RDataSource *GetDataSource() const { return fLoopManager->GetDataSource(); }
 
    ColumnNames_t GetValidatedColumnNames(const unsigned int nColumns, const ColumnNames_t &columns)
@@ -208,6 +207,8 @@ public:
    RInterfaceBase(std::shared_ptr<RDFDetail::RLoopManager> lm);
    RInterfaceBase(RDFDetail::RLoopManager &lm, const RDFInternal::RColumnRegister &colRegister);
 
+   RDFDetail::RLoopManager *GetLoopManager() const { return fLoopManager.get(); }
+
    ColumnNames_t GetColumnNames();
 
    std::string GetColumnType(std::string_view column);

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

@@ -55,6 +55,8 @@ ROOT::RLogChannel &RDFLogChannel();
 // fwd decl for ColumnName2ColumnTypeName
 class RDefineBase;
 
+class RLoopManager;
+
 // type used for tag dispatching
 struct RInferredType {
 };
@@ -339,6 +341,12 @@ auto MakeAliasedSharedPtr(T *rawPtr)
  */
 ROOT::RDF::Experimental::RDatasetSpec RetrieveSpecFromJson(const std::string &jsonFile);
 
+/**
+ * \brief Function to retrieve the entry ranges for each cluster in the dataset,
+ * across files, with a global offset.
+ */
+std::vector<std::pair<ULong64_t, ULong64_t>> GetClusterRanges(Detail::RDF::RLoopManager &lm);
+
 /**
  * Tag to let data sources use the native data type when creating a column reader.
  *

tree/dataframe/src/RDFUtils.cxx

@@ -16,13 +16,17 @@
 #include "ROOT/RDF/RSampleInfo.hxx"
 #include "ROOT/RDF/Utils.hxx"
 #include "ROOT/RLogger.hxx"
+#include "ROOT/RTTreeDS.hxx"
 #include "RtypesCore.h"
 #include "TBranch.h"
 #include "TBranchElement.h"
+#include <TChain.h>
+#include "TChainElement.h"
 #include "TClass.h"
 #include "TClassEdit.h"
 #include "TClassRef.h"
 #include "TError.h" // Info
+#include "TFile.h"
 #include "TInterpreter.h"
 #include "TLeaf.h"
 #include "TROOT.h" // IsImplicitMTEnabled, GetThreadPoolSize
@@ -632,9 +636,69 @@ ROOT::RDF::Experimental::RDatasetSpec RetrieveSpecFromJson(const std::string &js
    return spec;
 };
 
-} // end NS RDF
-} // end NS Internal
-} // end NS ROOT
+std::vector<std::pair<ULong64_t, ULong64_t>> GetClusterRanges(ROOT::Detail::RDF::RLoopManager &lm)
+{
+   std::vector<std::pair<ULong64_t, ULong64_t>> ranges{};
+
+   auto *ds = lm.GetDataSource();
+   if (!ds) {
+      throw std::runtime_error("Cannot retrieve cluster ranges, no data source available");
+   }
+
+   auto *ttreeds = dynamic_cast<RTTreeDS *>(ds);
+   if (!ttreeds) {
+      throw std::runtime_error("Cannot retrieve cluster ranges, data source is not a TTree");
+   }
+
+   TTree *tree = ttreeds->GetTree();
+   if (!tree) {
+      return ranges;
+   }
+
+   // Collect cluster boundaries from a single tree applying a global offset
+   auto collectClusters = [&ranges](TTree *t, Long64_t offset) -> Long64_t {
+      const Long64_t nEntries = t->GetEntries();
+      auto it = t->GetClusterIterator(0);
+      Long64_t start = 0;
+      while ((start = it()) < nEntries)
+         ranges.emplace_back(offset + start, offset + it.GetNextEntry());
+      return nEntries;
+   };
+
+   if (auto *chain = dynamic_cast<TChain *>(tree)) {
+      // If TChain, open each file and iterate over its clusters
+      TObjArray *files = chain->GetListOfFiles();
+      if (!files) {
+         return ranges;
+      }
+
+      Long64_t offset = 0;
+      for (int i = 0; i < files->GetEntries(); ++i) {
+         auto *el = static_cast<TChainElement *>(files->At(i));
+         if (!el) {
+            continue;
+         }
+         std::unique_ptr<TFile> f{TFile::Open(el->GetTitle(), "READ")};
+         if (!f || f->IsZombie()) {
+            continue;
+         }
+         auto *t = f->Get<TTree>(el->GetName());
+         if (!t) {
+            continue;
+         }
+         offset += collectClusters(t, offset);
+      }
+   } else {
+      // If single TTree, iterate over its clusters
+      collectClusters(tree, 0);
+   }
+
+   return ranges;
+}
+
+} // namespace RDF
+} // namespace Internal
+} // namespace ROOT
 
 std::string
 ROOT::Internal::RDF::GetTypeNameWithOpts(const ROOT::RDF::RDataSource &df, std::string_view colName, bool vector2RVec)