RActionSnapshot.hxx

Go to the documentation of this file.

// Author: Vincenzo Eduardo Padulano CERN 06/2025
 
/*************************************************************************
 * Copyright (C) 1995-2025, Rene Brun and Fons Rademakers.               *
 * All rights reserved.                                                  *
 *                                                                       *
 * For the licensing terms see $ROOTSYS/LICENSE.                         *
 * For the list of contributors see $ROOTSYS/README/CREDITS.             *
 *************************************************************************/
 
#ifndef ROOT_RACTIONSNAPSHOT
#define ROOT_RACTIONSNAPSHOT
 
#include "ROOT/RDF/ColumnReaderUtils.hxx"
#include "ROOT/RDF/GraphNode.hxx"
#include "ROOT/RDF/RActionBase.hxx"
#include "ROOT/RDF/RFilterBase.hxx"
#include "ROOT/RDF/RJittedFilter.hxx"
#include "ROOT/RDF/RLoopManager.hxx"
 
#include <cstddef> // std::size_t
#include <memory>
#include <string>
#include <vector>
 
namespace ROOT::Internal::RDF {
 
namespace GraphDrawing {
std::shared_ptr<GraphNode> AddDefinesToGraph(std::shared_ptr<GraphNode> node, const RColumnRegister &colRegister,
                                             const std::vector<std::string> &prevNodeDefines,
                                             std::unordered_map<void *, std::shared_ptr<GraphNode>> &visitedMap);
} // namespace GraphDrawing
 
class SnapshotHelperWithVariations;
 
template <typename Helper, typename PrevNode>
class R__CLING_PTRCHECK(off) RActionSnapshot final : public RActionBase {
 
   // Template needed to avoid dependency on ActionHelpers.hxx
   Helper fHelper;
 
   // If the PrevNode is a RJittedFilter, our collection of previous nodes will have to use the RFilterBase type:
   // we'll have a RJittedFilter for the nominal case, but the others will be concrete filters.
   using PrevNodeCommon_t = std::conditional_t<std::is_same_v<PrevNode, ROOT::Detail::RDF::RJittedFilter>,
                                               ROOT::Detail::RDF::RFilterBase, PrevNode>;
   /// Previous nodes in the computation graph. First element is nominal, others are varied.
   std::vector<std::shared_ptr<PrevNodeCommon_t>> fPrevNodes;
 
   /// Column readers per slot and per input column
   std::vector<std::vector<RColumnReaderBase *>> fValues;
 
   /// The nth flag signals whether the nth input column is a custom column or not.
   std::vector<bool> fIsDefine;
 
   /// Types of the columns to Snapshot
   std::vector<const std::type_info *> fColTypeIDs;
 
   ROOT::RDF::SampleCallback_t GetSampleCallback() final { return fHelper.GetSampleCallback(); }
 
   void AppendVariedPrevNodes()
   {
      // This method only makes sense if we're appending the varied filters to the list after the nominal
      assert(fPrevNodes.size() == 1);
      const auto &currentVariations = GetVariations();
 
      // If this node hangs from the RLoopManager itself, just use that as the upstream node for each variation
      auto nominalPrevNode = fPrevNodes.front();
      if (static_cast<ROOT::Detail::RDF::RNodeBase *>(nominalPrevNode.get()) == fLoopManager) {
         fPrevNodes.resize(1 + currentVariations.size(), nominalPrevNode);
         return;
      }
 
      // Otherwise, append one varied filter per variation
      const auto &prevVariations = nominalPrevNode->GetVariations();
      fPrevNodes.reserve(1 + currentVariations.size());
 
      for (const auto &variation : currentVariations) {
         if (IsStrInVec(variation, prevVariations)) {
            fPrevNodes.emplace_back(
               std::static_pointer_cast<PrevNodeCommon_t>(nominalPrevNode->GetVariedFilter(variation)));
         } else {
            fPrevNodes.push_back(nominalPrevNode);
         }
      }
   }
 
public:
   RActionSnapshot(Helper &&h, const std::vector<std::string> &columns,
                   const std::vector<const std::type_info *> &colTypeIDs, std::shared_ptr<PrevNode> pd,
                   const RColumnRegister &colRegister)
      : RActionBase(pd->GetLoopManagerUnchecked(), columns, colRegister, pd->GetVariations()),
        fHelper(std::move(h)),
        fPrevNodes{std::static_pointer_cast<PrevNodeCommon_t>(pd)},
        fValues(GetNSlots()),
        fColTypeIDs(colTypeIDs)
   {
      fLoopManager->Register(this);
 
      const auto nColumns = columns.size();
      fIsDefine.reserve(nColumns);
      for (auto i = 0u; i < nColumns; ++i)
         fIsDefine.push_back(colRegister.IsDefineOrAlias(columns[i]));
 
      if constexpr (std::is_same_v<Helper, SnapshotHelperWithVariations>) {
         // Need to populate parts of the computation graph for which we have empty shells, e.g. RJittedFilters and
         // varied Defines
         if (!GetVariations().empty())
            fLoopManager->Jit();
 
         AppendVariedPrevNodes();
 
         for (auto i = 0u; i < nColumns; ++i) {
            if (fIsDefine[i]) {
               auto define = colRegister.GetDefine(columns[i]);
               define->MakeVariations(GetVariations());
            }
         }
      }
   }
 
