DPL Analysis: move spawner caches outside of the spawner function (#14281)

aalkin · web-flow · commit 930d83707083 · 2025-05-20T19:27:52.000+02:00
diff --git a/Framework/Core/include/Framework/AnalysisHelpers.h b/Framework/Core/include/Framework/AnalysisHelpers.h
@@ -297,6 +297,7 @@ struct Spawns : decltype(transformBase<T>()) {
   using extension_t = typename metadata::extension_table_t;
   using base_table_t = typename metadata::base_table_t;
   using expression_pack_t = typename metadata::expression_pack_t;
+  static constexpr size_t N = framework::pack_size(expression_pack_t{});
 
   constexpr auto pack()
   {
@@ -318,7 +319,13 @@ struct Spawns : decltype(transformBase<T>()) {
   }
   std::shared_ptr<typename T::table_t> table = nullptr;
   std::shared_ptr<extension_t> extension = nullptr;
+  std::array<o2::framework::expressions::Projector, N> projectors = []<typename... C>(framework::pack<C...>) -> std::array<expressions::Projector, sizeof...(C)>
+  {
+    return {{std::move(C::Projector())...}};
+  }
+  (expression_pack_t{});
   std::shared_ptr<gandiva::Projector> projector = nullptr;
+  std::shared_ptr<arrow::Schema> schema = std::make_shared<arrow::Schema>(o2::soa::createFieldsFromColumns(expression_pack_t{}));
 };
 
 template <typename T>
@@ -365,6 +372,7 @@ struct Defines : decltype(transformBase<T>()) {
 
   std::array<o2::framework::expressions::Projector, N> projectors;
   std::shared_ptr<gandiva::Projector> projector = nullptr;
+  std::shared_ptr<arrow::Schema> schema = std::make_shared<arrow::Schema>(o2::soa::createFieldsFromColumns(placeholders_pack_t{}));
 };
 
 template <typename T>
@@ -828,8 +836,10 @@ template <soa::is_table T, soa::is_spawnable_column... Cs>
 auto Extend(T const& table)
 {
   using output_t = Join<T, soa::Table<o2::aod::Hash<"JOIN"_h>, o2::aod::Hash<"JOIN/0"_h>, o2::aod::Hash<"JOIN"_h>, Cs...>>;
+  static std::array<framework::expressions::Projector, sizeof...(Cs)> projectors{{std::move(Cs::Projector())...}};
   static std::shared_ptr<gandiva::Projector> projector = nullptr;
-  return output_t{{o2::framework::spawner(framework::pack<Cs...>{}, {table.asArrowTable()}, "dynamicExtension", projector), table.asArrowTable()}, 0};
+  static auto schema = std::make_shared<arrow::Schema>(o2::soa::createFieldsFromColumns(framework::pack<Cs...>{}));
+  return output_t{{o2::framework::spawner(framework::pack<Cs...>{}, {table.asArrowTable()}, "dynamicExtension", projectors.data(), projector, schema), table.asArrowTable()}, 0};
 }
 
 /// Template function to attach dynamic columns on-the-fly (e.g. inside
diff --git a/Framework/Core/include/Framework/AnalysisManagers.h b/Framework/Core/include/Framework/AnalysisManagers.h
@@ -287,8 +287,13 @@ bool prepareOutput(ProcessingContext& context, T& spawns)
     using base_table_t = typename T::base_table_t::table_t;
     originalTable = makeEmptyTable<base_table_t>(o2::aod::label<metadata::extension_table_t::ref>());
   }
+  using D = o2::aod::Hash<metadata::extension_table_t::ref.desc_hash>;
 
-  spawns.extension = std::make_shared<typename T::extension_t>(o2::framework::spawner<o2::aod::Hash<metadata::extension_table_t::ref.desc_hash>>(originalTable, o2::aod::label<metadata::extension_table_t::ref>(), spawns.projector));
+  spawns.extension = std::make_shared<typename T::extension_t>(o2::framework::spawner<D>(originalTable,
+                                                                                         o2::aod::label<metadata::extension_table_t::ref>(),
+                                                                                         spawns.projectors.data(),
+                                                                                         spawns.projector,
+                                                                                         spawns.schema));
   spawns.table = std::make_shared<typename T::spawnable_t::table_t>(soa::ArrowHelpers::joinTables({spawns.extension->asArrowTable(), originalTable}, std::span{T::spawnable_t::table_t::originalLabels}));
   return true;
 }
