From 631ae2dba7a09e348ac10dc9f524bdd3e4f639cb Mon Sep 17 00:00:00 2001
From: Maja Kabus <mkabus@cern.ch>
Date: Wed, 29 Jun 2022 11:22:54 +0200
Subject: [PATCH 1/6] Improve tests for persistent and index columns

---
 Framework/Core/include/Framework/ASoA.h | 43 +++++++++++++++++++------
 1 file changed, 34 insertions(+), 9 deletions(-)

diff --git a/Framework/Core/include/Framework/ASoA.h b/Framework/Core/include/Framework/ASoA.h
index 9030cf5fe7eec..542c6a88fb3db 100644
--- a/Framework/Core/include/Framework/ASoA.h
+++ b/Framework/Core/include/Framework/ASoA.h
@@ -541,8 +541,26 @@ struct Index : o2::soa::IndexColumn<Index<START, END>> {
 template <typename T>
 using is_dynamic_t = framework::is_specialization<typename T::base, DynamicColumn>;
 
+// template <typename T>
+// using is_persistent_t = typename std::decay_t<T>::persistent::type;
+
+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 <class T>
+typename T::persistent test(int T::persistent::*);
+
+template <class>
+std::false_type test(...);
+} // namespace persistent_type_helper
+
 template <typename T>
-using is_persistent_t = typename std::decay_t<T>::persistent::type;
+using is_persistent_t = decltype(persistent_type_helper::test<T>(nullptr));
+
+template <typename T>
+using is_persistent_v = typename is_persistent_t<T>::value;
 
 template <typename T>
 using is_external_index_t = typename std::conditional<is_index_column_v<T>, std::true_type, std::false_type>::type;
