From 20bebc0a0c6706dd0714baeaaf709d9195429438 Mon Sep 17 00:00:00 2001 From: Vincenzo Eduardo Padulano Date: Mon, 19 Oct 2026 00:11:17 +0200 Subject: [PATCH] [df] Rework RNTuple MT scheduling with splittable tasks The MT scheduling of RNTupleDS is changed. The data source adopts now a similar pattern as TTreeDS, implementing a ProcessMT(ROOT::Detail::RDF::RLoopManager &lm) method. This allows the data source to manage the creation and scheduling of multithreaded tasks, inside of which the loop manager is invoked to apply the computation graph to the range of entries in the task. A brief summary of the scheduling approach follows: 1. One task is launched in parallel per input file 2. At the beginning of the file-based task, a page source is created and the list of clusters is computed a. If the user has not specified a global entry range, it is the actual list of clusters of the page source. This happens in ScheduleTasksForFile. b. If the user has specified a global entry range, before launching the file-based tasks the data source had computed a full list of global entry offsets across all input files. The list of clusters of the page source is then adjusted according to the position of the global entry range w.r.t. the global offsets of this file. This happens in ScheduleTasksForFileWithGlobalOffset. 3. The task goes on within ScheduleTaskForPageSource, which invokes TBB's parallel_for with two arguments, one respecting the Range and the other respecting the Body requirements. See the docs respectively at https://oneapi-spec.uxlfoundation.org/specifications/oneapi/latest/elements/onetbb/source/named_requirements/algorithms/range and https://oneapi-spec.uxlfoundation.org/specifications/oneapi/latest/elements/onetbb/source/named_requirements/algorithms/par_for_body. a. The Range argument is of type RPageSourceSplittableRange, a struct representing a range of entries (N consecutive clusters) tied to a specific page source to be processed. b. The first time parallel_for is invoked, an RPageSourceSplittableRange instance is passed to it which includes the whole list of clusters for that page source. c. TBB immediately calls the Range object's SplitWork method, which will split the cluster range in roughly two equal parts, the first part is re-assigned to the same object, the second part is used to create a new RPageSourceSplittableRange. d. TBB continues to split the task recursively, until a list of all non-splittable tasks is available. e. Only then, the TBB parallel_for will schedule the various calls to the ParallelForBody function, each taking one of the cluster ranges generated before. f. Within ParallelForBody, the computation graph is executed by reading data from the page source, restricting it to the assigned cluster range, by calling RLoopManager::RNTupleThreadTask. This approach delegates to the scheduling engine (TBB in this case) the responsibility to optimise the task processing and split tasks that may be running for too long (e.g. some files may contain many more entries than others). The strategy is also documented in a markdown file for posterity. --- core/imt/CMakeLists.txt | 1 + core/imt/inc/ROOT/InternalIMTUtils.hxx | 42 ++ core/imt/inc/ROOT/TThreadExecutor.hxx | 6 + core/imt/src/TThreadExecutor.cxx | 50 ++ tree/dataframe/doc/parallel_scheduling.md | 82 ++++ tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx | 7 + tree/dataframe/inc/ROOT/RNTupleDS.hxx | 112 ++++- tree/dataframe/src/RLoopManager.cxx | 34 ++ tree/dataframe/src/RNTupleDS.cxx | 477 +++++++++++++++---- tree/dataframe/test/datasource_ntuple.cxx | 48 ++ 10 files changed, 739 insertions(+), 120 deletions(-) create mode 100644 core/imt/inc/ROOT/InternalIMTUtils.hxx create mode 100644 tree/dataframe/doc/parallel_scheduling.md 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..ca2feabf88862 --- /dev/null +++ b/core/imt/inc/ROOT/InternalIMTUtils.hxx @@ -0,0 +1,42 @@ +#ifndef ROOT_INTERNAL_IMTUTILS +#define ROOT_INTERNAL_IMTUTILS + +#include +#include + +namespace ROOT { +class TThreadExecutor; +} + +namespace ROOT::Internal::IMTUtils { + +/// \brief Abstract implementation of the oneAPI Range requirements +/// +/// For the specification, see +/// https://oneapi-spec.uxlfoundation.org/specifications/oneapi/latest/elements/onetbb/source/named_requirements/algorithms/range. +/// +/// The class is meant to pass implementations from other parts of ROOT to the +/// internal calls to the TBB parallel_for overload taking a Range object. +class RAbstractTBBRangeImpl { +public: + virtual std::unique_ptr SplitWork() = 0; + virtual bool Empty() const = 0; + virtual bool IsDivisible() const = 0; + RAbstractTBBRangeImpl() = default; + + // Destructor is required for dynamic polymorphism, so apply rule of five + virtual ~RAbstractTBBRangeImpl() = default; + RAbstractTBBRangeImpl(const RAbstractTBBRangeImpl &) = default; + RAbstractTBBRangeImpl &operator=(const RAbstractTBBRangeImpl &) = default; + RAbstractTBBRangeImpl(RAbstractTBBRangeImpl &&) = default; + RAbstractTBBRangeImpl &operator=(RAbstractTBBRangeImpl &&) = default; +}; + +// Entry point to a parallel_for that takes a Range and Body +// https://oneapi-spec.uxlfoundation.org/specifications/oneapi/latest/elements/onetbb/source/algorithms/functions/parallel_for_func +void ParallelFor(ROOT::TThreadExecutor &pool, std::unique_ptr rangeImpl, + const std::function &body); + +} // namespace ROOT::Internal::IMTUtils + +#endif diff --git a/core/imt/inc/ROOT/TThreadExecutor.hxx b/core/imt/inc/ROOT/TThreadExecutor.hxx index dd1494f0c1cf1..11ceb0f506c82 100644 --- a/core/imt/inc/ROOT/TThreadExecutor.hxx +++ b/core/imt/inc/ROOT/TThreadExecutor.hxx @@ -27,6 +27,7 @@ #include "ROOT/TypeTraits.hxx" // InvokeResult #include "RTaskArena.hxx" #include "TError.h" +#include #include //std::function #include @@ -40,6 +41,9 @@ namespace ROOT { class TThreadExecutor: public TExecutorCRTP { friend TExecutorCRTP; + friend void ROOT::Internal::IMTUtils::ParallelFor( + TThreadExecutor &pool, std::unique_ptr rangeImpl, + const std::function &body); public: @@ -124,6 +128,8 @@ namespace ROOT { // Functions that interface with the parallel library used as a backend void ParallelFor(unsigned start, unsigned end, unsigned step, const std::function &f); + void ParallelFor(std::unique_ptr rangeImpl, + const std::function &body); double ParallelReduce(const std::vector &objs, const std::function &redfunc); float ParallelReduce(const std::vector &objs, const std::function &redfunc); template diff --git a/core/imt/src/TThreadExecutor.cxx b/core/imt/src/TThreadExecutor.cxx index 6fd443cf6b66d..0becf7c652b45 100644 --- a/core/imt/src/TThreadExecutor.cxx +++ b/core/imt/src/TThreadExecutor.cxx @@ -214,3 +214,53 @@ unsigned TThreadExecutor::GetPoolSize() const } } // namespace ROOT + +namespace { +/// \brief The Range used in the TThreadExecutor::ParallelFor call below +/// +/// A simple struct in an anonymous namespace to be able to use the tbb::split tag in the basic splitting constructor. +/// https://oneapi-spec.uxlfoundation.org/specifications/oneapi/latest/elements/onetbb/source/named_requirements/algorithms/range +/// +/// The struct implements the Range requirements: +/// - Copy constructor and destructor (implicitly done by using rule of zero) +/// - "Basic splitting constructor", a constructor taking tbb::split tag which calls the SplitWork method of the +/// implementation to create a new task with roughly half the work of the original one. +/// - is_divisible to check if the task can be divided in two. +/// - empty to check if the task is empty. +struct TThreadExecutorTBBRange { + + std::shared_ptr fImpl; + + // Constructor to capture implementation from outside + TThreadExecutorTBBRange(std::unique_ptr impl) + : fImpl(std::move(impl)) + { + } + // Basic splitting constructor + TThreadExecutorTBBRange(TThreadExecutorTBBRange &r, tbb::split) : fImpl(r.fImpl->SplitWork()) {} + // True if range can be partitioned into two subranges + bool is_divisible() const { return fImpl->IsDivisible(); } + // True if range is empty + bool empty() const { return fImpl->Empty(); } +}; + +} // namespace + +void ROOT::TThreadExecutor::ParallelFor( + std::unique_ptr rangeImpl, + const std::function &body) +{ + fTaskArenaW->Access().execute([&] { + tbb::this_task_arena::isolate([&] { + tbb::parallel_for(TThreadExecutorTBBRange{std::move(rangeImpl)}, + [body](const TThreadExecutorTBBRange &r) { body(*r.fImpl); }); + }); + }); +} + +void ROOT::Internal::IMTUtils::ParallelFor( + ROOT::TThreadExecutor &pool, std::unique_ptr rangeImpl, + const std::function &body) +{ + pool.ParallelFor(std::move(rangeImpl), body); +} diff --git a/tree/dataframe/doc/parallel_scheduling.md b/tree/dataframe/doc/parallel_scheduling.md new file mode 100644 index 0000000000000..7b99470b37f4c --- /dev/null +++ b/tree/dataframe/doc/parallel_scheduling.md @@ -0,0 +1,82 @@ +# Parallel scheduling in RDataFrame + +RDataFrame allows to process input data in parallel in two ways: with implicit multi-threading and with distributed multiprocessing. This document explains the strategies for the implementations of the two modes. + +[//]: # (This should be a good way to write comments in markdown according to https://stackoverflow.com/questions/4823468/comments-in-markdown) +[//]: # (TODO: The document only features a description of the multithreading strategy for RNTuple, add the rest.) + +## Multi-threading + +### RNTuple + +[//]: # (TODO: The text is meant to document the current status of the MT strategy, knowing that in the future we may adopt a different, more optimised strategy.) + +RNTuple data reading is handled in `RNTupleDS`, an `RDataSource`-derived class. The reading is triggered by the +computation graph being run, and nodes requesting data from the data source. The computation graph is started and +managed by an `RLoopManager`, which delegates to the `RNTupleDS` for decisions related to scheduling of different MT +tasks, reading data and giving it back to the `RLoopManager` to run the computation graph until done. This is a +description of the high-level steps involved: + +1. When the user triggers the computation graph, `RLoopManager::RunDataSourceMT` prepares the computation graph nodes + and calls `RDataSource::ProcessMT`. This method is overridden by the RNTuple data source. This method, as well as all + other methods methods described from now on, are private to the `RNTupleDS` class. +2. In `RNTupleDS::ProcessMT`: + a. A `TThreadExecutor` is initialised with the number of threads given by `ROOT::GetThreadPoolSize()`. + b. An `RSlotStack` is initialised with the same number of slots. + c. Optionally, if the user has specified a global entry range to restrict processing to, a list of global entry + offsets per file is populated by getting the number of entries in each file upfront. + d. `TThreadExecutor::Foreach` is called, spawning one task per input file of the (primary) dataset. +3. The file-based tasks specialise depending on option 2c, respectively: + a. `RNTupleDS::ScheduleTasksForFile` will schedule a task reading all entries from the file. + b. `RNTupleDS::ScheduleTasksForFileWithGlobalOffset` will schedule a task with an adjusted list of clusters. The + user-provided global entry range is matched against the global offsets computed by 2c, and a new list of + clusters is generated to restrict processing of this page source according to the match. Consequently, some + tasks may be stopped immediately if the page source boundaries do not overlap with the global entry range. +4. Both file-based task implementations described above eventually call `RNTupleDS::ScheduleTaskForPageSource`. This + method redirects to the TBB parallel_for overload that takes two arguments, one respecting the + [Range](https://oneapi-spec.uxlfoundation.org/specifications/oneapi/latest/elements/onetbb/source/named_requirements/algorithms/range) and the other respecting the [Body](https://oneapi-spec.uxlfoundation.org/specifications/oneapi/latest/elements/onetbb/source/named_requirements/algorithms/par_for_body) requirements. Note that the parallel_for call is hidden by + `ROOT::Internal::IMTUtils::ParallelFor`, so that the actual TBB implementation is not leaked. + a. The Range argument is of type `RSplittableRangeOfClusterBoundaries`, a struct representing a range of entries (N + consecutive clusters) tied to a specific page source to be processed. `ScheduleTaskForPageSource` calls + parallel_for with an instance of `RSplittableRangeOfClusterBoundaries` that includes the full list of valid clusters of + the corresponding page source, as defined by step 3. + b. Immediately after, TBB calls `RSplittableRangeOfClusterBoundaries::SplitWork`, which will split the + cluster range in roughly two equal parts, recursively until there are only non-splittable tasks based on the + condition given by `RSplittableRangeOfClusterBoundaries::IsDivisible`. This is the core of the strategy: we cannot know a + priori if one of the input files will be much larger than the others, or if one thread finishes earlier than + others. For this reason, we allow tasks to split themselves, so that threads that finish a file-based task + earlier can pick up one of the smaller cluster-based tasks coming as a result of + `RSplittableRangeOfClusterBoundaries::SplitWork` of another file-based task. Also, the splitting done with this strategy + ensures all the ranges of entries are non-overlapping. `RSplittableRangeOfClusterBoundaries::IsDivisible` currently + dictates that a task can split itself only if it has three or more clusters to process. The reason for requiring + at least three clusters is that the page source will prefetch the next cluster while processing the current one. + While TBB allows to provide a proportional splitting, we decide to go for the balanced basic splitting and have + a task only split into two roughly equal parts, for lack of a motivation to do differently. + c. The Body argument calls `RNTupleDS::ParallelForBody`. This function is responsible to apply the computation graph + to a page source restricted to the given range of clusters of the `RSplittableRangeOfClusterBoundaries` argument by + calling `RLoopManager::RNTupleThreadTask`. + +One cause for concern in MT runs is the risk of having too many files open at the same time, if the data source keeps +spawning tasks that open a page source. The strategy described above also limits the number of concurrently open files +through the mechanism of the RDataFrame processing slot. Here is a description of how that works, using concepts +described above with a different perspective: + +1. In the beginning of the computation graph run, the pool spawns one task per file. The pool has `N` available threads, + which start and pick one of the file-based tasks. +2. The file-based task is a function. At the beginning of this function, a page source is created, so that the valid + list of clusters of step 3 above can be computed. +3. Within `ScheduleTaskForPageSource`, parallel_for is called with the Range and Body arguments. There is no creation + of new page sources or cloning of existing ones yet. +4. Once `ParallelForBody` starts running, it will request a slot from the `ROOT::Internal::RSlotStackRAII` + argument. This effectively locks that function to run and apply the whole computation graph on the given range before + ending. Only as many `ParallelForBody` functions can run as the number of available slots. +5. Only after a slot has been taken, then `ParallelForBody` clones the input page source and restricts the cloned page + source to the assigned cluster range of the `RSplittableRangeOfClusterBoundaries` argument, with `RPageSource::SetEntryRange`. + a. In order to save one clone, one of the `ParallelForBody` will just take the input (shared) page source without + cloning it. At the moment, this is chosen to be the function execution corresponding to a + `RSplittableRangeOfClusterBoundaries` that processes the first cluster of the page source. + +In summary, creation of a page source happens in one of the two methods `ScheduleTasksForFile` or +`ScheduleTasksForFileWithGlobalOffset`. A page source created at the beginning of a file-based task can also be cloned +within a `ParallelForBody`. There are only `N` available threads, each of which could have launched a file-based task, +then end up running a `ParallelForBody`. Thus, there can be at most `2*N` open page sources at the same time. diff --git a/tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx b/tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx index 6b6af72f84331..47e5e747ed76f 100644 --- a/tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx +++ b/tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx @@ -330,12 +330,19 @@ public: return fSuppressErrorsForMissingBranches; } + // TODO(vpadulan): it seems undesirable that we need different "thread tasks" depending on the data source. It would + // be nice to merge everything into one, potentially opening up the possibility to making the MT processing interface + // of the data source public. + /// The task run by every thread on the input entry range, for the generic RDataSource. void DataSourceThreadTask(const std::pair &entryRange, ROOT::Internal::RSlotStack &slotStack, std::atomic &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 &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..b9033fad61708 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 @@ -57,39 +87,43 @@ std::pair, ROOT::NTupleSiz GetClustersAndEntries(std::string_view ntupleName, std::string_view location); } // 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; +class TThreadExecutor; } -namespace ROOT::Internal::RDF { -class RNTupleColumnReader; -std::vector> -GetDatasetGlobalClusterBoundaries(const ROOT::RDF::RInterface &node); -} -namespace ROOT::Internal { -class RPageSource; +namespace ROOT::Internal::IMTUtils { +class RAbstractTBBRangeImpl; } +#endif namespace ROOT::RDF { class RNTupleDS final : public ROOT::RDF::RDataSource { friend class ROOT::Internal::RDF::RNTupleColumnReader; - /// The PrepareNextRanges() method populates the fNextRanges list with REntryRangeDS records. - /// 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. + /// \brief Page source information consumed while running the computation graph + /// + /// REntryRangeDS objects are populated differently in single-thread and multi-thread mode. + /// + /// In single-thread mode, the PrepareNextRanges() method populates the fNextRanges list with REntryRangeDS records. + /// 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. + /// + /// In multi-thread mode, the task being executed by a thread injects an REntryRangeDS into the fActiveRangesPerSlot + /// member, which is then retrievable for relevant operations e.g. connecting the column readers to the page source + /// and populating the sample information. struct REntryRangeDS { - std::unique_ptr fSource; + /// The page source. It can be shared between this object and other multithreaded tasks. + 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; + + // TODO(vpadulan): Having to have the entry offset is ugly, currently it's necessary because the implementation of + // RNTupleColumnReader::GetImpl interprets the 'entry' input argument as a "global" entry w.r.t. some offset. For + // lack of a better communication strategy at the moment, we carry this offset over with REntryRangeDS as it is + // also used by the RLoopManager e.g. to decide the values for rdfentry_ for a given slot + + 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 +149,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 +177,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 + + /// The page source and entry window seen by an active thread task in multi-threaded runs + 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 +256,36 @@ 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 + // Helpers for the MT scheduling + + void InsertActiveEntryRange(unsigned int slot, std::string_view fileName, + std::shared_ptr pageSource, std::uint64_t beginEntry, + std::uint64_t endEntry, std::uint64_t offset); + void ParallelForBody(std::shared_ptr pageSource, + const ROOT::Internal::IMTUtils::RAbstractTBBRangeImpl &rangeImpl, + ROOT::Detail::RDF::RLoopManager &lm, std::string_view fileName, + ROOT::Internal::RSlotStack &slotStack, std::atomic &processedEntries, + ULong64_t columnReaderOffset); + void ScheduleTaskForPageSource(ROOT::TThreadExecutor &pool, std::unique_ptr pageSource, + std::string_view fileName, + std::span clusterBoundaries, + std::uint64_t nEntries, ROOT::Detail::RDF::RLoopManager &lm, + ROOT::Internal::RSlotStack &slotStack, std::atomic &processedEntries, + std::atomic &globalEntries); + std::vector ComputeGlobalEntryOffsets(ROOT::TThreadExecutor &pool) const; + void ScheduleTasksForFileWithGlobalOffset(std::size_t fileIdx, const std::vector &globalFileOffsets, + ROOT::TThreadExecutor &pool, ROOT::Detail::RDF::RLoopManager &lm, + ROOT::Internal::RSlotStack &slotStack, + std::atomic &processedEntries, + std::atomic &globalEntries); + void ScheduleTasksForFile(std::size_t fileIdx, ROOT::TThreadExecutor &pool, ROOT::Detail::RDF::RLoopManager &lm, + ROOT::Internal::RSlotStack &slotStack, std::atomic &processedEntries, + std::atomic &globalEntries); + void ProcessMT(ROOT::Detail::RDF::RLoopManager &lm) final; + +#endif public: RNTupleDS(std::string_view ntupleName, std::string_view fileName); RNTupleDS(std::string_view ntupleName, const std::vector &fileNames); diff --git a/tree/dataframe/src/RLoopManager.cxx b/tree/dataframe/src/RLoopManager.cxx index 591b3790666e9..d750b4dfa8ddf 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 thread who owns 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 [begin, 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(), begin - columnReaderOffset, end - columnReaderOffset, slot}); + + try { + for (auto entry = begin; 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 135dd04f18ac8..1628b5d6b7eb3 100644 --- a/tree/dataframe/src/RNTupleDS.cxx +++ b/tree/dataframe/src/RNTupleDS.cxx @@ -2,6 +2,8 @@ /// \author Jakob Blomer /// \author Enrico Guiraud /// \date 2018-10-04 +/// \author Vincenzo Eduardo Padulano CERN +/// \date 2026-10-19 /************************************************************************* * Copyright (C) 1995-2020, Rene Brun and Fons Rademakers. * @@ -37,6 +39,12 @@ #include #include #include +#ifdef R__USE_IMT +#include +#include +#include +#include +#endif // clang-format off /** @@ -303,7 +311,9 @@ ROOT::RDF::RNTupleDS::RNTupleDS(std::unique_ptr pag { pageSource->Attach(); fPrincipalDescriptor = pageSource->GetSharedDescriptorGuard()->Clone(); - fStagingArea.emplace_back(std::move(pageSource)); + // TODO(vpadulan): file staging is disabled with the current MT strategy, but it would be good to re-enable it. + if (!ROOT::IsImplicitMTEnabled()) + fStagingArea.emplace_back(std::move(pageSource)); AddField(fPrincipalDescriptor, "", fPrincipalDescriptor.GetFieldZeroId(), std::vector()); @@ -537,70 +547,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(); - source->LoadAllPageLists(); - 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 - auto descGuard = source->GetSharedDescriptorGuard(); - return ROOT::Internal::GetClusterBoundaries(descGuard.GetConstRef()); - }(); - - 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(!ROOT::IsImplicitMTEnabled()); // MT scheduling doesn't use this function std::vector> ranges; // We need to distinguish between single threaded and multi-threaded runs. @@ -608,11 +559,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 @@ -724,7 +672,7 @@ std::vector> ROOT::RDF::RNTupleDS::GetEntryRange fSeenEntriesNoGlobalRange += nEntriesPerSource; - if ((fNSlots == 1) && (fCurrentRanges[0].fSource)) { + if (fCurrentRanges[0].fSource) { for (auto r : fActiveColumnReaders[0]) { r->Connect(*fCurrentRanges[0].fSource, fOriginalRanges[0].first); } @@ -733,32 +681,24 @@ 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; + // Only MT scheduling uses this function + if (!ROOT::IsImplicitMTEnabled()) 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); } } void ROOT::RDF::RNTupleDS::FinalizeSlot(unsigned int slot) { - if (fNSlots == 1) + // Only MT scheduling uses this function + if (!ROOT::IsImplicitMTEnabled()) return; for (auto r : fActiveColumnReaders[slot]) { @@ -789,7 +729,9 @@ void ROOT::RDF::RNTupleDS::Initialize() fSeenEntriesNoGlobalRange = 0; fNextFileIndex = 0; fIsReadyForStaging = fHasNextSources = fStagingThreadShouldTerminate = false; - fThreadStaging = std::thread(&RNTupleDS::ExecStaging, this); + // TODO(vpadulan): file staging is disabled with the current MT strategy, but it would be good to re-enable it. + if (!ROOT::IsImplicitMTEnabled()) + fThreadStaging = std::thread(&RNTupleDS::ExecStaging, this); assert(fNextRanges.empty()); if (fCurrentRanges.empty() || (fFileNames.size() > fNSlots)) { @@ -818,7 +760,9 @@ void ROOT::RDF::RNTupleDS::Finalize() fStagingThreadShouldTerminate = true; } fCvStaging.notify_one(); - fThreadStaging.join(); + // TODO(vpadulan): file staging is disabled with the current MT strategy, but it would be good to re-enable it. + if (!ROOT::IsImplicitMTEnabled()) + fThreadStaging.join(); // If we have a chain with more files than the number of slots, the files opened at the end of the // event loop won't be reused when the event loop restarts, so we can close them. if (fFileNames.size() > fNSlots) { @@ -835,7 +779,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) @@ -856,27 +800,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 = (!ROOT::IsImplicitMTEnabled()) + ? 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, @@ -897,3 +845,340 @@ ROOT::Internal::RDF::GetClustersAndEntries(std::string_view ntupleName, std::str auto descGuard = source->GetSharedDescriptorGuard(); return std::make_pair(ROOT::Internal::GetClusterBoundaries(descGuard.GetConstRef()), descGuard->GetNEntries()); } + +#ifdef R__USE_IMT +namespace { + +/// \brief Implementation of the Range type to be used in a parallel_for when processing RNTuple data with RDataFrame. +/// +/// The class represents one range of entries (N consecutive clusters) tied to a specific page source to be processed. +/// The current MT strategy launches an outer parallel_for with exactly one task per page source, then within each of +/// these launches another parallel_for that will take an instance of this class. The initial state of the class is to +/// include the whole list of clusters of the corresponding page source. TBB will eagerly split this range in as many +/// ranges as possible, until there is a full set of non-splittable ranges (based on IsDivisible). The split ranges +/// only have a different begin and end cluster, there is no page source copying or cloning done by this range, instead +/// that happens in the ParallelForBody so that it can be limited by the number of RDF slots used. +class RSplittableRangeOfClusterBoundaries final : public ROOT::Internal::IMTUtils::RAbstractTBBRangeImpl { + // A view of all the cluster boundaries in the page source + std::span fClusters; + std::size_t fClusterIdxBegin; // inclusive + std::size_t fClusterIdxEnd; // exclusive + +public: + RSplittableRangeOfClusterBoundaries(std::span clusters, + std::size_t clusterIdxBegin, std::size_t clusterIdxEnd) + : fClusters(clusters), fClusterIdxBegin(clusterIdxBegin), fClusterIdxEnd(clusterIdxEnd) + { + assert(fClusterIdxEnd > fClusterIdxBegin); + } + + std::unique_ptr SplitWork() final + { + // Split the cluster range in half + auto clusterIdxMid = fClusterIdxBegin + (fClusterIdxEnd - fClusterIdxBegin) / 2; + auto newClusterIdxBegin = clusterIdxMid; + auto newClusterIdxEnd = fClusterIdxEnd; + fClusterIdxEnd = clusterIdxMid; + + return std::make_unique(fClusters, newClusterIdxBegin, newClusterIdxEnd); + } + + bool IsDivisible() const final + { + // We choose two because the page source prefetches one cluster ahead of the current one. + return fClusterIdxEnd - fClusterIdxBegin > 2; + } + + bool Empty() const final + { + // empty is called once at the beginning of the parallel_for, to guard against empty tasks. We only pass tasks + // with actual work by construction + return false; + } + + std::uint64_t GetCurrentBeginEntry() const { return fClusters[fClusterIdxBegin].fFirstEntry; } + + std::uint64_t GetCurrentEndEntry() const { return fClusters[fClusterIdxEnd - 1].fLastEntryPlusOne; } + + bool IsProcessingFirstCluster() const { return fClusterIdxBegin == 0; } +}; +} // namespace + +// For now we need this function as a pass-through for the information contained in a +// RSplittableRangeOfClusterBoundaries to be available to the rest of the computation graph (e.g. to connect the column +// readers and update the sample info) +void ROOT::RDF::RNTupleDS::InsertActiveEntryRange(unsigned int slot, std::string_view fileName, + std::shared_ptr pageSource, + std::uint64_t beginEntry, std::uint64_t endEntry, + std::uint64_t offset) +{ + REntryRangeDS entryRangeDS; + entryRangeDS.fFileName = fileName; + entryRangeDS.fSource = std::move(pageSource); + entryRangeDS.fFirstEntry = beginEntry + offset; + entryRangeDS.fLastEntry = endEntry + offset; + entryRangeDS.fEntryOffset = offset; + fActiveRangesPerSlot.at(slot * ROOT::Internal::RDF::CacheLineStep()) = std::move(entryRangeDS); +} + +/// \brief Apply the computation graph to an input range. +/// \param pageSource The page source to read data from. +/// \param rangeImpl A Range type carrying information on the range of entries to process. +/// \param lm The RLoopManager which will run the computation graph. +/// \param fileName Path to the storage location of the page source. +/// \param slotStack A stack of slots, each thread needs to request one before running the computation graph. +/// \param processedEntries Cumulating number of processed entries by this computation graph run. +/// \param columnReaderOffset Offset value to pass to column readers used in this task. The column reader will request +/// entries from the page source receiving a global entry number, this offset is subtracted so that the column reader +/// requests the local entry from the page source. +void ROOT::RDF::RNTupleDS::ParallelForBody(std::shared_ptr pageSource, + const ROOT::Internal::IMTUtils::RAbstractTBBRangeImpl &rangeImpl, + ROOT::Detail::RDF::RLoopManager &lm, std::string_view fileName, + ROOT::Internal::RSlotStack &slotStack, + std::atomic &processedEntries, ULong64_t columnReaderOffset) +{ + ROOT::Internal::RSlotStackRAII slotRAII{slotStack}; + const auto &slot = slotRAII.fSlot; + + const auto &range = static_cast(rangeImpl); + const auto &beginEntry = range.GetCurrentBeginEntry(); + const auto &endEntry = range.GetCurrentEndEntry(); + + // We clone the page source only within this function scope, to limit the number of open page sources to 2*fNSlots + auto localPageSource = [&]() -> std::shared_ptr { + // In order to avoid an extra Clone, we need to choose one of the ranges that will take the original page source, + // we choose it to be the one processing the first valid cluster + if (range.IsProcessingFirstCluster()) + return pageSource; + + auto newPageSource = pageSource->Clone(); + newPageSource->Attach(); + return newPageSource; + }(); + + // Notify the current page source about the current entry range + localPageSource->SetEntryRange({beginEntry, endEntry - beginEntry}); + + processedEntries.fetch_add(endEntry - beginEntry); + + // At this point we let the RNTupleDS know about the range being processed by this task. This is used for relevant + // operations e.g. connecting the column readers as well as populating sample information. + InsertActiveEntryRange(slot, fileName, std::move(localPageSource), beginEntry, endEntry, columnReaderOffset); + + lm.RNTupleThreadTask({beginEntry + columnReaderOffset, endEntry + columnReaderOffset}, slot, columnReaderOffset); + + // Clear the active range after having processed it, to avoid keeping around the page source for longer than intended + fActiveRangesPerSlot.at(slot * ROOT::Internal::RDF::CacheLineStep()) = {}; +} + +/// \brief Call parallel_for to schedule a splittable task for a page source +/// \param pool The pool of threads used to schedule the task. +/// \param pageSource The RNTuple page source for the task. +/// \param fileName Path to the storage location of the page source. +/// \param clusterBoundaries A list of cluster boundaries valid for the page source. This may be computed differently, +/// see ProcessMT. +/// \param nEntries Number of entries in the page source. +/// \param lm The RLoopManager which will run the computation graph. +/// \param slotStack A stack of slots, each thread needs to request one before running the computation graph. +/// \param processedEntries Cumulating number of processed entries by this computation graph run. +/// \param globalEntries Cumulating total number of entries in the input dataset. +/// +/// This function calls the parallel_for overload that takes two arguments, one respecting the Range and the other +/// respecting the Body requirements. The Range argument is a RSplittableRangeOfClusterBoundaries, a struct with a range +/// of N cluster boundaries to process (initialised to the whole list of clusters). Immediately when calling +/// parallel_for, TBB will recursively ask the RSplittableRangeOfClusterBoundaries to split itself into two roughly +/// equal sub-tasks, until there are only leaf tasks that cannot be split further. This helps balancing cases where some +/// tasks may end up running much longer than others. +void ROOT::RDF::RNTupleDS::ScheduleTaskForPageSource( + ROOT::TThreadExecutor &pool, std::unique_ptr pageSource, std::string_view fileName, + std::span 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)}; + + // We need to share the same page source pointer across all splitted tasks, to decide if we clone it (in most cases) + // or not (one task will take the original page source). + std::shared_ptr sharedPageSource{std::move(pageSource)}; + + auto parallelForBody = [this, &sharedPageSource, &lm, &fileName, &slotStack, &processedEntries, + &columnReaderOffset](const ROOT::Internal::IMTUtils::RAbstractTBBRangeImpl &rangeImpl) { + this->ParallelForBody(sharedPageSource, rangeImpl, lm, fileName, slotStack, processedEntries, columnReaderOffset); + }; + + ROOT::Internal::IMTUtils::ParallelFor( + pool, std::make_unique(clusterBoundaries, 0, clusterBoundaries.size()), + parallelForBody); +} + +/// \brief Computes global entry offsets for the dataset +/// +/// Used in MT scheduling when the user selects a global entry range to restrict processing to. +std::vector ROOT::RDF::RNTupleDS::ComputeGlobalEntryOffsets(ROOT::TThreadExecutor &pool) const +{ + std::vector fileEntries(fFileNames.size()); + auto getEntries = [](std::string_view ntupleName, std::string_view fileName) { + auto source = ROOT::Internal::RPageSource::Create(ntupleName, fileName); + source->Attach(); + const auto descGuard = source->GetSharedDescriptorGuard(); + return descGuard->GetNEntries(); + }; + pool.Foreach([&fileEntries, this, + getEntries](std::size_t idx) { fileEntries[idx] = getEntries(fNTupleName, fFileNames[idx]); }, + ROOT::TSeq(fFileNames.size())); + + // We want the output list of offsets to start from zero: + // {0, n_entries_file_0, n_entries_file_1 + n_entries_file_0, ...} + std::uint64_t offset{}; + std::vector ret; + ret.reserve(fFileNames.size()); + for (auto nEntries : fileEntries) { + ret.push_back(offset); + offset += nEntries; + } + + return ret; +} + +/// \brief Launch the task(s) for the input file index, taking into account dataset global entry offsets. +/// +/// This function is used in MT scheduling to launch a task processing one input file. It is selected when +/// the user specifies a global entry range to restrict processing to. In this scenario we have precomputed the global +/// entry offsets of each file upfront so we can then adjust the range of entries in each file accordingly. The task may +/// be split into smaller tasks based on entry ranges aligned to cluster boundaries, see ScheduleTaskForPageSource. +void ROOT::RDF::RNTupleDS::ScheduleTasksForFileWithGlobalOffset( + std::size_t fileIdx, const std::vector &globalFileOffsets, ROOT::TThreadExecutor &pool, + ROOT::Detail::RDF::RLoopManager &lm, ROOT::Internal::RSlotStack &slotStack, std::atomic &processedEntries, + std::atomic &globalEntries) +{ + const auto &fileName = fFileNames[fileIdx]; + auto source = CreatePageSource(fNTupleName, fileName); + source->Attach(); + // TODO: improve MT scheduling such that not all page lists are required right away + source->LoadAllPageLists(); + + const std::uint64_t globalFileOffset = globalFileOffsets[fileIdx]; + const auto [globalBegin, globalEnd] = fGlobalEntryRange.value(); + + const auto clustersAndEntries = + [&]() -> std::pair, ROOT::NTupleSize_t> { + const auto descGuard = source->GetSharedDescriptorGuard(); + const auto nEntries = descGuard->GetNEntries(); + // Early return: + // - There are no entries in the file + // - The boundaries of this file are not touched by the global entry range + if (nEntries == 0 || (globalFileOffset + nEntries) < globalBegin || globalFileOffset > globalEnd) + return {{}, ROOT::kInvalidNTupleIndex}; + + const auto clustersInFile = ROOT::Internal::GetClusterBoundaries(descGuard.GetConstRef()); + + // The list of cluster boundaries in this file needs to be adjusted if either the beginning or the end of the + // global entry range fall within the boundaries of the entries of this file. + if ((globalBegin >= globalFileOffset) || (globalEnd < (globalFileOffset + nEntries))) { + std::vector adjustedClusters; + for (const auto &cluster : clustersInFile) { + const std::uint64_t localBegin = globalBegin < globalFileOffset ? 0 : globalBegin - globalFileOffset; + const std::uint64_t localEnd = + globalEnd > globalFileOffset + nEntries ? globalFileOffset + nEntries : globalEnd - globalFileOffset; + const auto currentBegin = std::max(cluster.fFirstEntry, localBegin); + const auto currentEnd = std::min(cluster.fLastEntryPlusOne, localEnd); + // This is not satified if the desired begin is larger than the last entry of some cluster + // In this case, this cluster is not going to be processed further + if (currentBegin < currentEnd) + adjustedClusters.push_back(ROOT::Internal::RNTupleClusterBoundaries{currentBegin, currentEnd}); + } + return std::make_pair(adjustedClusters, nEntries); + } + + return std::make_pair(clustersInFile, nEntries); + }(); + const ROOT::NTupleSize_t nEntries = clustersAndEntries.second; + if (nEntries == ROOT::kInvalidNTupleIndex) + return; + const auto &clusters = clustersAndEntries.first; + ScheduleTaskForPageSource(pool, std::move(source), fileName, clusters, nEntries, lm, slotStack, processedEntries, + globalEntries); +} + +/// \brief Launch the task(s) for the input file index +/// +/// This function is used in MT scheduling to launch a task processing one input file. It is selected when +/// there is no global entry range. The task may be split into smaller tasks based on entry ranges aligned to cluster +/// boundaries, see ScheduleTaskForPageSource. +void ROOT::RDF::RNTupleDS::ScheduleTasksForFile(std::size_t fileIdx, ROOT::TThreadExecutor &pool, + ROOT::Detail::RDF::RLoopManager &lm, + ROOT::Internal::RSlotStack &slotStack, + std::atomic &processedEntries, + std::atomic &globalEntries) +{ + // 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(); + // TODO: improve MT scheduling such that not all page lists are required right away + source->LoadAllPageLists(); + const auto clustersAndEntries = [&]() { + const auto descGuard = source->GetSharedDescriptorGuard(); + return std::make_pair(ROOT::Internal::GetClusterBoundaries(descGuard.GetConstRef()), descGuard->GetNEntries()); + }(); + const ROOT::NTupleSize_t nEntries = clustersAndEntries.second; + if (nEntries == 0) + return; + const auto &clusters = clustersAndEntries.first; + ScheduleTaskForPageSource(pool, std::move(source), fileName, clusters, nEntries, lm, slotStack, processedEntries, + globalEntries); +} + +/// \brief Run the computation graph by scheduling MT tasks +/// +/// The data source interacts with the loop manager and with a pool of threads (via TThreadExecutor) at the same time. +/// The function schedules one MT task per input file, eventually calling the parallel_for overload taking two input +/// arguments respectively respecting the Range and Body requirements. This means that each of the file-based tasks +/// can be split into various sub-tasks, each processing a set of consecutive clusters of the same file. Before +/// launching the file-based tasks, the function checks if the user has selected a global range of entries to restrict +/// processing to. If so, the function first computes the global entry offsets of all the input files, so they can be +/// matched against the user-provided global range. These will be used to adjust the list of clusters of the file before +/// spawning the corresponding task. +void ROOT::RDF::RNTupleDS::ProcessMT(ROOT::Detail::RDF::RLoopManager &lm) +{ + ROOT::TThreadExecutor pool; + ROOT::Internal::RSlotStack slotStack(fNSlots); + std::atomic processedEntries(0ull); + std::atomic globalEntries(0ull); + + if (fGlobalEntryRange.has_value()) { + const auto globalFileOffsets = ComputeGlobalEntryOffsets(pool); + pool.Foreach( + [&](std::size_t fileIdx) { + ScheduleTasksForFileWithGlobalOffset(fileIdx, globalFileOffsets, pool, lm, slotStack, processedEntries, + globalEntries); + }, + ROOT::TSeq(fFileNames.size())); + } else { + pool.Foreach( + [&](std::size_t fileIdx) { + ScheduleTasksForFile(fileIdx, pool, lm, slotStack, processedEntries, globalEntries); + }, + ROOT::TSeq(fFileNames.size())); + } + + if (fGlobalEntryRange.has_value()) { + const auto [begin, end] = fGlobalEntryRange.value(); + ULong64_t finalProcessedEntries = processedEntries.load(); + // This check serves two purposes: + // - The immediate purpose is for when the user specifies an ending entry of the range which is actually beyond + // the true last entry in the dataset. This check will inform the user they should restrict their input range. + // This case is also tested in dataframe_datasetspec.cxx. + // - The second purpose is as a sanity check that we indeed processed as many entries as the user wanted. This + // is what we expect, but we take the chance here as a last resort to double check that our expectation + // corresponds to reality. + 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 diff --git a/tree/dataframe/test/datasource_ntuple.cxx b/tree/dataframe/test/datasource_ntuple.cxx index e3b647305744a..7645f7026b85c 100644 --- a/tree/dataframe/test/datasource_ntuple.cxx +++ b/tree/dataframe/test/datasource_ntuple.cxx @@ -337,6 +337,54 @@ TEST_F(RNTupleDSTest, ChainTailScheduling) auto sumX = df.Aggregate([](int &acc, int x) { acc += x; }, [](int a, int b) { return a + b; }, "x"); EXPECT_EQ(56, sumX.GetValue()); } + +TEST(RNTupleDS, MoreFilesThanFileDescriptorLimit) +{ + // Test that MT runs keep a limited number of open page sources at the same time + struct DatasetRAII { + std::string fDatasetName{"events"}; + std::string fFileName{"datasource_ntuple_more_files_than_file_descriptor_limit.root"}; + unsigned fNEntriesPerFile{10}; + + DatasetRAII() + { + auto model = RNTupleModel::Create(); + auto ptrX = model->MakeField("x"); + *ptrX = 1; + auto writer = RNTupleWriter::Recreate(std::move(model), fDatasetName, fFileName); + for (unsigned j = 0; j < fNEntriesPerFile; j++) { + writer->Fill(); + writer->CommitCluster(); + } + } + + ~DatasetRAII() { std::remove(fFileName.c_str()); } + + } dataset; + + auto test_fn = [&](unsigned nFiles) { + std::vector fileNames(nFiles, dataset.fFileName); + ROOT::RDataFrame df{dataset.fDatasetName, fileNames}; + auto sum = df.Sum("x"); + EXPECT_EQ(*sum, nFiles * dataset.fNEntriesPerFile); + }; + + { + // Run with the number of threads decided by the RAII struct + IMTRAII _; + test_fn(10000); + } + + { + // Run with one thread, but with IMT turned on. This caught an issue + // with the implementation that was distinguishing the MT vs non-MT case + // based on the number of active slots (fNSlots == 1) + ROOT::EnableImplicitMT(1); + test_fn(5); + ROOT::DisableImplicitMT(); + } +} + #endif TEST_F(RNTupleDSTest, ModifyColumnValues)