RNTupleWriterHelper.h

Go to the documentation of this file.

/*
  Copyright (C) 2002-2024 CERN for the benefit of the ATLAS collaboration
*/
 
#ifndef RNTUPLEWRITERHELPER_H
#define RNTUPLEWRITERHELPER_H
 
#include "AthenaBaseComps/AthMessaging.h"
#include "ROOT/REntry.hxx"
#include "ROOT/RField.hxx"
#include "ROOT/RNTuple.hxx"
#if ROOT_VERSION_CODE >= ROOT_VERSION(6, 31, 0)
#include "ROOT/RNTupleWriteOptions.hxx"
#include "ROOT/RNTupleWriter.hxx"
#else
#include "ROOT/RNTupleOptions.hxx"
#endif
 
#include <tuple>
 
namespace ROOT::Experimental {
class RNTupleModel;
}
 
namespace RootStorageSvc {
#if ROOT_VERSION_CODE >= ROOT_VERSION(6, 31, 0)
using RFieldBase = ROOT::Experimental::RFieldBase;
#else
using RFieldBase = ROOT::Experimental::Detail::RFieldBase;
#endif
using RNTupleWriter = ROOT::Experimental::RNTupleWriter;
using RNTupleModel = ROOT::Experimental::RNTupleModel;
using REntry = ROOT::Experimental::REntry;
using RNTupleWriteOptions = ROOT::Experimental::RNTupleWriteOptions;
 
class RNTupleWriterHelper : public AthMessaging {
 public:
  RNTupleWriterHelper(TFile* file, const std::string& ntupleName,
                      bool enableBufferedWrite, bool enableMetrics);
 
  ~RNTupleWriterHelper() = default;
 
  void makeNewEntry();
 
  // The convention for the tuple is <name, type, data>
  typedef std::tuple<std::string, std::string, void*> attrDataTuple;
  void addAttribute(const attrDataTuple& in);
 
  void addField(const std::string& field_name, const std::string& attr_type);
 
  void addFieldValue(const std::string& field_name, void* attr_data);
 
  int commit();
 
  const std::string& getName() const { return m_ntupleName; }
 
  size_t size() const { return m_rowN; }
 
  bool needsCommit() const { return m_needsCommit; }
 
  bool isGrouped() const { return m_clients > 1; }
 
  void increaseClientCount() { m_clients++; }
 
  void close();
 
 private:
  std::map<std::string, void*> m_attrDataMap;
 
  std::unique_ptr<RNTupleModel> m_model;
 
  std::unique_ptr<REntry> m_entry;
 
  std::unique_ptr<RNTupleWriter> m_ntupleWriter;
 
  std::string m_ntupleName;
  TFile* m_tfile;
  RNTupleWriteOptions m_opts;
  int m_rowN = 0;
 
  int m_clients = 0;
  bool m_needsCommit = false;
 
  bool m_collectMetrics;
};
 
}  // namespace RootStorageSvc
#endif