root/ROOT/RVecDS.hxx

0001 // Author: Stefan Wunsch CERN  04/2019
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 #include <ROOT/RDataFrame.hxx>
0012 #include <ROOT/RDataSource.hxx>
0013 #include <ROOT/RVec.hxx>
0014 #include <ROOT/TSeq.hxx>
0015
0016 #include <algorithm>
0017 #include <any>
0018 #include <map>
0019 #include <memory>
0020 #include <string>
0021 #include <tuple>
0022 #include <typeinfo>
0023 #include <utility>
0024 #include <vector>
0025
0026 #ifndef ROOT_RVECDS
0027 #define ROOT_RVECDS
0028
0029 namespace ROOT::Internal::RDF {
0030
0031 class R__CLING_PTRCHECK(off) RVecDSColumnReader final : public ROOT::Detail::RDF::RColumnReaderBase {
0032    TPointerHolder *fPtrHolder;
0033    void *GetImpl(Long64_t) final { return fPtrHolder->GetPointer(); }
0034
0035 public:
0036    RVecDSColumnReader(TPointerHolder *ptrHolder) : fPtrHolder(ptrHolder) {}
0037 };
0038
0039 ////////////////////////////////////////////////////////////////////////////////////////////////
0040 /// \brief A RDataSource implementation which takes a collection of RVecs, which
0041 /// are able to adopt data from Numpy arrays
0042 ///
0043 /// This component allows to create a data source on a set of columns with data
0044 /// coming from RVecs. The adoption of externally provided data, e.g., via Numpy
0045 /// arrays, with RVecs allows to read arbitrary data from memory.
0046 /// In addition, the data source has to keep a reference on the Python owned data
0047 /// so that the lifetime of the data is tied to the datasource.
0048 template <typename... ColumnTypes>
0049 class RVecDS final : public ROOT::RDF::RDataSource {
0050    using PointerHolderPtrs_t = std::vector<ROOT::Internal::RDF::TPointerHolder *>;
0051
0052    std::tuple<ROOT::RVec<ColumnTypes>...> fColumns;
0053    std::vector<std::string> fColNames;
0054    std::unordered_map<std::string, std::string> fColTypesMap;
0055    // The role of the fPointerHoldersModels is to be initialised with the pack
0056    // of arguments in the constrcutor signature at construction time
0057    // Once the number of slots is known, the fPointerHolders are initialised
0058    // according to the models.
0059    PointerHolderPtrs_t fPointerHoldersModels;
0060    std::vector<PointerHolderPtrs_t> fPointerHolders;
0061    std::vector<std::pair<ULong64_t, ULong64_t>> fEntryRanges{};
0062    std::any fLifeline;
0063
0064    Record_t GetColumnReadersImpl(std::string_view, const std::type_info &) { return {}; }
0065
0066    size_t GetEntriesNumber() { return std::get<0>(fColumns).size(); }
0067    template <std::size_t... S>
0068    void SetEntryHelper(unsigned int slot, ULong64_t entry, std::index_sequence<S...>)
0069    {
0070       std::initializer_list<int> expander{
0071          (*static_cast<ColumnTypes *>(fPointerHolders[S][slot]->GetPointer()) = std::get<S>(fColumns)[entry], 0)...};
0072       (void)expander; // avoid unused variable warnings
0073    }
0074
0075    template <std::size_t... S>
0076    void ColLengthChecker(std::index_sequence<S...>)
0077    {
0078       if (sizeof...(S) < 2)
0079          return;
0080
0081       const std::vector<size_t> colLengths{std::get<S>(fColumns).size()...};
0082       const auto expectedLen = colLengths[0];
0083       std::string err;
0084       for (auto i : TSeqI(1, colLengths.size())) {
0085          if (expectedLen != colLengths[i]) {
0086             err += "Column \"" + fColNames[i] + "\" and column \"" + fColNames[0] +
0087                    "\" have different lengths: " + std::to_string(expectedLen) + " and " +
0088                    std::to_string(colLengths[i]);
0089          }
0090       }
0091       if (!err.empty()) {
0092          throw std::runtime_error(err);
0093       }
0094    }
0095
0096 protected:
0097    std::string AsString() { return "Numpy data source"; };
0098
0099 public:
0100    RVecDS(std::any lifeline, std::pair<std::string, ROOT::RVec<ColumnTypes>> const &...colsNameVals)
0101       : fColumns(colsNameVals.second...),
0102         fColNames{colsNameVals.first...},
0103         fColTypesMap({{colsNameVals.first, ROOT::Internal::RDF::TypeID2TypeName(typeid(ColumnTypes))}...}),
0104         fPointerHoldersModels({new ROOT::Internal::RDF::TTypedPointerHolder<ColumnTypes>(new ColumnTypes())...}),
0105         fLifeline{std::move(lifeline)}
0106    {
0107    }
0108
0109    // Rule of five
0110    RVecDS(const RVecDS &) = delete;
0111    RVecDS &operator=(const RVecDS &) = delete;
0112    RVecDS(RVecDS &&) = delete;
0113    RVecDS &operator=(RVecDS &&) = delete;
0114    ~RVecDS() final
0115    {
0116       for (auto &&ptrHolderv : fPointerHolders) {
0117          for (auto &&ptrHolder : ptrHolderv) {
0118             delete ptrHolder;
0119          }
0120       }
0121    }
0122
0123    std::unique_ptr<ROOT::Detail::RDF::RColumnReaderBase>
0124    GetColumnReaders(unsigned int slot, std::string_view colName, const std::type_info &id) final
0125    {
0126       auto colNameStr = std::string(colName);
0127
0128       auto it = fColTypesMap.find(colNameStr);
0129       if (fColTypesMap.end() == it) {
0130          std::string err = "The specified column name, \"" + colNameStr + "\" is not known to the data source.";
0131          throw std::runtime_error(err);
0132       }
0133
0134       const auto &colIdName = it->second;
0135       const auto idName = ROOT::Internal::RDF::TypeID2TypeName(id);
0136       if (colIdName != idName) {
0137          std::string err = "Column " + colNameStr + " has type " + colIdName +
0138                            " while the id specified is associated to type " + idName;
0139          throw std::runtime_error(err);
0140       }
0141
0142       if (auto colNameIt = std::find(fColNames.begin(), fColNames.end(), colNameStr); colNameIt != fColNames.end()) {
0143          const auto index = std::distance(fColNames.begin(), colNameIt);
0144          return std::make_unique<ROOT::Internal::RDF::RVecDSColumnReader>(fPointerHolders[index][slot]);
0145       }
0146
0147       throw std::runtime_error("Could not find column name \"" + colNameStr + "\" in available column names.");
0148    }
0149
0150    const std::vector<std::string> &GetColumnNames() const { return fColNames; }
0151
0152    std::vector<std::pair<ULong64_t, ULong64_t>> GetEntryRanges()
0153    {
0154       auto entryRanges(std::move(fEntryRanges)); // empty fEntryRanges
0155       return entryRanges;
0156    }
0157
0158    std::string GetTypeName(std::string_view colName) const
0159    {
0160       const auto key = std::string(colName);
0161       return fColTypesMap.at(key);
0162    }
0163
0164    bool HasColumn(std::string_view colName) const
0165    {
0166       const auto key = std::string(colName);
0167       const auto endIt = fColTypesMap.end();
0168       return endIt != fColTypesMap.find(key);
0169    }
0170
0171    bool SetEntry(unsigned int slot, ULong64_t entry)
0172    {
0173       SetEntryHelper(slot, entry, std::index_sequence_for<ColumnTypes...>());
0174       return true;
0175    }
0176
0177    void SetNSlots(unsigned int nSlots) final
0178    {
0179       fNSlots = nSlots;
0180       const auto nCols = fColNames.size();
0181       fPointerHolders.resize(nCols); // now we need to fill it with the slots, all of the same type
0182       auto colIndex = 0U;
0183       for (auto &&ptrHolderv : fPointerHolders) {
0184          for (auto slot : ROOT::TSeqI(fNSlots)) {
0185             auto ptrHolder = fPointerHoldersModels[colIndex]->GetDeepCopy();
0186             ptrHolderv.emplace_back(ptrHolder);
0187             (void)slot;
0188          }
0189          colIndex++;
0190       }
0191       for (auto &&ptrHolder : fPointerHoldersModels)
0192          delete ptrHolder;
0193    }
0194
0195    void Initialize()
0196    {
0197       ColLengthChecker(std::index_sequence_for<ColumnTypes...>());
0198       const auto nEntries = GetEntriesNumber();
0199       const auto nEntriesInRange = nEntries / fNSlots; // between integers. Should make smaller?
0200       auto reminder = 1U == fNSlots ? 0 : nEntries % fNSlots;
0201       fEntryRanges.resize(fNSlots);
0202       auto init = 0ULL;
0203       auto end = 0ULL;
0204       for (auto &&range : fEntryRanges) {
0205          end = init + nEntriesInRange;
0206          if (0 != reminder) { // Distribute the reminder among the first chunks
0207             reminder--;
0208             end += 1;
0209          }
0210          range.first = init;
0211          range.second = end;
0212          init = end;
0213       }
0214    }
0215
0216    std::string GetLabel() { return "RVecDS"; }
0217 };
0218
0219 // Factory to create datasource able to read Numpy arrays through RVecs.
0220 // \param pyRVecs Pointer to PyObject holding RVecs.
0221 //                The RVecs itself hold a reference to the associated Numpy arrays so that
0222 //                the data cannot go out of scope as long as the datasource survives.
0223 template <typename... ColumnTypes>
0224 std::unique_ptr<RDataFrame>
0225 MakeRVecDataFrame(std::any lifeline, std::pair<std::string, ROOT::RVec<ColumnTypes>> const &...colNameProxyPairs)
0226 {
0227    return std::make_unique<RDataFrame>(
0228       std::make_unique<RVecDS<ColumnTypes...>>(std::move(lifeline), colNameProxyPairs...));
0229 }
0230
0231 } // namespace ROOT::Internal::RDF
0232
0233 #endif // ROOT_RNUMPYDS