@@ -309,8 +314,13 @@ bool prepareOutput(ProcessingContext& context, T& defines)
     using base_table_t = typename T::base_table_t::table_t;
     originalTable = makeEmptyTable<base_table_t>(o2::aod::label<metadata::extension_table_t::ref>());
   }
+  using D = o2::aod::Hash<metadata::extension_table_t::ref.desc_hash>;
 
-  defines.extension = std::make_shared<typename T::extension_t>(o2::framework::spawner<o2::aod::Hash<metadata::extension_table_t::ref.desc_hash>>(originalTable, o2::aod::label<metadata::extension_table_t::ref>(), defines.projectors.data(), defines.projector));
+  defines.extension = std::make_shared<typename T::extension_t>(o2::framework::spawner<D>(originalTable,
+                                                                                          o2::aod::label<metadata::extension_table_t::ref>(),
+                                                                                          defines.projectors.data(),
+                                                                                          defines.projector,
+                                                                                          defines.schema));
   defines.table = std::make_shared<typename T::spawnable_t::table_t>(soa::ArrowHelpers::joinTables({defines.extension->asArrowTable(), originalTable}, std::span{T::spawnable_t::table_t::originalLabels}));
   return true;
 }
diff --git a/Framework/Core/include/Framework/TableBuilder.h b/Framework/Core/include/Framework/TableBuilder.h
@@ -768,80 +768,51 @@ std::shared_ptr<arrow::Table> spawnerHelper(std::shared_ptr<arrow::Table> const&
 /// Expression-based column generator to materialize columns
 template <aod::is_aod_hash D>
   requires(soa::has_configurable_extension<typename o2::aod::MetadataTrait<D>::metadata>)
-auto spawner(std::vector<std::shared_ptr<arrow::Table>>&& tables, const char* name, o2::framework::expressions::Projector* projectors, std::shared_ptr<gandiva::Projector>& projector)
+auto spawner(std::shared_ptr<arrow::Table> const& fullTable, const char* name, o2::framework::expressions::Projector* projectors, std::shared_ptr<gandiva::Projector>& projector, std::shared_ptr<arrow::Schema> const& schema)
 {
   using placeholders_pack_t = typename o2::aod::MetadataTrait<D>::metadata::placeholders_pack_t;
-  auto fullTable = soa::ArrowHelpers::joinTables(std::move(tables), std::span{o2::aod::MetadataTrait<D>::metadata::base_table_t::originalLabels});
   if (fullTable->num_rows() == 0) {
     return makeEmptyTable(name, placeholders_pack_t{});
   }
-  static auto new_schema = std::make_shared<arrow::Schema>(o2::soa::createFieldsFromColumns(placeholders_pack_t{}));
-
-  return spawnerHelper(fullTable, new_schema, framework::pack_size(placeholders_pack_t{}), projectors, name, projector);
+  return spawnerHelper(fullTable, schema, framework::pack_size(placeholders_pack_t{}), projectors, name, projector);
 }
 
 template <aod::is_aod_hash D>
   requires(soa::has_configurable_extension<typename o2::aod::MetadataTrait<D>::metadata>)
-auto spawner(std::shared_ptr<arrow::Table> const& fullTable, const char* name, o2::framework::expressions::Projector* projectors, std::shared_ptr<gandiva::Projector>& projector)
+auto spawner(std::vector<std::shared_ptr<arrow::Table>>&& tables, const char* name, o2::framework::expressions::Projector* projectors, std::shared_ptr<gandiva::Projector>& projector, std::shared_ptr<arrow::Schema> const& schema)
 {
-  using placeholders_pack_t = typename o2::aod::MetadataTrait<D>::metadata::placeholders_pack_t;
-  if (fullTable->num_rows() == 0) {
-    return makeEmptyTable(name, placeholders_pack_t{});
-  }
-  static auto new_schema = std::make_shared<arrow::Schema>(o2::soa::createFieldsFromColumns(placeholders_pack_t{}));
-
-  return spawnerHelper(fullTable, new_schema, framework::pack_size(placeholders_pack_t{}), projectors, name, projector);
+  auto fullTable = soa::ArrowHelpers::joinTables(std::move(tables), std::span{o2::aod::MetadataTrait<D>::metadata::base_table_t::originalLabels});
+  return spawner<D>(fullTable, name, projectors, projector, schema);
 }
 
 template <aod::is_aod_hash D>
   requires(soa::has_extension<typename o2::aod::MetadataTrait<D>::metadata> && !soa::has_configurable_extension<typename o2::aod::MetadataTrait<D>::metadata>)
-auto spawner(std::vector<std::shared_ptr<arrow::Table>>&& tables, const char* name, std::shared_ptr<gandiva::Projector>& projector)
+auto spawner(std::shared_ptr<arrow::Table> const& fullTable, const char* name, expressions::Projector* projectors, std::shared_ptr<gandiva::Projector>& projector, std::shared_ptr<arrow::Schema> const& schema)
 {
   using expression_pack_t = typename o2::aod::MetadataTrait<D>::metadata::expression_pack_t;
-  auto fullTable = soa::ArrowHelpers::joinTables(std::move(tables), std::span{o2::aod::MetadataTrait<D>::metadata::base_table_t::originalLabels});
   if (fullTable->num_rows() == 0) {
     return makeEmptyTable(name, expression_pack_t{});
   }
-  static auto new_schema = std::make_shared<arrow::Schema>(o2::soa::createFieldsFromColumns(expression_pack_t{}));
-
-  auto projectors = []<typename... C>(framework::pack<C...>) -> std::array<expressions::Projector, sizeof...(C)>
-  {
-    return {{std::move(C::Projector())...}};
-  }
-  (expression_pack_t{});
-
-  return spawnerHelper(fullTable, new_schema, framework::pack_size(expression_pack_t{}), projectors.data(), name, projector);
+  return spawnerHelper(fullTable, schema, framework::pack_size(expression_pack_t{}), projectors, name, projector);
 }
 
 template <aod::is_aod_hash D>
   requires(soa::has_extension<typename o2::aod::MetadataTrait<D>::metadata> && !soa::has_configurable_extension<typename o2::aod::MetadataTrait<D>::metadata>)
-auto spawner(std::shared_ptr<arrow::Table> const& fullTable, const char* name, std::shared_ptr<gandiva::Projector>& projector)
+auto spawner(std::vector<std::shared_ptr<arrow::Table>>&& tables, const char* name, expressions::Projector* projectors, std::shared_ptr<gandiva::Projector>& projector, std::shared_ptr<arrow::Schema> const& schema)
 {
-  using expression_pack_t = typename o2::aod::MetadataTrait<D>::metadata::expression_pack_t;
-  if (fullTable->num_rows() == 0) {
-    return makeEmptyTable(name, expression_pack_t{});
-  }
-  static auto new_schema = std::make_shared<arrow::Schema>(o2::soa::createFieldsFromColumns(expression_pack_t{}));
-  auto projectors = []<typename... C>(framework::pack<C...>) -> std::array<expressions::Projector, sizeof...(C)>
-  {
-    return {{std::move(C::Projector())...}};
-  }
-  (expression_pack_t{});
-
-  return spawnerHelper(fullTable, new_schema, framework::pack_size(expression_pack_t{}), projectors.data(), name, projector);
+  auto fullTable = soa::ArrowHelpers::joinTables(std::move(tables), std::span{o2::aod::MetadataTrait<D>::metadata::base_table_t::originalLabels});
+  return spawner<D>(fullTable, name, projectors, projector, schema);
 }
 
 template <typename... C>
-auto spawner(framework::pack<C...> columns, std::vector<std::shared_ptr<arrow::Table>>&& tables, const char* name, std::shared_ptr<gandiva::Projector>& projector)
+auto spawner(framework::pack<C...>, std::vector<std::shared_ptr<arrow::Table>>&& tables, const char* name, expressions::Projector* projectors, std::shared_ptr<gandiva::Projector>& projector, std::shared_ptr<arrow::Schema> const& schema)
 {
   std::array<const char*, 1> labels{"original"};
   auto fullTable = soa::ArrowHelpers::joinTables(std::move(tables), std::span<const char* const>{labels});
   if (fullTable->num_rows() == 0) {
     return makeEmptyTable(name, framework::pack<C...>{});
   }
-  static auto new_schema = std::make_shared<arrow::Schema>(o2::soa::createFieldsFromColumns(columns));
-  std::array<expressions::Projector, sizeof...(C)> projectors{{std::move(C::Projector())...}};
-  return spawnerHelper(fullTable, new_schema, sizeof...(C), projectors.data(), name, projector);
+  return spawnerHelper(fullTable, schema, sizeof...(C), projectors, name, projector);
 }
 
 template <typename... T>
diff --git a/Framework/Core/src/AODReaderHelpers.cxx b/Framework/Core/src/AODReaderHelpers.cxx
@@ -158,7 +158,13 @@ auto make_spawn(InputSpec const& input, ProcessingContext& pc)
   using metadata_t = o2::aod::MetadataTrait<D>::metadata;
   constexpr auto sources = metadata_t::sources;
   static std::shared_ptr<gandiva::Projector> projector = nullptr;
-  return o2::framework::spawner<D>(extractOriginals<sources.size(), sources>(pc), input.binding.c_str(), projector);
+  static std::shared_ptr<arrow::Schema> schema = std::make_shared<arrow::Schema>(o2::soa::createFieldsFromColumns(typename metadata_t::expression_pack_t{}));
+  static auto projectors = []<typename... C>(framework::pack<C...>) -> std::array<expressions::Projector, sizeof...(C)>
+  {
+    return {{std::move(C::Projector())...}};
+  }
+  (typename metadata_t::expression_pack_t{});
+  return o2::framework::spawner<D>(extractOriginals<sources.size(), sources>(pc), input.binding.c_str(), projectors.data(), projector, schema);
 }
 } // namespace
 
