Skip to content

Commit d507b6e

Browse files
aalkinaalkin
committed
improve caching behavior of projectors
1 parent add580c commit d507b6e

File tree

5 files changed

+76
-57
lines changed

5 files changed

+76
-57
lines changed

‎Framework/Core/include/Framework/AnalysisHelpers.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -795,7 +795,8 @@ template <soa::is_table T, soa::is_spawnable_column... Cs>
795795
auto Extend(T const& table)
796796
{
797797
using output_t = Join<T, soa::Table<o2::aod::Hash<"JOIN"_h>, o2::aod::Hash<"JOIN/0"_h>, o2::aod::Hash<"JOIN"_h>, Cs...>>;
798-
return output_t{{o2::framework::spawner(framework::pack<Cs...>{}, {table.asArrowTable()}, "dynamicExtension"), table.asArrowTable()}, 0};
798+
static std::shared_ptr<gandiva::Projector> projector = nullptr;
799+
return output_t{{o2::framework::spawner(framework::pack<Cs...>{}, {table.asArrowTable()}, "dynamicExtension", projector), table.asArrowTable()}, 0};
799800
}
800801

801802
/// Template function to attach dynamic columns on-the-fly (e.g. inside

‎Framework/Core/include/Framework/TableBuilder.h‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -855,12 +855,12 @@ auto makeEmptyTable(const char* name, framework::pack<Cs...> p)
855855
}
856856

857857
std::shared_ptr<arrow::Table> spawnerHelper(std::shared_ptr<arrow::Table> const& fullTable, std::shared_ptr<arrow::Schema> newSchema, size_t nColumns,
858-
expressions::Projector* projectors, std::vector<std::shared_ptr<arrow::Field>> const& fields, const char* name, std::shared_ptr<gandiva::Projector> projector = nullptr);
858+
expressions::Projector* projectors, std::vector<std::shared_ptr<arrow::Field>> const& fields, const char* name, std::shared_ptr<gandiva::Projector>& projector);
859859

