root/ROOT/RTTreeDS.hxx

0001 /**
0002  \file ROOT/RTTreeDS.hxx
0003  \ingroup dataframe
0004  \author Vincenzo Eduardo Padulano
0005  \date 2024-12
0006 */
0007
0008 /*************************************************************************
0009  * Copyright (C) 1995-2024, Rene Brun and Fons Rademakers.               *
0010  * All rights reserved.                                                  *
0011  *                                                                       *
0012  * For the licensing terms see $ROOTSYS/LICENSE.                         *
0013  * For the list of contributors see $ROOTSYS/README/CREDITS.             *
0014  *************************************************************************/
0015
0016 #ifndef ROOT_INTERNAL_RDF_RTTREEDS
0017 #define ROOT_INTERNAL_RDF_RTTREEDS
0018
0019 #include "ROOT/RDataSource.hxx"
0020
0021 #include <memory>
0022 #include <string>
0023 #include <vector>
0024 #include <stdexcept>
0025 #include <string_view>
0026
0027 // Begin forward decls
0028
0029 namespace ROOT {
0030 class RDataFrame;
0031 }
0032
0033 namespace ROOT::Detail::RDF {
0034 class RLoopManager;
0035 }
0036
0037 namespace ROOT::RDF {
0038 class RSampleInfo;
0039 }
0040
0041 namespace ROOT::RDF::Experimental {
0042 class RSample;
0043 }
0044
0045 namespace ROOT::TreeUtils {
0046 struct RFriendInfo;
0047 }
0048
0049 namespace ROOT::Internal::TreeUtils {
0050 class RNoCleanupNotifier;
0051 }
0052
0053 class TChain;
0054 class TDirectory;
0055 class TTree;
0056 class TTreeReader;
0057
0058 // End forward decls
0059
0060 namespace ROOT::Internal::RDF {
0061
0062 class RTTreeDS final : public ROOT::RDF::RDataSource {
0063    std::vector<std::string> fBranchNamesWithDuplicates{};
0064    std::vector<std::string> fBranchNamesWithoutDuplicates{};
0065    std::vector<std::string> fTopLevelBranchNames{};
0066
0067    std::shared_ptr<TTree> fTree;
0068
0069    std::unique_ptr<TTreeReader> fTreeReader;
0070
0071    std::vector<std::unique_ptr<TChain>> fFriends;
0072
0073    // Should be needed mostly for MT runs, but we keep it here to document and align the existing functionality
0074    // from RLoopManager. See https://github.com/root-project/root/pull/10729
0075    std::unique_ptr<ROOT::Internal::TreeUtils::RNoCleanupNotifier> fNoCleanupNotifier;
0076
0077    ROOT::RDF::RSampleInfo
0078    CreateSampleInfo(unsigned int,
0079                     const std::unordered_map<std::string, ROOT::RDF::Experimental::RSample *> &sampleMap) const final;
0080
0081    void RunFinalChecks(bool nodesLeftNotRun) const final;
0082
0083    void Setup(std::shared_ptr<TTree> &&tree, const ROOT::TreeUtils::RFriendInfo *friendInfo = nullptr);
0084
0085    std::vector<std::pair<ULong64_t, ULong64_t>> GetTTreeEntryRange(TTree &tree);
0086    std::vector<std::pair<ULong64_t, ULong64_t>> GetTChainEntryRange(TChain &chain);
0087
0088 public:
0089    RTTreeDS(std::shared_ptr<TTree> tree);
0090    RTTreeDS(std::shared_ptr<TTree> tree, const ROOT::TreeUtils::RFriendInfo &friendInfo);
0091    RTTreeDS(std::string_view treeName, TDirectory *dirPtr);
0092    RTTreeDS(std::string_view treeName, std::string_view fileNameGlob);
0093    RTTreeDS(std::string_view treeName, const std::vector<std::string> &fileNameGlobs);
0094
0095    // Rule of five
0096    RTTreeDS(const RTTreeDS &) = delete;
0097    RTTreeDS &operator=(const RTTreeDS &) = delete;
0098    RTTreeDS(RTTreeDS &&) = delete;
0099    RTTreeDS &operator=(RTTreeDS &&) = delete;
0100    ~RTTreeDS() final; // Define destructor where data member types are defined
0101
0102    void Initialize() final;
0103
0104    void Finalize() final;
0105
0106    std::vector<std::pair<ULong64_t, ULong64_t>> GetEntryRanges() final;
0107
0108    const std::vector<std::string> &GetColumnNames() const final { return fBranchNamesWithDuplicates; }
0109
0110    bool HasColumn(std::string_view colName) const final
0111    {
0112       return std::find(fBranchNamesWithDuplicates.begin(), fBranchNamesWithDuplicates.end(), colName) !=
0113              fBranchNamesWithDuplicates.end();
0114    }
0115
0116    std::string GetTypeName(std::string_view colName) const final;
0117
0118    std::string GetTypeNameWithOpts(std::string_view colName, bool vector2RVec) const final;
0119
0120    bool SetEntry(unsigned int, ULong64_t entry) final;
0121
0122    Record_t GetColumnReadersImpl(std::string_view /* name */, const std::type_info & /* ti */) final
0123    {
0124       // This datasource uses the newer GetColumnReaders() API
0125       return {};
0126    }
0127
0128    std::unique_ptr<ROOT::Detail::RDF::RColumnReaderBase>
0129    GetColumnReaders(unsigned int, std::string_view, const std::type_info &) final
0130    {
0131       // This data source creates column readers via CreateColumnReader
0132       throw std::runtime_error("GetColumnReaders should not be called on this data source, something wrong happened!");
0133    }
0134
0135    std::unique_ptr<ROOT::Detail::RDF::RColumnReaderBase> CreateColumnReader(unsigned int slot, std::string_view col,
0136                                                                             const std::type_info &tid,
0137                                                                             TTreeReader *treeReader) final;
0138
0139    std::string GetLabel() final { return "TTreeDS"; }
0140
0141    TTree *GetTree();
0142
0143    const std::vector<std::string> &GetTopLevelFieldNames() const final { return fTopLevelBranchNames; }
0144
0145    const std::vector<std::string> &GetColumnNamesNoDuplicates() const final { return fBranchNamesWithoutDuplicates; }
0146
0147    void InitializeWithOpts(const std::set<std::string> &suppressErrorsForMissingBranches) final;
0148
0149    std::string DescribeDataset() final;
0150
0151    std::string AsString() final { return "TTree data source"; }
0152
0153    std::size_t GetNFiles() const final;
0154
0155    void ProcessMT(ROOT::Detail::RDF::RLoopManager &lm) final;
0156 };
0157
0158 ROOT::RDataFrame FromTTree(std::string_view treeName, std::string_view fileNameGlob);
0159 ROOT::RDataFrame FromTTree(std::string_view treeName, const std::vector<std::string> &fileNameGlobs);
0160
0161 } // namespace ROOT::Internal::RDF
0162
0163 #endif