Podio/src/PodioWriter.cpp

0001 // This file is part of the ACTS project.
0002 //
0003 // Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008
0009 #include "ActsExamples/Io/Podio/PodioWriter.hpp"
0010
0011 #include "Acts/Plugins/Podio/PodioUtil.hpp"
0012 #include "Acts/Utilities/Logger.hpp"
0013 #include "ActsExamples/Framework/DataHandle.hpp"
0014
0015 #include <algorithm>
0016 #include <list>
0017
0018 #include <podio/CollectionBase.h>
0019 #include <podio/Frame.h>
0020
0021 namespace ActsExamples {
0022 namespace detail {
0023
0024 using CollectionHandle =
0025     ConsumeDataHandle<std::unique_ptr<podio::CollectionBase>>;
0026
0027 class PodioWriterImpl {
0028  public:
0029   PodioWriterImpl(const PodioWriter::Config& config, PodioWriter& parent)
0030       : m_cfg(config),
0031         m_inputPodioFrame(&parent, "InputPodioFrame"),
0032         m_writer(config.outputPath) {}
0033
0034   PodioWriter::Config m_cfg;
0035
0036   ConsumeDataHandle<podio::Frame> m_inputPodioFrame;
0037
0038   std::vector<std::unique_ptr<CollectionHandle>> m_collections;
0039
0040   std::mutex m_writeMutex;
0041   Acts::PodioUtil::ROOTWriter m_writer;
0042 };
0043 }  // namespace detail
0044
0045 PodioWriter::PodioWriter(const Config& config, Acts::Logging::Level level)
0046     : m_logger(Acts::getDefaultLogger("PodioWriter", level)),
0047       m_impl(std::make_unique<detail::PodioWriterImpl>(config, *this)) {
0048   ACTS_DEBUG("Creating output file " << config.outputPath);
0049
0050   if (m_impl->m_cfg.category.empty()) {
0051     throw std::invalid_argument("Category name is not set");
0052   }
0053   if (!m_impl->m_cfg.inputFrame.has_value()) {
0054     ACTS_DEBUG("No input frame name set, will create a new one");
0055   } else {
0056     m_impl->m_inputPodioFrame.initialize(m_impl->m_cfg.inputFrame.value());
0057   }
0058
0059   std::set<std::string> collections;
0060   std::ranges::copy(m_impl->m_cfg.collections,
0061                     std::inserter(collections, collections.end()));
0062   if (collections.size() != m_impl->m_cfg.collections.size()) {
0063     throw std::invalid_argument(
0064         "Duplicate collection names in config, please use check your "
0065         "configuration");
0066   }
0067   if (std::ranges::any_of(m_impl->m_cfg.collections,
0068                           [](const auto& c) { return c.empty(); })) {
0069     throw std::invalid_argument("Collection name is empty");
0070   }
0071
0072   ACTS_DEBUG("Adding " << m_impl->m_cfg.collections.size() << " collections");
0073   for (const auto& collection : m_impl->m_cfg.collections) {
0074     auto up = std::make_unique<detail::CollectionHandle>(this, collection);
0075     m_impl->m_collections.push_back(std::move(up));
0076
0077     m_impl->m_collections.back()->initialize(collection);
0078   }
0079   ACTS_DEBUG("PodioWriter::PodioWriter: " << m_impl->m_collections.size());
0080 }
0081
0082 std::string PodioWriter::name() const {
0083   return "PodioWriter";
0084 }
0085
0086 ProcessCode PodioWriter::write(const AlgorithmContext& ctx) {
0087   ACTS_DEBUG("PodioWriter::write");
0088   podio::Frame frame = [this, &ctx]() -> podio::Frame {
0089     if (m_impl->m_inputPodioFrame.isInitialized()) {
0090       ACTS_VERBOSE(
0091           "PodioWriter::write: taking inputPodioFrame from WhiteBoard");
0092       return m_impl->m_inputPodioFrame(ctx);
0093     } else {
0094       ACTS_VERBOSE("PodioWriter::write: creating new inputPodioFrame");
0095       return {};
0096     }
0097   }();
0098
0099   std::lock_guard guard(m_impl->m_writeMutex);
0100   for (const auto& handle : m_impl->m_collections) {
0101     auto collectionPtr = (*handle)(ctx);
0102     if (!collectionPtr) {
0103       ACTS_ERROR("PodioWriter::write: collection is not initialized");
0104       return ProcessCode::ABORT;
0105     }
0106     frame.put(std::move(collectionPtr), handle->name());
0107   }
0108   m_impl->m_writer.writeFrame(frame, m_impl->m_cfg.category);
0109
0110   return ProcessCode::SUCCESS;
0111 }
0112
0113 ProcessCode PodioWriter::finalize() {
0114   m_impl->m_writer.finish();
0115
0116   return ProcessCode::SUCCESS;
0117 }
0118
0119 PodioWriter::~PodioWriter() = default;
0120
0121 const PodioWriter::Config& PodioWriter::config() const {
0122   return m_impl->m_cfg;
0123 }
0124
0125 }  // namespace ActsExamples