diff --git a/core/imt/CMakeLists.txt b/core/imt/CMakeLists.txt index c99e8cd53ece9..e042fab222151 100644 --- a/core/imt/CMakeLists.txt +++ b/core/imt/CMakeLists.txt @@ -32,6 +32,7 @@ if(imt) ROOT/RSlotStack.hxx ROOT/TExecutor.hxx ROOT/TThreadExecutor.hxx + ROOT/InternalIMTUtils.hxx LINKDEF LinkDef.h MODULE diff --git a/core/imt/inc/ROOT/InternalIMTUtils.hxx b/core/imt/inc/ROOT/InternalIMTUtils.hxx new file mode 100644 index 0000000000000..b12b73185bd15 --- /dev/null +++ b/core/imt/inc/ROOT/InternalIMTUtils.hxx @@ -0,0 +1,23 @@ +#ifndef ROOT_INTERNAL_IMTUTILS +#define ROOT_INTERNAL_IMTUTILS + +#include + +namespace ROOT::Internal::IMTUtils { +class RParallelSplitFileProcessor { +public: + virtual std::unique_ptr SplitWork() = 0; + virtual bool Empty() const = 0; + virtual bool IsDivisible() const = 0; + RParallelSplitFileProcessor() = default; + + // Rule of five + virtual ~RParallelSplitFileProcessor() = default; + RParallelSplitFileProcessor(const RParallelSplitFileProcessor &) = delete; + RParallelSplitFileProcessor &operator=(const RParallelSplitFileProcessor &) = delete; + RParallelSplitFileProcessor(RParallelSplitFileProcessor &&) = delete; + RParallelSplitFileProcessor &operator=(RParallelSplitFileProcessor &&) = delete; +}; +} // namespace ROOT::Internal::IMTUtils + +#endif \ No newline at end of file diff --git a/core/imt/inc/ROOT/TThreadExecutor.hxx b/core/imt/inc/ROOT/TThreadExecutor.hxx index 7ab57cdce6628..76d6940a749b9 100644 --- a/core/imt/inc/ROOT/TThreadExecutor.hxx +++ b/core/imt/inc/ROOT/TThreadExecutor.hxx @@ -36,10 +36,24 @@ #include //std::move #include +namespace ROOT { +class TThreadExecutor; +} +namespace ROOT::Internal::IMTUtils { +class RParallelSplitFileProcessor; +void ParallelFor(TThreadExecutor &pool, std::shared_ptr proc, + const std::function &body); +} // namespace ROOT::Internal::IMTUtils + namespace ROOT { class TThreadExecutor: public TExecutorCRTP { friend TExecutorCRTP; + void ParallelFor(std::shared_ptr fileProcessor, + const std::function &body); + friend void ROOT::Internal::IMTUtils::ParallelFor( + TThreadExecutor &pool, std::shared_ptr proc, + const std::function &body); public: diff --git a/core/imt/src/TThreadExecutor.cxx b/core/imt/src/TThreadExecutor.cxx index 6fd443cf6b66d..32334f2cb2044 100644 --- a/core/imt/src/TThreadExecutor.cxx +++ b/core/imt/src/TThreadExecutor.cxx @@ -9,6 +9,7 @@ #if !defined(_MSC_VER) #pragma GCC diagnostic pop #endif +#include ////////////////////////////////////////////////////////////////////////// /// @@ -214,3 +215,39 @@ unsigned TThreadExecutor::GetPoolSize() const } } // namespace ROOT + +namespace { +struct RSplittableRange { + + // Needs to be a copyable type to comply with TBB's Range requirements + // See + // https://oneapi-spec.uxlfoundation.org/specifications/oneapi/latest/elements/onetbb/source/named_requirements/algorithms/range + std::shared_ptr fFileProcessor; + + RSplittableRange(std::shared_ptr fp) : fFileProcessor(fp) {} + + RSplittableRange(RSplittableRange &r, tbb::split) : fFileProcessor(r.fFileProcessor->SplitWork()) {} + bool is_divisible() const { return fFileProcessor->IsDivisible(); } + bool empty() const { return fFileProcessor->Empty(); } +}; + +} // namespace + +void ROOT::TThreadExecutor::ParallelFor( + std::shared_ptr fileProcessor, + const std::function &body) +{ + fTaskArenaW->Access().execute([&] { + tbb::this_task_arena::isolate([&] { + tbb::parallel_for(RSplittableRange{fileProcessor}, + [body](const RSplittableRange &r) { body(*r.fFileProcessor); }); + }); + }); +} + +void ROOT::Internal::IMTUtils::ParallelFor( + ROOT::TThreadExecutor &pool, std::shared_ptr proc, + const std::function &body) +{ + pool.ParallelFor(proc, body); +} diff --git a/tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx b/tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx index 6b6af72f84331..8e4c6c7ae8898 100644 --- a/tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx +++ b/tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx @@ -336,6 +336,9 @@ public: /// 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 &entryCount); + /// The task run by every thread on the input entry range, for the RNTuple data source. + void RNTupleThreadTask(const std::pair &entryRange, unsigned int slot, + std::uint64_t columnReaderOffset); }; /// \brief Create an RLoopManager that reads a TChain. diff --git a/tree/dataframe/inc/ROOT/RNTupleDS.hxx b/tree/dataframe/inc/ROOT/RNTupleDS.hxx index 8bddba044e01c..36ef6749fb891 100644 --- a/tree/dataframe/inc/ROOT/RNTupleDS.hxx +++ b/tree/dataframe/inc/ROOT/RNTupleDS.hxx @@ -35,6 +35,36 @@ namespace ROOT { class RDataFrame; } + +namespace ROOT::Detail::RDF { +class RLoopManager; +} + +namespace ROOT::Internal { +class RSlotStack; +} + +namespace ROOT { +class RFieldBase; +class RDataFrame; +class RNTuple; +} // namespace ROOT +namespace ROOT::Detail::RDF { +class RNodeBase; +} +namespace ROOT::RDF { +template +class RInterface; +} +namespace ROOT::Internal::RDF { +class RNTupleColumnReader; +std::vector> +GetDatasetGlobalClusterBoundaries(const ROOT::RDF::RInterface &node); +} // namespace ROOT::Internal::RDF +namespace ROOT::Internal { +class RPageSource; +} + namespace ROOT::Internal::RDF { /** * \brief Internal overload of the function that allows passing a range of entries @@ -55,28 +85,21 @@ ROOT::RDataFrame FromRNTuple(std::string_view ntupleName, const std::vector, ROOT::NTupleSize_t> GetClustersAndEntries(std::string_view ntupleName, std::string_view location); + +/** + * \brief Retrieves the cluster boundaries and the number of entries for the input RNTuple + * + * \param[in] pageSource the concrete page source + */ +std::pair, ROOT::NTupleSize_t> +GetClustersAndEntries(const ROOT::Internal::RPageSource &); } // namespace ROOT::Internal::RDF +#ifdef R__USE_IMT namespace ROOT { -class RFieldBase; -class RDataFrame; -class RNTuple; -} // namespace ROOT -namespace ROOT::Detail::RDF { -class RNodeBase; -} -namespace ROOT::RDF { -template -class RInterface; -} -namespace ROOT::Internal::RDF { -class RNTupleColumnReader; -std::vector> -GetDatasetGlobalClusterBoundaries(const ROOT::RDF::RInterface &node); -} -namespace ROOT::Internal { -class RPageSource; +class TThreadExecutor; } +#endif namespace ROOT::RDF { class RNTupleDS final : public ROOT::RDF::RDataSource { @@ -86,10 +109,11 @@ class RNTupleDS final : public ROOT::RDF::RDataSource { /// The GetEntryRanges() swaps fNextRanges and fCurrentRanges and uses the list of /// REntryRangeDS records to return the list of ranges ready to use by the RDF loop manager. struct REntryRangeDS { - std::unique_ptr fSource; + std::shared_ptr fSource; 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; + ULong64_t fEntryOffset = 0; /// Offset of both first and last entries w.r.t. the page source std::string_view fFileName; ///< Storage location of the current RNTuple }; @@ -115,7 +139,7 @@ class RNTupleDS final : public ROOT::RDF::RDataSource { /// and /// c) trigger staging of the next batch of files in the I/O background thread. /// 4. On `Finalize()`, the I/O background thread is stopped. - std::vector> fStagingArea; + std::vector> fStagingArea; std::size_t fNextFileIndex = 0; ///< Index into fFileNames to the next file to process /// We prepare a prototype field for every column. If a column reader is actually requested @@ -143,14 +167,15 @@ class RNTupleDS final : public ROOT::RDF::RDataSource { std::vector fCurrentRanges; ///< Basis for the ranges returned by the last GetEntryRanges() call std::vector fNextRanges; ///< Basis for the ranges populated by the PrepareNextRanges() call + + // During MT runs, the current window of entries seen by an active slot + std::vector fActiveRangesPerSlot; /// Maps the first entries from the ranges of the last GetEntryRanges() call to their corresponding index in /// 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 fFirstEntry2RangeIdx; // Keep track of the scheduled entries - necessary for processing of GlobalEntries std::vector> fOriginalRanges; - /// One element per slot, corresponding to the current range index for that slot, as filled by InitSlot - std::vector fSlotsToRangeIdxs; /// The background thread that runs StageNextSources() std::thread fThreadStaging; @@ -221,7 +246,14 @@ class RNTupleDS final : public ROOT::RDF::RDataSource { explicit RNTupleDS(std::string_view ntupleName, const std::vector &fileNames, const std::pair &range); - +#ifdef R__USE_IMT + void ProcessMTRange(ROOT::TThreadExecutor &pool, std::shared_ptr pageSource, + const std::string &fileName, + const std::vector &clusterBoundaries, + std::uint64_t nEntries, ROOT::Detail::RDF::RLoopManager &lm, + ROOT::Internal::RSlotStack &slotStack, std::atomic &processedEntries, + std::atomic &globalEntries); +#endif public: RNTupleDS(std::string_view ntupleName, std::string_view fileName); RNTupleDS(std::string_view ntupleName, const std::vector &fileNames); @@ -255,7 +287,12 @@ public: // Old API, unused bool SetEntry(unsigned int, ULong64_t) final { return true; } - +#ifdef R__USE_IMT + void ProcessMT(ROOT::Detail::RDF::RLoopManager &lm) final; + void InsertActiveEntryRange(unsigned int slot, const std::string &fileName, + std::shared_ptr pageSource, std::uint64_t beginEntry, + std::uint64_t endEntry, std::uint64_t offset); +#endif protected: Record_t GetColumnReadersImpl(std::string_view name, const std::type_info &) final; }; diff --git a/tree/dataframe/src/RLoopManager.cxx b/tree/dataframe/src/RLoopManager.cxx index 591b3790666e9..8418b82321ece 100644 --- a/tree/dataframe/src/RLoopManager.cxx +++ b/tree/dataframe/src/RLoopManager.cxx @@ -1343,6 +1343,40 @@ void ROOT::Detail::RDF::RLoopManager::DataSourceThreadTask(const std::pair &entryRange, + unsigned int slot, std::uint64_t columnReaderOffset) +{ +#ifdef R__USE_IMT + // These are begin and end entries of the current cluster in the file currently opened by the slot, offset by a + // global atomic counter with the value passed via columnReaderOffset. The only reason we use it for the moment is to + // provide a seed for a unique rdfentry_ sequence in the current slot task. + const auto &[start, end] = entryRange; + + RDSRangeRAII _{*this, slot, columnReaderOffset}; + RCallCleanUpTask cleanup(*this, slot); + + fSampleInfos[slot] = ROOT::Internal::RDF::CreateSampleInfo(*fDataSource, slot, fSampleMap); + + R__LOG_DEBUG(0, RDFLogChannel()) << LogRangeProcessing( + {fDataSource->GetLabel(), start - columnReaderOffset, end - columnReaderOffset, slot}); + + try { + for (auto entry = start; entry < end; ++entry) { + if (fDataSource->SetEntry(slot, entry)) { + RunAndCheckFilters(slot, entry); + } + } + } catch (...) { + std::cerr << "RDataFrame::Run: event loop was interrupted\n"; + throw; + } +#else + (void)entryRange; + (void)slot; + (void)columnReaderOffset; +#endif +} + void ROOT::Detail::RDF::RLoopManager::TTreeThreadTask(TTreeReader &treeReader, ROOT::Internal::RSlotStack &slotStack, std::atomic &entryCount) { diff --git a/tree/dataframe/src/RNTupleDS.cxx b/tree/dataframe/src/RNTupleDS.cxx index 82c72ee81102e..e9052169443f5 100644 --- a/tree/dataframe/src/RNTupleDS.cxx +++ b/tree/dataframe/src/RNTupleDS.cxx @@ -35,6 +35,11 @@ #include #include #include +#ifdef R__USE_IMT +#include +#include +#include +#endif // clang-format off /** @@ -678,69 +683,11 @@ void ROOT::RDF::RNTupleDS::PrepareNextRanges() } return; } - - // Work scheduling of the tail: multiple slots work on the same file. - // Every slot still has its own page source but these page sources may open the same file. - // Again, we need to skip empty files. - unsigned int nSlotsPerFile = fNSlots / nRemainingFiles; - for (std::size_t i = 0; (fNextRanges.size() < fNSlots) && (fNextFileIndex < nFiles); ++i) { - std::unique_ptr 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 - source = CreatePageSource(fNTupleName, fFileNames[fNextFileIndex]); - } - source->Attach(); - fNextFileIndex++; - - auto nEntries = source->GetNEntries(); - if (nEntries == 0) - continue; - - // If last file: use all remaining slots - if (i == (nRemainingFiles - 1)) - nSlotsPerFile = fNSlots - fNextRanges.size(); - - const auto rangesByCluster = [&source]() { - // Take the shared lock of the descriptor just for the time necessary - const auto descGuard = source->GetSharedDescriptorGuard(); - return ROOT::Internal::GetClusterBoundaries(descGuard.GetRef()); - }(); - - const unsigned int nRangesByCluster = rangesByCluster.size(); - - // Distribute slots equidistantly over the entry range, aligned on cluster boundaries - const auto nClustersPerSlot = nRangesByCluster / nSlotsPerFile; - const auto remainder = nRangesByCluster % nSlotsPerFile; - std::size_t iRange = 0; - unsigned int iSlot = 0; - const unsigned int N = std::min(nSlotsPerFile, nRangesByCluster); - for (; iSlot < N; ++iSlot) { - auto start = rangesByCluster[iRange].fFirstEntry; - iRange += nClustersPerSlot + static_cast(iSlot < remainder); - assert(iRange > 0); - auto end = rangesByCluster[iRange - 1].fLastEntryPlusOne; - - 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); - } else { - range.fSource = source->Clone(); - } - range.fSource->SetEntryRange({start, end - start}); - range.fFirstEntry = start; - range.fLastEntry = end; - fNextRanges.emplace_back(std::move(range)); - } - } // loop over tail of remaining files } std::vector> ROOT::RDF::RNTupleDS::GetEntryRanges() { + assert(fNSlots == 1); // MT scheduling doesn't use this function std::vector> ranges; // We need to distinguish between single threaded and multi-threaded runs. @@ -748,11 +695,8 @@ std::vector> ROOT::RDF::RNTupleDS::GetEntryRange // to new page sources of the chain in GetEntryRanges. In multi-threaded mode, on the other hand, // InitSlot is called for every returned range, thus rewiring the column readers takes place in // InitSlot and FinalizeSlot. - - if (fNSlots == 1) { - for (auto r : fActiveColumnReaders[0]) { - r->Disconnect(true /* keepValue */); - } + for (auto r : fActiveColumnReaders[0]) { + r->Disconnect(true /* keepValue */); } // If we have fewer files than slots and we run multiple event loops, we can reuse fCurrentRanges and don't need @@ -873,26 +817,17 @@ std::vector> ROOT::RDF::RNTupleDS::GetEntryRange return ranges; } -void ROOT::RDF::RNTupleDS::InitSlot(unsigned int slot, ULong64_t firstEntry) +void ROOT::RDF::RNTupleDS::InitSlot(unsigned int slot, ULong64_t columnReaderOffset) { 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()] = idxRange; + const auto &entryRangeDS = fActiveRangesPerSlot.at(slot * ROOT::Internal::RDF::CacheLineStep()); + assert(entryRangeDS.fSource != nullptr); for (auto r : fActiveColumnReaders[slot]) { - r->Connect(*fCurrentRanges[idxRange].fSource, - fOriginalRanges[idxRange].first - fCurrentRanges[idxRange].fFirstEntry); + r->Connect(*entryRangeDS.fSource, columnReaderOffset); } } @@ -975,7 +910,7 @@ void ROOT::RDF::RNTupleDS::SetNSlots(unsigned int nSlots) assert(nSlots > 0); fNSlots = nSlots; fActiveColumnReaders.resize(fNSlots); - fSlotsToRangeIdxs.resize(fNSlots * ROOT::Internal::RDF::CacheLineStep()); + fActiveRangesPerSlot.resize(fNSlots * ROOT::Internal::RDF::CacheLineStep()); } ROOT::RDataFrame ROOT::RDF::FromRNTuple(std::string_view ntupleName, std::string_view fileName) @@ -996,27 +931,31 @@ ROOT::RDF::RSampleInfo ROOT::Internal::RDF::RNTupleDS::CreateSampleInfo( // 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()); + const auto &entryRangeDS = fNSlots == 1 + ? fCurrentRanges[0] + : fActiveRangesPerSlot.at(slot * ROOT::Internal::RDF::CacheLineStep()); // Missing source if a file does not exist - if (!fCurrentRanges[rangeIdx].fSource) + if (!entryRangeDS.fSource) return ROOT::RDF::RSampleInfo{}; - const auto &ntupleName = fCurrentRanges[rangeIdx].fSource->GetNTupleName(); - const auto &ntuplePath = fCurrentRanges[rangeIdx].fFileName; + const auto &ntupleName = entryRangeDS.fSource->GetNTupleName(); + const auto &ntuplePath = entryRangeDS.fFileName; const auto ntupleID = std::string(ntuplePath) + '/' + ntupleName; if (sampleMap.empty()) - return ROOT::RDF::RSampleInfo( - ntupleID, std::make_pair(fCurrentRanges[rangeIdx].fFirstEntry, fCurrentRanges[rangeIdx].fLastEntry), nullptr, - fPrincipalDescriptor.GetNEntries()); + return ROOT::RDF::RSampleInfo(ntupleID, + std::make_pair(entryRangeDS.fFirstEntry - entryRangeDS.fEntryOffset, + entryRangeDS.fLastEntry - entryRangeDS.fEntryOffset), + nullptr, fPrincipalDescriptor.GetNEntries()); 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(ntupleID), fPrincipalDescriptor.GetNEntries()); + return ROOT::RDF::RSampleInfo(ntupleID, + std::make_pair(entryRangeDS.fFirstEntry - entryRangeDS.fEntryOffset, + entryRangeDS.fLastEntry - entryRangeDS.fEntryOffset), + sampleMap.at(ntupleID), fPrincipalDescriptor.GetNEntries()); } ROOT::RDataFrame ROOT::Internal::RDF::FromRNTuple(std::string_view ntupleName, @@ -1035,3 +974,244 @@ ROOT::Internal::RDF::GetClustersAndEntries(std::string_view ntupleName, std::str const auto descGuard = source->GetSharedDescriptorGuard(); return std::make_pair(ROOT::Internal::GetClusterBoundaries(descGuard.GetRef()), descGuard->GetNEntries()); } + +std::pair, ROOT::NTupleSize_t> +ROOT::Internal::RDF::GetClustersAndEntries(const ROOT::Internal::RPageSource &source) +{ + // We assume the source to be already attached + const auto descGuard = source.GetSharedDescriptorGuard(); + return std::make_pair(ROOT::Internal::GetClusterBoundaries(descGuard.GetRef()), descGuard->GetNEntries()); +} + +#ifdef R__USE_IMT + +namespace { + +class RNTupleSplitFileProcessor final : public ROOT::Internal::IMTUtils::RParallelSplitFileProcessor { + std::shared_ptr fPageSource; + const std::string &fFileName; + const std::vector &fClusters; + std::size_t fClusterIdxBegin; // inclusive + std::size_t fClusterIdxEnd; // exclusive + +public: + RNTupleSplitFileProcessor(std::shared_ptr pageSource, const std::string &fileName, + const std::vector &clusters, + std::size_t clusterIdxBegin, std::size_t clusterIdxEnd) + : fPageSource(pageSource), + fFileName(fileName), + fClusters(clusters), + fClusterIdxBegin(clusterIdxBegin), + fClusterIdxEnd(clusterIdxEnd) + { + fPageSource->SetEntryRange( + {fClusters[fClusterIdxBegin].fFirstEntry, + fClusters[fClusterIdxEnd - 1].fLastEntryPlusOne - fClusters[fClusterIdxBegin].fFirstEntry}); + } + + ~RNTupleSplitFileProcessor() final = default; + RNTupleSplitFileProcessor(const RNTupleSplitFileProcessor &) = delete; + RNTupleSplitFileProcessor &operator=(const RNTupleSplitFileProcessor &) = delete; + RNTupleSplitFileProcessor(RNTupleSplitFileProcessor &&) = delete; + RNTupleSplitFileProcessor &operator=(RNTupleSplitFileProcessor &&) = delete; + + std::unique_ptr SplitWork() final + { + // Split the cluster range in half and (re)assign as needed + auto clusterIdxMid = fClusterIdxBegin + (fClusterIdxEnd - fClusterIdxBegin) / 2; + + auto newClusterIdxBegin = clusterIdxMid; + auto newClusterIdxEnd = fClusterIdxEnd; + auto newPageSource = fPageSource->Clone(); + newPageSource->Attach(); + newPageSource->SetEntryRange( + {fClusters[newClusterIdxBegin].fFirstEntry, + fClusters[newClusterIdxEnd - 1].fLastEntryPlusOne - fClusters[newClusterIdxBegin].fFirstEntry}); + + fClusterIdxEnd = clusterIdxMid; + fPageSource->SetEntryRange( + {fClusters[fClusterIdxBegin].fFirstEntry, + fClusters[fClusterIdxEnd - 1].fLastEntryPlusOne - fClusters[fClusterIdxBegin].fFirstEntry}); + + return std::make_unique(std::move(newPageSource), fFileName, fClusters, + newClusterIdxBegin, newClusterIdxEnd); + } + + bool IsDivisible() const final + { + return fClusterIdxEnd - fClusterIdxBegin > 2; // The page source prefetches one cluster + } + + bool Empty() const final + { + // empty is called once at the beginning of the parallel_for by TBB, to guard against empty tasks + // We only pass tasks with actual work by construction + return false; + } + + std::shared_ptr GetPageSource() const { return fPageSource; } + + std::uint64_t GetCurrentBeginEntry() const { return fClusters[fClusterIdxBegin].fFirstEntry; } + + std::uint64_t GetCurrentEndEntry() const { return fClusters[fClusterIdxEnd - 1].fLastEntryPlusOne; } +}; +} // namespace + +void ROOT::RDF::RNTupleDS::InsertActiveEntryRange(unsigned int slot, const std::string &fileName, + std::shared_ptr pageSource, + std::uint64_t beginEntry, std::uint64_t endEntry, + std::uint64_t offset) +{ + REntryRangeDS entryRangeDS; + entryRangeDS.fFileName = fileName; + entryRangeDS.fSource = pageSource; + entryRangeDS.fFirstEntry = beginEntry + offset; + entryRangeDS.fLastEntry = endEntry + offset; + entryRangeDS.fEntryOffset = offset; + fActiveRangesPerSlot[slot * ROOT::Internal::RDF::CacheLineStep()] = std::move(entryRangeDS); +} + +void ROOT::RDF::RNTupleDS::ProcessMTRange( + ROOT::TThreadExecutor &pool, std::shared_ptr pageSource, const std::string &fileName, + const std::vector &clusterBoundaries, std::uint64_t nEntries, + ROOT::Detail::RDF::RLoopManager &lm, ROOT::Internal::RSlotStack &slotStack, std::atomic &processedEntries, + std::atomic &globalEntries) +{ + + const auto columnReaderOffset{globalEntries.fetch_add(nEntries)}; + + auto parallelForBody = [this, &lm, &fileName, &slotStack, &processedEntries, + &columnReaderOffset](const ROOT::Internal::IMTUtils::RParallelSplitFileProcessor &fp) { + ROOT::Internal::RSlotStackRAII slotRAII{slotStack}; + const auto &rntupleSplitFileProcessor = dynamic_cast(fp); + const auto &slot = slotRAII.fSlot; + const auto &beginEntry = rntupleSplitFileProcessor.GetCurrentBeginEntry(); + const auto &endEntry = rntupleSplitFileProcessor.GetCurrentEndEntry(); + + processedEntries.fetch_add(endEntry - beginEntry); + + this->InsertActiveEntryRange(slot, fileName, rntupleSplitFileProcessor.GetPageSource(), beginEntry, endEntry, + columnReaderOffset); + + lm.RNTupleThreadTask({beginEntry + columnReaderOffset, endEntry + columnReaderOffset}, slot, columnReaderOffset); + }; + + ROOT::Internal::IMTUtils::ParallelFor( + pool, + std::make_shared(pageSource, fileName, clusterBoundaries, 0, clusterBoundaries.size()), + parallelForBody); +} + +void ROOT::RDF::RNTupleDS::ProcessMT(ROOT::Detail::RDF::RLoopManager &lm) +{ + ROOT::Internal::RSlotStack slotStack(fNSlots); + std::atomic processedEntries(0ull); + std::atomic globalEntries(0ull); + ROOT::TThreadExecutor pool; + + std::vector globalFileOffsets; + + if (fGlobalEntryRange.has_value()) { + // If the user explicitly request a global entry range, we pay the cost of + // computing the global entry offsets upfront + std::vector fileEntries(fFileNames.size()); + auto processFile = [](const std::string &ntupleName, const std::string &fileName) { + // RNTuple + auto source = ROOT::Internal::RPageSource::Create(ntupleName, fileName); + source->Attach(); + const auto descGuard = source->GetSharedDescriptorGuard(); + return descGuard->GetNEntries(); + }; + pool.Foreach([&fileEntries, &processFile, + this](std::size_t idx) { fileEntries[idx] = processFile(fNTupleName, fFileNames[idx]); }, + ROOT::TSeq(fFileNames.size())); + + std::uint64_t offset{}; + globalFileOffsets.reserve(fFileNames.size()); + for (auto entries : fileEntries) { + globalFileOffsets.push_back(offset); + offset += entries; + } + } + + auto processFileWithGlobalOffset = [&](std::size_t fileIdx) { + // This function is called when the user has explicitly requested to process only a global entry range. In this + // scenario we have precomputed the entry offsets of each file upfront (aligned with the global dataset entry + // index) so we can then adjust the range of entries in each file accordingly. + + // Evaluate clusters (with local entry numbers) and number of entries for this file + const auto &fileName = fFileNames[fileIdx]; + auto source = CreatePageSource(fNTupleName, fileName); + source->Attach(); + const auto clustersAndEntries = ROOT::Internal::RDF::GetClustersAndEntries(*source); + const auto &nEntries = clustersAndEntries.second; + if (nEntries == 0) + return; + + const auto &clustersInFile = clustersAndEntries.first; + const auto &globalFileOffset = globalFileOffsets[fileIdx]; + const auto &[globalBegin, globalEnd] = fGlobalEntryRange.value(); + + if ((globalFileOffset + nEntries) < globalBegin || globalFileOffset > globalEnd) + return; + + if (((globalBegin >= globalFileOffset) && (globalBegin < (globalFileOffset + nEntries))) || + ((globalEnd >= globalFileOffset) && (globalEnd < (globalFileOffset + nEntries)))) { + std::vector adjustedClusters; + for (const auto &cluster : clustersInFile) { + // If either the global begin or global end fit within the current file, we adjust the cluster ranges + // accordingly + const std::uint64_t localBegin = globalBegin < globalFileOffset ? 0 : globalBegin - globalFileOffset; + const std::uint64_t localEnd = + globalEnd > globalFileOffset + nEntries ? globalFileOffset + nEntries : globalEnd - globalFileOffset; + const auto currentStart = std::max(cluster.fFirstEntry, localBegin); + const auto currentEnd = std::min(cluster.fLastEntryPlusOne, localEnd); + // This is not satified if the desired start is larger than the last entry of some cluster + // In this case, this cluster is not going to be processes further + if (currentStart < currentEnd) + adjustedClusters.push_back(ROOT::Internal::RNTupleClusterBoundaries{currentStart, currentEnd}); + } + + this->ProcessMTRange(pool, std::move(source), fileName, adjustedClusters, nEntries, lm, slotStack, + processedEntries, globalEntries); + } else { + // Otherwise, we just use the cluster boundaries of the whole file + this->ProcessMTRange(pool, std::move(source), fileName, clustersInFile, nEntries, lm, slotStack, + processedEntries, globalEntries); + } + }; + + // Per-file processing that also retrieves cluster info for a file + auto processFileRetrievingClusters = [&](std::size_t fileIdx) { + // Evaluate clusters (with local entry numbers) and number of entries for this file + const auto &fileName = fFileNames[fileIdx]; + auto source = CreatePageSource(fNTupleName, fFileNames[fileIdx]); + source->Attach(); + const auto clustersAndEntries = ROOT::Internal::RDF::GetClustersAndEntries(*source); + const auto &nEntries = clustersAndEntries.second; + if (nEntries == 0) + return; + const auto &clusters = clustersAndEntries.first; + this->ProcessMTRange(pool, std::move(source), fileName, clusters, nEntries, lm, slotStack, processedEntries, + globalEntries); + }; + + std::vector fileIdxs(fFileNames.size()); + std::iota(fileIdxs.begin(), fileIdxs.end(), 0); + if (fGlobalEntryRange.has_value()) + pool.Foreach(processFileWithGlobalOffset, fileIdxs); + else + pool.Foreach(processFileRetrievingClusters, fileIdxs); + + if (fGlobalEntryRange.has_value()) { + auto &&[begin, end] = fGlobalEntryRange.value(); + auto &&finalProcessedEntries = processedEntries.load(); + if ((end - begin) > finalProcessedEntries) { + Warning("RDataFrame::Run", + "RDataFrame stopped processing after %lld entries, whereas an entry range (begin=%lld,end=%lld) was " + "requested. Consider adjusting the end value of the entry range to a maximum of %lld.", + finalProcessedEntries, begin, end, begin + finalProcessedEntries); + } + } +} +#endif