fix filtered grouping for merged AODs

Author: aalkin

Framework/Core/include/Framework/AnalysisTask.h

@@ -282,7 +282,8 @@ struct AnalysisDataProcessorBuilder {
                                                        x.asArrowTable(),
                                                        static_cast<int32_t>(gt.tableSize()),
                                                        &groups[index],
-                                                       &offsets[index]);
+                                                       &offsets[index],
+                                                       &sizes[index]);
             if (result.ok() == false) {
               throw runtime_error("Cannot split collection");
             }
@@ -394,7 +395,7 @@ struct AnalysisDataProcessorBuilder {
 
             // for each grouping element we need to slice the selection vector
             auto start_iterator = std::lower_bound(starts[index], selections[index]->end(), (offsets[index])[pos]);
-            auto stop_iterator = std::lower_bound(start_iterator, selections[index]->end(), (offsets[index])[pos + 1]);
+            auto stop_iterator = std::lower_bound(start_iterator, selections[index]->end(), (offsets[index])[pos] + (sizes[index])[pos]);
             starts[index] = stop_iterator;
             soa::SelectionVector slicedSelection{start_iterator, stop_iterator};
             std::transform(slicedSelection.begin(), slicedSelection.end(), slicedSelection.begin(),
@@ -421,6 +422,7 @@ struct AnalysisDataProcessorBuilder {
       soa::SelectionVector const* groupSelection = nullptr;
       std::array<std::vector<arrow::Datum>, sizeof...(A)> groups;
       std::array<std::vector<uint64_t>, sizeof...(A)> offsets;
+      std::array<std::vector<int>, sizeof...(A)> sizes;
       std::array<soa::SelectionVector const*, sizeof...(A)> selections;
       std::array<soa::SelectionVector::const_iterator, sizeof...(A)> starts;
     };

Framework/Core/include/Framework/Kernels.h

@@ -36,6 +36,7 @@ auto sliceByColumn(
   T fullSize,
   std::vector<arrow::Datum>* slices,
   std::vector<uint64_t>* offsets = nullptr,
+  std::vector<int>* sizes = nullptr,
   std::vector<arrow::Datum>* unassignedSlices = nullptr,
   std::vector<uint64_t>* unassignedOffsets = nullptr)
 {
@@ -50,7 +51,6 @@ auto sliceByColumn(
 
   // create slices and offsets
   uint64_t offset = 0;
-  uint64_t unassignedOffset = 0;
   auto count = 0;
   auto size = values.length();
 
@@ -59,6 +59,9 @@ auto sliceByColumn(
     if (offsets) {
       offsets->emplace_back(offset_);
     }
+    if (sizes) {
+      sizes->emplace_back(count_);
+    }
   };
 
   auto makeUnassignedSlice = [&](uint64_t offset_, T count_) {
@@ -85,7 +88,7 @@ auto sliceByColumn(
       offset += count;
       continue;
     }
-    nzeros = v - vprev - ((i == 0) ? 0 : 1);
+    nzeros = v - vprev - ((i == 0 || slices->empty() == true) ? 0 : 1);
     for (auto z = 0; z < nzeros; ++z) {
       makeSlice(offset, 0);
     }