@@ -1112,10 +1130,12 @@ class Table
     auto getId() const
     {
       using decayed = std::decay_t<TI>;
-      if constexpr (framework::has_type_v<decayed, bindings_pack_t>) {
+      if constexpr (framework::has_type_v<decayed, bindings_pack_t>) { // index to another table
         constexpr auto idx = framework::has_type_at_v<decayed>(bindings_pack_t{});
         return framework::pack_element_t<idx, external_index_columns_t>::getId();
-      } else if constexpr (std::is_same_v<decayed, Parent>) {
+      } else if constexpr (std::is_same_v<decayed, Parent>) { // self index
+        return this->globalIndex();
+      } else if constexpr (is_index_t<decayed>::value && decayed::mLabel == "Index") { // soa::Index<>
         return this->globalIndex();
       } else {
         return static_cast<int32_t>(-1);
@@ -1505,14 +1525,14 @@ namespace row_helpers
 template <typename... Cs>
 std::array<arrow::ChunkedArray*, sizeof...(Cs)> getArrowColumns(arrow::Table* table, framework::pack<Cs...>)
 {
-  static_assert(std::conjunction_v<typename Cs::persistent...>, "BinningPolicy: only persistent columns accepted (not dynamic and not index ones");
+  static_assert(std::conjunction_v<typename Cs::persistent...>, "Arrow columns: only persistent columns accepted (not dynamic and not index ones");
   return std::array<arrow::ChunkedArray*, sizeof...(Cs)>{o2::soa::getIndexFromLabel(table, Cs::columnLabel())...};
 }
 
 template <typename... Cs>
 std::array<std::shared_ptr<arrow::Array>, sizeof...(Cs)> getChunks(arrow::Table* table, framework::pack<Cs...>, uint64_t ci)
 {
-  static_assert(std::conjunction_v<typename Cs::persistent...>, "BinningPolicy: only persistent columns accepted (not dynamic and not index ones");
+  static_assert(std::conjunction_v<typename Cs::persistent...>, "Arrow chunks: only persistent columns accepted (not dynamic and not index ones");
   return std::array<std::shared_ptr<arrow::Array>, sizeof...(Cs)>{o2::soa::getIndexFromLabel(table, Cs::columnLabel())->chunk(ci)...};
 }
 
@@ -1541,11 +1561,16 @@ typename C::type getSingleRowData(arrow::Table* table, T& rowIterator, uint64_t
 {
   using decayed = std::decay_t<C>;
   if constexpr (decayed::persistent::value) {
-    return getSingleRowPersistentData<C>(table, ci, ai);
+    auto val = getSingleRowPersistentData<C>(table, ci, ai);
+    return val;
   } else if constexpr (o2::soa::is_dynamic_t<decayed>()) {
-    return getSingleRowDynamicData<T, C>(rowIterator, globalIndex);
-  } else if constexpr (o2::soa::is_index_column_v<decayed>) {
-    return getSingleRowIndexData<T, C>(rowIterator, globalIndex);
+    auto val = getSingleRowDynamicData<T, C>(rowIterator, globalIndex);
+    return val;
+  } else if constexpr (o2::soa::is_index_t<decayed>::value) {
+    auto val = getSingleRowIndexData<T, C>(rowIterator, globalIndex);
+    return val;
+  } else {
+    static_assert(!sizeof(decayed*), "Unrecognized column kind"); // A trick to delay static_assert until we actually instantiate this branch
   }
 }
 

From 47c5d624229aa54596ef228740b8143dc2200ea4 Mon Sep 17 00:00:00 2001
From: Maja Kabus <mkabus@cern.ch>
Date: Wed, 29 Jun 2022 11:29:27 +0200
Subject: [PATCH 2/6] Refactor binning policy to accept lambdas

---
 Framework/Core/include/Framework/ASoA.h       |   2 +-
 .../Core/include/Framework/ASoAHelpers.h      |  11 +-
 .../Core/include/Framework/BinningPolicy.h    | 308 ++++++++++++------
 Framework/Core/test/benchmark_EventMixing.cxx |   4 +-
 Framework/Core/test/test_ASoAHelpers.cxx      |  16 +-
 5 files changed, 218 insertions(+), 123 deletions(-)

diff --git a/Framework/Core/include/Framework/ASoA.h b/Framework/Core/include/Framework/ASoA.h
index 542c6a88fb3db..b4308f3334b17 100644
--- a/Framework/Core/include/Framework/ASoA.h
+++ b/Framework/Core/include/Framework/ASoA.h
@@ -1575,7 +1575,7 @@ typename C::type getSingleRowData(arrow::Table* table, T& rowIterator, uint64_t
 }
 
 template <typename T, typename... Cs>
-std::tuple<typename Cs::type...> getRowData(arrow::Table* table, T rowIterator, uint64_t ci, uint64_t ai, uint64_t globalIndex)
+std::tuple<typename Cs::type...> getRowData(arrow::Table* table, T rowIterator, framework::pack<Cs...>, uint64_t ci, uint64_t ai, uint64_t globalIndex)
 {
   return std::make_tuple(getSingleRowData<T, Cs>(table, rowIterator, ci, ai, globalIndex)...);
 }
diff --git a/Framework/Core/include/Framework/ASoAHelpers.h b/Framework/Core/include/Framework/ASoAHelpers.h
index 19b1cd1408865..01ddaf1989246 100644
--- a/Framework/Core/include/Framework/ASoAHelpers.h
+++ b/Framework/Core/include/Framework/ASoAHelpers.h
@@ -72,8 +72,8 @@ inline bool diffCategory(BinningIndex const& a, BinningIndex const& b)
   return a.bin >= b.bin;
 }
 
-template <template <typename... Cs> typename BP, typename T, typename... Cs>
-std::vector<BinningIndex> groupTable(const T& table, const BP<Cs...>& binningPolicy, int minCatSize, int outsider)
+template <template <typename C, typename... Cs> typename BP, typename T, typename C, typename... Cs>
+std::vector<BinningIndex> groupTable(const T& table, const BP<C, Cs...>& binningPolicy, int minCatSize, int outsider)
 {
   arrow::Table* arrowTable = table.asArrowTable().get();
   auto rowIterator = table.begin();
@@ -92,7 +92,7 @@ std::vector<BinningIndex> groupTable(const T& table, const BP<Cs...>& binningPol
     selectedRows = table.getSelectedRows(); // vector<int64_t>
   }
 
-  auto persistentColumns = typename BP<Cs...>::persistent_columns_t{};
+  auto persistentColumns = typename BP<C, Cs...>::persistent_columns_t{};
   constexpr auto persistentColumnsCount = pack_size(persistentColumns);
   auto arrowColumns = o2::soa::row_helpers::getArrowColumns(arrowTable, persistentColumns);
   auto chunksCount = arrowColumns[0]->num_chunks();
@@ -125,9 +125,8 @@ std::vector<BinningIndex> groupTable(const T& table, const BP<Cs...>& binningPol
         selInd = selectedRows[ind];
       }
 
-      auto rowData = o2::soa::row_helpers::getRowData<decltype(rowIterator), Cs...>(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..8dda91c664144 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 <optional>
@@ -21,110 +20,186 @@
 namespace o2::framework
 {
 
-template <typename C, typename... Cs>
-struct BinningPolicy {
-  BinningPolicy(std::array<std::vector<double>, sizeof...(Cs) + 1> bins, bool ignoreOverflows = true) : mBins(bins), mIgnoreOverflows(ignoreOverflows)
+namespace binning_helpers
+{
+void expandConstantBinning(std::vector<double> const& bins, std::vector<double>& expanded)
+{
+  if (bins[0] != VARIABLE_WIDTH) {
+    int nBins = static_cast<int>(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;
+    }
+  }
+}
+
+int getOverflowShift(bool ignoreOverflows)
+{
+  return ignoreOverflows ? -1 : 1;
+}
+
+// Note: Overflow / underflow bin -1 is not included
+int getBinsCount(std::vector<double> const& bins, bool ignoreOverflows)
+{
+  if (bins.size() == 0) {
+    return 0;
+  }
+  if (bins[0] == VARIABLE_WIDTH) {
+    return bins.size() - 1 + getOverflowShift(ignoreOverflows);
+  }
+  return bins[0] + getOverflowShift(ignoreOverflows);
+}
+
+} // namespace binning_helpers
+
+template <std::size_t N>
+struct BinningPolicyBase {
+  BinningPolicyBase(std::array<std::vector<double>, 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<typename C::type, typename Cs::type...> const& data) const
+  // Note: Overflow / underflow bin -1 is not included
+  int getAllBinsCount() const
+  {
+    if constexpr (N == 1) {
+      return binning_helpers::getBinsCount(mBins[0], mIgnoreOverflows);
+    }
+    if constexpr (N == 2) {
+      return binning_helpers::getBinsCount(mBins[0], mIgnoreOverflows) * binning_helpers::getBinsCount(mBins[1], mIgnoreOverflows);
+    }
+    if constexpr (N == 2) {
+      return binning_helpers::getBinsCount(mBins[0], mIgnoreOverflows) * binning_helpers::getBinsCount(mBins[1], mIgnoreOverflows) * binning_helpers::getBinsCount(mBins[2], mIgnoreOverflows);
+    }
+    return -1;
+  }
+
+  // Note: Overflow / underflow bin -1 is not included
+  int getXBinsCount() const
   {
+    return binning_helpers::getBinsCount(mBins[0], mIgnoreOverflows);
+  }
+
+  // Note: Overflow / underflow bin -1 is not included
+  int getYBinsCount() const
+  {
+    if constexpr (N == 1) {
+      return 0;
+    }
+    return binning_helpers::getBinsCount(mBins[1], mIgnoreOverflows);
+  }
+
+  // Note: Overflow / underflow bin -1 is not included
+  int getZBinsCount() const
+  {
+    if constexpr (N < 3) {
+      return 0;
+    }
+    return binning_helpers::getBinsCount(mBins[2], mIgnoreOverflows);
+  }
+
+  template <typename... Ts>
+  int getBin(std::tuple<Ts...> 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) {
+    if (mIgnoreOverflows) {
       // underflow
-      if (std::get<0>(data) < this->mBins[0][1]) { // xBins[0] is a dummy VARIABLE_WIDTH
+      if (std::get<0>(data) < 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) < 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) < mBins[2][1]) { // mBins[2][0] is a dummy VARIABLE_WIDTH
           return -1;
         }
+      } else {
+        i = 1;
+        j = 1;
+        k = 1;
       }
-    } else {
-      i = 1;
-      j = 1;
-      k = 1;
     }
 
-    for (; i < this->mBins[0].size(); i++) {
-      if (std::get<0>(data) < this->mBins[0][i]) {
+    for (; i < mBins[0].size(); i++) {
+      if (std::get<0>(data) < mBins[0][i]) {
 
-        if constexpr (sizeof...(Cs) > 0) {
-          for (; j < this->mBins[1].size(); j++) {
-            if (std::get<1>(data) < this->mBins[1][j]) {
+        if constexpr (N > 1) {
+          for (; j < mBins[1].size(); j++) {
+            if (std::get<1>(data) < mBins[1][j]) {
 
-              if constexpr (sizeof...(Cs) > 1) {
-                for (; k < this->mBins[2].size(); k++) {
-                  if (std::get<2>(data) < this->mBins[2][k]) {
+              if constexpr (N > 2) {
+                for (; k < mBins[2].size(); k++) {
+                  if (std::get<2>(data) < mBins[2][k]) {
                     return getBinAt(i, j, k);
                   }
                 }
-                if (this->mIgnoreOverflows) {
+                if (mIgnoreOverflows) {
                   return -1;
                 }
               }
 
-              // overflow for this->mBins[2] only
+              // overflow for mBins[2] only
               return getBinAt(i, j, k);
             }
           }
 
-          if (this->mIgnoreOverflows) {
+          if (mIgnoreOverflows) {
             return -1;
           }
 
-          // overflow for this->mBins[1] only
-          if constexpr (sizeof...(Cs) > 1) {
-            for (k = 2; k < this->mBins[2].size(); k++) {
-              if (std::get<2>(data) < this->mBins[2][k]) {
+          // overflow for mBins[1] only
+          if constexpr (N > 2) {
+            for (k = 2; k < mBins[2].size(); k++) {
+              if (std::get<2>(data) < mBins[2][k]) {
                 return getBinAt(i, j, k);
               }
             }
           }
         }
 
-        // overflow for this->mBins[2] and this->mBins[1]
+        // overflow for mBins[2] and mBins[1]
         return getBinAt(i, j, k);
       }
     }
 
-    if (this->mIgnoreOverflows) {
+    if (mIgnoreOverflows) {
       // overflow
       return -1;
     }
 
-    // overflow for this->mBins[0] only
-    if constexpr (sizeof...(Cs) > 0) {
-      for (j = 2; j < this->mBins[1].size(); j++) {
-        if (std::get<1>(data) < this->mBins[1][j]) {
+    // overflow for mBins[0] only
+    if constexpr (N > 1) {
+      for (j = 2; j < mBins[1].size(); j++) {
+        if (std::get<1>(data) < mBins[1][j]) {
 
-          if constexpr (sizeof...(Cs) > 1) {
-            for (k = 2; k < this->mBins[2].size(); k++) {
-              if (std::get<2>(data) < this->mBins[2][k]) {
+          if constexpr (N > 2) {
+            for (k = 2; k < mBins[2].size(); k++) {
+              if (std::get<2>(data) < mBins[2][k]) {
                 return getBinAt(i, j, k);
               }
             }
           }
 
-          // 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) {
-      for (k = 2; k < this->mBins[2].size(); k++) {
-        if (std::get<2>(data) < this->mBins[2][k]) {
+    // overflow for mBins[0] and mBins[1]
+    if constexpr (N > 2) {
+      for (k = 2; k < mBins[2].size(); k++) {
+        if (std::get<2>(data) < mBins[2][k]) {
           return getBinAt(i, j, k);
         }
       }
@@ -134,89 +209,104 @@ struct BinningPolicy {
     return getBinAt(i, j, k);
   }
 
-  // Note: Overflow / underflow bin -1 is not included
-  int getXBinsCount() const
-  {
-    return this->mBins[0].size() - 1 - getOverflowShift();
-  }
+  std::array<std::vector<double>, N> mBins;
+  bool mIgnoreOverflows;
 
-  // Note: Overflow / underflow bin -1 is not included
-  int getYBinsCount() const
+ private:
+  // We substract 1 to account for VARIABLE_WIDTH in the bins vector
+  // We substract second 1 if we omit values below minima (underflow, mapped to -1)
+  // Otherwise we add 1 and we get the number of bins including those below and over the outer edges
+  int getBinAt(unsigned int iRaw, unsigned int jRaw, unsigned int kRaw) const
   {
-    if constexpr (sizeof...(Cs) == 0) {
-      return 0;
+    int shiftBinsWithoutOverflow = binning_helpers::getOverflowShift(mIgnoreOverflows);
+    unsigned int i = iRaw - 1 + shiftBinsWithoutOverflow;
+    unsigned int j = jRaw - 1 + shiftBinsWithoutOverflow;
+    unsigned int k = kRaw - 1 + shiftBinsWithoutOverflow;
+    auto xBinsCount = binning_helpers::getBinsCount(mBins[0], mIgnoreOverflows);
+    if constexpr (N == 1) {
+      return i;
+    } else if constexpr (N == 2) {
+      return i + j * xBinsCount;
+    } else if constexpr (N == 3) {
+      return i + j * xBinsCount + k * xBinsCount * binning_helpers::getBinsCount(mBins[1], mIgnoreOverflows);
+    } else {
+      return -1;
     }
-    return this->mBins[1].size() - 1 - getOverflowShift();
   }
+};
 
-  // Note: Overflow / underflow bin -1 is not included
-  int getZBinsCount() const
+template <typename, typename...>
+struct BinningPolicy;
+
+template <typename... Ts, typename... Ls>
+struct BinningPolicy<std::tuple<Ls...>, Ts...> : BinningPolicyBase<sizeof...(Ts)> {
+  BinningPolicy(std::tuple<Ls...> const& lambdaPtrs, std::array<std::vector<double>, sizeof...(Ts)> bins, bool ignoreOverflows = true) : BinningPolicyBase<sizeof...(Ts)>(bins, ignoreOverflows), mBinningFunctions{lambdaPtrs}
   {
-    if constexpr (sizeof...(Cs) < 2) {
-      return 0;
-    }
-    return this->mBins[2].size() - 1 - getOverflowShift();
   }
 
-  // Note: Overflow / underflow bin -1 is not included
-  int getAllBinsCount() const
+  template <typename T, typename T2>
+  auto getBinningValue(T& rowIterator, arrow::Table* table, uint64_t globalIndex = -1, uint64_t ci = -1, uint64_t ai = -1) const
   {
-    if constexpr (sizeof...(Cs) == 0) {
-      return getXBinsCount();
-    }
-    if constexpr (sizeof...(Cs) == 1) {
-      return getXBinsCount() * getYBinsCount();
+    using decayed = std::decay_t<T2>;
+    if (globalIndex == -1) {
+      globalIndex = *(std::get<0>(rowIterator.getIndices()));
     }
-    if constexpr (sizeof...(Cs) == 2) {
-      return getXBinsCount() * getYBinsCount() * getZBinsCount();
+
+    if constexpr (has_type_v<T2, pack<Ls...>>) {
+      rowIterator.setCursor(globalIndex);
+      return std::get<T2>(mBinningFunctions)(rowIterator);
+    } else {
+      if (ci == -1 && ai == -1) {
+        auto colIterator = static_cast<decayed>(rowIterator).mColumnIterator;
+        ci = colIterator.mCurrentChunk;
+        ai = *(colIterator.mCurrentPos) - colIterator.mFirstIndex;
+      }
+      return soa::row_helpers::getSingleRowData<T, T2>(table, rowIterator, ci, ai, globalIndex);
     }
-    return -1;
   }
 
-  using persistent_columns_t = framework::selected_pack<o2::soa::is_persistent_t, C, Cs...>;
+  template <typename T>
+  auto getBinningValues(T& rowIterator, arrow::Table* table, uint64_t globalIndex = -1, uint64_t ci = -1, uint64_t ai = -1) const
+  {
+    return std::make_tuple(getBinningValue<T, Ts>(rowIterator, table, globalIndex, ci, ai)...);
+  }
+
+  template <typename T>
+  auto getBinningValues(typename T::iterator rowIterator, T& table, uint64_t globalIndex = -1, uint64_t ci = -1, uint64_t ai = -1) const
+  {
+    return getBinningValues(rowIterator, table.asArrowTable().get(), globalIndex, ci, ai);
+  }
+
+  template <typename... T2s>
+  int getBin(std::tuple<T2s...> const& data) const
+  {
+    return BinningPolicyBase<sizeof...(Ts)>::template getBin<T2s...>(data);
+  }
+
+  using persistent_columns_t = framework::selected_pack<o2::soa::is_persistent_t, Ts...>;
 
  private:
-  // We substract 1 to account for VARIABLE_WIDTH in the bins vector
-  // We substract second 1 if we omit values below minima (underflow, mapped to -1)
-  // Otherwise we add 1 and we get the number of bins including those below and over the outer edges
-  int getBinAt(unsigned int iRaw, unsigned int jRaw, unsigned int kRaw) const
+  std::tuple<Ls...> mBinningFunctions;
+};
+
+template <typename... Ts>
+struct ColumnBinningPolicy : BinningPolicyBase<sizeof...(Ts)> {
+  ColumnBinningPolicy(std::array<std::vector<double>, sizeof...(Ts)> bins, bool ignoreOverflows = true) : BinningPolicyBase<sizeof...(Ts)>(bins, ignoreOverflows)
   {
-    int shiftBinsWithoutOverflow = getOverflowShift();
-    unsigned int i = iRaw - 1 - shiftBinsWithoutOverflow;
-    unsigned int j = jRaw - 1 - shiftBinsWithoutOverflow;
-    unsigned int k = kRaw - 1 - shiftBinsWithoutOverflow;
-    auto xBinsCount = getXBinsCount();
-    if constexpr (sizeof...(Cs) == 0) {
-      return i;
-    } else if constexpr (sizeof...(Cs) == 1) {
-      return i + j * xBinsCount;
-    } else if constexpr (sizeof...(Cs) == 2) {
-      return i + j * xBinsCount + k * xBinsCount * (this->mBins[1].size() - 1 - shiftBinsWithoutOverflow);
-    } else {
-      return -1;
-    }
   }
 
-  int getOverflowShift() const
+  template <typename T>
+  auto getBinningValues(T& rowIterator, arrow::Table* table, uint64_t ci, uint64_t ai, uint64_t globalIndex) const
   {
-    return mIgnoreOverflows ? 1 : -1;
+    return std::make_tuple(soa::row_helpers::getSingleRowData<T, Ts>(table, rowIterator, ci, ai, globalIndex)...);
   }
 
-  void expandConstantBinning(std::vector<double> const& bins, int ind)
+  int getBin(std::tuple<typename Ts::type...> const& data) const
   {
-    if (bins[0] != VARIABLE_WIDTH) {
-      int nBins = static_cast<int>(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;
-      }
-    }
+    return BinningPolicyBase<sizeof...(Ts)>::template getBin<typename Ts::type...>(data);
   }
 
-  std::array<std::vector<double>, sizeof...(Cs) + 1> mBins;
-  bool mIgnoreOverflows;
+  using persistent_columns_t = framework::selected_pack<o2::soa::is_persistent_t, Ts...>;
 };
 
 template <typename C>
@@ -224,6 +314,12 @@ struct NoBinningPolicy {
   // Just take the bin number from the column data
   NoBinningPolicy() = default;
 
+  template <typename T>
+  auto getBinningValues(T& rowIterator, arrow::Table* table, uint64_t ci, uint64_t ai, uint64_t globalIndex) const
+  {
+    return std::make_tuple(soa::row_helpers::getSingleRowData<T, C>(table, rowIterator, ci, ai, globalIndex));
+  }
+
   int getBin(std::tuple<typename C::type> const& data) const
   {
     return std::get<0>(data);
diff --git a/Framework/Core/test/benchmark_EventMixing.cxx b/Framework/Core/test/benchmark_EventMixing.cxx
index ddf6dbc20bdde..71f8d6ab2e0c7 100644
--- a/Framework/Core/test/benchmark_EventMixing.cxx
+++ b/Framework/Core/test/benchmark_EventMixing.cxx
@@ -48,7 +48,7 @@ static void BM_EventMixingTraditional(benchmark::State& state)
 
   std::vector<double> xBins{VARIABLE_WIDTH, -0.064, -0.062, -0.060, 0.066, 0.068, 0.070, 0.072};
   std::vector<double> yBins{VARIABLE_WIDTH, -0.320, -0.301, -0.300, 0.330, 0.340, 0.350, 0.360};
-  using BinningType = BinningPolicy<o2::aod::collision::PosX, o2::aod::collision::PosY>;
+  using BinningType = ColumnBinningPolicy<o2::aod::collision::PosX, o2::aod::collision::PosY>;
   BinningType binningOnPositions{{xBins, yBins}, true}; // true is for 'ignore overflows' (true by default)
 
   TableBuilder colBuilder, trackBuilder;
@@ -135,7 +135,7 @@ static void BM_EventMixingCombinations(benchmark::State& state)
 
   std::vector<double> xBins{VARIABLE_WIDTH, -0.064, -0.062, -0.060, 0.066, 0.068, 0.070, 0.072};
   std::vector<double> yBins{VARIABLE_WIDTH, -0.320, -0.301, -0.300, 0.330, 0.340, 0.350, 0.360};
-  using BinningType = BinningPolicy<o2::aod::collision::PosX, o2::aod::collision::PosY>;
+  using BinningType = ColumnBinningPolicy<o2::aod::collision::PosX, o2::aod::collision::PosY>;
   BinningType binningOnPositions{{xBins, yBins}, true}; // true is for 'ignore overflows' (true by default)
 
   TableBuilder colBuilder, trackBuilder;
diff --git a/Framework/Core/test/test_ASoAHelpers.cxx b/Framework/Core/test/test_ASoAHelpers.cxx
index 1e4c963c44011..936f38ce248ad 100644
--- a/Framework/Core/test/test_ASoAHelpers.cxx
+++ b/Framework/Core/test/test_ASoAHelpers.cxx
@@ -133,7 +133,7 @@ BOOST_AUTO_TEST_CASE(CombinationsGeneratorConstruction)
   std::vector<double> yBins{VARIABLE_WIDTH, 0, 5, 10, 20, 30, 40, 50, 101};
   std::vector<double> zBins{VARIABLE_WIDTH, -7.0, -5.0, -3.0, -1.0, 1.0, 3.0, 5.0, 7.0};
 
-  BinningPolicy<test::Y, test::FloatZ> pairBinning{{yBins, zBins}, false};
+  ColumnBinningPolicy<test::Y, test::FloatZ> pairBinning{{yBins, zBins}, false};
 
   CombinationsGenerator<CombinationsStrictlyUpperIndexPolicy<TestA, TestA>>::CombinationsIterator combIt(CombinationsStrictlyUpperIndexPolicy(testsA, testsA));
   BOOST_REQUIRE_NE(static_cast<test::X>(std::get<0>(*(combIt))).getIterator().mCurrentPos, nullptr);
@@ -293,8 +293,8 @@ BOOST_AUTO_TEST_CASE(CombinationsGeneratorConstruction)
 
   auto combBlock = combinations(CombinationsBlockStrictlyUpperSameIndexPolicy(pairBinning, 2, -1, testsA, testsA));
 
-  static_assert(std::is_same_v<decltype(combBlock.begin()), CombinationsGenerator<CombinationsBlockStrictlyUpperSameIndexPolicy<BinningPolicy<test::Y, test::FloatZ>, int32_t, TestA, TestA>>::CombinationsIterator>, "Wrong iterator type");
-  static_assert(std::is_same_v<decltype(*(combBlock.begin())), CombinationsBlockStrictlyUpperSameIndexPolicy<BinningPolicy<test::Y, test::FloatZ>, int32_t, TestA, TestA>::CombinationType&>, "Wrong combination type");
+  static_assert(std::is_same_v<decltype(combBlock.begin()), CombinationsGenerator<CombinationsBlockStrictlyUpperSameIndexPolicy<ColumnBinningPolicy<test::Y, test::FloatZ>, int32_t, TestA, TestA>>::CombinationsIterator>, "Wrong iterator type");
+  static_assert(std::is_same_v<decltype(*(combBlock.begin())), CombinationsBlockStrictlyUpperSameIndexPolicy<ColumnBinningPolicy<test::Y, test::FloatZ>, int32_t, TestA, TestA>::CombinationType&>, "Wrong combination type");
 
   auto beginBlockCombination = *(combBlock.begin());
   BOOST_REQUIRE_NE(static_cast<test::X>(std::get<0>(beginBlockCombination)).getIterator().mCurrentPos, nullptr);
@@ -948,8 +948,8 @@ BOOST_AUTO_TEST_CASE(BlockCombinations)
   std::vector<double> yBins{VARIABLE_WIDTH, 0, 5, 10, 20, 30, 40, 50, 101};
   std::vector<double> zBins{VARIABLE_WIDTH, -7.0, -5.0, -3.0, -1.0, 1.0, 3.0, 5.0, 7.0};
 
-  BinningPolicy<test::Y, test::FloatZ> pairBinning{{yBins, zBins}, false};
-  BinningPolicy<test::Y, test::FloatZ> pairBinningNoOverflows{{yBins, zBins}, true};
+  ColumnBinningPolicy<test::Y, test::FloatZ> pairBinning{{yBins, zBins}, false};
+  ColumnBinningPolicy<test::Y, test::FloatZ> pairBinningNoOverflows{{yBins, zBins}, true};
 
   // 2, 3, 5, 8, 9 have overflows in testA
   std::vector<std::tuple<int32_t, int32_t>> expectedFullPairsNoOverflows{
@@ -1168,8 +1168,8 @@ BOOST_AUTO_TEST_CASE(BlockCombinations)
   // [3, 5] [0, 4], [7], [1, 6], [2], [8, 9]
   // Assuming bins intervals: [ , )
   std::vector<double> xBins{VARIABLE_WIDTH, 0, 7, 10};
-  BinningPolicy<test::X, test::Y, test::FloatZ> tripleBinning{{xBins, yBins, zBins}, false};
-  BinningPolicy<test::X, test::Y, test::FloatZ> tripleBinningNoOverflows{{xBins, yBins, zBins}, true};
+  ColumnBinningPolicy<test::X, test::Y, test::FloatZ> tripleBinning{{xBins, yBins, zBins}, false};
+  ColumnBinningPolicy<test::X, test::Y, test::FloatZ> tripleBinningNoOverflows{{xBins, yBins, zBins}, true};
 
   // 2, 3, 5, 8, 9 have overflows in testA
   std::vector<std::tuple<int32_t, int32_t>> expectedFullPairsTripleBinningNoOverflows{
@@ -1276,7 +1276,7 @@ BOOST_AUTO_TEST_CASE(CombinationsHelpers)
   std::vector<double> yBins{VARIABLE_WIDTH, 0, 5, 10, 20, 30, 40, 50, 101};
   std::vector<double> zBins{VARIABLE_WIDTH, -7.0, -5.0, -3.0, -1.0, 1.0, 3.0, 5.0, 7.0};
 
-  BinningPolicy<test::Y, test::FloatZ> pairBinning{{yBins, zBins}, false};
+  ColumnBinningPolicy<test::Y, test::FloatZ> pairBinning{{yBins, zBins}, false};
 
   std::vector<std::tuple<int32_t, int32_t>> expectedStrictlyUpperPairs{
     {0, 4}, {0, 7}, {4, 7}, {1, 6}, {3, 5}, {2, 8}, {2, 9}, {8, 9}};

From bec24ae554cb32a2cb57fc2ca3869715e6e3a569 Mon Sep 17 00:00:00 2001
From: Maja Kabus <mkabus@cern.ch>
Date: Wed, 29 Jun 2022 11:42:17 +0200
Subject: [PATCH 3/6] Nitpick with bin calculation, simplify some codes

---
 .../Core/include/Framework/BinningPolicy.h    | 55 ++++++++-----------
 1 file changed, 24 insertions(+), 31 deletions(-)

diff --git a/Framework/Core/include/Framework/BinningPolicy.h b/Framework/Core/include/Framework/BinningPolicy.h
index 8dda91c664144..d5a7034961ded 100644
--- a/Framework/Core/include/Framework/BinningPolicy.h
+++ b/Framework/Core/include/Framework/BinningPolicy.h
@@ -34,24 +34,6 @@ void expandConstantBinning(std::vector<double> const& bins, std::vector<double>&
     }
   }
 }
-
-int getOverflowShift(bool ignoreOverflows)
-{
-  return ignoreOverflows ? -1 : 1;
-}
-
-// Note: Overflow / underflow bin -1 is not included
-int getBinsCount(std::vector<double> const& bins, bool ignoreOverflows)
-{
-  if (bins.size() == 0) {
-    return 0;
-  }
-  if (bins[0] == VARIABLE_WIDTH) {
-    return bins.size() - 1 + getOverflowShift(ignoreOverflows);
-  }
-  return bins[0] + getOverflowShift(ignoreOverflows);
-}
-
 } // namespace binning_helpers
 
 template <std::size_t N>
@@ -68,13 +50,13 @@ struct BinningPolicyBase {
   int getAllBinsCount() const
   {
     if constexpr (N == 1) {
-      return binning_helpers::getBinsCount(mBins[0], mIgnoreOverflows);
+      return getBinsCount(mBins[0]);
     }
     if constexpr (N == 2) {
-      return binning_helpers::getBinsCount(mBins[0], mIgnoreOverflows) * binning_helpers::getBinsCount(mBins[1], mIgnoreOverflows);
+      return getBinsCount(mBins[0]) * getBinsCount(mBins[1]);
     }
     if constexpr (N == 2) {
-      return binning_helpers::getBinsCount(mBins[0], mIgnoreOverflows) * binning_helpers::getBinsCount(mBins[1], mIgnoreOverflows) * binning_helpers::getBinsCount(mBins[2], mIgnoreOverflows);
+      return getBinsCount(mBins[0]) * getBinsCount(mBins[1]) * getBinsCount(mBins[2]);
     }
     return -1;
   }
@@ -82,7 +64,7 @@ struct BinningPolicyBase {
   // Note: Overflow / underflow bin -1 is not included
   int getXBinsCount() const
   {
-    return binning_helpers::getBinsCount(mBins[0], mIgnoreOverflows);
+    return getBinsCount(mBins[0]);
   }
 
   // Note: Overflow / underflow bin -1 is not included
@@ -91,7 +73,7 @@ struct BinningPolicyBase {
     if constexpr (N == 1) {
       return 0;
     }
-    return binning_helpers::getBinsCount(mBins[1], mIgnoreOverflows);
+    return getBinsCount(mBins[1]);
   }
 
   // Note: Overflow / underflow bin -1 is not included
@@ -100,7 +82,7 @@ struct BinningPolicyBase {
     if constexpr (N < 3) {
       return 0;
     }
-    return binning_helpers::getBinsCount(mBins[2], mIgnoreOverflows);
+    return getBinsCount(mBins[2]);
   }
 
   template <typename... Ts>
@@ -123,11 +105,11 @@ struct BinningPolicyBase {
         if (std::get<2>(data) < mBins[2][1]) { // mBins[2][0] is a dummy VARIABLE_WIDTH
           return -1;
         }
-      } else {
-        i = 1;
-        j = 1;
-        k = 1;
       }
+    } else {
+      i = 1;
+      j = 1;
+      k = 1;
     }
 
     for (; i < mBins[0].size(); i++) {
@@ -218,21 +200,32 @@ struct BinningPolicyBase {
   // Otherwise we add 1 and we get the number of bins including those below and over the outer edges
   int getBinAt(unsigned int iRaw, unsigned int jRaw, unsigned int kRaw) const
   {
-    int shiftBinsWithoutOverflow = binning_helpers::getOverflowShift(mIgnoreOverflows);
+    int shiftBinsWithoutOverflow = getOverflowShift();
     unsigned int i = iRaw - 1 + shiftBinsWithoutOverflow;
     unsigned int j = jRaw - 1 + shiftBinsWithoutOverflow;
     unsigned int k = kRaw - 1 + shiftBinsWithoutOverflow;
-    auto xBinsCount = binning_helpers::getBinsCount(mBins[0], mIgnoreOverflows);
+    auto xBinsCount = getXBinsCount();
     if constexpr (N == 1) {
       return i;
     } else if constexpr (N == 2) {
       return i + j * xBinsCount;
     } else if constexpr (N == 3) {
-      return i + j * xBinsCount + k * xBinsCount * binning_helpers::getBinsCount(mBins[1], mIgnoreOverflows);
+      return i + j * xBinsCount + k * xBinsCount * getYBinsCount();
     } else {
       return -1;
     }
   }
+
+  int getOverflowShift()
+  {
+    return mIgnoreOverflows ? -1 : 1;
+  }
+
+  // Note: Overflow / underflow bin -1 is not included
+  int getBinsCount(std::vector<double> const& bins)
+  {
+    return bins.size() - 1 + getOverflowShift();
+  }
 };
 
 template <typename, typename...>

From 1a4842c6db11e0a902202b97b6b9b6d6308ed622 Mon Sep 17 00:00:00 2001
From: Maja Kabus <mkabus@cern.ch>
Date: Wed, 29 Jun 2022 13:37:14 +0200
Subject: [PATCH 4/6] Make all methods const

---
 Framework/Core/include/Framework/BinningPolicy.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/Framework/Core/include/Framework/BinningPolicy.h b/Framework/Core/include/Framework/BinningPolicy.h
index d5a7034961ded..2769f3f399801 100644
--- a/Framework/Core/include/Framework/BinningPolicy.h
+++ b/Framework/Core/include/Framework/BinningPolicy.h
@@ -216,13 +216,13 @@ struct BinningPolicyBase {
     }
   }
 
-  int getOverflowShift()
+  int getOverflowShift() const
   {
     return mIgnoreOverflows ? -1 : 1;
   }
 
   // Note: Overflow / underflow bin -1 is not included
-  int getBinsCount(std::vector<double> const& bins)
+  int getBinsCount(std::vector<double> const& bins) const
   {
     return bins.size() - 1 + getOverflowShift();
   }

From 994c1af23b3c989d7b45bcde016250a141a7b4f7 Mon Sep 17 00:00:00 2001
From: Maja Kabus <mkabus@cern.ch>
Date: Wed, 29 Jun 2022 13:40:23 +0200
Subject: [PATCH 5/6] Remove commented lines

---
 Framework/Core/include/Framework/ASoA.h | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/Framework/Core/include/Framework/ASoA.h b/Framework/Core/include/Framework/ASoA.h
index b4308f3334b17..888e30f317570 100644
--- a/Framework/Core/include/Framework/ASoA.h
+++ b/Framework/Core/include/Framework/ASoA.h
@@ -541,9 +541,6 @@ struct Index : o2::soa::IndexColumn<Index<START, END>> {
 template <typename T>
 using is_dynamic_t = framework::is_specialization<typename T::base, DynamicColumn>;
 
-// template <typename T>
-// using is_persistent_t = typename std::decay_t<T>::persistent::type;
-
 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.

From 98ef753d11de2cf486753b5a4fbd4cfeb7dcee7e Mon Sep 17 00:00:00 2001
From: Maja Kabus <mkabus@cern.ch>
Date: Wed, 29 Jun 2022 16:27:34 +0200
Subject: [PATCH 6/6] Adjust template argument in BP call

---
 Framework/Core/include/Framework/ASoAHelpers.h | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/Framework/Core/include/Framework/ASoAHelpers.h b/Framework/Core/include/Framework/ASoAHelpers.h
index 01ddaf1989246..40ccda09fc4b3 100644
--- a/Framework/Core/include/Framework/ASoAHelpers.h
+++ b/Framework/Core/include/Framework/ASoAHelpers.h
@@ -72,8 +72,8 @@ inline bool diffCategory(BinningIndex const& a, BinningIndex const& b)
   return a.bin >= b.bin;
 }
 
-template <template <typename C, typename... Cs> typename BP, typename T, typename C, typename... Cs>
-std::vector<BinningIndex> groupTable(const T& table, const BP<C, Cs...>& binningPolicy, int minCatSize, int outsider)
+template <template <typename... Cs> typename BP, typename T, typename... Cs>
+std::vector<BinningIndex> groupTable(const T& table, const BP<Cs...>& binningPolicy, int minCatSize, int outsider)
 {
   arrow::Table* arrowTable = table.asArrowTable().get();
   auto rowIterator = table.begin();
@@ -92,7 +92,7 @@ std::vector<BinningIndex> groupTable(const T& table, const BP<C, Cs...>& binning
     selectedRows = table.getSelectedRows(); // vector<int64_t>
   }
 
-  auto persistentColumns = typename BP<C, Cs...>::persistent_columns_t{};
+  auto persistentColumns = typename BP<Cs...>::persistent_columns_t{};
   constexpr auto persistentColumnsCount = pack_size(persistentColumns);
   auto arrowColumns = o2::soa::row_helpers::getArrowColumns(arrowTable, persistentColumns);
   auto chunksCount = arrowColumns[0]->num_chunks();