Skip to content

Commit 8616240

Browse files
aalkinaalkin
committed
update IndexBuilder
1 parent 6af3449 commit 8616240

File tree

4 files changed

+346
-26
lines changed

4 files changed

+346
-26
lines changed

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

Lines changed: 75 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -253,6 +253,35 @@ struct TableMetadataNG {
253253
using persistent_columns_t = framework::selected_pack<soa::is_persistent_t, Cs...>;
254254
using external_index_columns_t = framework::selected_pack<soa::is_external_index_t, Cs...>;
255255
using internal_index_columns_t = framework::selected_pack<soa::is_self_index_t, Cs...>;
256+
257+
template <typename Key, typename... PCs>
258+
static consteval std::array<bool, sizeof...(PCs)> getMap(framework::pack<PCs...>)
259+
{
260+
return std::array<bool, sizeof...(PCs)>{[]() {
261+
if constexpr (requires { PCs::index_targets.size(); }) {
262+
return Key::template isIndexTargetOf<PCs::index_targets.size(), PCs::index_targets>();
263+
} else {
264+
return false;
265+
}
266+
}()...};
267+
}
268+
269+
template <typename Key>
270+
static consteval int getIndexPosToKey()
271+
{
272+
return getIndexPosToKey_impl<Key, framework::pack_size(persistent_columns_t{}), getMap<Key>(persistent_columns_t{})>();
273+
}
274+
275+
template <typename Key, size_t N, std::array<bool, N> map>
276+
static consteval int getIndexPosToKey_impl()
277+
{
278+
constexpr const auto pos = std::find(map.begin(), map.end(), true);
279+
if constexpr (pos != map.end()) {
280+
return std::distance(map.begin(), pos);
281+
} else {
282+
return -1;
283+
}
284+
}
256285
};
257286

258287
template <typename T>
@@ -2069,7 +2098,7 @@ class TableNG
20692098
}();
20702099

20712100
template <size_t N, std::array<TableRef, N> bindings>
2072-
static constexpr auto isIndexTargetOf()
2101+
static consteval auto isIndexTargetOf()
20732102
{
20742103
if constexpr (std::same_as<O, o2::aod::Hash<"CONC"_h>>) {
20752104
return false;
@@ -2082,6 +2111,12 @@ class TableNG
20822111
}
20832112
}
20842113