   RActionSnapshot(const RActionSnapshot &) = delete;
   RActionSnapshot &operator=(const RActionSnapshot &) = delete;
   RActionSnapshot(RActionSnapshot &&) = delete;
   RActionSnapshot &operator=(RActionSnapshot &&) = delete;
 
   ~RActionSnapshot() final { fLoopManager->Deregister(this); }
 
   /**
      Retrieve a wrapper to the result of the action that knows how to merge
      with others of the same type.
   */
   std::unique_ptr<ROOT::Detail::RDF::RMergeableValueBase> GetMergeableValue() const final
   {
      return fHelper.GetMergeableValue();
   }
 
   void Initialize() final { fHelper.Initialize(); }
 
   void InitSlot(TTreeReader *r, unsigned int slot) final
   {
      fValues[slot] = GetUntypedColumnReaders(slot, r, RActionBase::GetColRegister(), *fLoopManager,
                                              RActionBase::GetColumnNames(), fColTypeIDs);
 
      if constexpr (std::is_same_v<Helper, SnapshotHelperWithVariations>) {
         // In case of systematic variations, append also the varied column readers to the values
         // that get passed to the helpers
         auto const &variations = GetVariations();
         for (unsigned int variationIndex = 0; variationIndex < variations.size(); ++variationIndex) {
            auto const &readers =
               GetUntypedColumnReaders(slot, r, RActionBase::GetColRegister(), *fLoopManager,
                                       RActionBase::GetColumnNames(), fColTypeIDs, variations[variationIndex]);
            for (unsigned int i = 0; i < readers.size(); ++i) {
               if (fValues[slot][i] != readers[i]) {
                  // The reader with variations differs from nominal, so this column needs to be added to the output
                  fValues[slot].push_back(readers[i]);
                  // Both the original and the varied column need to be registered for masking
                  fHelper.RegisterVariedColumn(slot, i, i, 0,
                                               "nominal"); // (No harm flagging the nominal multiple times)
                  fHelper.RegisterVariedColumn(slot, fValues[slot].size() - 1, i, variationIndex + 1,
                                               variations[variationIndex]);
               }
            }
         }
      }
 
      fHelper.InitTask(r, slot);
   }
 
   void *GetValue(unsigned int slot, std::size_t readerIdx, Long64_t entry)
   {
      assert(slot < fValues.size());
      assert(readerIdx < fValues[slot].size());
      if (auto *val = fValues[slot][readerIdx]->template TryGet<void>(entry))
         return val;
 
      throw std::out_of_range{"RDataFrame: Action (" + fHelper.GetActionName() +
                              ") could not retrieve value for column '" + fColumnNames[readerIdx] + "' for entry " +
                              std::to_string(entry) +
                              ". You can use the DefaultValueFor operation to provide a default value, or "
                              "FilterAvailable/FilterMissing to discard/keep entries with missing values instead."};
   }
 
   void CallExec(unsigned int slot, Long64_t entry)
   {
      std::vector<void *> untypedValues;
      auto nReaders = fValues[slot].size();
      untypedValues.reserve(nReaders);
      for (decltype(nReaders) readerIdx{}; readerIdx < nReaders; readerIdx++)
         untypedValues.push_back(GetValue(slot, readerIdx, entry));
 
      fHelper.Exec(slot, untypedValues);
   }
 
