diff --git a/Framework/Core/include/Framework/ASoA.h b/Framework/Core/include/Framework/ASoA.h index 9030cf5fe7eec..b4f941c3ab128 100644 --- a/Framework/Core/include/Framework/ASoA.h +++ b/Framework/Core/include/Framework/ASoA.h @@ -541,8 +541,23 @@ struct Index : o2::soa::IndexColumn> { template using is_dynamic_t = framework::is_specialization; +namespace persistent_type_helper +{ +// This checks both for the existence of the ::persistent member in the class T as well as the value returned stored in it. +// Hack: a pointer to any field of type int inside persistent. Both true_type and false_type do not have any int field, but anyways we pass nullptr. +// The compiler picks the version with exact number of arguments when only it can, i.e., when T::persistent is defined. +template +typename T::persistent test(int T::persistent::*); + +template +std::false_type test(...); +} // namespace persistent_type_helper + template -using is_persistent_t = typename std::decay_t::persistent::type; +using is_persistent_t = decltype(persistent_type_helper::test(nullptr)); + +template +using is_persistent_v = typename is_persistent_t::value; template using is_external_index_t = typename std::conditional, std::true_type, std::false_type>::type; @@ -1112,10 +1127,12 @@ class Table auto getId() const { using decayed = std::decay_t; - if constexpr (framework::has_type_v) { + if constexpr (framework::has_type_v) { // index to another table constexpr auto idx = framework::has_type_at_v(bindings_pack_t{}); return framework::pack_element_t::getId(); - } else if constexpr (std::is_same_v) { + } else if constexpr (std::is_same_v) { // self index + return this->globalIndex(); + } else if constexpr (is_index_t::value && decayed::mLabel == "Index") { // soa::Index<> return this->globalIndex(); } else { return static_cast(-1); @@ -1505,52 +1522,63 @@ namespace row_helpers template std::array getArrowColumns(arrow::Table* table, framework::pack) { - static_assert(std::conjunction_v, "BinningPolicy: only persistent columns accepted (not dynamic and not index ones"); + static_assert(std::conjunction_v, "Arrow columns: only persistent columns accepted (not dynamic and not index ones"); return std::array{o2::soa::getIndexFromLabel(table, Cs::columnLabel())...}; } template std::array, sizeof...(Cs)> getChunks(arrow::Table* table, framework::pack, uint64_t ci) { - static_assert(std::conjunction_v, "BinningPolicy: only persistent columns accepted (not dynamic and not index ones"); + static_assert(std::conjunction_v, "Arrow chunks: only persistent columns accepted (not dynamic and not index ones"); return std::array, sizeof...(Cs)>{o2::soa::getIndexFromLabel(table, Cs::columnLabel())->chunk(ci)...}; } -template -typename C::type getSingleRowPersistentData(arrow::Table* table, uint64_t ci, uint64_t ai) +template +typename C::type getSingleRowPersistentData(arrow::Table* table, T& rowIterator, uint64_t ci = -1, uint64_t ai = -1) { + if (ci == -1 || ai == -1) { + auto colIterator = static_cast(rowIterator).getIterator(); + ci = colIterator.mCurrentChunk; + ai = *(colIterator.mCurrentPos) - colIterator.mFirstIndex; + } return std::static_pointer_cast>(o2::soa::getIndexFromLabel(table, C::columnLabel())->chunk(ci))->raw_values()[ai]; } template -typename C::type getSingleRowDynamicData(T& rowIterator, uint64_t globalIndex) +typename C::type getSingleRowDynamicData(T& rowIterator, uint64_t globalIndex = -1) { - rowIterator.setCursor(globalIndex); + if (globalIndex != -1 && globalIndex != *std::get<0>(rowIterator.getIndices())) { + rowIterator.setCursor(globalIndex); + } return rowIterator.template getDynamicColumn(); } template -typename C::type getSingleRowIndexData(T& rowIterator, uint64_t globalIndex) +typename C::type getSingleRowIndexData(T& rowIterator, uint64_t globalIndex = -1) { - rowIterator.setCursor(globalIndex); + if (globalIndex != -1 && globalIndex != *std::get<0>(rowIterator.getIndices())) { + rowIterator.setCursor(globalIndex); + } return rowIterator.template getId(); } template -typename C::type getSingleRowData(arrow::Table* table, T& rowIterator, uint64_t ci, uint64_t ai, uint64_t globalIndex) +typename C::type getSingleRowData(arrow::Table* table, T& rowIterator, uint64_t ci = -1, uint64_t ai = -1, uint64_t globalIndex = -1) { using decayed = std::decay_t; if constexpr (decayed::persistent::value) { - return getSingleRowPersistentData(table, ci, ai); + return getSingleRowPersistentData(table, rowIterator, ci, ai); } else if constexpr (o2::soa::is_dynamic_t()) { return getSingleRowDynamicData(rowIterator, globalIndex); - } else if constexpr (o2::soa::is_index_column_v) { + } else if constexpr (o2::soa::is_index_t::value) { return getSingleRowIndexData(rowIterator, globalIndex); + } else { + static_assert(!sizeof(decayed*), "Unrecognized column kind"); // A trick to delay static_assert until we actually instantiate this branch } } template -std::tuple getRowData(arrow::Table* table, T rowIterator, uint64_t ci, uint64_t ai, uint64_t globalIndex) +std::tuple getRowData(arrow::Table* table, T rowIterator, uint64_t ci = -1, uint64_t ai = -1, uint64_t globalIndex = -1) { return std::make_tuple(getSingleRowData(table, rowIterator, ci, ai, globalIndex)...); } diff --git a/Framework/Core/include/Framework/ASoAHelpers.h b/Framework/Core/include/Framework/ASoAHelpers.h index 19b1cd1408865..40ccda09fc4b3 100644 --- a/Framework/Core/include/Framework/ASoAHelpers.h +++ b/Framework/Core/include/Framework/ASoAHelpers.h @@ -125,9 +125,8 @@ std::vector groupTable(const T& table, const BP& binningPol selInd = selectedRows[ind]; } - auto rowData = o2::soa::row_helpers::getRowData(arrowTable, rowIterator, ci, ai, ind); - - int val = binningPolicy.getBin(rowData); + auto values = binningPolicy.getBinningValues(rowIterator, arrowTable, ci, ai, ind); + auto val = binningPolicy.getBin(values); if (val != outsider) { groupedIndices.emplace_back(val, ind); } diff --git a/Framework/Core/include/Framework/BinningPolicy.h b/Framework/Core/include/Framework/BinningPolicy.h index eb398e44b15cd..d4da8db8d627f 100644 --- a/Framework/Core/include/Framework/BinningPolicy.h +++ b/Framework/Core/include/Framework/BinningPolicy.h @@ -13,7 +13,6 @@ #define FRAMEWORK_BINNINGPOLICY_H #include "Framework/HistogramSpec.h" // only for VARIABLE_WIDTH -#include "Framework/ASoAHelpers.h" #include "Framework/Pack.h" #include "Framework/ArrowTypes.h" #include @@ -21,31 +20,50 @@ namespace o2::framework { -template -struct BinningPolicy { - BinningPolicy(std::array, sizeof...(Cs) + 1> bins, bool ignoreOverflows = true) : mBins(bins), mIgnoreOverflows(ignoreOverflows) +namespace binning_helpers +{ +void expandConstantBinning(std::vector const& bins, std::vector& expanded) +{ + if (bins[0] != VARIABLE_WIDTH) { + int nBins = static_cast(bins[0]); + expanded.clear(); + expanded.resize(nBins + 2); + expanded[0] = VARIABLE_WIDTH; + for (int i = 0; i <= nBins; i++) { + expanded[i + 1] = bins[1] + i * (bins[2] - bins[1]) / nBins; + } + } +} +} // namespace binning_helpers + +template +struct BinningPolicyBase { + BinningPolicyBase(std::array, N> bins, bool ignoreOverflows = true) : mBins(bins), mIgnoreOverflows(ignoreOverflows) { - static_assert(sizeof...(Cs) < 3, "No default binning for more than 3 columns, you need to implement a binning class yourself"); - for (int i = 0; i < sizeof...(Cs) + 1; i++) { - expandConstantBinning(bins[i], i); + static_assert(N <= 3, "No default binning for more than 3 columns, you need to implement a binning class yourself"); + for (int i = 0; i < N; i++) { + binning_helpers::expandConstantBinning(bins[i], mBins[i]); } } - int getBin(std::tuple const& data) const + template + int getBin(std::tuple const& data) const { + static_assert(sizeof...(Ts) == N, "There must be the same number of binning axes and data values/columns"); + unsigned int i = 2, j = 2, k = 2; if (this->mIgnoreOverflows) { // underflow - if (std::get<0>(data) < this->mBins[0][1]) { // xBins[0] is a dummy VARIABLE_WIDTH + if (std::get<0>(data) < this->mBins[0][1]) { // mBins[0][0] is a dummy VARIABLE_WIDTH return -1; } - if constexpr (sizeof...(Cs) > 0) { - if (std::get<1>(data) < this->mBins[1][1]) { // this->mBins[1][0] is a dummy VARIABLE_WIDTH + if constexpr (N > 1) { + if (std::get<1>(data) < this->mBins[1][1]) { // mBins[1][0] is a dummy VARIABLE_WIDTH return -1; } } - if constexpr (sizeof...(Cs) > 1) { - if (std::get<2>(data) < this->mBins[2][1]) { // this->mBins[2][0] is a dummy VARIABLE_WIDTH + if constexpr (N > 2) { + if (std::get<2>(data) < this->mBins[2][1]) { // mBins[2][0] is a dummy VARIABLE_WIDTH return -1; } } @@ -58,11 +76,11 @@ struct BinningPolicy { for (; i < this->mBins[0].size(); i++) { if (std::get<0>(data) < this->mBins[0][i]) { - if constexpr (sizeof...(Cs) > 0) { + if constexpr (N > 1) { for (; j < this->mBins[1].size(); j++) { if (std::get<1>(data) < this->mBins[1][j]) { - if constexpr (sizeof...(Cs) > 1) { + if constexpr (N > 2) { for (; k < this->mBins[2].size(); k++) { if (std::get<2>(data) < this->mBins[2][k]) { return getBinAt(i, j, k); @@ -73,7 +91,7 @@ struct BinningPolicy { } } - // overflow for this->mBins[2] only + // overflow for mBins[2] only return getBinAt(i, j, k); } } @@ -82,8 +100,8 @@ struct BinningPolicy { return -1; } - // overflow for this->mBins[1] only - if constexpr (sizeof...(Cs) > 1) { + // overflow for mBins[1] only + if constexpr (N > 2) { for (k = 2; k < this->mBins[2].size(); k++) { if (std::get<2>(data) < this->mBins[2][k]) { return getBinAt(i, j, k); @@ -92,7 +110,7 @@ struct BinningPolicy { } } - // overflow for this->mBins[2] and this->mBins[1] + // overflow for mBins[2] and mBins[1] return getBinAt(i, j, k); } } @@ -102,12 +120,12 @@ struct BinningPolicy { return -1; } - // overflow for this->mBins[0] only - if constexpr (sizeof...(Cs) > 0) { + // overflow for mBins[0] only + if constexpr (N > 1) { for (j = 2; j < this->mBins[1].size(); j++) { if (std::get<1>(data) < this->mBins[1][j]) { - if constexpr (sizeof...(Cs) > 1) { + if constexpr (N > 2) { for (k = 2; k < this->mBins[2].size(); k++) { if (std::get<2>(data) < this->mBins[2][k]) { return getBinAt(i, j, k); @@ -115,14 +133,14 @@ struct BinningPolicy { } } - // overflow for this->mBins[0] and this->mBins[2] + // overflow for mBins[0] and mBins[2] return getBinAt(i, j, k); } } } - // overflow for this->mBins[0] and this->mBins[1] - if constexpr (sizeof...(Cs) > 1) { + // overflow for mBins[0] and mBins[1] + if constexpr (N > 2) { for (k = 2; k < this->mBins[2].size(); k++) { if (std::get<2>(data) < this->mBins[2][k]) { return getBinAt(i, j, k); @@ -137,43 +155,44 @@ struct BinningPolicy { // Note: Overflow / underflow bin -1 is not included int getXBinsCount() const { - return this->mBins[0].size() - 1 - getOverflowShift(); + return getBinsCount(mBins[0]); } // Note: Overflow / underflow bin -1 is not included int getYBinsCount() const { - if constexpr (sizeof...(Cs) == 0) { + if constexpr (N == 1) { return 0; } - return this->mBins[1].size() - 1 - getOverflowShift(); + return getBinsCount(mBins[1]); } // Note: Overflow / underflow bin -1 is not included int getZBinsCount() const { - if constexpr (sizeof...(Cs) < 2) { + if constexpr (N < 3) { return 0; } - return this->mBins[2].size() - 1 - getOverflowShift(); + return getBinsCount(mBins[2]); } // Note: Overflow / underflow bin -1 is not included int getAllBinsCount() const { - if constexpr (sizeof...(Cs) == 0) { + if constexpr (N == 1) { return getXBinsCount(); } - if constexpr (sizeof...(Cs) == 1) { + if constexpr (N == 2) { return getXBinsCount() * getYBinsCount(); } - if constexpr (sizeof...(Cs) == 2) { + if constexpr (N == 2) { return getXBinsCount() * getYBinsCount() * getZBinsCount(); } return -1; } - using persistent_columns_t = framework::selected_pack; + std::array, N> mBins; + bool mIgnoreOverflows; private: // We substract 1 to account for VARIABLE_WIDTH in the bins vector @@ -186,12 +205,12 @@ struct BinningPolicy { unsigned int j = jRaw - 1 - shiftBinsWithoutOverflow; unsigned int k = kRaw - 1 - shiftBinsWithoutOverflow; auto xBinsCount = getXBinsCount(); - if constexpr (sizeof...(Cs) == 0) { + if constexpr (N == 1) { return i; - } else if constexpr (sizeof...(Cs) == 1) { + } else if constexpr (N == 2) { return i + j * xBinsCount; - } else if constexpr (sizeof...(Cs) == 2) { - return i + j * xBinsCount + k * xBinsCount * (this->mBins[1].size() - 1 - shiftBinsWithoutOverflow); + } else if constexpr (N == 3) { + return i + j * xBinsCount + k * xBinsCount * getYBinsCount(); } else { return -1; } @@ -202,28 +221,105 @@ struct BinningPolicy { return mIgnoreOverflows ? 1 : -1; } - void expandConstantBinning(std::vector const& bins, int ind) + // Note: Overflow / underflow bin -1 is not included + int getBinsCount(std::vector const& bins) const + { + return bins.size() - 1 - getOverflowShift(); + } +}; + +template +struct FlexibleBinningPolicy; + +template +struct FlexibleBinningPolicy, Ts...> : BinningPolicyBase { + FlexibleBinningPolicy(std::tuple const& lambdaPtrs, std::array, sizeof...(Ts)> bins, bool ignoreOverflows = true) : BinningPolicyBase(bins, ignoreOverflows), mBinningFunctions{lambdaPtrs} { - if (bins[0] != VARIABLE_WIDTH) { - int nBins = static_cast(bins[0]); - this->mBins[ind].clear(); - this->mBins[ind].resize(nBins + 2); - this->mBins[ind][0] = VARIABLE_WIDTH; - for (int i = 0; i <= nBins; i++) { - this->mBins[ind][i + 1] = bins[1] + i * (bins[2] - bins[1]) / nBins; + } + + template + auto getBinningValue(T& rowIterator, arrow::Table* table, uint64_t ci = -1, uint64_t ai = -1, uint64_t globalIndex = -1) const + { + if constexpr (has_type_v>) { + if (globalIndex != -1) { + rowIterator.setCursor(globalIndex); } + return std::get(mBinningFunctions)(rowIterator); + } else { + return soa::row_helpers::getSingleRowData(table, rowIterator, ci, ai, globalIndex); } } - std::array, sizeof...(Cs) + 1> mBins; - bool mIgnoreOverflows; + template + auto getBinningValues(T& rowIterator, arrow::Table* table, uint64_t ci = -1, uint64_t ai = -1, uint64_t globalIndex = -1) const + { + return std::make_tuple(getBinningValue(rowIterator, table, ci, ai, globalIndex)...); + } + + template + auto getBinningValues(typename T::iterator rowIterator, T& table, uint64_t ci = -1, uint64_t ai = -1, uint64_t globalIndex = -1) const + { + return getBinningValues(rowIterator, table.asArrowTable().get(), ci, ai, globalIndex); + } + + template + int getBin(std::tuple const& data) const + { + return BinningPolicyBase::template getBin(data); + } + + using persistent_columns_t = framework::selected_pack; + + private: + std::tuple mBinningFunctions; }; +template +struct ColumnBinningPolicy : BinningPolicyBase { + ColumnBinningPolicy(std::array, sizeof...(Ts)> bins, bool ignoreOverflows = true) : BinningPolicyBase(bins, ignoreOverflows) + { + } + + template + auto getBinningValues(T& rowIterator, arrow::Table* table, uint64_t ci = -1, uint64_t ai = -1, uint64_t globalIndex = -1) const + { + return std::make_tuple(soa::row_helpers::getSingleRowData(table, rowIterator, ci, ai, globalIndex)...); + } + + template + auto getBinningValues(typename T::iterator rowIterator, T& table, uint64_t ci = -1, uint64_t ai = -1, uint64_t globalIndex = -1) const + { + return getBinningValues(rowIterator, table.asArrowTable().get(), ci, ai, globalIndex); + } + + int getBin(std::tuple const& data) const + { + return BinningPolicyBase::template getBin(data); + } + + using persistent_columns_t = framework::selected_pack; +}; + +template +using BinningPolicy = ColumnBinningPolicy; + template struct NoBinningPolicy { // Just take the bin number from the column data NoBinningPolicy() = default; + template + auto getBinningValues(T& rowIterator, arrow::Table* table, uint64_t ci = -1, uint64_t ai = -1, uint64_t globalIndex = -1) const + { + return std::make_tuple(soa::row_helpers::getSingleRowData(table, rowIterator, ci, ai, globalIndex)); + } + + template + auto getBinningValues(typename T::iterator rowIterator, T& table, uint64_t ci = -1, uint64_t ai = -1, uint64_t globalIndex = -1) const + { + return getBinningValues(rowIterator, table.asArrowTable().get(), ci, ai, globalIndex); + } + int getBin(std::tuple const& data) const { return std::get<0>(data);