fix: fix progress bar for RDataFrame with Range limits

tree/dataframe/inc/ROOT/RDFHelpers.hxx

@@ -17,6 +17,7 @@
 #include <ROOT/RDF/RActionBase.hxx>
 #include <ROOT/RDF/RResultMap.hxx>
 #include <ROOT/RResultHandle.hxx> // users of RunGraphs might rely on this transitive include
+#include <ROOT/RResultPtr.hxx>
 #include <ROOT/TypeTraits.hxx>
 
 #include <array>
@@ -276,8 +277,9 @@ SnapshotPtr_t VariationsFor(SnapshotPtr_t resPtr);
 /// \brief Add ProgressBar to a ROOT::RDF::RNode
 /// \param[in] df RDataFrame node at which ProgressBar is called.
 ///
-/// The ProgressBar can be added not only at the RDataFrame head node, but also at any any computational node,
-/// such as Filter or Define.
+/// The ProgressBar can be added not only at the RDataFrame head node, but also at any computational node,
+/// such as Filter or Define. To correctly account for the entries processed, place the progress bar
+/// after transformations that reduce the number of events (e.g. `Range`).
 /// ###Example usage:
 /// ~~~{.cpp}
 /// ROOT::RDataFrame df("tree", "file.root");
@@ -360,6 +362,8 @@ private:
    bool fUseShellColours;
 
    std::shared_ptr<TTree> fTree{nullptr};
+   ROOT::RDF::RResultPtr<ULong64_t> fTotalEntries{};
+
 
 public:
    /// Create a progress helper.
@@ -370,7 +374,8 @@ public:
    /// \param useColors Use shell colour codes to colour the output. Automatically disabled when
    /// we are not writing to a tty.
    ProgressHelper(std::size_t increment, unsigned int totalFiles = 1, unsigned int progressBarWidth = 40,
-                  unsigned int printInterval = 1, bool useColors = true);
+                  unsigned int printInterval = 1, bool useColors = true,
+                  ROOT::RDF::RResultPtr<ULong64_t> totalEntries = {});
 
    ~ProgressHelper() = default;
 

tree/dataframe/src/RDFHelpers.cxx

@@ -15,6 +15,7 @@
 #include "ROOT/RDF/RLoopManager.hxx" // for RLoopManager
 #include "ROOT/RDF/Utils.hxx"
 #include "ROOT/RResultHandle.hxx" // for RResultHandle, RunGraphs
+#include "ROOT/RResultPtr.hxx"
 
 #include "TROOT.h"      // IsImplicitMTEnabled
 #include "TError.h"     // Warning
@@ -157,7 +158,8 @@ void ThreadsPerTH3(unsigned int N)
 }
 
 ProgressHelper::ProgressHelper(std::size_t increment, unsigned int totalFiles, unsigned int progressBarWidth,
-                               unsigned int printInterval, bool useColors)
+                              unsigned int printInterval, bool useColors,
+                              ROOT::RDF::RResultPtr<ULong64_t> totalEntries)
    : fPrintInterval(printInterval),
      fIncrement{increment},
      fBarWidth{progressBarWidth = int(get_tty_size() / 4)},
@@ -169,6 +171,7 @@ ProgressHelper::ProgressHelper(std::size_t increment, unsigned int totalFiles, u
      fIsTTY{isatty(fileno(stdout)) == 1},
      fUseShellColours{useColors && fIsTTY} // Control characters only with terminals.
 #endif
+      , fTotalEntries(std::move(totalEntries))
 {
 }
 
@@ -241,7 +244,7 @@ void ProgressHelper::PrintStats(std::ostream &stream, std::size_t currentEventCo
 {
    RestoreStreamState restore(stream);
    auto evtpersec = EvtPerSec();
-   auto GetNEventsOfCurrentFile = ComputeNEventsSoFar();
+   auto GetNEventsOfCurrentFile = fTotalEntries && fTotalEntries.IsReady() ? const_cast<ROOT::RDF::RResultPtr<ULong64_t>&>(fTotalEntries).GetValue() : ComputeNEventsSoFar();
    auto currentFileIdx = ComputeCurrentFileIdx();
    auto totalFiles = fTotalFiles;
 
@@ -287,7 +290,7 @@ void ProgressHelper::PrintStats(std::ostream &stream, std::size_t currentEventCo
 void ProgressHelper::PrintStatsFinal(std::ostream &stream, std::chrono::seconds elapsedSeconds) const
 {
    RestoreStreamState restore(stream);
-   auto totalEvents = ComputeNEventsSoFar();
+   auto totalEvents = fTotalEntries && fTotalEntries.IsReady() ? const_cast<ROOT::RDF::RResultPtr<ULong64_t>&>(fTotalEntries).GetValue() : ComputeNEventsSoFar();
    auto totalFiles = fTotalFiles;
 
    if (fUseShellColours)
@@ -316,7 +319,7 @@ void ProgressHelper::PrintStatsFinal(std::ostream &stream, std::chrono::seconds
 /// Print a progress bar of width `ProgressHelper::fBarWidth` if `fGetNEventsOfCurrentFile` is known.
 void ProgressHelper::PrintProgressBar(std::ostream &stream, std::size_t currentEventCount) const
 {
-   auto GetNEventsOfCurrentFile = ComputeNEventsSoFar();
+   auto GetNEventsOfCurrentFile = fTotalEntries && fTotalEntries.IsReady() ? const_cast<ROOT::RDF::RResultPtr<ULong64_t>&>(fTotalEntries).GetValue() : ComputeNEventsSoFar();
    if (GetNEventsOfCurrentFile == 0)
       return;
 
@@ -387,7 +390,8 @@ class ProgressBarAction final : public ROOT::Detail::RDF::RActionImpl<ProgressBa
 void AddProgressBar(ROOT::RDF::RNode node)
 {
    auto total_files = node.GetNFiles();
-   auto progress = std::make_shared<ProgressHelper>(1000, total_files);
+   auto totalEntries = node.Count();
+   auto progress = std::make_shared<ProgressHelper>(1000, total_files, 40, 1, true, totalEntries);
    ProgressBarAction c(progress);
    auto r = node.Book<>(c);
    r.OnPartialResultSlot(1000, [progress](unsigned int slot, auto &&arg) { (*progress)(slot, arg); });

tutorials/analysis/dataframe/df108_ProgressRange.C

@@ -0,0 +1,15 @@
+/// Minimal example showing how a progress bar interacts with `Range`.
+///
+/// The progress bar must be attached to the node *after* the range is applied
+/// so the total number of entries displayed corresponds to the ranged dataset
+/// and not the original input.
+///
+/// \macro_code
+void df108_ProgressRange()
+{
+   ROOT::RDataFrame df(100);
+   auto ranged = df.Range(0, 10);
+   ROOT::RDF::Experimental::AddProgressBar(ranged);
+   auto h = ranged.Define("x", []() { return gRandom->Rndm(); }).Histo1D("x");
+   std::cout << h->GetEntries() << std::endl;
+}