DPL Analysis: implemented grouping for Filtered argument (#3003)

Author: aalkin

Framework/Core/include/Framework/ASoA.h

@@ -844,7 +844,8 @@ class Table
   }
   std::shared_ptr<arrow::Table> mTable;
   /// This is a cached lookup of the column index in a given
-  std::tuple<std::pair<C*, arrow::Column*>...> mColumnIndex;
+  std::tuple<std::pair<C*, arrow::Column*>...>
+    mColumnIndex;
   /// Cached begin iterator for this table.
   unfiltered_iterator mBegin;
   /// Cached end iterator for this table.
@@ -1196,6 +1197,11 @@ class Filtered : public T
     return mSelection->GetNumSlots();
   }
 
+  int64_t tableSize() const
+  {
+    return table_t::asArrowTable()->num_rows();
+  }
+
   framework::expressions::Selection getSelection() const
   {
     return mSelection;

Framework/Core/include/Framework/AnalysisTask.h

@@ -398,10 +398,50 @@ struct AnalysisDataProcessorBuilder {
               ++const_cast<std::decay_t<Grouping>&>(groupingElement);
               ++oi;
             }
-          }
-          if constexpr (is_specialization<std::decay_t<AssociatedType>, o2::soa::Filtered>::value) {
-            // FIXME: we need to implement the case for the grouped filtered case.
-            static_assert(always_static_assert_v<std::decay_t<AssociatedType>>, "Grouping Filtered is not yet supported");
+          } else if constexpr (is_specialization<std::decay_t<AssociatedType>, o2::soa::Filtered>::value) {
+            auto& fullFiltered = std::get<0>(associatedTables);
+            auto selectionBuffer = std::shared_ptr<arrow::Buffer>(&(fullFiltered.getSelection()->GetBuffer()));
+            auto selectionArray = fullFiltered.getSelection()->ToArray();
+            offsets.push_back(fullFiltered.tableSize());
+            uint64_t selectionIndex = 0;
+            uint64_t sliceStart = 0;
+            uint64_t sliceStop = 0;
+
+            auto findSliceBounds = [&](int64_t l, int64_t h) {
+              size_t s = 0;
+              for (auto i = selectionIndex; i < selectionArray->length(); ++i) {
+                auto value = selectionArray->data()->template GetValues<uint64_t>(i);
+                if (*value == l) {
+                  sliceStart = i;
+                  s = i;
+                  break;
+                }
+              }
+              for (auto i = s + 1; i < selectionArray->length(); ++i) {
+                auto value = selectionArray->data()->template GetValues<uint64_t>(i);
+                if (*value == h) {
+                  sliceStop = i;
+                  selectionIndex = i;
+                  break;
+                }
+              }
+            };
+
+            for (auto& groupedDatum : groupsCollection) {
+              auto groupedElementsTable = arrow::util::get<std::shared_ptr<arrow::Table>>(groupedDatum.value);
+              // for each grouping element we need to slice the selection vector
+              findSliceBounds(offsets[oi], offsets[oi + 1]);
+              auto slicedBuffer = arrow::SliceBuffer(selectionBuffer, sliceStart, sliceStop - sliceStart + 1);
+              expressions::Selection slicedSelection;
+              if (!gandiva::SelectionVector::MakeInt64(sliceStop - sliceStart + 1, slicedBuffer, &slicedSelection).ok()) {
+                throw std::runtime_error("Cannot create sliced selection");
+              }
+              std::decay_t<AssociatedType> typedTable{{groupedElementsTable}, slicedSelection, offsets[oi]};
+              typedTable.bindExternalIndices(&groupingTable);
+              task.process(groupingElement, typedTable);
+              ++const_cast<std::decay_t<Grouping>&>(groupingElement);
+              ++oi;
+            }
           } else if constexpr (is_specialization<std::decay_t<AssociatedType>, o2::soa::Join>::value || is_specialization<std::decay_t<AssociatedType>, o2::soa::Concat>::value) {
             for (auto& groupedDatum : groupsCollection) {
               auto groupedElementsTable = arrow::util::get<std::shared_ptr<arrow::Table>>(groupedDatum.value);

Framework/Core/test/test_AnalysisTask.cxx

@@ -123,7 +123,7 @@ struct GTask {
 //};
 
 struct ITask {
-  void process(o2::soa::Filtered<o2::soa::Join<o2::aod::Foos, o2::aod::Bars, o2::aod::XYZ>> const& foobars)
+  void process(o2::aod::Collision const&, o2::soa::Filtered<o2::soa::Join<o2::aod::Foos, o2::aod::Bars, o2::aod::XYZ>> const& foobars)
   {
     for (auto foobar : foobars) {
       foobar.x();
@@ -170,5 +170,5 @@ BOOST_AUTO_TEST_CASE(AdaptorCompilation)
   //  BOOST_CHECK_EQUAL(task8.inputs.size(), 3);
 
   auto task9 = adaptAnalysisTask<ITask>("test9");
-  BOOST_CHECK_EQUAL(task9.inputs.size(), 3);
+  BOOST_CHECK_EQUAL(task9.inputs.size(), 4);
 }