2114+
template <TableRef r>
2115+
static consteval bool hasOriginal()
2116+
{
2117+
return std::find_if(originals.begin(), originals.end(), [](TableRef const& o){ return o.desc_hash == r.desc_hash;}) != originals.end();
2118+
}
2119+
20852120
using columns_t = decltype(
20862121
[]() {
20872122
if constexpr (sizeof...(Ts) == 0) {
@@ -2323,13 +2358,20 @@ class TableNG
23232358
template <typename Key>
23242359
inline arrow::ChunkedArray* getIndexToKey()
23252360
{
2326-
if constexpr (framework::has_type_conditional<is_binding_compatible, Key>(external_index_columns_t{})) {
2327-
using IC = framework::pack_element_t<framework::has_type_at_conditional_v<is_binding_compatible, Key>(external_index_columns_t{}), external_index_columns_t>;
2328-
return mColumnChunks[framework::has_type_at_v<IC>(persistent_columns_t{})];
2329-
} else if constexpr (std::is_same_v<table_t, Key>) {
2330-
return nullptr;
2361+
constexpr auto map = []<typename... Cs>(framework::pack<Cs...>) {
2362+
return std::array<bool, sizeof...(Cs)>{[]() {
2363+
if constexpr (requires { Cs::index_targets.size(); }) {
2364+
return Key::template isIndexTargetOf<Cs::index_targets.size(), Cs::index_targets>();
2365+
} else {
2366+
return false;
2367+
}
2368+
}()...};
2369+
}(persistent_columns_t{});
2370+
constexpr auto pos = std::find(map.begin(), map.end(), true);
2371+
if constexpr (pos != map.end()) {
2372+
return mColumnChunks[std::distance(map.begin(), pos)];
23312373
} else {
2332-
static_assert(framework::always_static_assert_v<Key>, "This table does not have an index to this type");
2374+
static_assert(framework::always_static_assert_v<Key>, "This table does not have an index to given Key");
23332375
}
23342376
}
23352377

@@ -3949,19 +3991,39 @@ consteval auto getIndexTargets()
39493991
DECLARE_SOA_INDEX_TABLE_FULL(_Name_, _Key_, "AOD", _Description_, true, __VA_ARGS__)
39503992

39513993
#define DECLARE_SOA_INDEX_TABLE_NG_FULL(_Name_, _Key_, _Origin_, _Version_, _Desc_, _Exclusive_, ...) \
3994+
O2HASH(#_Name_); \
3995+
O2HASH(_Desc_ "/" #_Version_); \
39523996
template <typename O = o2::aod::Hash<_Origin_ ""_h>> \
39533997
using _Name_##From = o2::soa::IndexTableNG<o2::aod::Hash<#_Name_ ""_h>, o2::aod::Hash<_Desc_ "/" #_Version_ ""_h>, O, _Key_, __VA_ARGS__>; \
39543998
using _Name_ = _Name_##From<o2::aod::Hash<_Origin_ ""_h>>; \
39553999
\
39564000
template <typename O = o2::aod::Hash<_Origin_ ""_h>> \
3957-
struct _Name_##Metadata : o2::soa::TableMetadataNG<o2::aod::Hash<_Desc_ "/" #_Version_ ""_h>, __VA_ARGS__> { \
4001+
struct _Name_##MetadataFrom : o2::aod::TableMetadataNG<o2::aod::Hash<_Desc_ "/" #_Version_ ""_h>, Index<>, __VA_ARGS__> { \
39584002
static constexpr bool exclusive = _Exclusive_; \
39594003
using table_t = _Name_##From<O>; \
39604004
using Key = _Key_; \
39614005
using index_pack_t = framework::pack<__VA_ARGS__>; \
3962-
static constexpr auto sources = table_t::soriginals; \
4006+
static constexpr auto sources = table_t::originals; \
4007+
}; \
4008+
using _Name_##Metadata = _Name_##MetadataFrom<o2::aod::Hash<_Origin_ ""_h>>; \
4009+
\
4010+
template <> \
4011+
struct MetadataTraitNG<o2::aod::Hash<_Desc_ "/" #_Version_ ""_h>> { \
4012+
using metadata = _Name_##Metadata; \
39634013
};
39644014

4015+
#define DECLARE_SOA_INDEX_TABLE_NG(_Name_, _Key_, _Description_, ...) \
4016+
DECLARE_SOA_INDEX_TABLE_NG_FULL(_Name_, _Key_, "IDX", 0, _Description_, false, __VA_ARGS__)
4017+
4018+
#define DECLARE_SOA_INDEX_TABLE_NG_EXCLUSIVE(_Name_, _Key_, _Description_, ...) \
4019+
DECLARE_SOA_INDEX_TABLE_NG_FULL(_Name_, _Key_, "IDX", 0, _Description_, true, __VA_ARGS__)
4020+
4021+
#define DECLARE_SOA_INDEX_TABLE_NG_USER(_Name_, _Key_, _Description_, ...) \
4022+
DECLARE_SOA_INDEX_TABLE_NG_FULL(_Name_, _Key_, "AOD", 0, _Description_, false, __VA_ARGS__)
4023+
4024+
#define DECLARE_SOA_INDEX_TABLE_NG_EXCLUSIVE_USER(_Name_, _Key_, _Description_, ...) \
4025+
DECLARE_SOA_INDEX_TABLE_NG_FULL(_Name_, _Key_, "AOD", 0, _Description_, true, __VA_ARGS__)
4026+
39654027
namespace o2::soa
39664028
{
39674029
template <typename T>
@@ -5422,15 +5484,15 @@ class Filtered<Filtered<T>> : public FilteredBase<typename T::table_t>
54225484
/// are index columns defined for the required tables.
54235485
/// First index will be used by process() as the grouping
54245486
template <typename L, typename D, typename O, typename Key, typename H, typename... Ts>
5425-
struct IndexTableNG : TableNG<L, D, O, soa::Index<>, typename H::binding_t, typename Ts::binding_t...> {
5487+
struct IndexTableNG : TableNG<L, D, O> {
54265488
using self_t = IndexTableNG<L, D, O, Key, H, Ts...>;
5427-
using base_t = TableNG<L, D, O, soa::Index<>, typename H::binding_t, typename Ts::binding_t...>;
5489+
using base_t = TableNG<L, D, O>;
54285490
using table_t = base_t;
5429-
using safe_base_t = TableNG<L, D, O, typename H::binding_t, typename Ts::binding_t...>;
5491+
using safe_base_t = TableNG<L, D, O>;
54305492
using indexing_t = Key;
54315493
using first_t = typename H::binding_t;
54325494
using rest_t = framework::pack<typename Ts::binding_t...>;
5433-
using base_t::originals;
5495+
static constexpr const auto originals = o2::soa::mergeOriginals<typename H::binding_t, typename Ts::binding_t...>();
54345496

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

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

Lines changed: 80 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -495,7 +495,6 @@ struct IndexBuilder {
495495
auto pool = arrow::default_memory_pool();
496496
SelfIndexColumnBuilder self{C1::columnLabel(), pool};
497497
std::unique_ptr<ChunkedArrayIterator> keyIndex = nullptr;
498-
int64_t counter = 0;
499498
if constexpr (!std::is_same_v<T1, Key>) {
500499
keyIndex = std::make_unique<ChunkedArrayIterator>(getIndexToKey<T1, Key>(tables[0].get()));
501500
}
@@ -505,7 +504,7 @@ struct IndexBuilder {
505504
pool)...};
506505
std::array<bool, sizeof...(Cs)> finds;
507506

508-
for (counter = 0; counter < tables[0]->num_rows(); ++counter) {
507+
for (int64_t counter = 0; counter < tables[0]->num_rows(); ++counter) {
509508
auto idx = -1;
510509
if constexpr (std::is_same_v<T1, Key>) {
511510
idx = counter;
@@ -541,6 +540,85 @@ struct IndexBuilder {
541540
}
542541
};
543542

543+
template <typename Kind>
544+
struct IndexBuilderNG {
545+
template <typename Key, size_t N, std::array<soa::TableRef, N> refs, typename C1, typename... Cs>
546+
static auto indexBuilder(const char* label, std::vector<std::shared_ptr<arrow::Table>>&& tables, framework::pack<C1, Cs...>)
547+
{
548+
auto pool = arrow::default_memory_pool();
549+
SelfIndexColumnBuilder self{C1::columnLabel(), pool};
550+
std::unique_ptr<ChunkedArrayIterator> keyIndex = nullptr;
551+
if constexpr (!Key::template hasOriginal<refs[0]>()) {
552+
keyIndex = std::make_unique<ChunkedArrayIterator>(tables[0]->column(o2::aod::MetadataTraitNG<o2::aod::Hash<refs[0].desc_hash>>::metadata::template getIndexPosToKey<Key>()));
553+
}
554+
555+
auto sq = std::make_index_sequence<sizeof...(Cs)>();
556+
557+
auto columnBuilders = [&tables, &pool]<size_t... Is>(std::index_sequence<Is...>) -> std::array<std::shared_ptr<framework::SelfIndexColumnBuilder>, sizeof...(Cs)> {
558+
return {[](arrow::Table* table, arrow::MemoryPool* pool) {
559+
using T = framework::pack_element_t<Is, framework::pack<Cs...>>;
560+
if constexpr (!Key::template hasOriginal<refs[Is + 1]>()) {
561+
constexpr auto pos = o2::aod::MetadataTraitNG<o2::aod::Hash<refs[Is + 1].desc_hash>>::metadata::template getIndexPosToKey<Key>();
562+
return std::make_shared<IndexColumnBuilder>(table->column(pos), T::columnLabel(), ColumnTrait<T>::listSize(), pool);
563+
} else {
564+
return std::make_shared<SelfIndexColumnBuilder>(T::columnLabel(), pool);
565+
}
566+
}(tables[Is + 1].get(), pool)...};
567+
}(sq);
568+
569+
std::array<bool, sizeof...(Cs)> finds;
570+
571+
for (int64_t counter = 0; counter < tables[0]->num_rows(); ++counter) {
572+
int64_t idx = -1;
573+
if constexpr (Key::template hasOriginal<refs[0]>()) {
574+
idx = counter;
575+
} else {
576+
idx = keyIndex->valueAt(counter);
577+
}
578+
finds = [&idx, &columnBuilders]<size_t... Is>(std::index_sequence<Is...>){
579+
return std::array{
580+
[&idx, &columnBuilders](){
581+
using T = typename framework::pack_element_t<Is, framework::pack<Cs...>>;
582+
return std::static_pointer_cast<typename Reduction<Key, T>::type>(columnBuilders[Is])->template find<T>(idx);
583+
}()
584+
...};
585+
}(sq);
586+
if constexpr (std::is_same_v<Kind, Sparse>) {
587+
[&idx, &columnBuilders]<size_t... Is>(std::index_sequence<Is...>){
588+
([&idx, &columnBuilders](){
589+
using T = typename framework::pack_element_t<Is, framework::pack<Cs...>>;
590+
return std::static_pointer_cast<typename Reduction<Key, T>::type>(columnBuilders[Is])->template fill<T>(idx);}()
591+
, ...);
592+
}(sq);
593+
self.fill<C1>(counter);
594+
} else if constexpr (std::is_same_v<Kind, Exclusive>) {
595+
if (std::none_of(finds.begin(), finds.end(), [](bool const x) { return x == false; })) {
596+
[&idx, &columnBuilders]<size_t... Is>(std::index_sequence<Is...>){
597+
([&idx, &columnBuilders](){
598+
using T = typename framework::pack_element_t<Is, framework::pack<Cs...>>;
599+
return std::static_pointer_cast<typename Reduction<Key, T>::type>(columnBuilders[Is])->template fill<T>(idx);
600+
}()
601+
, ...);
602+
}(sq);
603+
self.fill<C1>(counter);
604+
}
605+
}
606+
}
607+
608+
return [&label,&columnBuilders,&self]<size_t... Is>(std::index_sequence<Is...>){
609+
return makeArrowTable(label,
610+
{self.template result<C1>(), [&columnBuilders](){
611+
using T = typename framework::pack_element_t<Is, framework::pack<Cs...>>;
612+
return std::static_pointer_cast<typename Reduction<Key, T>::type>(columnBuilders[Is])->template result<T>();
613+
}()...},
614+
{self.field(), [&columnBuilders](){
615+
using T = typename framework::pack_element_t<Is, framework::pack<Cs...>>;
616+
return std::static_pointer_cast<typename Reduction<Key, T>::type>(columnBuilders[Is])->field();
617+
}()...});
618+
}(sq);
619+
}
620+
};
621+
544622
/// This helper struct allows you to declare index tables to be created in a task
545623
template <typename T>
546624
struct Builds : TableTransform<typename aod::MetadataTrait<T>::metadata> {

‎Framework/Core/src/IndexBuilderHelpers.cxx‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,8 +43,8 @@ std::shared_ptr<arrow::Field> SelfIndexColumnBuilder::field() const
4343
}
4444

4545
IndexColumnBuilder::IndexColumnBuilder(std::shared_ptr<arrow::ChunkedArray> source, const char* name, int listSize, arrow::MemoryPool* pool)
46-
: ChunkedArrayIterator{source},
47-
SelfIndexColumnBuilder{name, pool},
46+
: SelfIndexColumnBuilder{name, pool},
47+
ChunkedArrayIterator{source},
4848
mListSize{listSize},
4949
mSourceSize{(size_t)source->length()}
5050
{

0 commit comments

Comments
 (0)