From 5d8b64bed9b073090962fb4c50e2cfb7b3ba6874 Mon Sep 17 00:00:00 2001
From: Anton Alkin <anton.alkin@cern.ch>
Date: Thu, 26 Nov 2020 12:53:54 +0100
Subject: [PATCH] DPL Analysis: easy way to filter a table on the fly

---
 Analysis/Tutorials/src/associatedExample.cxx  | 34 +++++++++++++++++--
 .../Tutorials/src/histogramsFullTracks.cxx    |  2 +-
 Framework/Core/include/Framework/ASoA.h       | 22 ++++++++----
 3 files changed, 49 insertions(+), 9 deletions(-)

diff --git a/Analysis/Tutorials/src/associatedExample.cxx b/Analysis/Tutorials/src/associatedExample.cxx
index 6ac8e1225467c..19c171fb6fba7 100644
--- a/Analysis/Tutorials/src/associatedExample.cxx
+++ b/Analysis/Tutorials/src/associatedExample.cxx
@@ -75,7 +75,7 @@ struct BTask {
 struct TTask {
   using myCol = soa::Join<aod::Collisions, aod::CollisionsExtra>;
   expressions::Filter multfilter = aod::collision::mult > 10;
-  void process(soa::Filtered<soa::Join<aod::Collisions, aod::CollisionsExtra>>::iterator const& col, aod::Tracks const& tracks)
+  void process(soa::Filtered<myCol>::iterator const& col, aod::Tracks const& tracks)
   {
     LOGF(INFO, "[direct] ID: %d; %d == %d", col.globalIndex(), col.mult(), tracks.size());
     if (tracks.size() > 0) {
@@ -85,11 +85,41 @@ struct TTask {
   }
 };
 
+struct ZTask {
+  using myCol = soa::Join<aod::Collisions, aod::CollisionsExtra>;
+
+  void process(myCol const& collisions, aod::Tracks const& tracks)
+  {
+    auto multbin0_10 = collisions.select(aod::collision::mult >= 0 && aod::collision::mult < 10);
+    auto multbin10_30 = collisions.select(aod::collision::mult >= 10 && aod::collision::mult < 30);
+    auto multbin30_100 = collisions.select(aod::collision::mult >= 30 && aod::collision::mult < 100);
+
+    LOGF(INFO, "Bin 0-10");
+    for (auto& col : multbin0_10) {
+      auto groupedTracks = tracks.select(aod::track::collisionId == col.globalIndex());
+      LOGF(INFO, "Collision %d; Ntrk = %d vs %d", col.globalIndex(), col.mult(), groupedTracks.size());
+    }
+
+    LOGF(INFO, "Bin 10-30");
+    for (auto& col : multbin10_30) {
+      auto groupedTracks = tracks.select(aod::track::collisionId == col.globalIndex());
+      LOGF(INFO, "Collision %d; Ntrk = %d vs %d", col.globalIndex(), col.mult(), groupedTracks.size());
+    }
+
+    LOGF(INFO, "Bin 30-100");
+    for (auto& col : multbin30_100) {
+      auto groupedTracks = tracks.select(aod::track::collisionId == col.globalIndex());
+      LOGF(INFO, "Collision %d; Ntrk = %d vs %d", col.globalIndex(), col.mult(), groupedTracks.size());
+    }
+  }
+};
+
 WorkflowSpec defineDataProcessing(ConfigContext const&)
 {
   return WorkflowSpec{
     adaptAnalysisTask<ATask>("produce-etaphi"),
     adaptAnalysisTask<BTask>("consume-etaphi"),
     adaptAnalysisTask<MTask>("produce-mult"),
-    adaptAnalysisTask<TTask>("consume-mult")};
+    adaptAnalysisTask<TTask>("consume-mult"),
+    adaptAnalysisTask<ZTask>("partition-mult")};
 }
diff --git a/Analysis/Tutorials/src/histogramsFullTracks.cxx b/Analysis/Tutorials/src/histogramsFullTracks.cxx
index 3dab7398facdf..f264ce8cf7361 100644
--- a/Analysis/Tutorials/src/histogramsFullTracks.cxx
+++ b/Analysis/Tutorials/src/histogramsFullTracks.cxx
@@ -19,7 +19,7 @@ using namespace o2::framework;
 using namespace o2::framework::expressions;
 
 struct ATask {
-  OutputObj<TH2F> etaClsH{TH2F("eta_vs_cls", "eta_vs_cls", 102, -2.01, 2.01, 160, -0.5, 159.5)};
+  OutputObj<TH2F> etaClsH{TH2F("eta_vs_cls", "#eta vs N_{cls}", 102, -2.01, 2.01, 160, -0.5, 159.5)};
 
   void process(aod::FullTracks const& tracks)
   {
diff --git a/Framework/Core/include/Framework/ASoA.h b/Framework/Core/include/Framework/ASoA.h
index a31588e003489..3a9e89267416a 100644
--- a/Framework/Core/include/Framework/ASoA.h
+++ b/Framework/Core/include/Framework/ASoA.h
@@ -808,6 +808,17 @@ static constexpr auto extractBindings(framework::pack<Is...>)
   return framework::pack<typename Is::binding_t...>{};
 }
 
+template <typename T>
+class Filtered;
+
+template <typename T>
+auto select(T const& t, framework::expressions::Filter&& f)
+{
+  return Filtered<T>({t.asArrowTable()}, framework::expressions::createExpressionTree(
+                                           framework::expressions::createOperations(f),
+                                           t.asArrowTable()->schema()));
+}
+
 /// A Table class which observes an arrow::Table and provides
 /// It is templated on a set of Column / DynamicColumn types.
 template <typename... C>
@@ -1000,6 +1011,11 @@ class Table
     mBegin.bindExternalIndices(current...);
   }
 
+  auto select(framework::expressions::Filter&& f) const
+  {
+    return o2::soa::select(*this, std::forward<framework::expressions::Filter>(f));
+  }
+
  private:
   template <typename T>
   arrow::ChunkedArray* lookupColumn()
@@ -1762,12 +1778,6 @@ class Filtered<Filtered<T>> : public FilteredPolicy<typename T::table_t>
 template <typename T>
 using is_soa_filtered_t = typename framework::is_base_of_template<soa::FilteredPolicy, T>;
 
-template <typename T>
-auto filter(T&& t, framework::expressions::Filter const& expr)
-{
-  return Filtered<T>(t.asArrowTable(), expr);
-}
-
 /// Template for building an index table to access matching rows from non-
 /// joinable, but compatible tables, e.g. Collisions and ZDCs.
 /// First argument is the key table (BCs for the Collisions+ZDCs case), the rest