Binning policy changes minimized, backward compatible

Author: saganatt

Framework/Core/include/Framework/ASoAHelpers.h

@@ -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);
       }

Framework/Core/include/Framework/BinningPolicy.h

@@ -13,39 +13,57 @@
 #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>
 
 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;
+    }
+  }
+}
+} // 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
+  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) {
       // 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,27 +120,27 @@ 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);
               }
             }
           }
 
-          // 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<o2::soa::is_persistent_t, C, Cs...>;
+  std::array<std::vector<double>, 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,101 @@ struct BinningPolicy {
     return mIgnoreOverflows ? 1 : -1;
   }
 
-  void expandConstantBinning(std::vector<double> const& bins, int ind)
+  // Note: Overflow / underflow bin -1 is not included
+  int getBinsCount(std::vector<double> const& bins) 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 bins.size() - 1 - getOverflowShift();
+  }
+};
+
+template <typename, typename...>
+struct FlexibleBinningPolicy;
+
+template <typename... Ts, typename... Ls>
+struct FlexibleBinningPolicy<std::tuple<Ls...>, Ts...> : BinningPolicyBase<sizeof...(Ts)> {
+  FlexibleBinningPolicy(std::tuple<Ls...> const& lambdaPtrs, std::array<std::vector<double>, sizeof...(Ts)> bins, bool ignoreOverflows = true) : BinningPolicyBase<sizeof...(Ts)>(bins, ignoreOverflows), mBinningFunctions{lambdaPtrs}
+  {
+  }
+
+  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
+  {
+    using decayed = std::decay_t<T2>;
+    if (globalIndex == -1) {
+      globalIndex = *(std::get<0>(rowIterator.getIndices()));
+    }
+
+    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);
     }
   }
 
-  std::array<std::vector<double>, sizeof...(Cs) + 1> mBins;
-  bool mIgnoreOverflows;
+  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:
+  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)
+  {
+  }
+
+  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, Ts>(table, rowIterator, ci, ai, globalIndex)...);
+  }
+
+  int getBin(std::tuple<typename Ts::type...> const& data) const
+  {
+    return BinningPolicyBase<sizeof...(Ts)>::template getBin<typename Ts::type...>(data);
+  }
+
+  using persistent_columns_t = framework::selected_pack<o2::soa::is_persistent_t, Ts...>;
+};
+
+template <typename... Ts>
+using BinningPolicy = ColumnBinningPolicy<Ts...>;
+
 template <typename C>
 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);