diff --git a/Framework/Core/test/test_TableSpawner.cxx b/Framework/Core/test/test_TableSpawner.cxx
@@ -50,10 +50,9 @@ TEST_CASE("TestTableSpawner")
   auto t1 = b1.finalize();
   Points st1{t1};
 
-  std::shared_ptr<gandiva::Projector> projector = nullptr;
-
   auto expoints_a = o2::soa::Extend<o2::aod::Points, test::Rsq, test::Sin>(st1);
-  auto extension = ExPointsExtension{o2::framework::spawner<o2::aod::Hash<"EXPTSNG/0"_h>>(t1, o2::aod::Hash<"ExPoints"_h>::str, projector)};
+  Spawns<ExPoints> s;
+  auto extension = ExPointsExtension{o2::framework::spawner<o2::aod::Hash<"EXPTSNG/0"_h>>(t1, o2::aod::Hash<"ExPoints"_h>::str, s.projectors.data(), s.projector, s.schema)};
   auto expoints = ExPoints{{t1, extension.asArrowTable()}, 0};
 
   REQUIRE(expoints_a.size() == 9);
@@ -81,7 +80,7 @@ TEST_CASE("TestTableSpawner")
   Defines<ExcPoints> excpts;
   excpts.projectors[0] = test::x * test::x + test::y * test::y + test::z * test::z;
 
-  auto extension_2 = ExcPointsCfgExtension{o2::framework::spawner<o2::aod::Hash<"EXCFGPTS/0"_h>>({t1}, o2::aod::Hash<"ExcPoints"_h>::str, excpts.projectors.data(), excpts.projector)};
+  auto extension_2 = ExcPointsCfgExtension{o2::framework::spawner<o2::aod::Hash<"EXCFGPTS/0"_h>>({t1}, o2::aod::Hash<"ExcPoints"_h>::str, excpts.projectors.data(), excpts.projector, excpts.schema)};
   auto excpoints = ExcPoints{{t1, extension_2.asArrowTable()}, 0};
 
   rex = extension.begin();