   void Run(unsigned int slot, Long64_t entry) final
   {
      if constexpr (std::is_same_v<Helper, SnapshotHelperWithVariations>) {
         // check if entry passes all filters
         std::vector<bool> filterPassed(fPrevNodes.size(), false);
         for (unsigned int variation = 0; variation < fPrevNodes.size(); ++variation) {
            filterPassed[variation] = fPrevNodes[variation]->CheckFilters(slot, entry);
         }
 
         // Currently, every event where any of nominal or variations pass gets written to the output.
         // This logic could be extended for different use cases if the need arises.
         if (std::any_of(filterPassed.begin(), filterPassed.end(), [](bool val) { return val; })) {
            // TODO: Don't allocate
            std::vector<void *> untypedValues;
            auto nReaders = fValues[slot].size();
            untypedValues.reserve(nReaders);
            for (decltype(nReaders) readerIdx{}; readerIdx < nReaders; readerIdx++)
               untypedValues.push_back(GetValue(slot, readerIdx, entry));
 
            fHelper.Exec(slot, untypedValues, filterPassed);
         }
      } else {
         if (fPrevNodes.front()->CheckFilters(slot, entry))
            CallExec(slot, entry);
      }
   }
 
   void TriggerChildrenCount() final
   {
      for (auto const &node : fPrevNodes)
         node->IncrChildrenCount();
   }
 
   /// Clean-up operations to be performed at the end of a task.
   void FinalizeSlot(unsigned int slot) final
   {
      fValues[slot].clear();
      fHelper.CallFinalizeTask(slot);
   }
 
   /// Clean-up and finalize the action result (e.g. merging slot-local results).
   /// It invokes the helper's Finalize method.
   void Finalize() final
   {
      fHelper.Finalize();
      SetHasRun();
   }
 
   std::shared_ptr<GraphDrawing::GraphNode>
   GetGraph(std::unordered_map<void *, std::shared_ptr<GraphDrawing::GraphNode>> &visitedMap) final
   {
      // Action nodes do not need to go through CreateFilterNode: they are never common nodes between multiple branches
      const auto nodeType = HasRun() ? GraphDrawing::ENodeType::kUsedAction : GraphDrawing::ENodeType::kAction;
      auto thisNode = std::make_shared<GraphDrawing::GraphNode>(fHelper.GetActionName(), visitedMap.size(), nodeType);
      visitedMap[(void *)this] = thisNode;
 
      for (auto const &node : fPrevNodes) {
         auto prevNode = node->GetGraph(visitedMap);
         const auto &prevColumns = prevNode->GetDefinedColumns();
         auto upmostNode = AddDefinesToGraph(thisNode, GetColRegister(), prevColumns, visitedMap);
 
         thisNode->AddDefinedColumns(GetColRegister().GenerateColumnNames());
         upmostNode->SetPrevNode(prevNode);
      }
      return thisNode;
   }
 
   /// Forwards to the action helpers; will throw since PartialUpdate not supported for most snapshot helpers.
   void *PartialUpdate(unsigned int slot) final { return fHelper.CallPartialUpdate(slot); }
 
   /// Will throw, since varied actions are unsupported. Instead, set a flag in RSnapshotOptions.
   [[maybe_unused]] std::unique_ptr<RActionBase> MakeVariedAction(std::vector<void *> && /*results*/) final
   {
      throw std::logic_error("RDataFrame::Snapshot: The snapshot action cannot be varied. Instead, switch on "
                             "variations in RSnapshotOptions.");
   }
 
   /**
    * \brief Returns a new action with a cloned helper.
    *
    * \param[in] newResult The result to be filled by the new action (needed to clone the helper).
    * \return A unique pointer to the new action.
    */
   std::unique_ptr<RActionBase> CloneAction(void *newResult) final
   {
      return std::make_unique<RActionSnapshot>(fHelper.CallMakeNew(newResult), GetColumnNames(), fColTypeIDs,
                                               std::static_pointer_cast<PrevNode>(fPrevNodes.front()),
                                               GetColRegister());
   }
};
 
} // namespace ROOT::Internal::RDF
 
#endif // ROOT_RACTIONSNAPSHOT