[df] Support DefinePerSample for RNTuple

tree/dataframe/inc/ROOT/RDataSource.hxx

@@ -94,7 +94,7 @@ const std::vector<std::string> &GetColumnNamesNoDuplicates(const ROOT::RDF::RDat
 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,
+CreateSampleInfo(const ROOT::RDF::RDataSource &ds, unsigned int slot,
                  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);
@@ -168,10 +168,10 @@ protected:
    virtual std::string DescribeDataset() { return "Dataframe from datasource " + GetLabel(); }
 
    friend ROOT::RDF::RSampleInfo
-   ROOT::Internal::RDF::CreateSampleInfo(const ROOT::RDF::RDataSource &,
+   ROOT::Internal::RDF::CreateSampleInfo(const ROOT::RDF::RDataSource &, unsigned int,
                                          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;
+   CreateSampleInfo(unsigned int, 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 {}

tree/dataframe/inc/ROOT/RNTupleDS.hxx

@@ -55,6 +55,7 @@ class RNTupleDS final : public ROOT::RDF::RDataSource {
       ULong64_t fFirstEntry = 0; ///< First entry index in fSource
       /// End entry index in fSource, e.g. the number of entries in the range is fLastEntry - fFirstEntry
       ULong64_t fLastEntry = 0;
+      std::string_view fFileName; ///< Storage location of the current RNTuple
    };
 
    /// A clone of the first pages source's descriptor.
@@ -109,6 +110,8 @@ class RNTupleDS final : public ROOT::RDF::RDataSource {
    /// the fCurrentRanges vectors.  This is necessary because the returned ranges get distributed arbitrarily
    /// onto slots.  In the InitSlot method, the column readers use this map to find the correct range to connect to.
    std::unordered_map<ULong64_t, std::size_t> fFirstEntry2RangeIdx;
+   /// One element per slot, corresponding to the current range index for that slot, as filled by InitSlot
+   std::vector<std::size_t> fSlotsToRangeIdxs;
 
    /// The background thread that runs StageNextSources()
    std::thread fThreadStaging;
@@ -194,6 +197,10 @@ public:
    std::unique_ptr<ROOT::Detail::RDF::RColumnReaderBase>
    GetColumnReaders(unsigned int /*slot*/, std::string_view /*name*/, const std::type_info &) final;
 
+   ROOT::RDF::RSampleInfo
+   CreateSampleInfo(unsigned int,
+                    const std::unordered_map<std::string, ROOT::RDF::Experimental::RSample *> &) const final;
+
    // Old API, unused
    bool SetEntry(unsigned int, ULong64_t) final { return true; }
 

tree/dataframe/inc/ROOT/RTTreeDS.hxx

@@ -75,7 +75,8 @@ class RTTreeDS final : public ROOT::RDF::RDataSource {
    std::unique_ptr<ROOT::Internal::TreeUtils::RNoCleanupNotifier> fNoCleanupNotifier;
 
    ROOT::RDF::RSampleInfo
-   CreateSampleInfo(const std::unordered_map<std::string, ROOT::RDF::Experimental::RSample *> &sampleMap) const final;
+   CreateSampleInfo(unsigned int,
+                    const std::unordered_map<std::string, ROOT::RDF::Experimental::RSample *> &sampleMap) const final;
 
    void RunFinalChecks(bool nodesLeftNotRun) const final;
 

tree/dataframe/src/RDFUtils.cxx

@@ -572,10 +572,10 @@ std::string ROOT::Internal::RDF::DescribeDataset(ROOT::RDF::RDataSource &ds)
 }
 
 ROOT::RDF::RSampleInfo ROOT::Internal::RDF::CreateSampleInfo(
-   const ROOT::RDF::RDataSource &ds,
+   const ROOT::RDF::RDataSource &ds, unsigned int slot,
    const std::unordered_map<std::string, ROOT::RDF::Experimental::RSample *> &sampleMap)
 {
-   return ds.CreateSampleInfo(sampleMap);
+   return ds.CreateSampleInfo(slot, sampleMap);
 }
 
 void ROOT::Internal::RDF::RunFinalChecks(const ROOT::RDF::RDataSource &ds, bool nodesLeftNotRun)

tree/dataframe/src/RDataSource.cxx

@@ -9,7 +9,7 @@
 #endif
 
 ROOT::RDF::RSampleInfo ROOT::RDF::RDataSource::CreateSampleInfo(
-   const std::unordered_map<std::string, ROOT::RDF::Experimental::RSample *> &) const
+   unsigned int, const std::unordered_map<std::string, ROOT::RDF::Experimental::RSample *> &) const
 {
    // Currently not implemented for the generic data source, only works correctly for TTree.
    // TODO: Implement the feature also for the generic data source.

tree/dataframe/src/RLoopManager.cxx

@@ -614,7 +614,7 @@ void RLoopManager::RunDataSource()
 
       ranges = fDataSource->GetEntryRanges();
 
-      fSampleInfos[0] = ROOT::Internal::RDF::CreateSampleInfo(*fDataSource, fSampleMap);
+      fSampleInfos[0] = ROOT::Internal::RDF::CreateSampleInfo(*fDataSource, /*slot*/ 0, fSampleMap);
 
       try {
          for (const auto &range : ranges) {
@@ -1326,6 +1326,8 @@ void ROOT::Detail::RDF::RLoopManager::DataSourceThreadTask(const std::pair<ULong
    RCallCleanUpTask cleanup(*this, slot);
    RDSRangeRAII _{*this, slot, start};
 
+   fSampleInfos[slot] = ROOT::Internal::RDF::CreateSampleInfo(*fDataSource, slot, fSampleMap);
+
    R__LOG_DEBUG(0, RDFLogChannel()) << LogRangeProcessing({fDataSource->GetLabel(), start, end, slot});
 
    try {

tree/dataframe/src/RNTupleDS.cxx

@@ -13,6 +13,7 @@
 
 #include <ROOT/RDF/RColumnReaderBase.hxx>
 #include <ROOT/RDataFrame.hxx>
+#include <ROOT/RDF/Utils.hxx>
 #include <ROOT/RField.hxx>
 #include <ROOT/RFieldUtils.hxx>
 #include <ROOT/RPageStorageFile.hxx>
@@ -411,6 +412,7 @@ std::unique_ptr<ROOT::Internal::RPageSource> CreatePageSource(std::string_view n
 ROOT::RDF::RNTupleDS::RNTupleDS(std::string_view ntupleName, std::string_view fileName)
    : RNTupleDS(CreatePageSource(ntupleName, fileName))
 {
+   fFileNames = std::vector<std::string>{std::string{fileName}};
 }
 
 ROOT::RDF::RNTupleDS::RNTupleDS(RNTuple *ntuple) : RNTupleDS(ROOT::Internal::RPageSourceFile::CreateFromAnchor(*ntuple))
@@ -534,6 +536,7 @@ void ROOT::RDF::RNTupleDS::PrepareNextRanges()
             // to open and attach files here.
             range.fSource = CreatePageSource(fNTupleName, fFileNames[fNextFileIndex]);
          }
+         range.fFileName = fFileNames[fNextFileIndex];
          range.fSource->Attach();
          fNextFileIndex++;
 
@@ -553,6 +556,8 @@ void ROOT::RDF::RNTupleDS::PrepareNextRanges()
    unsigned int nSlotsPerFile = fNSlots / nRemainingFiles;
    for (std::size_t i = 0; (fNextRanges.size() < fNSlots) && (fNextFileIndex < nFiles); ++i) {
       std::unique_ptr<ROOT::Internal::RPageSource> source;
+      // Need to look for the file name to populate the sample info later
+      const auto &sourceFileName = fFileNames[fNextFileIndex];
       std::swap(fStagingArea[fNextFileIndex], source);
       if (!source) {
          // Empty files trigger this condition
@@ -595,6 +600,7 @@ void ROOT::RDF::RNTupleDS::PrepareNextRanges()
          auto end = rangesByCluster[iRange - 1].second;
 
          REntryRangeDS range;
+         range.fFileName = sourceFileName;
          // The last range for this file just takes the already opened page source. All previous ranges clone.
          if (iSlot == N - 1) {
             range.fSource = std::move(source);
@@ -689,10 +695,20 @@ std::vector<std::pair<ULong64_t, ULong64_t>> ROOT::RDF::RNTupleDS::GetEntryRange
 
 void ROOT::RDF::RNTupleDS::InitSlot(unsigned int slot, ULong64_t firstEntry)
 {
-   if (fNSlots == 1)
+   if (fNSlots == 1) {
+      // Ensure the connection between slot and range is valid also in single-thread mode
+      fSlotsToRangeIdxs[0] = 0;
       return;
+   }
 
+   // The same slot ID could be picked multiple times in the same execution, thus
+   // ending up processing different page sources. Here we re-establish the
+   // connection between the slot and the correct page source by finding which
+   // range index corresponds to the first entry passed.
    auto idxRange = fFirstEntry2RangeIdx.at(firstEntry);
+   // We also remember this connection so it can later be retrieved in CreateSampleInfo
+   fSlotsToRangeIdxs[slot * ROOT::Internal::RDF::CacheLineStep<std::size_t>()] = idxRange;
+
    for (auto r : fActiveColumnReaders[slot]) {
       r->Connect(*fCurrentRanges[idxRange].fSource, firstEntry - fCurrentRanges[idxRange].fFirstEntry);
    }
@@ -777,6 +793,7 @@ void ROOT::RDF::RNTupleDS::SetNSlots(unsigned int nSlots)
    assert(nSlots > 0);
    fNSlots = nSlots;
    fActiveColumnReaders.resize(fNSlots);
+   fSlotsToRangeIdxs.resize(fNSlots * ROOT::Internal::RDF::CacheLineStep<std::size_t>());
 }
 
 ROOT::RDataFrame ROOT::RDF::FromRNTuple(std::string_view ntupleName, std::string_view fileName)
@@ -788,3 +805,34 @@ ROOT::RDataFrame ROOT::RDF::FromRNTuple(std::string_view ntupleName, const std::
 {
    return ROOT::RDataFrame(std::make_unique<ROOT::RDF::RNTupleDS>(ntupleName, fileNames));
 }
+
+ROOT::RDF::RSampleInfo ROOT::Internal::RDF::RNTupleDS::CreateSampleInfo(
+   unsigned int slot, const std::unordered_map<std::string, ROOT::RDF::Experimental::RSample *> &sampleMap) const
+{
+   // The same slot ID could be picked multiple times in the same execution, thus
+   // ending up processing different page sources. Here we re-establish the
+   // connection between the slot and the correct page source by retrieving
+   // which range is connected currently to the slot
+   const auto &rangeIdx = fSlotsToRangeIdxs.at(slot * ROOT::Internal::RDF::CacheLineStep<std::size_t>());
+
+   // Missing source if a file does not exist
+   if (!fCurrentRanges[rangeIdx].fSource)
+      return ROOT::RDF::RSampleInfo{};
+
+   const auto &ntupleName = fCurrentRanges[rangeIdx].fSource->GetNTupleName();
+   const auto &ntuplePath = fCurrentRanges[rangeIdx].fFileName;
+   const auto ntupleID = std::string(ntuplePath) + '/' + ntupleName;
+
+   // TODO: There is no support for RNTuple in RDatasetSpec, thus the sample map
+   // is always empty at the moment.
+   if (sampleMap.empty())
+      return ROOT::RDF::RSampleInfo(
+         ntupleID, std::make_pair(fCurrentRanges[rangeIdx].fFirstEntry, fCurrentRanges[rangeIdx].fLastEntry));
+
+   if (sampleMap.find(ntupleID) == sampleMap.end())
+      throw std::runtime_error("Full sample identifier '" + ntupleID + "' cannot be found in the available samples.");
+
+   return ROOT::RDF::RSampleInfo(
+      ntupleID, std::make_pair(fCurrentRanges[rangeIdx].fFirstEntry, fCurrentRanges[rangeIdx].fLastEntry),
+      sampleMap.at(ntupleName));
+}

tree/dataframe/src/RTTreeDS.cxx

@@ -155,7 +155,7 @@ ROOT::Internal::RDF::FromTTree(std::string_view treeName, const std::vector<std:
 }
 
 ROOT::RDF::RSampleInfo ROOT::Internal::RDF::RTTreeDS::CreateSampleInfo(
-   const std::unordered_map<std::string, ROOT::RDF::Experimental::RSample *> &sampleMap) const
+   unsigned int, const std::unordered_map<std::string, ROOT::RDF::Experimental::RSample *> &sampleMap) const
 {
    // one GetTree to retrieve the TChain, another to retrieve the underlying TTree
    auto *tree = fTreeReader->GetTree()->GetTree();