860860
/// Expression-based column generator to materialize columns
861861
template <aod::is_aod_hash D>
862862
requires(soa::has_configurable_extension<typename o2::aod::MetadataTrait<D>::metadata>)
863-
auto spawner(std::vector<std::shared_ptr<arrow::Table>>&& tables, const char* name, o2::framework::expressions::Projector* projectors, std::shared_ptr<gandiva::Projector> projector)
863+
auto spawner(std::vector<std::shared_ptr<arrow::Table>>&& tables, const char* name, o2::framework::expressions::Projector* projectors, std::shared_ptr<gandiva::Projector>& projector)
864864
{
865865
using placeholders_pack_t = typename o2::aod::MetadataTrait<D>::metadata::placeholders_pack_t;
866866
auto fullTable = soa::ArrowHelpers::joinTables(std::move(tables));
@@ -875,7 +875,7 @@ auto spawner(std::vector<std::shared_ptr<arrow::Table>>&& tables, const char* na
875875

876876
template <aod::is_aod_hash D>
877877
requires(soa::has_configurable_extension<typename o2::aod::MetadataTrait<D>::metadata>)
878-
auto spawner(std::shared_ptr<arrow::Table> const& fullTable, const char* name, o2::framework::expressions::Projector* projectors, std::shared_ptr<gandiva::Projector> projector)
878+
auto spawner(std::shared_ptr<arrow::Table> const& fullTable, const char* name, o2::framework::expressions::Projector* projectors, std::shared_ptr<gandiva::Projector>& projector)
879879
{
880880
using placeholders_pack_t = typename o2::aod::MetadataTrait<D>::metadata::placeholders_pack_t;
881881
if (fullTable->num_rows() == 0) {
@@ -889,7 +889,7 @@ auto spawner(std::shared_ptr<arrow::Table> const& fullTable, const char* name, o
889889

890890
template <aod::is_aod_hash D>
891891
requires(soa::has_extension<typename o2::aod::MetadataTrait<D>::metadata> && !soa::has_configurable_extension<typename o2::aod::MetadataTrait<D>::metadata>)
892-
auto spawner(std::vector<std::shared_ptr<arrow::Table>>&& tables, const char* name, std::shared_ptr<gandiva::Projector> projector = nullptr)
892+
auto spawner(std::vector<std::shared_ptr<arrow::Table>>&& tables, const char* name, std::shared_ptr<gandiva::Projector>& projector)
893893
{
894894
using expression_pack_t = typename o2::aod::MetadataTrait<D>::metadata::expression_pack_t;
895895
auto fullTable = soa::ArrowHelpers::joinTables(std::move(tables));
@@ -909,7 +909,7 @@ auto spawner(std::vector<std::shared_ptr<arrow::Table>>&& tables, const char* na
909909

910910
template <aod::is_aod_hash D>
911911
requires(soa::has_extension<typename o2::aod::MetadataTrait<D>::metadata> && !soa::has_configurable_extension<typename o2::aod::MetadataTrait<D>::metadata>)
912-
auto spawner(std::shared_ptr<arrow::Table> const& fullTable, const char* name, std::shared_ptr<gandiva::Projector> projector = nullptr)
912+
auto spawner(std::shared_ptr<arrow::Table> const& fullTable, const char* name, std::shared_ptr<gandiva::Projector>& projector)
913913
{
914914
using expression_pack_t = typename o2::aod::MetadataTrait<D>::metadata::expression_pack_t;
915915
if (fullTable->num_rows() == 0) {
@@ -927,7 +927,7 @@ auto spawner(std::shared_ptr<arrow::Table> const& fullTable, const char* name, s
927927
}
928928

929929
template <typename... C>
930-
auto spawner(framework::pack<C...> columns, std::vector<std::shared_ptr<arrow::Table>>&& tables, const char* name)
930+
auto spawner(framework::pack<C...> columns, std::vector<std::shared_ptr<arrow::Table>>&& tables, const char* name, std::shared_ptr<gandiva::Projector>& projector)
931931
{
932932
auto fullTable = soa::ArrowHelpers::joinTables(std::move(tables));
933933
if (fullTable->num_rows() == 0) {
@@ -936,7 +936,7 @@ auto spawner(framework::pack<C...> columns, std::vector<std::shared_ptr<arrow::T
936936
static auto fields = o2::soa::createFieldsFromColumns(columns);
937937
static auto new_schema = std::make_shared<arrow::Schema>(fields);
938938
std::array<expressions::Projector, sizeof...(C)> projectors{{std::move(C::Projector())...}};
939-
return spawnerHelper(fullTable, new_schema, sizeof...(C), projectors.data(), fields, name);
939+
return spawnerHelper(fullTable, new_schema, sizeof...(C), projectors.data(), fields, name, projector);
940940
}
941941

942942
template <typename... T>

‎Framework/Core/src/AODReaderHelpers.cxx‎

Lines changed: 63 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -83,53 +83,64 @@ static inline auto extractOriginals(ProcessingContext& pc)
8383
return {pc.inputs().get<TableConsumer>(o2::aod::label<refs[Is]>())->asArrowTable()...};
8484
}(std::make_index_sequence<refs.size()>());
8585
}
86+
namespace {
87+
template <typename D>
88+
requires(D::exclusive)
89+
auto make_build(D metadata, InputSpec const& input, ProcessingContext& pc)
90+
{
91+
using metadata_t = decltype(metadata);
92+
using Key = typename metadata_t::Key;
93+
using index_pack_t = typename metadata_t::index_pack_t;
94+
constexpr auto sources = metadata_t::sources;
95+
return o2::framework::IndexBuilder<o2::framework::Exclusive>::indexBuilder<Key, sources.size(), sources>(input.binding.c_str(),
96+
extractOriginals<sources.size(), sources>(pc),
97+
index_pack_t{});
98+
}
99+
100+
template <typename D>
101+
requires(!D::exclusive)
102+
auto make_build(D metadata, InputSpec const& input, ProcessingContext& pc)
103+
{
104+
using metadata_t = decltype(metadata);
105+
using Key = typename metadata_t::Key;
106+
using index_pack_t = typename metadata_t::index_pack_t;
107+
constexpr auto sources = metadata_t::sources;
108+
return o2::framework::IndexBuilder<o2::framework::Sparse>::indexBuilder<Key, sources.size(), sources>(input.binding.c_str(),
109+
extractOriginals<sources.size(), sources>(pc),
110+
index_pack_t{});
111+
}
112+
}
86113

87114
AlgorithmSpec AODReaderHelpers::indexBuilderCallback(std::vector<InputSpec>& requested)
88115
{
89-
return AlgorithmSpec::InitCallback{[requested](InitContext& ic) {
116+
return AlgorithmSpec::InitCallback{[requested](InitContext& /*ic*/) {
90117
return [requested](ProcessingContext& pc) {
91118
auto outputs = pc.outputs();
92119
// spawn tables
93120
for (auto& input : requested) {
94121
auto&& [origin, description, version] = DataSpecUtils::asConcreteDataMatcher(input);
95-
auto maker = [&](auto metadata) {
96-
using metadata_t = decltype(metadata);
97-
using Key = typename metadata_t::Key;
98-
using index_pack_t = typename metadata_t::index_pack_t;
99-
constexpr auto sources = metadata_t::sources;
100-
if constexpr (metadata_t::exclusive == true) {
101-
return o2::framework::IndexBuilder<o2::framework::Exclusive>::indexBuilder<Key, sources.size(), sources>(input.binding.c_str(),
102-
extractOriginals<sources.size(), sources>(pc),
103-
index_pack_t{});
104-
} else {
105-
return o2::framework::IndexBuilder<o2::framework::Sparse>::indexBuilder<Key, sources.size(), sources>(input.binding.c_str(),
106-
extractOriginals<sources.size(), sources>(pc),
107-
index_pack_t{});
108-
}
109-
};
110-
111122
if (description == header::DataDescription{"MA_RN2_EX"}) {
112-
outputs.adopt(Output{origin, description, version}, maker(o2::aod::Run2MatchedExclusiveMetadata{}));
123+
outputs.adopt(Output{origin, description, version}, make_build(o2::aod::Run2MatchedExclusiveMetadata{}, input, pc));
113124
} else if (description == header::DataDescription{"MA_RN2_SP"}) {
114-
outputs.adopt(Output{origin, description, version}, maker(o2::aod::Run2MatchedSparseMetadata{}));
125+
outputs.adopt(Output{origin, description, version}, make_build(o2::aod::Run2MatchedSparseMetadata{}, input, pc));
115126
} else if (description == header::DataDescription{"MA_RN3_EX"}) {
116-
outputs.adopt(Output{origin, description, version}, maker(o2::aod::Run3MatchedExclusiveMetadata{}));
127+
outputs.adopt(Output{origin, description, version}, make_build(o2::aod::Run3MatchedExclusiveMetadata{}, input, pc));
117128
} else if (description == header::DataDescription{"MA_RN3_SP"}) {
118-
outputs.adopt(Output{origin, description, version}, maker(o2::aod::Run3MatchedSparseMetadata{}));
129+
outputs.adopt(Output{origin, description, version}, make_build(o2::aod::Run3MatchedSparseMetadata{}, input, pc));
119130
} else if (description == header::DataDescription{"MA_BCCOL_EX"}) {
120-
outputs.adopt(Output{origin, description, version}, maker(o2::aod::MatchedBCCollisionsExclusiveMetadata{}));
131+
outputs.adopt(Output{origin, description, version}, make_build(o2::aod::MatchedBCCollisionsExclusiveMetadata{}, input, pc));
121132
} else if (description == header::DataDescription{"MA_BCCOL_SP"}) {
122-
outputs.adopt(Output{origin, description, version}, maker(o2::aod::MatchedBCCollisionsSparseMetadata{}));
133+
outputs.adopt(Output{origin, description, version}, make_build(o2::aod::MatchedBCCollisionsSparseMetadata{}, input, pc));
123134
} else if (description == header::DataDescription{"MA_BCCOLS_EX"}) {
124-
outputs.adopt(Output{origin, description, version}, maker(o2::aod::MatchedBCCollisionsExclusiveMultiMetadata{}));
135+
outputs.adopt(Output{origin, description, version}, make_build(o2::aod::MatchedBCCollisionsExclusiveMultiMetadata{}, input, pc));
125136
} else if (description == header::DataDescription{"MA_BCCOLS_SP"}) {
126-
outputs.adopt(Output{origin, description, version}, maker(o2::aod::MatchedBCCollisionsSparseMultiMetadata{}));
137+
outputs.adopt(Output{origin, description, version}, make_build(o2::aod::MatchedBCCollisionsSparseMultiMetadata{}, input, pc));
127138
} else if (description == header::DataDescription{"MA_RN3_BC_SP"}) {
128-
outputs.adopt(Output{origin, description, version}, maker(o2::aod::Run3MatchedToBCSparseMetadata{}));
139+
outputs.adopt(Output{origin, description, version}, make_build(o2::aod::Run3MatchedToBCSparseMetadata{}, input, pc));
129140
} else if (description == header::DataDescription{"MA_RN3_BC_EX"}) {
130-
outputs.adopt(Output{origin, description, version}, maker(o2::aod::Run3MatchedToBCExclusiveMetadata{}));
141+
outputs.adopt(Output{origin, description, version}, make_build(o2::aod::Run3MatchedToBCExclusiveMetadata{}, input, pc));
131142
} else if (description == header::DataDescription{"MA_RN2_BC_SP"}) {
132-
outputs.adopt(Output{origin, description, version}, maker(o2::aod::Run2MatchedToBCSparseMetadata{}));
143+
outputs.adopt(Output{origin, description, version}, make_build(o2::aod::Run2MatchedToBCSparseMetadata{}, input, pc));
133144
} else {
134145
throw std::runtime_error("Not an index table");
135146
}
@@ -138,6 +149,17 @@ AlgorithmSpec AODReaderHelpers::indexBuilderCallback(std::vector<InputSpec>& req
138149
}};
139150
}
140151

152+
namespace{
153+
template <o2::aod::is_aod_hash D>
154+
auto make_spawn(InputSpec const& input, ProcessingContext& pc)
155+
{
156+
using metadata_t = o2::aod::MetadataTrait<D>::metadata;
157+
constexpr auto sources = metadata_t::sources;
158+
static std::shared_ptr<gandiva::Projector> projector = nullptr;
159+
return o2::framework::spawner<D>(extractOriginals<sources.size(), sources>(pc), input.binding.c_str(), projector);
160+
}
161+
}
162+
141163
AlgorithmSpec AODReaderHelpers::aodSpawnerCallback(std::vector<InputSpec>& requested)
142164
{
143165
return AlgorithmSpec::InitCallback{[requested](InitContext& /*ic*/) {
@@ -146,43 +168,37 @@ AlgorithmSpec AODReaderHelpers::aodSpawnerCallback(std::vector<InputSpec>& reque
146168
// spawn tables
147169
for (auto& input : requested) {
148170
auto&& [origin, description, version] = DataSpecUtils::asConcreteDataMatcher(input);
149-
auto maker = [&]<o2::aod::is_aod_hash D>() {
150-
using metadata_t = o2::aod::MetadataTrait<D>::metadata;
151-
constexpr auto sources = metadata_t::sources;
152-
return o2::framework::spawner<D>(extractOriginals<sources.size(), sources>(pc), input.binding.c_str());
153-
};
154-
155171
if (description == header::DataDescription{"EXTRACK"}) {
156-
outputs.adopt(Output{origin, description, version}, maker.template operator()<o2::aod::Hash<"EXTRACK/0"_h>>());
172+
outputs.adopt(Output{origin, description, version}, make_spawn<o2::aod::Hash<"EXTRACK/0"_h>>(input, pc));
157173
} else if (description == header::DataDescription{"EXTRACK_IU"}) {
158-
outputs.adopt(Output{origin, description, version}, maker.template operator()<o2::aod::Hash<"EXTRACK_IU/0"_h>>());
174+
outputs.adopt(Output{origin, description, version}, make_spawn<o2::aod::Hash<"EXTRACK_IU/0"_h>>(input, pc));
159175
} else if (description == header::DataDescription{"EXTRACKCOV"}) {
160-
outputs.adopt(Output{origin, description, version}, maker.template operator()<o2::aod::Hash<"EXTRACKCOV/0"_h>>());
176+
outputs.adopt(Output{origin, description, version}, make_spawn<o2::aod::Hash<"EXTRACKCOV/0"_h>>(input, pc));
161177
} else if (description == header::DataDescription{"EXTRACKCOV_IU"}) {
162-
outputs.adopt(Output{origin, description, version}, maker.template operator()<o2::aod::Hash<"EXTRACKCOV_IU/0"_h>>());
178+
outputs.adopt(Output{origin, description, version}, make_spawn<o2::aod::Hash<"EXTRACKCOV_IU/0"_h>>(input, pc));
163179
} else if (description == header::DataDescription{"EXTRACKEXTRA"}) {
164180
if (version == 0U) {
165-
outputs.adopt(Output{origin, description, version}, maker.template operator()<o2::aod::Hash<"EXTRACKEXTRA/0"_h>>());
181+
outputs.adopt(Output{origin, description, version}, make_spawn<o2::aod::Hash<"EXTRACKEXTRA/0"_h>>(input, pc));
166182
} else if (version == 1U) {
167-
outputs.adopt(Output{origin, description, version}, maker.template operator()<o2::aod::Hash<"EXTRACKEXTRA/1"_h>>());
183+
outputs.adopt(Output{origin, description, version}, make_spawn<o2::aod::Hash<"EXTRACKEXTRA/1"_h>>(input, pc));
168184
} else if (version == 2U) {
169-
outputs.adopt(Output{origin, description, version}, maker.template operator()<o2::aod::Hash<"EXTRACKEXTRA/2"_h>>());
185+
outputs.adopt(Output{origin, description, version}, make_spawn<o2::aod::Hash<"EXTRACKEXTRA/2"_h>>(input, pc));
170186
}
171187
} else if (description == header::DataDescription{"EXMFTTRACK"}) {
172188
if (version == 0U) {
173-
outputs.adopt(Output{origin, description, version}, maker.template operator()<o2::aod::Hash<"EXMFTTRACK/0"_h>>());
189+
outputs.adopt(Output{origin, description, version}, make_spawn<o2::aod::Hash<"EXMFTTRACK/0"_h>>(input, pc));
174190
} else if (version == 1U) {
175-
outputs.adopt(Output{origin, description, version}, maker.template operator()<o2::aod::Hash<"EXMFTTRACK/1"_h>>());
191+
outputs.adopt(Output{origin, description, version}, make_spawn<o2::aod::Hash<"EXMFTTRACK/1"_h>>(input, pc));
176192
}
177193
} else if (description == header::DataDescription{"EXFWDTRACK"}) {
178-
outputs.adopt(Output{origin, description, version}, maker.template operator()<o2::aod::Hash<"EXFWDTRACK/0"_h>>());
194+
outputs.adopt(Output{origin, description, version}, make_spawn<o2::aod::Hash<"EXFWDTRACK/0"_h>>(input, pc));
179195
} else if (description == header::DataDescription{"EXFWDTRACKCOV"}) {
180-
outputs.adopt(Output{origin, description, version}, maker.template operator()<o2::aod::Hash<"EXFWDTRACKCOV/0"_h>>());
196+
outputs.adopt(Output{origin, description, version}, make_spawn<o2::aod::Hash<"EXFWDTRACKCOV/0"_h>>(input, pc));
181197
} else if (description == header::DataDescription{"EXMCPARTICLE"}) {
182198
if (version == 0U) {
183-
outputs.adopt(Output{origin, description, version}, maker.template operator()<o2::aod::Hash<"EXMCPARTICLE/0"_h>>());
199+
outputs.adopt(Output{origin, description, version}, make_spawn<o2::aod::Hash<"EXMCPARTICLE/0"_h>>(input, pc));
184200
} else if (version == 1U) {
185-
outputs.adopt(Output{origin, description, version}, maker.template operator()<o2::aod::Hash<"EXMCPARTICLE/1"_h>>());
201+
outputs.adopt(Output{origin, description, version}, make_spawn<o2::aod::Hash<"EXMCPARTICLE/1"_h>>(input, pc));
186202
}
187203
} else {
188204
throw runtime_error("Not an extended table");

‎Framework/Core/src/TableBuilder.cxx‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ void TableBuilder::setLabel(const char* label)
8686

8787
std::shared_ptr<arrow::Table> spawnerHelper(std::shared_ptr<arrow::Table> const& fullTable, std::shared_ptr<arrow::Schema> newSchema, size_t nColumns,
8888
expressions::Projector* projectors, std::vector<std::shared_ptr<arrow::Field>> const& fields, const char* name,
89-
std::shared_ptr<gandiva::Projector> projector)
89+
std::shared_ptr<gandiva::Projector>& projector)
9090
{
9191
if (projector == nullptr) {
9292
projector = framework::expressions::createProjectorHelper(nColumns, projectors, fullTable->schema(), fields);

‎Framework/Core/test/test_TableSpawner.cxx‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,8 +50,10 @@ TEST_CASE("TestTableSpawner")
5050
auto t1 = b1.finalize();
5151
Points st1{t1};
5252

53+
std::shared_ptr<gandiva::Projector> projector = nullptr;
54+
5355
auto expoints_a = o2::soa::Extend<o2::aod::Points, test::Rsq, test::Sin>(st1);
54-
auto extension = ExPointsExtension{o2::framework::spawner<o2::aod::Hash<"EXPTSNG/0"_h>>(t1, o2::aod::Hash<"ExPoints"_h>::str)};
56+
auto extension = ExPointsExtension{o2::framework::spawner<o2::aod::Hash<"EXPTSNG/0"_h>>(t1, o2::aod::Hash<"ExPoints"_h>::str, projector)};
5557
auto expoints = ExPoints{{t1, extension.asArrowTable()}, 0};
5658

5759
REQUIRE(expoints_a.size() == 9);

0 commit comments

Comments
 (0)