diff --git a/Framework/Core/include/Framework/ASoA.h b/Framework/Core/include/Framework/ASoA.h index 7219812cc11a5..e894afcd67e1b 100644 --- a/Framework/Core/include/Framework/ASoA.h +++ b/Framework/Core/include/Framework/ASoA.h @@ -1045,7 +1045,12 @@ class Table return filtered_iterator(mColumnChunks, {selection, mOffset}); } - unfiltered_iterator iteratorAt(uint64_t i) const + iterator iteratorAt(uint64_t i) const + { + return rawIteratorAt(i); + } + + unfiltered_iterator rawIteratorAt(uint64_t i) const { return mBegin + (i - mOffset); } @@ -1120,6 +1125,16 @@ class Table return t; } + auto slice(uint64_t start, uint64_t end) + { + return rawSlice(start, end); + } + + auto rawSlice(uint64_t start, uint64_t end) + { + return table_t{mTable->Slice(start, end - start + 1), start}; + } + protected: /// Offset of the table within a larger table. uint64_t mOffset; @@ -1340,153 +1355,153 @@ constexpr auto is_binding_compatible_v() /// Array index: return an array of iterators, defined by values in its elements /// SLICE -#define DECLARE_SOA_SLICE_INDEX_COLUMN_FULL(_Name_, _Getter_, _Type_, _Table_, _Suffix_) \ - struct _Name_##IdSlice : o2::soa::Column<_Type_[2], _Name_##IdSlice> { \ - static_assert(std::is_integral_v<_Type_>, "Index type must be integral"); \ - static_assert((*_Suffix_ == '\0') || (*_Suffix_ == '_'), "Suffix has to begin with _"); \ - static constexpr const char* mLabel = "fIndexSlice" #_Table_ _Suffix_; \ - using base = o2::soa::Column<_Type_[2], _Name_##IdSlice>; \ - using type = _Type_[2]; \ - using column_t = _Name_##IdSlice; \ - using binding_t = _Table_; \ - _Name_##IdSlice(arrow::ChunkedArray const* column) \ - : o2::soa::Column<_Type_[2], _Name_##IdSlice>(o2::soa::ColumnIterator(column)) \ - { \ - } \ - \ - _Name_##IdSlice() = default; \ - _Name_##IdSlice(_Name_##IdSlice const& other) = default; \ - _Name_##IdSlice& operator=(_Name_##IdSlice const& other) = default; \ - std::array<_Type_, 2> inline getIds() const \ - { \ - return _Getter_##Ids(); \ - } \ - \ - std::array<_Type_, 2> _Getter_##Ids() const \ - { \ - return std::array{(*mColumnIterator)[0], (*mColumnIterator)[1]}; \ - } \ - \ - bool has_##_Getter_() const \ - { \ - return (*mColumnIterator)[0] >= 0; \ - } \ - \ - template \ - auto _Getter_##_as() const \ - { \ - assert(mBinding != nullptr); \ - auto t = T{static_cast(mBinding)->asArrowTable()->Slice((*mColumnIterator)[0], (*mColumnIterator)[1] - (*mColumnIterator)[0] + 1u), static_cast((*mColumnIterator)[0])}; \ - static_cast(mBinding)->copyIndexBindings(t); \ - return t; \ - } \ - \ - auto _Getter_() const \ - { \ - return _Getter_##_as(); \ - } \ - \ - template \ - bool setCurrent(T* current) \ - { \ - if constexpr (o2::soa::is_binding_compatible_v()) { \ - assert(current != nullptr); \ - this->mBinding = current; \ - return true; \ - } \ - return false; \ - } \ - \ - bool setCurrentRaw(void* current) \ - { \ - this->mBinding = current; \ - return true; \ - } \ - binding_t* getCurrent() const { return static_cast(mBinding); } \ - void* getCurrentRaw() const { return mBinding; } \ - void* mBinding = nullptr; \ +#define DECLARE_SOA_SLICE_INDEX_COLUMN_FULL(_Name_, _Getter_, _Type_, _Table_, _Suffix_) \ + struct _Name_##IdSlice : o2::soa::Column<_Type_[2], _Name_##IdSlice> { \ + static_assert(std::is_integral_v<_Type_>, "Index type must be integral"); \ + static_assert((*_Suffix_ == '\0') || (*_Suffix_ == '_'), "Suffix has to begin with _"); \ + static constexpr const char* mLabel = "fIndexSlice" #_Table_ _Suffix_; \ + using base = o2::soa::Column<_Type_[2], _Name_##IdSlice>; \ + using type = _Type_[2]; \ + using column_t = _Name_##IdSlice; \ + using binding_t = _Table_; \ + _Name_##IdSlice(arrow::ChunkedArray const* column) \ + : o2::soa::Column<_Type_[2], _Name_##IdSlice>(o2::soa::ColumnIterator(column)) \ + { \ + } \ + \ + _Name_##IdSlice() = default; \ + _Name_##IdSlice(_Name_##IdSlice const& other) = default; \ + _Name_##IdSlice& operator=(_Name_##IdSlice const& other) = default; \ + std::array<_Type_, 2> inline getIds() const \ + { \ + return _Getter_##Ids(); \ + } \ + \ + std::array<_Type_, 2> _Getter_##Ids() const \ + { \ + return std::array{(*mColumnIterator)[0], (*mColumnIterator)[1]}; \ + } \ + \ + bool has_##_Getter_() const \ + { \ + return (*mColumnIterator)[0] >= 0; \ + } \ + \ + template \ + auto _Getter_##_as() const \ + { \ + assert(mBinding != nullptr); \ + auto t = static_cast(mBinding)->rawSlice((*mColumnIterator)[0], (*mColumnIterator)[1]); \ + static_cast(mBinding)->copyIndexBindings(t); \ + return t; \ + } \ + \ + auto _Getter_() const \ + { \ + return _Getter_##_as(); \ + } \ + \ + template \ + bool setCurrent(T* current) \ + { \ + if constexpr (o2::soa::is_binding_compatible_v()) { \ + assert(current != nullptr); \ + this->mBinding = current; \ + return true; \ + } \ + return false; \ + } \ + \ + bool setCurrentRaw(void* current) \ + { \ + this->mBinding = current; \ + return true; \ + } \ + binding_t* getCurrent() const { return static_cast(mBinding); } \ + void* getCurrentRaw() const { return mBinding; } \ + void* mBinding = nullptr; \ }; #define DECLARE_SOA_SLICE_INDEX_COLUMN(_Name_, _Getter_) DECLARE_SOA_SLICE_INDEX_COLUMN_FULL(_Name_, _Getter_, int32_t, _Name_##s, "") ///ARRAY -#define DECLARE_SOA_ARRAY_INDEX_COLUMN_FULL(_Name_, _Getter_, _Type_, _N_, _Table_, _Suffix_) \ - struct _Name_##Ids : o2::soa::Column<_Type_[_N_], _Name_##Ids> { \ - static_assert(std::is_integral_v<_Type_>, "Index type must be integral"); \ - static_assert((*_Suffix_ == '\0') || (*_Suffix_ == '_'), "Suffix has to begin with _"); \ - static constexpr const char* mLabel = "fIndexArray" #_Table_ _Suffix_; \ - using base = o2::soa::Column<_Type_[_N_], _Name_##Ids>; \ - using type = _Type_[_N_]; \ - using column_t = _Name_##Ids; \ - using binding_t = _Table_; \ - _Name_##Ids(arrow::ChunkedArray const* column) \ - : o2::soa::Column<_Type_[_N_], _Name_##Ids>(o2::soa::ColumnIterator(column)) \ - { \ - } \ - \ - _Name_##Ids() = default; \ - _Name_##Ids(_Name_##Ids const& other) = default; \ - _Name_##Ids& operator=(_Name_##Ids const& other) = default; \ - \ - std::array<_Type_, _N_> inline getIds() const \ - { \ - return _Getter_##Ids(); \ - } \ - \ - std::array<_Type_, _N_> _Getter_##Ids() const \ - { \ - return getIds(std::make_index_sequence<_N_>{}); \ - } \ - \ - template \ - std::array<_Type_, _N_> getIds(std::index_sequence) const \ - { \ - return std::array<_Type_, _N_>{(*mColumnIterator)[N]...}; \ - } \ - \ - template \ - bool has_##_Getter_() const \ - { \ - static_assert(N < _N_, "Out-of-bounds"); \ - return (*mColumnIterator)[N] >= 0; \ - } \ - \ - template \ - auto _Getter_##_as() const \ - { \ - assert(mBinding != nullptr); \ - return getIterators(std::make_index_sequence<_N_>{}); \ - } \ - template \ - auto getIterators(std::index_sequence) const \ - { \ - return std::array{(static_cast(mBinding)->begin() + (*mColumnIterator)[N])...}; \ - } \ - \ - auto _Getter_() const \ - { \ - return _Getter_##_as(); \ - } \ - \ - template \ - bool setCurrent(T* current) \ - { \ - if constexpr (o2::soa::is_binding_compatible_v()) { \ - assert(current != nullptr); \ - this->mBinding = current; \ - return true; \ - } \ - return false; \ - } \ - \ - bool setCurrentRaw(void* current) \ - { \ - this->mBinding = current; \ - return true; \ - } \ - binding_t* getCurrent() const { return static_cast(mBinding); } \ - void* getCurrentRaw() const { return mBinding; } \ - void* mBinding = nullptr; \ +#define DECLARE_SOA_ARRAY_INDEX_COLUMN_FULL(_Name_, _Getter_, _Type_, _N_, _Table_, _Suffix_) \ + struct _Name_##Ids : o2::soa::Column<_Type_[_N_], _Name_##Ids> { \ + static_assert(std::is_integral_v<_Type_>, "Index type must be integral"); \ + static_assert((*_Suffix_ == '\0') || (*_Suffix_ == '_'), "Suffix has to begin with _"); \ + static constexpr const char* mLabel = "fIndexArray" #_Table_ _Suffix_; \ + using base = o2::soa::Column<_Type_[_N_], _Name_##Ids>; \ + using type = _Type_[_N_]; \ + using column_t = _Name_##Ids; \ + using binding_t = _Table_; \ + _Name_##Ids(arrow::ChunkedArray const* column) \ + : o2::soa::Column<_Type_[_N_], _Name_##Ids>(o2::soa::ColumnIterator(column)) \ + { \ + } \ + \ + _Name_##Ids() = default; \ + _Name_##Ids(_Name_##Ids const& other) = default; \ + _Name_##Ids& operator=(_Name_##Ids const& other) = default; \ + \ + std::array<_Type_, _N_> inline getIds() const \ + { \ + return _Getter_##Ids(); \ + } \ + \ + std::array<_Type_, _N_> _Getter_##Ids() const \ + { \ + return getIds(std::make_index_sequence<_N_>{}); \ + } \ + \ + template \ + std::array<_Type_, _N_> getIds(std::index_sequence) const \ + { \ + return std::array<_Type_, _N_>{(*mColumnIterator)[N]...}; \ + } \ + \ + template \ + bool has_##_Getter_() const \ + { \ + static_assert(N < _N_, "Out-of-bounds"); \ + return (*mColumnIterator)[N] >= 0; \ + } \ + \ + template \ + auto _Getter_##_as() const \ + { \ + assert(mBinding != nullptr); \ + return getIterators(std::make_index_sequence<_N_>{}); \ + } \ + template \ + auto getIterators(std::index_sequence) const \ + { \ + return std::array{(static_cast(mBinding)->rawIteratorAt((*mColumnIterator)[N]))...}; \ + } \ + \ + auto _Getter_() const \ + { \ + return _Getter_##_as(); \ + } \ + \ + template \ + bool setCurrent(T* current) \ + { \ + if constexpr (o2::soa::is_binding_compatible_v()) { \ + assert(current != nullptr); \ + this->mBinding = current; \ + return true; \ + } \ + return false; \ + } \ + \ + bool setCurrentRaw(void* current) \ + { \ + this->mBinding = current; \ + return true; \ + } \ + binding_t* getCurrent() const { return static_cast(mBinding); } \ + void* getCurrentRaw() const { return mBinding; } \ + void* mBinding = nullptr; \ }; #define DECLARE_SOA_ARRAY_INDEX_COLUMN(_Name_, _Getter_, _Size_) DECLARE_SOA_ARRAY_INDEX_COLUMN_FULL(_Name_, _Getter_, int32_t, _Size_, _Name_##s, "") @@ -1528,7 +1543,7 @@ constexpr auto is_binding_compatible_v() auto _Getter_##_as() const \ { \ assert(mBinding != nullptr); \ - return static_cast(mBinding)->begin() + *mColumnIterator; \ + return static_cast(mBinding)->rawIteratorAt(*mColumnIterator); \ } \ \ auto _Getter_() const \ @@ -1807,6 +1822,7 @@ template class FilteredPolicy : public T { public: + using self_t = FilteredPolicy; using originals = originals_pack_t; using table_t = typename T::table_t; using persistent_columns_t = typename T::persistent_columns_t; @@ -1918,6 +1934,18 @@ class FilteredPolicy : public T doCopyIndexBindings(external_index_columns_t{}, dest); } + auto slice(uint64_t start, uint64_t end) + { + auto start_iterator = std::lower_bound(mSelectedRows.begin(), mSelectedRows.end(), start); + auto stop_iterator = std::lower_bound(start_iterator, mSelectedRows.end(), end); + SelectionVector slicedSelection{start_iterator, stop_iterator}; + std::transform(slicedSelection.begin(), slicedSelection.end(), slicedSelection.begin(), + [&](int64_t idx) { + return idx - static_cast(start); + }); + return self_t{{this->asArrowTable()->Slice(start, end - start + 1)}, std::move(slicedSelection), start}; + } + protected: void sumWithSelection(SelectionVector const& selection) { diff --git a/Framework/Core/test/test_ASoA.cxx b/Framework/Core/test/test_ASoA.cxx index 093fbbc58dfda..7c0471fb821a5 100644 --- a/Framework/Core/test/test_ASoA.cxx +++ b/Framework/Core/test/test_ASoA.cxx @@ -649,6 +649,43 @@ BOOST_AUTO_TEST_CASE(TestEmptyTables) BOOST_CHECK_EQUAL(spawned.size(), 0); } +DECLARE_SOA_TABLE(Origins, "TST", "ORIG", o2::soa::Index<>, test::X, test::SomeBool); +namespace test +{ +DECLARE_SOA_INDEX_COLUMN(Origin, origin); +} +DECLARE_SOA_TABLE(References, "TST", "REFS", o2::soa::Index<>, test::OriginId); + +BOOST_AUTO_TEST_CASE(TestIndexToFiltered) +{ + TableBuilder b; + auto writer = b.cursor(); + for (auto i = 0; i < 20; ++i) { + writer(0, i, i % 3 == 0); + } + auto origins = b.finalize(); + Origins o{origins}; + + TableBuilder w; + auto writer_w = w.cursor(); + for (auto i = 0; i < 5 * 20; ++i) { + writer_w(0, i % 20); + } + auto refs = w.finalize(); + References r{refs}; + expressions::Filter flt = test::someBool == true; + using Flt = o2::soa::Filtered; + Flt f{{o.asArrowTable()}, expressions::createSelection(o.asArrowTable(), flt)}; + r.bindExternalIndices(&f); + auto it = r.begin(); + it.moveByIndex(23); + BOOST_CHECK_EQUAL(it.origin().globalIndex(), 3); + it++; + BOOST_CHECK_EQUAL(it.origin().globalIndex(), 4); + it++; + BOOST_CHECK_EQUAL(it.origin().globalIndex(), 5); +} + namespace test { DECLARE_SOA_ARRAY_INDEX_COLUMN(Points3D, pointGroup, 3); @@ -684,10 +721,8 @@ BOOST_AUTO_TEST_CASE(TestAdvancedIndices) auto s1 = it.pointSlice(); auto g1 = it.pointGroup(); auto bb = std::is_same_v; - BOOST_CHECK(bb); BOOST_CHECK_EQUAL(s1.size(), 2); - aa = {2, 3, 4}; for (int i = 0; i < 3; ++i) { BOOST_CHECK_EQUAL(g1[i].globalIndex(), aa[i]); @@ -697,9 +732,31 @@ BOOST_AUTO_TEST_CASE(TestAdvancedIndices) auto s2 = it.pointSlice(); auto g2 = it.pointGroup(); BOOST_CHECK_EQUAL(s2.size(), 7); - aa = {12, 2, 19}; for (int i = 0; i < 3; ++i) { BOOST_CHECK_EQUAL(g2[i].globalIndex(), aa[i]); } + + using Flt = o2::soa::Filtered; + expressions::Filter flt = test::x <= -6; + Flt f{{pt.asArrowTable()}, expressions::createSelection(pt.asArrowTable(), flt)}; + prt.bindExternalIndices(&f); + + auto it2 = prt.begin(); + auto s1f = it2.pointSlice_as(); + auto g1f = it2.pointGroup_as(); + BOOST_CHECK_EQUAL(s1f.size(), 2); + aa = {2, 3, 4}; + for (int i = 0; i < 3; ++i) { + BOOST_CHECK_EQUAL(g1f[i].globalIndex(), aa[i]); + } + + ++it2; + auto s2f = it2.pointSlice_as(); + auto g2f = it2.pointGroup_as(); + BOOST_CHECK_EQUAL(s2f.size(), 7); + aa = {12, 2, 19}; + for (int i = 0; i < 3; ++i) { + BOOST_CHECK_EQUAL(g2f[i].globalIndex(), aa[i]); + } }