ROOT/RDF/SnapshotHelpers.hxx

0001 /**
0002  \file ROOT/RDF/SnapshotHelpers.hxx
0003  \ingroup dataframe
0004  \author Enrico Guiraud, CERN
0005  \author Danilo Piparo, CERN
0006  \date 2016-12
0007  \author Vincenzo Eduardo Padulano
0008  \author Stephan Hageboeck
0009  \date 2025-06
0010 */
0011
0012 /*************************************************************************
0013  * Copyright (C) 1995-2025, Rene Brun and Fons Rademakers.               *
0014  * All rights reserved.                                                  *
0015  *                                                                       *
0016  * For the licensing terms see $ROOTSYS/LICENSE.                         *
0017  * For the list of contributors see $ROOTSYS/README/CREDITS.             *
0018  *************************************************************************/
0019
0020 #ifndef RDF_SNAPSHOTHELPERS
0021 #define RDF_SNAPSHOTHELPERS
0022
0023 #include <ROOT/RSnapshotOptions.hxx>
0024
0025 #include <ROOT/RDF/RActionImpl.hxx>
0026 #include <ROOT/RDF/RLoopManager.hxx>
0027 #include <ROOT/RDF/Utils.hxx>
0028
0029 #include <array>
0030 #include <memory>
0031 #include <variant>
0032
0033 class TBranch;
0034 class TFile;
0035
0036 namespace ROOT {
0037 class REntry;
0038 class RFieldToken;
0039 class RNTupleFillContext;
0040 class RNTupleParallelWriter;
0041 class TBufferMerger;
0042 class TBufferMergerFile;
0043 } // namespace ROOT
0044
0045 namespace ROOT::Internal::RDF {
0046
0047 class R__CLING_PTRCHECK(off) UntypedSnapshotRNTupleHelper final : public RActionImpl<UntypedSnapshotRNTupleHelper> {
0048    std::string fFileName;
0049    std::string fDirName;
0050    std::string fNTupleName;
0051
0052    std::unique_ptr<TFile> fOutputFile;
0053
0054    RSnapshotOptions fOptions;
0055    ROOT::Detail::RDF::RLoopManager *fInputLoopManager;
0056    ROOT::Detail::RDF::RLoopManager *fOutputLoopManager;
0057    ColumnNames_t fInputFieldNames; // This contains the resolved aliases
0058    ColumnNames_t fOutputFieldNames;
0059    std::unique_ptr<ROOT::RNTupleParallelWriter> fWriter;
0060    std::vector<ROOT::RFieldToken> fFieldTokens;
0061
0062    unsigned int fNSlots;
0063    std::vector<std::shared_ptr<ROOT::RNTupleFillContext>> fFillContexts;
0064    std::vector<std::unique_ptr<ROOT::REntry>> fEntries;
0065
0066    std::vector<const std::type_info *> fInputColumnTypeIDs; // Types for the input columns
0067
0068 public:
0069    UntypedSnapshotRNTupleHelper(unsigned int nSlots, std::string_view filename, std::string_view dirname,
0070                                 std::string_view ntuplename, const ColumnNames_t &vfnames, const ColumnNames_t &fnames,
0071                                 const RSnapshotOptions &options, ROOT::Detail::RDF::RLoopManager *inputLM,
0072                                 ROOT::Detail::RDF::RLoopManager *outputLM,
0073                                 const std::vector<const std::type_info *> &colTypeIDs);
0074
0075    UntypedSnapshotRNTupleHelper(const UntypedSnapshotRNTupleHelper &) = delete;
0076    UntypedSnapshotRNTupleHelper &operator=(const UntypedSnapshotRNTupleHelper &) = delete;
0077    UntypedSnapshotRNTupleHelper(UntypedSnapshotRNTupleHelper &&) noexcept;
0078    UntypedSnapshotRNTupleHelper &operator=(UntypedSnapshotRNTupleHelper &&) noexcept;
0079    ~UntypedSnapshotRNTupleHelper() final;
0080
0081    void Initialize();
0082
0083    void Exec(unsigned int slot, const std::vector<void *> &values);
0084
0085    void InitTask(TTreeReader *, unsigned int slot);
0086
0087    void FinalizeTask(unsigned int slot);
0088
0089    void Finalize();
0090
0091    std::string GetActionName() { return "Snapshot"; }
0092
0093    ROOT::RDF::SampleCallback_t GetSampleCallback() final
0094    {
0095       return [](unsigned int, const RSampleInfo &) mutable {};
0096    }
0097
0098    UntypedSnapshotRNTupleHelper MakeNew(void *newName);
0099 };
0100
0101 /// Stores properties of each output branch in a Snapshot.
0102 struct RBranchData {
0103    /// Stores variations of a fundamental type.
0104    /// The bytes hold anything up to double or 64-bit numbers, and are cleared for every event.
0105    /// This allows for binding the branches directly to these bytes.
0106    struct FundamentalType {
0107       static constexpr std::size_t fNBytes = 8;
0108       alignas(8) std::array<std::byte, fNBytes> fBytes{std::byte{0}}; // 8 bytes to store any fundamental type
0109       unsigned short fSize = 0;
0110       FundamentalType(unsigned short size) : fSize(size) { assert(size <= fNBytes); }
0111    };
0112    /// Stores empty instances of classes, so a dummy object can be written when a systematic variation
0113    /// doesn't pass a selection cut.
0114    struct EmptyDynamicType {
0115       const TClass *fTClass = nullptr;
0116       std::shared_ptr<void> fEmptyInstance = nullptr;
0117       void *fRawPtrToEmptyInstance = nullptr; // Needed because TTree expects pointer to pointer
0118    };
0119
0120    std::string fInputBranchName; // This contains resolved aliases
0121    std::string fOutputBranchName;
0122    const std::type_info *fInputTypeID = nullptr;
0123    TBranch *fOutputBranch = nullptr;
0124    void *fBranchAddressForCArrays = nullptr; // Used to detect if branch addresses need to be updated
0125
0126    int fVariationIndex = -1; // For branches that are only valid if a specific filter passed
0127    std::variant<FundamentalType, EmptyDynamicType> fTypeData = FundamentalType{0};
0128    bool fIsCArray = false;
0129    bool fIsDefine = false;
0130
0131    RBranchData() = default;
0132    RBranchData(std::string inputBranchName, std::string outputBranchName, bool isDefine, const std::type_info *typeID);
0133
0134    void ClearBranchPointers()
0135    {
0136       fOutputBranch = nullptr;
0137       fBranchAddressForCArrays = nullptr;
0138    }
0139    void *EmptyInstance(bool pointerToPointer);
0140    void ClearBranchContents();
0141    /// For fundamental types represented by TDataType, fetch a value from the pointer into the local branch buffer.
0142    /// If the branch holds a class type, nothing happens.
0143    /// \return true if the branch holds a fundamental type, false if it holds a class type.
0144    bool WriteValueIfFundamental(void *valuePtr)
0145    {
0146       if (auto fundamentalType = std::get_if<FundamentalType>(&fTypeData); fundamentalType) {
0147          std::memcpy(fundamentalType->fBytes.data(), valuePtr, fundamentalType->fSize);
0148          return true;
0149       }
0150       return false;
0151    }
0152 };
0153
0154 class R__CLING_PTRCHECK(off) UntypedSnapshotTTreeHelper final : public RActionImpl<UntypedSnapshotTTreeHelper> {
0155    std::string fFileName;
0156    std::string fDirName;
0157    std::string fTreeName;
0158    RSnapshotOptions fOptions;
0159    std::unique_ptr<TFile> fOutputFile;
0160    std::unique_ptr<TTree> fOutputTree; // must be a ptr because TTrees are not copy/move constructible
0161    bool fBranchAddressesNeedReset{true};
0162    TTree *fInputTree = nullptr; // Current input tree. Set at initialization time (`InitTask`)
0163    std::vector<RBranchData> fBranchData; // Information for all output branches
0164    ROOT::Detail::RDF::RLoopManager *fOutputLoopManager;
0165    ROOT::Detail::RDF::RLoopManager *fInputLoopManager;
0166
0167 public:
0168    UntypedSnapshotTTreeHelper(std::string_view filename, std::string_view dirname, std::string_view treename,
0169                               const ColumnNames_t &vbnames, const ColumnNames_t &bnames,
0170                               const RSnapshotOptions &options, std::vector<bool> &&isDefine,
0171                               ROOT::Detail::RDF::RLoopManager *loopManager, ROOT::Detail::RDF::RLoopManager *inputLM,
0172                               const std::vector<const std::type_info *> &colTypeIDs);
0173
0174    UntypedSnapshotTTreeHelper(const UntypedSnapshotTTreeHelper &) = delete;
0175    UntypedSnapshotTTreeHelper &operator=(const UntypedSnapshotTTreeHelper &) = delete;
0176    UntypedSnapshotTTreeHelper(UntypedSnapshotTTreeHelper &&) noexcept;
0177    UntypedSnapshotTTreeHelper &operator=(UntypedSnapshotTTreeHelper &&) noexcept;
0178    ~UntypedSnapshotTTreeHelper() final;
0179
0180    void InitTask(TTreeReader *, unsigned int);
0181
0182    void Exec(unsigned int, const std::vector<void *> &values);
0183
0184    void UpdateCArraysPtrs(const std::vector<void *> &values);
0185
0186    void SetBranches(const std::vector<void *> &values);
0187
0188    void SetEmptyBranches(TTree *inputTree, TTree &outputTree);
0189
0190    void Initialize();
0191
0192    void Finalize();
0193
0194    std::string GetActionName() { return "Snapshot"; }
0195
0196    ROOT::RDF::SampleCallback_t GetSampleCallback() final
0197    {
0198       return [this](unsigned int, const RSampleInfo &) mutable { fBranchAddressesNeedReset = true; };
0199    }
0200
0201    UntypedSnapshotTTreeHelper MakeNew(void *newName, std::string_view /*variation*/ = "nominal");
0202 };
0203
0204 class R__CLING_PTRCHECK(off) UntypedSnapshotTTreeHelperMT final : public RActionImpl<UntypedSnapshotTTreeHelperMT> {
0205
0206    // IMT-specific data members
0207
0208    unsigned int fNSlots;
0209    std::unique_ptr<ROOT::TBufferMerger> fMerger; // must use a ptr because TBufferMerger is not movable
0210    std::vector<std::shared_ptr<ROOT::TBufferMergerFile>> fOutputFiles;
0211    std::vector<std::unique_ptr<TTree>> fOutputTrees;
0212    std::vector<int> fBranchAddressesNeedReset; // vector<bool> does not allow concurrent writing of different elements
0213    std::vector<TTree *> fInputTrees; // Current input trees, one per slot. Set at initialization time (`InitTask`)
0214    std::vector<std::vector<RBranchData>> fBranchData; // Information for all output branches of each slot
0215
0216    // Attributes of the output TTree
0217
0218    std::string fFileName;
0219    std::string fDirName;
0220    std::string fTreeName;
0221    TFile *fOutputFile; // Non-owning view on the output file
0222    RSnapshotOptions fOptions;
0223
0224    // Attributes related to the computation graph
0225
0226    ROOT::Detail::RDF::RLoopManager *fOutputLoopManager;
0227    ROOT::Detail::RDF::RLoopManager *fInputLoopManager;
0228
0229 public:
0230    UntypedSnapshotTTreeHelperMT(unsigned int nSlots, std::string_view filename, std::string_view dirname,
0231                                 std::string_view treename, const ColumnNames_t &vbnames, const ColumnNames_t &bnames,
0232                                 const RSnapshotOptions &options, std::vector<bool> &&isDefine,
0233                                 ROOT::Detail::RDF::RLoopManager *loopManager, ROOT::Detail::RDF::RLoopManager *inputLM,
0234                                 const std::vector<const std::type_info *> &colTypeIDs);
0235
0236    UntypedSnapshotTTreeHelperMT(const UntypedSnapshotTTreeHelperMT &) = delete;
0237    UntypedSnapshotTTreeHelperMT &operator=(const UntypedSnapshotTTreeHelperMT &) = delete;
0238    UntypedSnapshotTTreeHelperMT(UntypedSnapshotTTreeHelperMT &&) noexcept;
0239    UntypedSnapshotTTreeHelperMT &operator=(UntypedSnapshotTTreeHelperMT &&) noexcept;
0240    ~UntypedSnapshotTTreeHelperMT() final;
0241
0242    void InitTask(TTreeReader *r, unsigned int slot);
0243
0244    void FinalizeTask(unsigned int slot);
0245
0246    void Exec(unsigned int slot, const std::vector<void *> &values);
0247
0248    void UpdateCArraysPtrs(unsigned int slot, const std::vector<void *> &values);
0249
0250    void SetBranches(unsigned int slot, const std::vector<void *> &values);
0251
0252    void SetEmptyBranches(TTree *inputTree, TTree &outputTree);
0253
0254    void Initialize();
0255
0256    void Finalize();
0257
0258    std::string GetActionName() { return "Snapshot"; }
0259
0260    ROOT::RDF::SampleCallback_t GetSampleCallback() final
0261    {
0262       return [this](unsigned int slot, const RSampleInfo &) mutable { fBranchAddressesNeedReset[slot] = 1; };
0263    }
0264
0265    UntypedSnapshotTTreeHelperMT MakeNew(void *newName, std::string_view /*variation*/ = "nominal");
0266 };
0267
0268 struct SnapshotOutputWriter;
0269
0270 /// TTree snapshot helper with systematic variations.
0271 class R__CLING_PTRCHECK(off) SnapshotHelperWithVariations
0272    : public ROOT::Detail::RDF::RActionImpl<SnapshotHelperWithVariations> {
0273    RSnapshotOptions fOptions;
0274    std::shared_ptr<SnapshotOutputWriter> fOutputHandle;
0275    TTree *fInputTree = nullptr; // Current input tree. Set at initialization time (`InitTask`)
0276    std::vector<RBranchData> fBranchData;
0277    ROOT::Detail::RDF::RLoopManager *fInputLoopManager = nullptr;
0278    ROOT::Detail::RDF::RLoopManager *fOutputLoopManager = nullptr;
0279
0280    void ClearOutputBranches();
0281
0282 public:
0283    SnapshotHelperWithVariations(std::string_view filename, std::string_view dirname, std::string_view treename,
0284                                 const ColumnNames_t & /*vbnames*/, const ColumnNames_t &bnames,
0285                                 const RSnapshotOptions &options, std::vector<bool> && /*isDefine*/,
0286                                 ROOT::Detail::RDF::RLoopManager *outputLoopMgr,
0287                                 ROOT::Detail::RDF::RLoopManager *inputLoopMgr,
0288                                 const std::vector<const std::type_info *> &colTypeIDs);
0289
0290    SnapshotHelperWithVariations(SnapshotHelperWithVariations const &) noexcept = delete;
0291    SnapshotHelperWithVariations(SnapshotHelperWithVariations &&) noexcept = default;
0292    ~SnapshotHelperWithVariations() = default;
0293    SnapshotHelperWithVariations &operator=(SnapshotHelperWithVariations const &) = delete;
0294    SnapshotHelperWithVariations &operator=(SnapshotHelperWithVariations &&) noexcept = default;
0295
0296    void RegisterVariedColumn(unsigned int slot, unsigned int columnIndex, unsigned int originalColumnIndex,
0297                              unsigned int varationIndex, std::string const &variationName);
0298
0299    void InitTask(TTreeReader *, unsigned int slot);
0300
0301    void Exec(unsigned int /*slot*/, const std::vector<void *> &values, std::vector<bool> const &filterPassed);
0302
0303    /// Nothing to do. All initialisations run in the constructor or InitTask().
0304    void Initialize() {}
0305
0306    void Finalize();
0307
0308    std::string GetActionName() { return "SnapshotWithVariations"; }
0309 };
0310
0311 } // namespace ROOT::Internal::RDF
0312
0313 #endif