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 class TChain;
0050 class TDirectory;
0051 class TTree;
0052 class TTreeReader;
0053
0054 // End forward decls
0055
0056 namespace ROOT::Internal::RDF {
0057
0058 class RTTreeDS final : public ROOT::RDF::RDataSource {
0059    std::vector<std::string> fBranchNamesWithDuplicates{};
0060    std::vector<std::string> fBranchNamesWithoutDuplicates{};
0061    std::vector<std::string> fTopLevelBranchNames{};
0062
0063    std::shared_ptr<TTree> fTree;
0064
0065    std::unique_ptr<TTreeReader> fTreeReader;
0066
0067    std::vector<std::unique_ptr<TChain>> fFriends;
0068
0069    ROOT::RDF::RSampleInfo
0070    CreateSampleInfo(const std::unordered_map<std::string, ROOT::RDF::Experimental::RSample *> &sampleMap) const final;
0071
0072    void RunFinalChecks(bool nodesLeftNotRun) const final;
0073
0074    void Setup(std::shared_ptr<TTree> &&tree, const ROOT::TreeUtils::RFriendInfo *friendInfo = nullptr);
0075
0076    std::vector<std::pair<ULong64_t, ULong64_t>> GetTTreeEntryRange(TTree &tree);
0077    std::vector<std::pair<ULong64_t, ULong64_t>> GetTChainEntryRange(TChain &chain);
0078
0079 public:
0080    RTTreeDS(std::shared_ptr<TTree> tree);
0081    RTTreeDS(std::shared_ptr<TTree> tree, const ROOT::TreeUtils::RFriendInfo &friendInfo);
0082    RTTreeDS(std::string_view treeName, TDirectory *dirPtr);
0083    RTTreeDS(std::string_view treeName, std::string_view fileNameGlob);
0084    RTTreeDS(std::string_view treeName, const std::vector<std::string> &fileNameGlobs);
0085
0086    // Rule of five
0087    RTTreeDS(const RTTreeDS &) = delete;
0088    RTTreeDS &operator=(const RTTreeDS &) = delete;
0089    RTTreeDS(RTTreeDS &&) = delete;
0090    RTTreeDS &operator=(RTTreeDS &&) = delete;
0091    ~RTTreeDS() final; // Define destructor where data member types are defined
0092
0093    void Initialize() final;
0094
0095    void Finalize() final;
0096
0097    std::vector<std::pair<ULong64_t, ULong64_t>> GetEntryRanges() final;
0098
0099    const std::vector<std::string> &GetColumnNames() const final { return fBranchNamesWithDuplicates; }
0100
0101    bool HasColumn(std::string_view colName) const final
0102    {
0103       return std::find(fBranchNamesWithDuplicates.begin(), fBranchNamesWithDuplicates.end(), colName) !=
0104              fBranchNamesWithDuplicates.end();
0105    }
0106
0107    std::string GetTypeName(std::string_view colName) const final;
0108
0109    std::string GetTypeNameWithOpts(std::string_view colName, bool vector2RVec) const final;
0110
0111    bool SetEntry(unsigned int, ULong64_t entry) final;
0112
0113    Record_t GetColumnReadersImpl(std::string_view /* name */, const std::type_info & /* ti */) final
0114    {
0115       // This datasource uses the newer GetColumnReaders() API
0116       return {};
0117    }
0118
0119    std::unique_ptr<ROOT::Detail::RDF::RColumnReaderBase>
0120    GetColumnReaders(unsigned int, std::string_view, const std::type_info &) final
0121    {
0122       // This data source creates column readers via CreateColumnReader
0123       throw std::runtime_error("GetColumnReaders should not be called on this data source, something wrong happened!");
0124    }
0125
0126    std::unique_ptr<ROOT::Detail::RDF::RColumnReaderBase> CreateColumnReader(unsigned int slot, std::string_view col,
0127                                                                             const std::type_info &tid,
0128                                                                             TTreeReader *treeReader) final;
0129
0130    std::string GetLabel() final { return "TTreeDS"; }
0131
0132    TTree *GetTree();
0133
0134    const std::vector<std::string> &GetTopLevelFieldNames() const final { return fTopLevelBranchNames; }
0135
0136    const std::vector<std::string> &GetColumnNamesNoDuplicates() const final { return fBranchNamesWithoutDuplicates; }
0137
0138    void InitializeWithOpts(const std::set<std::string> &suppressErrorsForMissingBranches) final;
0139
0140    std::string DescribeDataset() final;
0141
0142    std::string AsString() final { return "TTree data source"; }
0143
0144    std::size_t GetNFiles() const final;
0145
0146    void ProcessMT(ROOT::Detail::RDF::RLoopManager &lm) final;
0147 };
0148
0149 ROOT::RDataFrame FromTTree(std::string_view treeName, std::string_view fileNameGlob);
0150 ROOT::RDataFrame FromTTree(std::string_view treeName, const std::vector<std::string> &fileNameGlobs);
0151
0152 } // namespace ROOT::Internal::RDF
0153
0154 #endif