diff --git a/src/mana_plugin_root_rntuple_writer.cc b/src/mana_plugin_root_rntuple_writer.cc index 6e56ace..69125f9 100644 --- a/src/mana_plugin_root_rntuple_writer.cc +++ b/src/mana_plugin_root_rntuple_writer.cc @@ -3,6 +3,7 @@ #include #include #include +#include #include #include #include @@ -25,9 +26,8 @@ using RNTupleWriter = ROOT::Experimental::RNTupleWriter; struct Context { - std::unique_ptr outputFile; - std::unique_ptr model; - std::vector>> eventRVecs; + std::vector>> rvecs; + std::unique_ptr writer; }; static Context *g_ctx = nullptr; @@ -56,48 +56,48 @@ MANA_DEFINE_PLUGIN_SHUTDOWN(shutdown) g_ctx = nullptr; } -inline void make_things(const nlohmann::json &jRun) -{ - std::vector>> result; - for (const auto &jEvent: jRun["events"]) - { - auto dir = gDirectory->mkdir(jEvent["name"].get().c_str(), "", true); - auto model = RNTupleModel::Create(); - std::vector> rvecs; - for (const auto &jOutput: jEvent["outputs"]) - { - auto rvec = model->MakeField(jOutput["name"].get()); - rvec->resize(jOutput["size"].get()); - rvecs.emplace_back(std::move(rvec)); - } - result.emplace_back(std::move(rvecs)); - } -} - MANA_DEFINE_PLUGIN_BEGIN_RUN(begin_run) { + assert(context); + assert(context == g_ctx); log_debug("begin_run: context=%p, descriptor_json=%s", context, descriptor_json); auto jRun = nlohmann::json::parse(descriptor_json); + + if (jRun["events"].size() > 1) + { + log_error("root-rntuple-writer: currently only supports a single event/trigger"); + return; + } + + auto ctx = reinterpret_cast(context); + ctx->rvecs.clear(); + auto jEvent = jRun["events"][0]; + auto model = RNTupleModel::Create(); + for (const auto &jOutput: jEvent["outputs"]) + { + auto name = jOutput["name"].get(); + std::replace(std::begin(name), std::end(name), '.', '_'); + auto rvec = model->MakeField>(name); + rvec->resize(jOutput["size"].get()); + ctx->rvecs.emplace_back(std::move(rvec)); + } + std::filesystem::path rp(jRun["name"].get()); auto filename = fmt::format("{}_rntuple_events.root", rp.filename().replace_extension().string()); - auto ctx = reinterpret_cast(context); - ctx->outputFile = std::make_unique(filename.c_str(), "RECREATE"); - ctx->outputFile->cd(); - - make_things(jRun); - - ctx->model = ROOT::Experimental::RNTupleModel::Create(); - auto field = ctx->model->MakeField("value"); log_info("writing rntuple into: %s", filename.c_str()); + + ctx->writer = + RNTupleWriter::Recreate(std::move(model), jEvent["name"].get(), filename); } MANA_DEFINE_PLUGIN_END_RUN(end_run) { + assert(context); + assert(context == g_ctx); log_debug("end: context=%p, descriptor_json=%s", context, descriptor_json); auto ctx = reinterpret_cast(context); - ctx->outputFile->Write(); *ctx = {}; } @@ -105,7 +105,20 @@ MANA_DEFINE_PLUGIN_EVENT_DATA(process_event) { log_trace("event: ctx=%p, eventIndex=%d, arrayCount=%zu, totalBytes=%zu", context, eventIndex, arrayCount, totalBytes); + if (eventIndex != 0) + return; + auto ctx = reinterpret_cast(context); + + for (size_t ai = 0; ai < arrayCount; ++ai) + { + auto &rvec = ctx->rvecs.at(ai); + auto input = mana::get_span(arrays[ai]); + + rvec->resize(input.size()); + std::copy(std::begin(input), std::end(input), std::begin(*rvec)); + } + ctx->writer->Fill(); } MANA_DEFINE_PLUGIN_SYSTEM_EVENT(process_system_event)