From b8482643fadaaf1c3fd642e3db1ed8163f26e04e Mon Sep 17 00:00:00 2001
From: Maja Kabus <mkabus@cern.ch>
Date: Mon, 23 Nov 2020 22:39:21 +0100
Subject: [PATCH] Example of partitions inside process

---
 Analysis/Tutorials/src/partitions.cxx          | 18 +++++++++++++++++-
 .../Core/include/Framework/AnalysisHelpers.h   |  7 +++++++
 Framework/Core/include/Framework/Expressions.h |  4 +++-
 Framework/Core/src/Expressions.cxx             | 11 ++++++++---
 4 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/Analysis/Tutorials/src/partitions.cxx b/Analysis/Tutorials/src/partitions.cxx
index 6479356a48cfa..e80fd07da5b58 100644
--- a/Analysis/Tutorials/src/partitions.cxx
+++ b/Analysis/Tutorials/src/partitions.cxx
@@ -59,8 +59,24 @@ struct ATask {
   }
 };
 
+// Partition inside process
+// Caveat: partitioned table cannot be passed as const& to process()
+struct BTask {
+  void process(aod::Collisions const& collisions, aod::Tracks& tracks)
+  {
+    for (auto& c : collisions) {
+      Partition<aod::Tracks> groupedTracks = aod::track::collisionId == c.globalIndex();
+      groupedTracks.bindTable(tracks);
+      for (auto& t : groupedTracks) {
+        LOGF(INFO, "collision global index: %d grouped track collision id: %d", c.globalIndex(), t.collisionId());
+      }
+    }
+  }
+};
+
 WorkflowSpec defineDataProcessing(ConfigContext const&)
 {
   return WorkflowSpec{
-    adaptAnalysisTask<ATask>("consume-tracks")};
+    adaptAnalysisTask<ATask>("consume-tracks"),
+    adaptAnalysisTask<BTask>("partition-in-process")};
 }
diff --git a/Framework/Core/include/Framework/AnalysisHelpers.h b/Framework/Core/include/Framework/AnalysisHelpers.h
index 0745170ebf804..a6376e57e0cf1 100644
--- a/Framework/Core/include/Framework/AnalysisHelpers.h
+++ b/Framework/Core/include/Framework/AnalysisHelpers.h
@@ -486,6 +486,13 @@ struct Partition {
   {
   }
 
+  void bindTable(T& table)
+  {
+    mFiltered.reset(getTableFromFilter(table, filter));
+    bindExternalIndices(&table);
+    getBoundToExternalIndices(table);
+  }
+
   void setTable(const T& table)
   {
     mFiltered.reset(getTableFromFilter(table, filter));
diff --git a/Framework/Core/include/Framework/Expressions.h b/Framework/Core/include/Framework/Expressions.h
index 743ae75ba974a..fc13daea14375 100644
--- a/Framework/Core/include/Framework/Expressions.h
+++ b/Framework/Core/include/Framework/Expressions.h
@@ -56,7 +56,7 @@ struct LiteralStorage {
   using stored_pack = framework::pack<T...>;
 };
 
-using LiteralValue = LiteralStorage<int, bool, float, double, uint8_t>;
+using LiteralValue = LiteralStorage<int, bool, float, double, uint8_t, int64_t>;
 
 template <typename T>
 constexpr auto selectArrowType()
@@ -73,6 +73,8 @@ constexpr auto selectArrowType()
     return atype::INT8;
   } else if constexpr (std::is_same_v<T, uint16_t>) {
     return atype::INT16;
+  } else if constexpr (std::is_same_v<T, int64_t>) {
+    return atype::INT64;
   } else if constexpr (std::is_same_v<T, uint8_t>) {
     return atype::UINT8;
   } else {
diff --git a/Framework/Core/src/Expressions.cxx b/Framework/Core/src/Expressions.cxx
index 94f63dd75353b..60323a1305d95 100644
--- a/Framework/Core/src/Expressions.cxx
+++ b/Framework/Core/src/Expressions.cxx
@@ -67,6 +67,8 @@ std::shared_ptr<arrow::DataType> concreteArrowType(atype::type type)
       return arrow::int16();
     case atype::INT32:
       return arrow::int32();
+    case atype::INT64:
+      return arrow::int64();
     case atype::FLOAT:
       return arrow::float32();
     case atype::DOUBLE:
@@ -268,8 +270,8 @@ Operations createOperations(Filter const& expression)
       return t1;
     }
 
-    if (t1 == atype::INT32 || t1 == atype::INT8 || t1 == atype::INT16 || t1 == atype::UINT8) {
-      if (t2 == atype::INT32 || t2 == atype::INT8 || t2 == atype::INT16 || t2 == atype::UINT8) {
+    if (t1 == atype::INT32 || t1 == atype::INT8 || t1 == atype::INT16 || t1 == atype::INT64 || t1 == atype::UINT8) {
+      if (t2 == atype::INT32 || t2 == atype::INT8 || t2 == atype::INT16 || t2 == atype::INT64 || t2 == atype::UINT8) {
         return atype::FLOAT;
       }
       if (t2 == atype::FLOAT) {
@@ -280,7 +282,7 @@ Operations createOperations(Filter const& expression)
       }
     }
     if (t1 == atype::FLOAT) {
-      if (t2 == atype::INT32 || t2 == atype::INT8 || t2 == atype::INT16 || t2 == atype::UINT8) {
+      if (t2 == atype::INT32 || t2 == atype::INT8 || t2 == atype::INT16 || t2 == atype::INT64 || t2 == atype::UINT8) {
         return atype::FLOAT;
       }
       if (t2 == atype::DOUBLE) {
@@ -451,6 +453,9 @@ gandiva::NodePtr createExpressionTree(Operations const& opSpecs,
       if (content.index() == 4) {
         return gandiva::TreeExprBuilder::MakeLiteral(std::get<uint8_t>(content));
       }
+      if (content.index() == 5) {
+        return gandiva::TreeExprBuilder::MakeLiteral(std::get<int64_t>(content));
+      }
       throw runtime_error("Malformed LiteralNode");
     }