ROOT/RDF/RInterfaceBase.hxx

0001 // Author: Enrico Guiraud CERN 08/2022
0002
0003 /*************************************************************************
0004  * Copyright (C) 1995-2022, 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_RINTERFACEBASE
0012 #define ROOT_RDF_RINTERFACEBASE
0013
0014 #include "ROOT/RVec.hxx"
0015 #include <ROOT/RDF/InterfaceUtils.hxx>
0016 #include <ROOT/RDF/RColumnRegister.hxx>
0017 #include <ROOT/RDF/RDisplay.hxx>
0018 #include <ROOT/RDF/RLoopManager.hxx>
0019 #include <ROOT/RDataSource.hxx>
0020 #include <ROOT/RResultPtr.hxx>
0021 #include <string_view>
0022 #include <TError.h> // R__ASSERT
0023
0024 #include <memory>
0025 #include <set>
0026 #include <string>
0027 #include <vector>
0028
0029 namespace ROOT {
0030 namespace RDF {
0031
0032 class RDFDescription;
0033 class RVariationsDescription;
0034
0035 using ColumnNames_t = std::vector<std::string>;
0036
0037 namespace RDFDetail = ROOT::Detail::RDF;
0038 namespace RDFInternal = ROOT::Internal::RDF;
0039
0040 // clang-format off
0041 /**
0042  * \class ROOT::RDF::RInterfaceBase
0043  * \ingroup dataframe
0044  * \brief The base public interface to the RDataFrame federation of classes.
0045  *
0046  * This class contains common methods for all RInterface instantiations.
0047  */
0048 // clang-format on
0049 class RInterfaceBase {
0050 protected:
0051    ///< The RLoopManager at the root of this computation graph. Never null.
0052    std::shared_ptr<ROOT::Detail::RDF::RLoopManager> fLoopManager;
0053
0054    /// Contains the columns defined up to this node.
0055    RDFInternal::RColumnRegister fColRegister;
0056
0057    std::string DescribeDataset() const;
0058
0059    ColumnNames_t GetColumnTypeNamesList(const ColumnNames_t &columnList);
0060
0061    void CheckIMTDisabled(std::string_view callerName);
0062
0063    void AddDefaultColumns();
0064
0065    template <typename RetType>
0066    void SanityChecksForVary(const std::vector<std::string> &colNames, const std::vector<std::string> &variationTags,
0067                             std::string_view variationName)
0068    {
0069       R__ASSERT(!variationTags.empty() && "Must have at least one variation.");
0070       R__ASSERT(!colNames.empty() && "Must have at least one varied column.");
0071       R__ASSERT(!variationName.empty() && "Must provide a variation name.");
0072
0073       for (auto &colName : colNames) {
0074          RDFInternal::CheckForDefinition("Vary", colName, fColRegister,
0075                                          GetDataSource() ? GetDataSource()->GetColumnNames() : ColumnNames_t{});
0076       }
0077       RDFInternal::CheckValidCppVarName(variationName, "Vary");
0078
0079       static_assert(ROOT::Internal::VecOps::IsRVec<RetType>::value, "Vary expressions must return an RVec.");
0080
0081       if (colNames.size() > 1) { // we are varying multiple columns simultaneously, RetType is RVec<RVec<T>>
0082          constexpr bool hasInnerRVec = ROOT::Internal::VecOps::IsRVec<typename RetType::value_type>::value;
0083          if (!hasInnerRVec)
0084             throw std::runtime_error("This Vary call is varying multiple columns simultaneously but the expression "
0085                                      "does not return an RVec of RVecs.");
0086
0087          // Check for type mismatches. We are interested in two cases:
0088          // - All columns that are going to be varied must be of the same type
0089          // - The return type of the expression must match the type of the nominal column
0090          auto colTypes = GetColumnTypeNamesList(colNames);
0091          auto &&nColTypes = colTypes.size();
0092          // Cache type_info when requested
0093          std::vector<const std::type_info *> colTypeIDs(nColTypes);
0094          const auto &innerTypeID = typeid(RDFInternal::InnerValueType_t<RetType>);
0095          for (decltype(nColTypes) i{}; i < nColTypes; ++i) {
0096             // Need to retrieve the type_info for each column. We start with
0097             // checking if the column comes from a Define, in which case the
0098             // type_info is cached already. Otherwise, we need to retrieve it
0099             // via TypeName2TypeID, which eventually might call the interpreter.
0100             const auto *define = fColRegister.GetDefine(colNames[i]);
0101             colTypeIDs[i] = define ? &define->GetTypeId() : &RDFInternal::TypeName2TypeID(colTypes[i]);
0102             // First check: whether the current column type is the same as the first one.
0103             if (*colTypeIDs[i] != *colTypeIDs[0]) {
0104                throw std::runtime_error("Cannot simultaneously vary multiple columns of different types.");
0105             }
0106             // Second check: mismatch between varied type and nominal type
0107             if (innerTypeID != *colTypeIDs[i])
0108                throw std::runtime_error("Varied values for column \"" + colNames[i] + "\" have a different type (" +
0109                                         RDFInternal::TypeID2TypeName(innerTypeID) + ") than the nominal value (" +
0110                                         colTypes[i] + ").");
0111          }
0112
0113       } else { // we are varying a single column, RetType is RVec<T>
0114          const auto &retTypeID = typeid(typename RetType::value_type);
0115          const auto &colName = colNames[0]; // we have only one element in there
0116          const auto *define = fColRegister.GetDefine(colName);
0117          const auto *expectedTypeID =
0118             define ? &define->GetTypeId() : &RDFInternal::TypeName2TypeID(GetColumnType(colName));
0119          if (retTypeID != *expectedTypeID)
0120             throw std::runtime_error("Varied values for column \"" + colName + "\" have a different type (" +
0121                                      RDFInternal::TypeID2TypeName(retTypeID) + ") than the nominal value (" +
0122                                      GetColumnType(colName) + ").");
0123       }
0124
0125       // when varying multiple columns, they must be different columns
0126       if (colNames.size() > 1) {
0127          std::set<std::string> uniqueCols(colNames.begin(), colNames.end());
0128          if (uniqueCols.size() != colNames.size())
0129             throw std::logic_error("A column name was passed to the same Vary invocation multiple times.");
0130       }
0131    }
0132
0133    RDFDetail::RLoopManager *GetLoopManager() const { return fLoopManager.get(); }
0134    RDataSource *GetDataSource() const { return fLoopManager->GetDataSource(); }
0135
0136    ColumnNames_t GetValidatedColumnNames(const unsigned int nColumns, const ColumnNames_t &columns)
0137    {
0138       return RDFInternal::GetValidatedColumnNames(*fLoopManager, nColumns, columns, fColRegister, GetDataSource());
0139    }
0140
0141    template <typename... ColumnTypes>
0142    void CheckAndFillDSColumns(ColumnNames_t validCols, TTraits::TypeList<ColumnTypes...> typeList)
0143    {
0144       if (auto dataSource = GetDataSource())
0145          RDFInternal::AddDSColumns(validCols, *fLoopManager, *dataSource, typeList, fColRegister);
0146    }
0147
0148    void CheckAndFillDSColumns(const std::vector<std::string> &colNames,
0149                               const std::vector<const std::type_info *> &colTypeIDs)
0150    {
0151       if (auto dataSource = GetDataSource())
0152          RDFInternal::AddDSColumns(colNames, *fLoopManager, *dataSource, colTypeIDs, fColRegister);
0153    }
0154
0155    /// Create RAction object, return RResultPtr for the action
0156    /// Overload for the case in which all column types were specified (no jitting).
0157    /// For most actions, `r` and `helperArg` will refer to the same object, because the only argument to forward to
0158    /// the action helper is the result value itself. We need the distinction for actions such as Snapshot or Cache,
0159    /// for which the constructor arguments of the action helper are different from the returned value.
0160    template <typename ActionTag, typename... ColTypes, typename ActionResultType, typename RDFNode,
0161              typename HelperArgType = ActionResultType,
0162              std::enable_if_t<!RDFInternal::RNeedJitting<ColTypes...>::value, int> = 0>
0163    RResultPtr<ActionResultType> CreateAction(const ColumnNames_t &columns, const std::shared_ptr<ActionResultType> &r,
0164                                              const std::shared_ptr<HelperArgType> &helperArg,
0165                                              const std::shared_ptr<RDFNode> &proxiedPtr, const int /*nColumns*/ = -1)
0166    {
0167       constexpr auto nColumns = sizeof...(ColTypes);
0168
0169       const auto validColumnNames = GetValidatedColumnNames(nColumns, columns);
0170       CheckAndFillDSColumns(validColumnNames, RDFInternal::TypeList<ColTypes...>());
0171
0172       const auto nSlots = fLoopManager->GetNSlots();
0173
0174       auto action = RDFInternal::BuildAction<ColTypes...>(validColumnNames, helperArg, nSlots, proxiedPtr, ActionTag{},
0175                                                           fColRegister);
0176       return MakeResultPtr(r, *fLoopManager, std::move(action));
0177    }
0178
0179    /// Create RAction object, return RResultPtr for the action
0180    /// Overload for the case in which one or more column types were not specified (RTTI + jitting).
0181    /// This overload has a `nColumns` optional argument. If present, the number of required columns for
0182    /// this action is taken equal to nColumns, otherwise it is assumed to be sizeof...(ColTypes).
0183    template <typename ActionTag, typename... ColTypes, typename ActionResultType, typename RDFNode,
0184              typename HelperArgType = ActionResultType,
0185              std::enable_if_t<RDFInternal::RNeedJitting<ColTypes...>::value, int> = 0>
0186    RResultPtr<ActionResultType>
0187    CreateAction(const ColumnNames_t &columns, const std::shared_ptr<ActionResultType> &r,
0188                 const std::shared_ptr<HelperArgType> &helperArg, const std::shared_ptr<RDFNode> &proxiedPtr,
0189                 const int nColumns = -1, const bool vector2RVec = true)
0190    {
0191       auto realNColumns = (nColumns > -1 ? nColumns : sizeof...(ColTypes));
0192
0193       const auto validColumnNames = GetValidatedColumnNames(realNColumns, columns);
0194       const unsigned int nSlots = fLoopManager->GetNSlots();
0195
0196       const auto jittedAction = std::make_shared<RDFInternal::RJittedAction>(
0197          *fLoopManager, validColumnNames, fColRegister, proxiedPtr->GetVariations(), proxiedPtr);
0198
0199       auto funcBody = RDFInternal::JitBuildAction(validColumnNames, typeid(HelperArgType), typeid(ActionTag), nullptr,
0200                                                   nSlots, fColRegister, GetDataSource(), vector2RVec);
0201       fLoopManager->RegisterJitHelperCall(funcBody,
0202                                           std::make_unique<ROOT::Internal::RDF::RColumnRegister>(fColRegister),
0203                                           validColumnNames, jittedAction, helperArg);
0204       return MakeResultPtr(r, *fLoopManager, std::move(jittedAction));
0205    }
0206
0207 public:
0208    RInterfaceBase(std::shared_ptr<RDFDetail::RLoopManager> lm);
0209    RInterfaceBase(RDFDetail::RLoopManager &lm, const RDFInternal::RColumnRegister &colRegister);
0210
0211    ColumnNames_t GetColumnNames();
0212    ColumnNames_t GetDatasetTopLevelFieldNames();
0213
0214    std::string GetColumnType(std::string_view column);
0215
0216    RDFDescription Describe();
0217
0218    RVariationsDescription GetVariations() const;
0219    bool HasColumn(std::string_view columnName);
0220    ColumnNames_t GetDefinedColumnNames();
0221    unsigned int GetNSlots() const;
0222    unsigned int GetNRuns() const;
0223    unsigned int GetNFiles();
0224 };
0225 } // namespace RDF
0226 } // namespace ROOT
0227
0228 #endif