Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-09-15 09:06:27

0001 #ifndef PODIO_ROOTWRITER_H
0002 #define PODIO_ROOTWRITER_H
0003 
0004 #include "podio/utilities/DatamodelRegistryIOHelpers.h"
0005 #include "podio/utilities/RootHelpers.h"
0006 
0007 #include "TFile.h"
0008 
0009 #include <memory>
0010 #include <string>
0011 #include <tuple>
0012 #include <unordered_map>
0013 #include <vector>
0014 
0015 // forward declarations
0016 class TTree;
0017 
0018 namespace podio {
0019 class Frame;
0020 class CollectionBase;
0021 class GenericParameters;
0022 
0023 /// The ROOTWriter writes podio files into ROOT files using TTrees.
0024 ///
0025 /// Each category gets its own TTree. Additionally, there is a podio_metadata
0026 /// TTree that contains metadata that is necessary for interpreting the files
0027 /// for reading.
0028 ///
0029 /// Files written with the ROOTWriter can be read with the ROOTReader.
0030 class ROOTWriter {
0031 public:
0032   /// Create a ROOTWriter to write to a file.
0033   ///
0034   /// @note Existing files will be overwritten without warning.
0035   ///
0036   /// @param filename The path to the file that will be created.
0037   ROOTWriter(const std::string& filename);
0038 
0039   /// ROOTWriter destructor
0040   ///
0041   /// This also takes care of writing all the necessary metadata to read files back again.
0042   ~ROOTWriter();
0043 
0044   /// The ROOTWriter is not copy-able
0045   ROOTWriter(const ROOTWriter&) = delete;
0046   /// The ROOTWriter is not copy-able
0047   ROOTWriter& operator=(const ROOTWriter&) = delete;
0048 
0049   /// Store the given frame with the given category.
0050   ///
0051   /// This stores all available collections from the Frame.
0052   ///
0053   /// @note The contents of the first Frame that is written in this way
0054   /// determines the contents that will be written for all subsequent Frames.
0055   ///
0056   /// @param frame    The Frame to store
0057   /// @param category The category name under which this Frame should be stored
0058   void writeFrame(const podio::Frame& frame, const std::string& category);
0059 
0060   /// Store the given Frame with the given category.
0061   ///
0062   /// This stores only the desired collections and not the complete frame.
0063   ///
0064   /// @note The contents of the first Frame that is written in this way
0065   /// determines the contents that will be written for all subsequent Frames.
0066   ///
0067   /// @param frame        The Frame to store
0068   /// @param category     The category name under which this Frame should be
0069   ///                     stored
0070   /// @param collsToWrite The collection names that should be written
0071   void writeFrame(const podio::Frame& frame, const std::string& category, const std::vector<std::string>& collsToWrite);
0072 
0073   /// Write the current file, including all the necessary metadata to read it
0074   /// again.
0075   ///
0076   /// @note The destructor will also call this, so letting a ROOTWriter go out
0077   /// of scope is also a viable way to write a readable file
0078   void finish();
0079 
0080   /// Check whether the collsToWrite are consistent with the state of the passed
0081   /// category.
0082   ///
0083   /// @note This will only be a meaningful check if the first Frame of the passed
0084   /// category has already been written. Also, this check is rather expensive as
0085   /// it has to effectively do two set differences.
0086   ///
0087   ///
0088   /// @param collsToWrite The collection names that should be checked for
0089   ///                     consistency
0090   /// @param category     The category name for which consistency should be
0091   ///                     checked
0092   ///
0093   /// @returns two vectors of collection names. The first one contains all the
0094   /// names that were missing from the collsToWrite but were present in the
0095   /// category. The second one contains the names that are present in the
0096   /// collsToWrite only. If both vectors are empty the category and the passed
0097   /// collsToWrite are consistent.
0098   std::tuple<std::vector<std::string>, std::vector<std::string>>
0099   checkConsistency(const std::vector<std::string>& collsToWrite, const std::string& category) const;
0100 
0101 private:
0102   /// Helper struct to group together all necessary state to write / process a
0103   /// given category. Created during the first writing of a category
0104   struct CategoryInfo {
0105     TTree* tree{nullptr};                                    ///< The TTree to which this category is written
0106     std::vector<root_utils::CollectionBranches> branches{};  ///< The branches for this category
0107     std::vector<root_utils::CollectionWriteInfo> collInfo{}; ///< Collection info for this category
0108     std::vector<std::string> collsToWrite{};                 ///< The collections to write for this category
0109 
0110     // Storage for the keys & values of all the parameters of this category
0111     // (resp. at least the current entry)
0112     root_utils::ParamStorage<int> intParams{};
0113     root_utils::ParamStorage<float> floatParams{};
0114     root_utils::ParamStorage<double> doubleParams{};
0115     root_utils::ParamStorage<std::string> stringParams{};
0116   };
0117 
0118   /// Initialize the branches for this category
0119   void initBranches(CategoryInfo& catInfo, const std::vector<root_utils::StoreCollection>& collections,
0120                     /*const*/ podio::GenericParameters& parameters);
0121 
0122   /// Get the (potentially uninitialized category information for this category)
0123   CategoryInfo& getCategoryInfo(const std::string& category);
0124 
0125   static void resetBranches(CategoryInfo& categoryInfo, const std::vector<root_utils::StoreCollection>& collections);
0126 
0127   /// Fill the parameter keys and values into the CategoryInfo storage
0128   static void fillParams(CategoryInfo& catInfo, const GenericParameters& params);
0129 
0130   std::unique_ptr<TFile> m_file{nullptr};                       ///< The storage file
0131   std::unordered_map<std::string, CategoryInfo> m_categories{}; ///< All categories
0132 
0133   DatamodelDefinitionCollector m_datamodelCollector{};
0134 
0135   bool m_finished{false}; ///< Whether writing has been actually done
0136 };
0137 
0138 } // namespace podio
0139 
0140 #endif // PODIO_ROOTWRITER_H