ROOT/RDF/RSampleInfo.hxx

0001 // Author: Enrico Guiraud, 2021
0002
0003 /*************************************************************************
0004  * Copyright (C) 1995-2021, 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_RSAMPLEINFO
0012 #define ROOT_RDF_RSAMPLEINFO
0013
0014 #include <ROOT/RDF/RSample.hxx>
0015 #include <string_view>
0016 #include <Rtypes.h>
0017
0018 #include <functional>
0019 #include <stdexcept>
0020 #include <string>
0021
0022 #include <tuple>
0023
0024 namespace ROOT {
0025 namespace RDF {
0026
0027 /// This type represents a sample identifier, to be used in conjunction with RDataFrame features such as
0028 /// \ref ROOT::RDF::RInterface<Proxied>::DefinePerSample "DefinePerSample()" and per-sample callbacks.
0029 ///
0030 /// When the input data comes from a TTree, the string representation of RSampleInfo (which is returned by AsString()
0031 /// and that can be queried e.g. with Contains()) is of the form "<filename>/<treename>".
0032 ///
0033 /// In multi-thread runs, different tasks might process different entry ranges of the same sample,
0034 /// so RSampleInfo also provides methods to inspect which part of a sample is being taken into consideration.
0035 class RSampleInfo {
0036    std::string fID;
0037    std::pair<ULong64_t, ULong64_t> fEntryRange;
0038
0039    const ROOT::RDF::Experimental::RSample *fSample = nullptr; // non-owning
0040    ULong64_t fTotalEntries = 0;                               // Number of entries in current file (if known).
0041
0042    void ThrowIfNoSample() const
0043    {
0044       if (fSample == nullptr) {
0045          const auto msg = "RSampleInfo: sample data was requested but no samples are available.";
0046          throw std::logic_error(msg);
0047       }
0048    }
0049
0050 public:
0051    RSampleInfo(std::string_view id, std::pair<ULong64_t, ULong64_t> entryRange,
0052                const ROOT::RDF::Experimental::RSample *sample = nullptr, ULong64_t totalEntries = 0)
0053       : fID(id), fEntryRange(entryRange), fSample(sample), fTotalEntries{totalEntries}
0054    {
0055    }
0056    RSampleInfo() = default;
0057    RSampleInfo(const RSampleInfo &) = default;
0058    RSampleInfo &operator=(const RSampleInfo &) = default;
0059    RSampleInfo(RSampleInfo &&) = default;
0060    RSampleInfo &operator=(RSampleInfo &&) = default;
0061    ~RSampleInfo() = default;
0062
0063    /// @brief Get the name of the sample as a string.
0064    const std::string &GetSampleName() const
0065    {
0066       ThrowIfNoSample();
0067       return fSample->GetSampleName();
0068    }
0069
0070    /// @brief Get the sample id as an int.
0071    unsigned int GetSampleId() const
0072    {
0073       ThrowIfNoSample();
0074       return fSample->GetSampleId();
0075    }
0076
0077    /// @brief Return the metadata value of type int given the key.
0078    int GetI(const std::string &key) const
0079    {
0080       ThrowIfNoSample();
0081       return fSample->GetMetaData().GetI(key);
0082    }
0083
0084    /// @brief Return the metadata value of type double given the key.
0085    double GetD(const std::string &key) const
0086    {
0087       ThrowIfNoSample();
0088       return fSample->GetMetaData().GetD(key);
0089    }
0090
0091    /// @brief Return the metadata value of type string given the key.
0092    std::string GetS(const std::string &key) const
0093    {
0094       ThrowIfNoSample();
0095       return fSample->GetMetaData().GetS(key);
0096    }
0097
0098    /// @brief Check whether the sample name contains the given substring.
0099    bool Contains(std::string_view substr) const
0100    {
0101       // C++14 needs the conversion from std::string_view to std::string
0102       return fID.find(std::string(substr)) != std::string::npos;
0103    }
0104
0105    /// @brief Check whether the sample name is empty.
0106    ///
0107    /// This is the case e.g. when using a RDataFrame with no input data, constructed as `RDataFrame(nEntries)`.
0108    bool Empty() const {
0109       return fID.empty();
0110    }
0111
0112    /// @brief Return a string representation of the sample name.
0113    ///
0114    /// The representation is of the form "<filename>/<treename>" if the input data comes from a TTree or a TChain.
0115    const std::string &AsString() const
0116    {
0117       return fID;
0118    }
0119
0120    /// @brief Return the entry range in the sample that is being taken into consideration.
0121    ///
0122    /// Multiple multi-threading tasks might process different entry ranges of the same sample.
0123    std::pair<ULong64_t, ULong64_t> EntryRange() const { return fEntryRange; }
0124
0125    /// @brief Return the number of entries of this range of the sample.
0126    ULong64_t NEntries() const { return fEntryRange.second - fEntryRange.first; }
0127
0128    /// Return the total number of entries in the underlying dataset.
0129    /// If the total number of entries is not known, the end of the current range is returned.
0130    /// This can be larger than NEntries() if the sample is split in multiple ranges.
0131    ULong64_t NEntriesTotal() const { return std::max(fTotalEntries, fEntryRange.second); }
0132
0133    bool operator==(const RSampleInfo &other) const { return fID == other.fID; }
0134    bool operator!=(const RSampleInfo &other) const { return !(*this == other); }
0135 };
0136
0137 /// The type of a data-block callback, registered with an RDataFrame computation graph via e.g.  \ref
0138 /// ROOT::RDF::RInterface<Proxied>::DefinePerSample "DefinePerSample()" or by certain actions (e.g. \ref
0139 /// ROOT::RDF::RInterface<Proxied>::Snapshot "Snapshot()").
0140 using SampleCallback_t = std::function<void(unsigned int, const ROOT::RDF::RSampleInfo &)>;
0141
0142 } // namespace RDF
0143 } // namespace ROOT
0144
0145 #endif