ROOT/RDF/RDefine.hxx

0001 // Author: Enrico Guiraud, Danilo Piparo CERN  09/2018
0002
0003 /*************************************************************************
0004  * Copyright (C) 1995-2018, Rene Brun and Fons Rademakers.               *
0005  * All rights reserved.                                                  *
0006  *                                                                       *
0007  * For the licensing terms see $ROOTSYS/LICENSE.                         *
0008  * For the list of contributors see $ROOTSYS/README/CREDITS.             *
0009  *************************************************************************/
0010
0011 #ifndef ROOT_RDF_RDEFINE
0012 #define ROOT_RDF_RDEFINE
0013
0014 #include "ROOT/RDF/ColumnReaderUtils.hxx"
0015 #include "ROOT/RDF/RColumnReaderBase.hxx"
0016 #include "ROOT/RDF/RDefineBase.hxx"
0017 #include "ROOT/RDF/RLoopManager.hxx"
0018 #include "ROOT/RDF/Utils.hxx"
0019 #include <string_view>
0020 #include "ROOT/TypeTraits.hxx"
0021 #include "RtypesCore.h"
0022
0023 #include <array>
0024 #include <deque>
0025 #include <type_traits>
0026 #include <utility> // std::index_sequence
0027 #include <vector>
0028
0029 class TTreeReader;
0030
0031 namespace ROOT {
0032 namespace Detail {
0033 namespace RDF {
0034
0035 using namespace ROOT::TypeTraits;
0036
0037 // clang-format off
0038 namespace ExtraArgsForDefine {
0039 struct None{};
0040 struct Slot{};
0041 struct SlotAndEntry{};
0042 }
0043 // clang-format on
0044
0045 template <typename F, typename ExtraArgsTag = ExtraArgsForDefine::None>
0046 class R__CLING_PTRCHECK(off) RDefine final : public RDefineBase {
0047    // shortcuts
0048    using NoneTag = ExtraArgsForDefine::None;
0049    using SlotTag = ExtraArgsForDefine::Slot;
0050    using SlotAndEntryTag = ExtraArgsForDefine::SlotAndEntry;
0051    // other types
0052    using FunParamTypes_t = typename CallableTraits<F>::arg_types;
0053    using ColumnTypesTmp_t =
0054       RDFInternal::RemoveFirstParameterIf_t<std::is_same<ExtraArgsTag, SlotTag>::value, FunParamTypes_t>;
0055    using ColumnTypes_t =
0056       RDFInternal::RemoveFirstTwoParametersIf_t<std::is_same<ExtraArgsTag, SlotAndEntryTag>::value, ColumnTypesTmp_t>;
0057    using TypeInd_t = std::make_index_sequence<ColumnTypes_t::list_size>;
0058    using ret_type = typename CallableTraits<F>::ret_type;
0059    // Avoid instantiating vector<bool> as `operator[]` returns temporaries in that case. Use std::deque instead.
0060    using ValuesPerSlot_t =
0061       std::conditional_t<std::is_same<ret_type, bool>::value, std::deque<ret_type>, std::vector<ret_type>>;
0062
0063    F fExpression;
0064    ValuesPerSlot_t fLastResults;
0065
0066    /// Column readers per slot and per input column
0067    std::vector<std::array<RColumnReaderBase *, ColumnTypes_t::list_size>> fValues;
0068
0069    /// Define objects corresponding to systematic variations other than nominal for this defined column.
0070    /// The map key is the full variation name, e.g. "pt:up".
0071    std::unordered_map<std::string, std::unique_ptr<RDefineBase>> fVariedDefines;
0072
0073    template <typename ColType>
0074    auto GetValueChecked(unsigned int slot, std::size_t readerIdx, Long64_t entry) -> ColType &
0075    {
0076       if (auto *val = fValues[slot][readerIdx]->template TryGet<ColType>(entry))
0077          return *val;
0078
0079       throw std::out_of_range{"RDataFrame: Define could not retrieve value for column '" + fColumnNames[readerIdx] +
0080                               "' for entry " + std::to_string(entry) +
0081                               ". You can use the DefaultValueFor operation to provide a default value, or "
0082                               "FilterAvailable/FilterMissing to discard/keep entries with missing values instead."};
0083    }
0084
0085    template <typename... ColTypes, std::size_t... S>
0086    void UpdateHelper(unsigned int slot, Long64_t entry, TypeList<ColTypes...>, std::index_sequence<S...>, NoneTag)
0087    {
0088       fLastResults[slot * RDFInternal::CacheLineStep<ret_type>()] =
0089          fExpression(GetValueChecked<ColTypes>(slot, S, entry)...);
0090       (void)entry; // avoid unused parameter warning (gcc 12.1)
0091    }
0092
0093    template <typename... ColTypes, std::size_t... S>
0094    void UpdateHelper(unsigned int slot, Long64_t entry, TypeList<ColTypes...>, std::index_sequence<S...>, SlotTag)
0095    {
0096       fLastResults[slot * RDFInternal::CacheLineStep<ret_type>()] =
0097          fExpression(slot, GetValueChecked<ColTypes>(slot, S, entry)...);
0098       (void)entry; // avoid unused parameter warning (gcc 12.1)
0099    }
0100
0101    template <typename... ColTypes, std::size_t... S>
0102    void
0103    UpdateHelper(unsigned int slot, Long64_t entry, TypeList<ColTypes...>, std::index_sequence<S...>, SlotAndEntryTag)
0104    {
0105       fLastResults[slot * RDFInternal::CacheLineStep<ret_type>()] =
0106          fExpression(slot, entry, GetValueChecked<ColTypes>(slot, S, entry)...);
0107    }
0108
0109 public:
0110    RDefine(std::string_view name, std::string_view type, F expression, const ROOT::RDF::ColumnNames_t &columns,
0111            const RDFInternal::RColumnRegister &colRegister, RLoopManager &lm,
0112            const std::string &variationName = "nominal")
0113       : RDefineBase(name, type, colRegister, lm, columns, variationName), fExpression(std::move(expression)),
0114         fLastResults(lm.GetNSlots() * RDFInternal::CacheLineStep<ret_type>()), fValues(lm.GetNSlots())
0115    {
0116       fLoopManager->Register(this);
0117    }
0118
0119    RDefine(const RDefine &) = delete;
0120    RDefine &operator=(const RDefine &) = delete;
0121    ~RDefine() { fLoopManager->Deregister(this); }
0122
0123    void InitSlot(TTreeReader *r, unsigned int slot) final
0124    {
0125       RDFInternal::RColumnReadersInfo info{fColumnNames, fColRegister, fIsDefine.data(), *fLoopManager};
0126       fValues[slot] = RDFInternal::GetColumnReaders(slot, r, ColumnTypes_t{}, info, fVariation);
0127       fLastCheckedEntry[slot * RDFInternal::CacheLineStep<Long64_t>()] = -1;
0128    }
0129
0130    /// Return the (type-erased) address of the Define'd value for the given processing slot.
0131    void *GetValuePtr(unsigned int slot) final
0132    {
0133       return static_cast<void *>(&fLastResults[slot * RDFInternal::CacheLineStep<ret_type>()]);
0134    }
0135
0136    /// Update the value at the address returned by GetValuePtr with the content corresponding to the given entry
0137    void Update(unsigned int slot, Long64_t entry) final
0138    {
0139       if (entry != fLastCheckedEntry[slot * RDFInternal::CacheLineStep<Long64_t>()]) {
0140          // evaluate this define expression, cache the result
0141          UpdateHelper(slot, entry, ColumnTypes_t{}, TypeInd_t{}, ExtraArgsTag{});
0142          fLastCheckedEntry[slot * RDFInternal::CacheLineStep<Long64_t>()] = entry;
0143       }
0144    }
0145
0146    void Update(unsigned int /*slot*/, const ROOT::RDF::RSampleInfo &/*id*/) final {}
0147
0148    const std::type_info &GetTypeId() const final { return typeid(ret_type); }
0149
0150    /// Clean-up operations to be performed at the end of a task.
0151    void FinalizeSlot(unsigned int slot) final
0152    {
0153       fValues[slot].fill(nullptr);
0154
0155       for (auto &e : fVariedDefines)
0156          e.second->FinalizeSlot(slot);
0157    }
0158
0159    /// Create clones of this Define that work with values in varied "universes".
0160    void MakeVariations(const std::vector<std::string> &variations) final
0161    {
0162       for (const auto &variation : variations) {
0163          if (std::find(fVariationDeps.begin(), fVariationDeps.end(), variation) == fVariationDeps.end()) {
0164             // this Defined quantity does not depend on this variation, so no need to create a varied RDefine
0165             continue;
0166          }
0167          if (fVariedDefines.find(variation) != fVariedDefines.end())
0168             continue; // we already have this variation stored
0169
0170          // the varied defines get a copy of the callable object.
0171          // TODO document this
0172          auto variedDefine = std::unique_ptr<RDefineBase>(
0173             new RDefine(fName, fType, fExpression, fColumnNames, fColRegister, *fLoopManager, variation));
0174          // TODO switch to fVariedDefines.insert({variationName, std::move(variedDefine)}) when we drop gcc 5
0175          fVariedDefines[variation] = std::move(variedDefine);
0176       }
0177    }
0178
0179    /// Return a clone of this Define that works with values in the variationName "universe".
0180    RDefineBase &GetVariedDefine(const std::string &variationName) final
0181    {
0182       auto it = fVariedDefines.find(variationName);
0183       if (it == fVariedDefines.end()) {
0184          // We don't have a varied RDefine for this variation.
0185          // This means we don't depend on it and we can return ourselves, i.e. the RDefine for the nominal universe.
0186          assert(std::find(fVariationDeps.begin(), fVariationDeps.end(), variationName) == fVariationDeps.end());
0187          return *this;
0188       }
0189
0190       return *(it->second);
0191    }
0192 };
0193
0194 } // ns RDF
0195 } // ns Detail
0196 } // ns ROOT
0197
0198 #endif // ROOT_RDF_RDEFINE