Add a reader and writer interface (#522)

* Use std::unique_ptr<T> * Push some changes needed for tests * Return the reader and writer by value * Use final, make error messages a bit more verbose * Add a new podioIO library * Add sio files to gitignore * Link with podioRootIO * Make sure test exclusion matches new test names * Split RTNtuple based interface tests off to make sanitizers run on them * Move unique_ptr instead of releasing and acquiring it * Add read_interface_sio to be excluded when using the Address sanitizer * Don't check type for SIO --------- Co-authored-by: jmcarcell <[email protected]> Co-authored-by: tmadlener <[email protected]>
AIDASoft · May 22, 2024 · 6eefdc4 · 6eefdc4
1 parent d0e9958
commit 6eefdc4
Show file tree

Hide file tree

Showing 18 changed files with 525 additions and 5 deletions.
diff --git a/.gitignore b/.gitignore
@@ -19,6 +19,7 @@ tests/unittests/Manifest.toml
 # Data Files
 *.root
 *.dat
+*.sio
 
 # Spack build folders
 spack*

diff --git a/include/podio/Reader.h b/include/podio/Reader.h
@@ -0,0 +1,110 @@
+#ifndef PODIO_READER_H
+#define PODIO_READER_H
+
+#include "podio/Frame.h"
+#include "podio/podioVersion.h"
+
+namespace podio {
+
+class Reader {
+public:
+  struct ReaderConcept {
+    virtual ~ReaderConcept() = default;
+
+    virtual podio::Frame readNextFrame(const std::string& name) = 0;
+    virtual podio::Frame readFrame(const std::string& name, size_t index) = 0;
+    virtual size_t getEntries(const std::string& name) = 0;
+    virtual podio::version::Version currentFileVersion() const = 0;
+    virtual std::vector<std::string_view> getAvailableCategories() const = 0;
+    virtual const std::string_view getDatamodelDefinition(const std::string& name) const = 0;
+    virtual std::vector<std::string> getAvailableDatamodels() const = 0;
+  };
+
+  template <typename T>
+  struct ReaderModel final : public ReaderConcept {
+    ReaderModel(std::unique_ptr<T> reader) : m_reader(std::move(reader)) {
+    }
+    ReaderModel(const ReaderModel&) = delete;
+    ReaderModel& operator=(const ReaderModel&) = delete;
+
+    podio::Frame readNextFrame(const std::string& name) override {
+      auto maybeFrame = m_reader->readNextEntry(name);
+      if (maybeFrame) {
+        return maybeFrame;
+      }
+      throw std::runtime_error("Failed reading category " + name + " (reading beyond bounds?)");
+    }
+
+    podio::Frame readFrame(const std::string& name, size_t index) override {
+      auto maybeFrame = m_reader->readEntry(name, index);
+      if (maybeFrame) {
+        return maybeFrame;
+      }
+      throw std::runtime_error("Failed reading category " + name + " at frame " + std::to_string(index) +
+                               " (reading beyond bounds?)");
+    }
+    size_t getEntries(const std::string& name) override {
+      return m_reader->getEntries(name);
+    }
+    podio::version::Version currentFileVersion() const override {
+      return m_reader->currentFileVersion();
+    }
+
+    std::vector<std::string_view> getAvailableCategories() const override {
+      return m_reader->getAvailableCategories();
+    }
+
+    const std::string_view getDatamodelDefinition(const std::string& name) const override {
+      return m_reader->getDatamodelDefinition(name);
+    }
+
+    std::vector<std::string> getAvailableDatamodels() const override {
+      return m_reader->getAvailableDatamodels();
+    }
+
+    std::unique_ptr<T> m_reader;
+  };
+
+  std::unique_ptr<ReaderConcept> m_self{nullptr};
+
+  template <typename T>
+  Reader(std::unique_ptr<T>);
+
+  podio::Frame readNextFrame(const std::string& name) {
+    return m_self->readNextFrame(name);
+  }
+  podio::Frame readNextEvent() {
+    return readNextFrame(podio::Category::Event);
+  }
+  podio::Frame readFrame(const std::string& name, size_t index) {
+    return m_self->readFrame(name, index);
+  }
+  podio::Frame readEvent(size_t index) {
+    return readFrame(podio::Category::Event, index);
+  }
+  size_t getEntries(const std::string& name) {
+    return m_self->getEntries(name);
+  }
+  size_t getEvents() {
+    return getEntries(podio::Category::Event);
+  }
+  podio::version::Version currentFileVersion() const {
+    return m_self->currentFileVersion();
+  }
+  std::vector<std::string_view> getAvailableCategories() const {
+    return m_self->getAvailableCategories();
+  }
+  const std::string_view getDatamodelDefinition(const std::string& name) const {
+    return m_self->getDatamodelDefinition(name);
+  }
+  std::vector<std::string> getAvailableDatamodels() const {
+    return m_self->getAvailableDatamodels();
+  }
+};
+
+Reader makeReader(const std::string& filename);
+Reader makeReader(const std::vector<std::string>& filename);
+
+} // namespace podio
+
+#endif // PODIO_READER_H
diff --git a/include/podio/Writer.h b/include/podio/Writer.h
@@ -0,0 +1,86 @@
+#ifndef PODIO_WRITER_H
+#define PODIO_WRITER_H
+
+#include "podio/Frame.h"
+#include "podio/podioVersion.h"
+
+namespace podio {
+
+class Writer {
+public:
+  struct WriterConcept {
+    virtual ~WriterConcept() = default;
+
+    virtual void writeFrame(const podio::Frame& frame, const std::string& category) = 0;
+    virtual void writeFrame(const podio::Frame& frame, const std::string& category,
+                            const std::vector<std::string>& collections) = 0;
+    virtual void writeEvent(const podio::Frame& frame) = 0;
+    virtual void writeEvent(const podio::Frame& frame, const std::vector<std::string>& collections) = 0;
+    virtual void finish() = 0;
+  };
+
+  template <typename T>
+  struct WriterModel final : public WriterConcept {
+    WriterModel(std::unique_ptr<T> writer) : m_writer(std::move(writer)) {
+    }
+    WriterModel(const WriterModel&) = delete;
+    WriterModel& operator=(const WriterModel&) = delete;
+    WriterModel(WriterModel&&) = default;
+    WriterModel& operator=(WriterModel&&) = default;
+
+    ~WriterModel() = default;
+
+    void writeFrame(const podio::Frame& frame, const std::string& category) override {
+      return m_writer->writeFrame(frame, category);
+    }
+    void writeFrame(const podio::Frame& frame, const std::string& category,
+                    const std::vector<std::string>& collections) override {
+      return m_writer->writeFrame(frame, category, collections);
+    }
+    void writeEvent(const podio::Frame& frame) override {
+      return writeFrame(frame, podio::Category::Event);
+    }
+    void writeEvent(const podio::Frame& frame, const std::vector<std::string>& collections) override {
+      return writeFrame(frame, podio::Category::Event, collections);
+    }
+    void finish() override {
+      return m_writer->finish();
+    }
+    std::unique_ptr<T> m_writer{nullptr};
+  };
+
+  std::unique_ptr<WriterConcept> m_self{nullptr};
+
+  template <typename T>
+  Writer(std::unique_ptr<T> reader) : m_self(std::make_unique<WriterModel<T>>(std::move(reader))) {
+  }
+
+  Writer(const Writer&) = delete;
+  Writer& operator=(const Writer&) = delete;
+  Writer(Writer&&) = default;
+  Writer& operator=(Writer&&) = default;
+
+  ~Writer() = default;
+
+  void writeFrame(const podio::Frame& frame, const std::string& category) {
+    return m_self->writeFrame(frame, category);
+  }
+  void writeFrame(const podio::Frame& frame, const std::string& category, const std::vector<std::string>& collections) {
+    return m_self->writeFrame(frame, category, collections);
+  }
+  void writeEvent(const podio::Frame& frame) {
+    return writeFrame(frame, podio::Category::Event);
+  }
+  void writeEvent(const podio::Frame& frame, const std::vector<std::string>& collections) {
+    return writeFrame(frame, podio::Category::Event, collections);
+  }
+  void finish() {
+    return m_self->finish();
+  }
+};
+
+Writer makeWriter(const std::string& filename, const std::string& type = "default");
+
+} // namespace podio
+
+#endif // PODIO_WRITER_H
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -135,8 +135,25 @@ if(ENABLE_SIO)
   LIST(APPEND INSTALL_LIBRARIES podioSioIO podioSioIODict)
 endif()
 
+# --- IO
+set(io_sources
+  Writer.cc
+  Reader.cc
+  )
+
+set(io_headers
+  ${PROJECT_SOURCE_DIR}/include/podio/Writer.h
+  ${PROJECT_SOURCE_DIR}/include/podio/Reader.h
+  )
+
+PODIO_ADD_LIB_AND_DICT(podioIO "${io_headers}" "${io_sources}" io_selection.xml)
+target_link_libraries(podioIO PUBLIC podio::podio podio::podioRootIO)
+if(ENABLE_SIO)
+  target_link_libraries(podioIO PUBLIC podio::podioSioIO)
+endif()
+
 # --- Install everything
-install(TARGETS podio podioDict podioRootIO podioRootIODict ${INSTALL_LIBRARIES}
+install(TARGETS podio podioDict podioRootIO podioRootIODict podioIO ${INSTALL_LIBRARIES}
   EXPORT podioTargets
   DESTINATION "${CMAKE_INSTALL_LIBDIR}")
 
@@ -153,6 +170,8 @@ install(FILES
   ${CMAKE_CURRENT_BINARY_DIR}/libpodioDict_rdict.pcm
   ${CMAKE_CURRENT_BINARY_DIR}/podioRootIODictDict.rootmap
   ${CMAKE_CURRENT_BINARY_DIR}/libpodioRootIODict_rdict.pcm
+  ${CMAKE_CURRENT_BINARY_DIR}/podioIODictDict.rootmap
+  ${CMAKE_CURRENT_BINARY_DIR}/libpodioIODict_rdict.pcm
   DESTINATION "${CMAKE_INSTALL_LIBDIR}")
 
 if (ENABLE_SIO)

diff --git a/src/Reader.cc b/src/Reader.cc
@@ -0,0 +1,75 @@
+#include "podio/Reader.h"
+
+#include "podio/ROOTReader.h"
+#if PODIO_ENABLE_RNTUPLE
+  #include "podio/RNTupleReader.h"
+#endif
+#if PODIO_ENABLE_SIO
+  #include "podio/SIOReader.h"
+#endif
+
+#include "TFile.h"
+#include "TKey.h"
+#include <memory>
+
+namespace podio {
+
+template <typename T>
+Reader::Reader(std::unique_ptr<T> reader) : m_self(std::make_unique<ReaderModel<T>>(std::move(reader))) {
+}
+
+Reader makeReader(const std::string& filename) {
+  return makeReader(std::vector<std::string>{filename});
+}
+
+Reader makeReader(const std::vector<std::string>& filenames) {
+
+  auto suffix = filenames[0].substr(filenames[0].find_last_of(".") + 1);
+  for (size_t i = 1; i < filenames.size(); ++i) {
+    if (filenames[i].substr(filenames[i].find_last_of(".") + 1) != suffix) {
+      throw std::runtime_error("All files must have the same extension");
+    }
+  }
+
+  if (suffix == "root") {
+    TFile* file = TFile::Open(filenames[0].c_str());
+    bool hasRNTuple = false;
+
+    for (auto key : *file->GetListOfKeys()) {
+      auto tkey = dynamic_cast<TKey*>(key);
+
+      if (tkey && std::string(tkey->GetClassName()) == "ROOT::Experimental::RNTuple") {
+        hasRNTuple = true;
+        break;
+      }
+    }
+    if (hasRNTuple) {
+#if PODIO_ENABLE_RNTUPLE
+      auto actualReader = std::make_unique<RNTupleReader>();
+      actualReader->openFiles(filenames);
+      Reader reader{std::move(actualReader)};
+      return reader;
+#else
+      throw std::runtime_error("ROOT RNTuple reader not available. Please recompile with ROOT RNTuple support.");
+#endif
+    } else {
+      auto actualReader = std::make_unique<ROOTReader>();
+      actualReader->openFiles(filenames);
+      Reader reader{std::move(actualReader)};
+      return reader;
+    }
+  } else if (suffix == "sio") {
+#if PODIO_ENABLE_SIO
+    auto actualReader = std::make_unique<SIOReader>();
+    actualReader->openFile(filenames[0]);
+    Reader reader{std::move(actualReader)};
+    return reader;
+#else
+    throw std::runtime_error("SIO reader not available. Please recompile with SIO support.");
+#endif
+  }
+
+  throw std::runtime_error("Unknown file extension: " + suffix);
+}
+
+} // namespace podio
diff --git a/src/Writer.cc b/src/Writer.cc
@@ -0,0 +1,44 @@
+#include "podio/Writer.h"
+
+#include "podio/ROOTWriter.h"
+#if PODIO_ENABLE_RNTUPLE
+  #include "podio/RNTupleWriter.h"
+#endif
+#if PODIO_ENABLE_SIO
+  #include "podio/SIOWriter.h"
+#endif
+
+#include <memory>
+
+namespace podio {
+
+Writer makeWriter(const std::string& filename, const std::string& type) {
+
+  auto endsWith = [](const std::string& str, const std::string& suffix) {
+    return str.size() >= suffix.size() && 0 == str.compare(str.size() - suffix.size(), suffix.size(), suffix);
+  };
+
+  auto lower = [](std::string str) {
+    std::transform(str.begin(), str.end(), str.begin(), [](unsigned char c) { return std::tolower(c); });
+    return str;
+  };
+
+  if ((type == "default" && endsWith(filename, ".root")) || lower(type) == "root") {
+    return Writer{std::make_unique<ROOTWriter>(filename)};
+  } else if (lower(type) == "rntuple") {
+#if PODIO_ENABLE_RNTUPLE
+    return Writer{std::make_unique<RNTupleWriter>(filename)};
+#else
+    throw std::runtime_error("ROOT RNTuple writer not available. Please recompile with ROOT RNTuple support.");
+#endif
+  } else if (endsWith(filename, ".sio")) {
+#if PODIO_ENABLE_SIO
+    return Writer{std::make_unique<SIOWriter>(filename)};
+#else
+    throw std::runtime_error("SIO writer not available. Please recompile with SIO support.");
+#endif
+  }
+  throw std::runtime_error("Unknown file type for file " + filename + " with type " + type);
+}
+
+} // namespace podio
diff --git a/src/io_selection.xml b/src/io_selection.xml
@@ -0,0 +1,6 @@
+<lcgdict>
+  <selection>
+    <class name="podio::Reader"/>
+    <class name="podio::Writer"/>
+  </selection>
+</lcgdict>
diff --git a/tests/CTestCustom.cmake b/tests/CTestCustom.cmake
@@ -22,6 +22,9 @@ if ((NOT "@FORCE_RUN_ALL_TESTS@" STREQUAL "ON") AND (NOT "@USE_SANITIZER@" STREQ
     write_frame_root
     read_frame_root
 
+    write_interface_root
+    read_interface_root
+
     write_python_frame_sio
     read_python_frame_sio
 
@@ -71,13 +74,12 @@ if ((NOT "@FORCE_RUN_ALL_TESTS@" STREQUAL "ON") AND (NOT "@USE_SANITIZER@" STREQ
     set(CTEST_CUSTOM_TESTS_IGNORE
       ${CTEST_CUSTOM_TESTS_IGNORE}
 
-      write_sio
       read_sio
       read_and_write_sio
       write_timed_sio
       read_timed_sio
-      write_frame_sio
       read_frame_sio
+      read_interface_sio
       read_frame_legacy_sio
       read_and_write_frame_sio
       )