Skip to content

Commit 591e23b

Browse files
aalkinaalkin
committed
improve handling of index tables
1 parent 6766ad6 commit 591e23b

File tree

5 files changed

+121
-49
lines changed

5 files changed

+121
-49
lines changed

‎Framework/Core/include/Framework/ASoA.h‎

Lines changed: 48 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,27 @@ consteval auto merge()
100100
return out;
101101
}
102102

103+
template <size_t N1, size_t N2, std::array<TableRef, N1> ar1, std::array<TableRef, N2> ar2, typename L>
104+
consteval auto merge_if(L l)
105+
{
106+
constexpr const int to_remove = std::ranges::count_if(ar1.begin(), ar1.end(), [&](TableRef const& a){ return !l(a); });
107+
constexpr const int duplicates = std::ranges::count_if(ar2.begin(), ar2.end(), [&](TableRef const& a) { return std::any_of(ar1.begin(), ar1.end(), [&](TableRef const& e) { return e == a; }) || !l(a); });
108+
std::array<TableRef, N1 + N2 - duplicates - to_remove> out;
109+
110+
auto pos = std::copy_if(ar1.begin(), ar1.end(), out.begin(), [&](TableRef const& a){ return l(a); });
111+
std::copy_if(ar2.begin(), ar2.end(), pos, [&](TableRef const& a) { return std::none_of(ar1.begin(), ar1.end(), [&](TableRef const& e) { return e == a; }) && l(a); });
112+
return out;
113+
}
114+
115+
template <size_t N, std::array<TableRef, N> ar, typename L>
116+
consteval auto remove_if(L l)
117+
{
118+
constexpr const int to_remove = std::ranges::count_if(ar.begin(), ar.end(), [&l](TableRef const& e){ return l(e);});
119+
std::array<TableRef, N - to_remove> out;
120+
std::copy_if(ar.begin(), ar.end(), out.begin(), [&l](TableRef const& e){ return !l(e); });
121+
return out;
122+
}
123+
103124
template <size_t N1, size_t N2, std::array<TableRef, N1> ar1, std::array<TableRef, N2> ar2>
104125
consteval auto intersect()
105126
{
@@ -298,6 +319,18 @@ struct Hash {
298319
static constexpr char const* const str{""};
299320
};
300321

322+
template <size_t N, std::array<soa::TableRef, N> ar, typename Key>
323+
consteval auto filterForKey()
324+
{
325+
constexpr std::array<bool, N> test = []<size_t... Is>(std::index_sequence<Is...>) {
326+
return std::array<bool, N>{(Key::template hasOriginal<ar[Is]>() || (o2::aod::MetadataTraitNG<o2::aod::Hash<ar[Is].desc_hash>>::metadata::template getIndexPosToKey<Key>() >= 0))...};
327+
}(std::make_index_sequence<N>());
328+
constexpr int correct = std::ranges::count(test.begin(), test.end(), true);
329+
std::array<soa::TableRef, correct> out;
330+
std::ranges::copy_if(ar.begin(), ar.end(), out.begin(), [&test](soa::TableRef const& r){ return test[std::distance(ar.begin(), std::find(ar.begin(), ar.end(), r))]; });
331+
return out;
332+
}
333+
301334
#define O2HASH(_Str_) \
302335
template <> \
303336
struct Hash<_Str_ ""_h> { \
@@ -2176,6 +2209,11 @@ class TableNG
21762209
return std::find_if(originals.begin(), originals.end(), [](TableRef const& o){ return o.desc_hash == r.desc_hash;}) != originals.end();
21772210
}
21782211

2212+
static consteval bool hasOriginal_alt(TableRef const& r)
2213+
{
2214+
return std::ranges::find_if(originals.begin(), originals.end(), [&](TableRef const& o){ return o.desc_hash == r.desc_hash;}) != originals.end();
2215+
}
2216+
21792217
using columns_t = decltype(
21802218
[]() {
21812219
if constexpr (sizeof...(Ts) == 0) {
@@ -3989,7 +4027,7 @@ consteval auto getIndexTargets()
39894027
O2HASH(_Label_); \
39904028
O2HASH(_Desc_ "/" #_Version_); \
39914029
template <typename O> \
3992-
using _Name_##ExtensionFrom = TableNG<o2::aod::Hash<_Label_ ""_h>, o2::aod::Hash<_Desc_ "/" #_Version_ ""_h>, O>; \
4030+
using _Name_##ExtensionFrom = soa::TableNG<o2::aod::Hash<_Label_ ""_h>, o2::aod::Hash<_Desc_ "/" #_Version_ ""_h>, O>; \
39934031
using _Name_##Extension = _Name_##ExtensionFrom<o2::aod::Hash<_Origin_ ""_h>>; \
39944032
template <typename O = o2::aod::Hash<_Origin_ ""_h>> \
39954033
struct _Name_##ExtensionMetadataFrom : TableMetadataNG<o2::aod::Hash<_Desc_ "/" #_Version_ ""_h>, __VA_ARGS__> { \
@@ -4059,19 +4097,21 @@ consteval auto getIndexTargets()
40594097
O2HASH(#_Name_); \
40604098
O2HASH(_Desc_ "/" #_Version_); \
40614099
template <typename O = o2::aod::Hash<_Origin_ ""_h>> \
4062-
using _Name_##From = o2::soa::IndexTableNG<o2::aod::Hash<#_Name_ ""_h>, o2::aod::Hash<_Desc_ "/" #_Version_ ""_h>, O, _Key_, __VA_ARGS__>; \
4063-
using _Name_ = _Name_##From<o2::aod::Hash<_Origin_ ""_h>>; \
4064-
\
4065-
template <typename O = o2::aod::Hash<_Origin_ ""_h>> \
4066-
struct _Name_##MetadataFrom : o2::aod::TableMetadataNG<o2::aod::Hash<_Desc_ "/" #_Version_ ""_h>, Index<>, __VA_ARGS__> { \
4100+
struct _Name_##MetadataFrom : o2::aod::TableMetadataNG<o2::aod::Hash<_Desc_ "/" #_Version_ ""_h>, soa::Index<>, __VA_ARGS__> { \
40674101
static constexpr bool exclusive = _Exclusive_; \
4068-
using table_t = _Name_##From<O>; \
40694102
using Key = _Key_; \
40704103
using index_pack_t = framework::pack<__VA_ARGS__>; \
4071-
static constexpr auto sources = table_t::originals; \
4104+
static constexpr const auto sources = []<typename... Cs>(framework::pack<Cs...>){ \
4105+
constexpr auto a = o2::soa::mergeOriginals<typename Cs::binding_t...>(); \
4106+
return o2::aod::filterForKey<a.size(), a, Key>(); \
4107+
}(framework::pack<__VA_ARGS__>{}); \
40724108
}; \
40734109
using _Name_##Metadata = _Name_##MetadataFrom<o2::aod::Hash<_Origin_ ""_h>>; \
40744110
\
4111+
template <typename O = o2::aod::Hash<_Origin_ ""_h>> \
4112+
using _Name_##From = o2::soa::IndexTableNG<o2::aod::Hash<#_Name_ ""_h>, o2::aod::Hash<_Desc_ "/" #_Version_ ""_h>, O, _Key_, __VA_ARGS__>; \
4113+
using _Name_ = _Name_##From<o2::aod::Hash<_Origin_ ""_h>>; \
4114+
\
40754115
template <> \
40764116
struct MetadataTraitNG<o2::aod::Hash<_Desc_ "/" #_Version_ ""_h>> { \
40774117
using metadata = _Name_##Metadata; \
@@ -5557,7 +5597,6 @@ struct IndexTableNG : TableNG<L, D, O> {
55575597
using indexing_t = Key;
55585598
using first_t = typename H::binding_t;
55595599
using rest_t = framework::pack<typename Ts::binding_t...>;
5560-
static constexpr const auto originals = o2::soa::mergeOriginals<typename H::binding_t, typename Ts::binding_t...>();
55615600

55625601
IndexTableNG(std::shared_ptr<arrow::Table> table, uint64_t offset = 0)
55635602
: base_t{table, offset}

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

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ constexpr auto tableRef2InputSpec()
3333
{
3434
return framework::InputSpec{
3535
o2::aod::Hash<R.label_hash>::str,
36-
o2::aod::Hash<R.origin_hash>::str,
36+
o2::aod::Hash<R.origin_hash>::origin,
3737
o2::aod::description(o2::aod::Hash<R.desc_hash>::str),
3838
R.version
3939
};
@@ -44,7 +44,7 @@ constexpr auto tableRef2OutputSpec()
4444
{
4545
return framework::OutputSpec{
4646
framework::OutputLabel{o2::aod::Hash<R.label_hash>::str},
47-
o2::aod::Hash<R.origin_hash>::str,
47+
o2::aod::Hash<R.origin_hash>::origin,
4848
o2::aod::description(o2::aod::Hash<R.desc_hash>::str),
4949
R.version
5050
};
@@ -54,7 +54,7 @@ template <TableRef R>
5454
constexpr auto tableRef2Output()
5555
{
5656
return framework::Output{
57-
o2::aod::Hash<R.origin_hash>::str,
57+
o2::aod::Hash<R.origin_hash>::origin,
5858
o2::aod::description(o2::aod::Hash<R.desc_hash>::str),
5959
R.version
6060
};
@@ -342,12 +342,12 @@ template <o2::aod::NGMetadata M, soa::TableRef Ref>
342342
struct TableTransformNG
343343
{
344344
using metadata = M;
345-
using M::sources;
345+
constexpr static auto sources = M::sources;
346346

347347
template <soa::TableRef R>
348348
static constexpr auto base_spec()
349349
{
350-
return tableRef2InputSpec<R>();
350+
return soa::tableRef2InputSpec<R>();
351351
}
352352

353353
static auto base_specs()
@@ -359,17 +359,17 @@ struct TableTransformNG
359359

360360
constexpr auto spec() const
361361
{
362-
return tableRef2OutputSpec<Ref>();
362+
return soa::tableRef2OutputSpec<Ref>();
363363
}
364364

365365
constexpr auto output() const
366366
{
367-
return tableRef2Output<Ref>();
367+
return soa::tableRef2Output<Ref>();
368368
}
369369

370370
constexpr auto ref() const
371371
{
372-
return tableRef2OutputRef<Ref>();
372+
return soa::tableRef2OutputRef<Ref>();
373373
}
374374
};
375375

@@ -404,9 +404,9 @@ struct Spawns : TableTransform<typename aod::MetadataTrait<framework::pack_head_
404404
};
405405

406406
template <typename T>
407-
struct SpawnsNG : TableTransformNG<typename aod::MetadataTraitNG<o2::aod::Hash<T::ref.desc_hash>>, T::ref>
407+
struct SpawnsNG : TableTransformNG<typename aod::MetadataTraitNG<o2::aod::Hash<T::ref.desc_hash>>::metadata, T::ref>
408408
{
409-
using metadata = TableTransformNG<typename aod::MetadataTraitNG<o2::aod::Hash<T::ref.desc_hash>>, T::ref>::metadata;
409+
using metadata = TableTransformNG<typename aod::MetadataTraitNG<o2::aod::Hash<T::ref.desc_hash>>::metadata, T::ref>::metadata;
410410
using extension_t = typename metadata::extension_table_t;
411411
using base_table_t = typename metadata::base_table_t;
412412
using expression_pack_t = typename metadata::expression_pack_t;
@@ -657,9 +657,9 @@ struct Builds : TableTransform<typename aod::MetadataTrait<T>::metadata> {
657657
};
658658

659659
template <typename T>
660-
struct BuildsNG : TableTransformNG<aod::MetadataTraitNG<aod::Hash<T::ref.desc_hash>>, T::ref> {
661-
using metadata = TableTransformNG<aod::MetadataTraitNG<aod::Hash<T::ref.desc_hash>>, T::ref>::metadata;
662-
using IP = std::conditional_t<metadata::exclusive, IndexBuilder<Exclusive>, IndexBuilder<Sparse>>;
660+
struct BuildsNG : TableTransformNG<typename aod::MetadataTraitNG<aod::Hash<T::ref.desc_hash>>::metadata, T::ref> {
661+
using metadata = TableTransformNG<typename aod::MetadataTraitNG<aod::Hash<T::ref.desc_hash>>::metadata, T::ref>::metadata;
662+
using IP = std::conditional_t<metadata::exclusive, IndexBuilderNG<Exclusive>, IndexBuilderNG<Sparse>>;
663663
using Key = metadata::Key;
664664
using H = typename T::first_t;
665665
using Ts = typename T::rest_t;
@@ -685,10 +685,10 @@ struct BuildsNG : TableTransformNG<aod::MetadataTraitNG<aod::Hash<T::ref.desc_ha
685685
return index_pack_t{};
686686
}
687687

688-
template <typename Key, typename... Cs, typename... Ts>
689-
auto build(framework::pack<Cs...>, framework::pack<Ts...>, std::vector<std::shared_ptr<arrow::Table>>&& tables)
688+
template <typename Key, typename... Cs>
689+
auto build(framework::pack<Cs...>, std::vector<std::shared_ptr<arrow::Table>>&& tables)
690690
{
691-
this->table = std::make_shared<T>(IP::template indexBuilder<Key>(o2::aod::Hash<T::ref.label_hash>::str, std::forward<std::vector<std::shared_ptr<arrow::Table>>>(tables), framework::pack<Cs...>{}, framework::pack<Ts...>{}));
691+
this->table = std::make_shared<T>(IP::template indexBuilder<Key, metadata::sources.size(), metadata::sources>(o2::aod::Hash<T::ref.label_hash>::str, std::forward<std::vector<std::shared_ptr<arrow::Table>>>(tables), framework::pack<Cs...>{}));
692692
return (this->table != nullptr);
693693
}
694694
};

‎Framework/Core/include/Framework/AnalysisManagers.h‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -449,13 +449,13 @@ struct OutputManager<SpawnsNG<T>> {
449449
static bool prepare(ProcessingContext& pc, SpawnsNG<T>& what)
450450
{
451451
using metadata = o2::aod::MetadataTraitNG<o2::aod::Hash<T::ref.desc_hash>>::metadata;
452-
auto originalTable = soa::ArrowHelpers::joinTables(extractOriginals<metadata::sources>(pc));
452+
auto originalTable = soa::ArrowHelpers::joinTables(extractOriginals<metadata::sources.size(), metadata::sources>(pc));
453453
if (originalTable->schema()->fields().empty() == true) {
454454
using base_table_t = typename SpawnsNG<T>::base_table_t::table_t;
455455
originalTable = makeEmptyTable<base_table_t>(o2::aod::Hash<metadata::extension_table_t::ref.label_hash>::str);
456456
}
457457

458-
what.extension = std::make_shared<typename SpawnsNG<T>::extension_t>(o2::framework::spawner<o2::aod::Hash<metadata::extension_t::ref.desc_hash>>(originalTable, o2::aod::Hash<metadata::extension_table_t::ref.label_hash>::str));
458+
what.extension = std::make_shared<typename SpawnsNG<T>::extension_t>(o2::framework::spawner<o2::aod::Hash<metadata::extension_table_t::ref.desc_hash>>(originalTable, o2::aod::Hash<metadata::extension_table_t::ref.label_hash>::str));
459459
what.table = std::make_shared<typename T::table_t>(soa::ArrowHelpers::joinTables({what.extension->asArrowTable(), originalTable}));
460460
return true;
461461
}
@@ -523,25 +523,25 @@ struct OutputManager<Builds<T>> {
523523

524524
template <typename T>
525525
struct OutputManager<BuildsNG<T>> {
526-
static bool appendOutput(std::vector<OutputSpec>& outputs, Builds<T>& what, uint32_t)
526+
static bool appendOutput(std::vector<OutputSpec>& outputs, BuildsNG<T>& what, uint32_t)
527527
{
528528
outputs.emplace_back(what.spec());
529529
return true;
530530
}
531531

532-
static bool prepare(ProcessingContext& pc, Builds<T>& what)
532+
static bool prepare(ProcessingContext& pc, BuildsNG<T>& what)
533533
{
534-
return what.template build<typename T::indexing_t>(what.pack(), what.originals_pack(),
535-
extractOriginalsVector(what.originals_pack(), pc));
534+
using metadata = o2::aod::MetadataTraitNG<o2::aod::Hash<T::ref.desc_hash>>::metadata;
535+
return what.template build<typename T::indexing_t>(what.pack(), extractOriginals<metadata::sources.size(), metadata::sources>(pc));
536536
}
537537

538-
static bool finalize(ProcessingContext& pc, Builds<T>& what)
538+
static bool finalize(ProcessingContext& pc, BuildsNG<T>& what)
539539
{
540540
pc.outputs().adopt(what.output(), what.asArrowTable());
541541
return true;
542542
}
543543

544-
static bool postRun(EndOfStreamContext&, Builds<T>&)
544+
static bool postRun(EndOfStreamContext&, BuildsNG<T>&)
545545
{
546546
return true;
547547
}
@@ -745,7 +745,7 @@ struct SpawnManager<Spawns<TABLE>> {
745745

746746
template <soa::ngTable TABLE>
747747
struct SpawnManager<SpawnsNG<TABLE>> {
748-
static bool requestInputs(std::vector<InputSpec>& inputs, Spawns<TABLE>& spawns)
748+
static bool requestInputs(std::vector<InputSpec>& inputs, SpawnsNG<TABLE>& spawns)
749749
{
750750
auto base_specs = spawns.base_specs();
751751
for (auto base_spec : base_specs) {
@@ -777,7 +777,7 @@ struct IndexManager<Builds<IDX>> {
777777

778778
template <typename IDX>
779779
struct IndexManager<BuildsNG<IDX>> {
780-
static bool requestInputs(std::vector<InputSpec>& inputs, Builds<IDX>& builds)
780+
static bool requestInputs(std::vector<InputSpec>& inputs, BuildsNG<IDX>& builds)
781781
{
782782
auto base_specs = builds.base_specs();
783783
for (auto base_spec : base_specs) {

0 commit comments

Comments
 (0)