improve partition handling; update test

Author: aalkin

Framework/Core/include/Framework/AnalysisHelpers.h

@@ -805,14 +805,28 @@ struct Service {
   }
 };
 
-template <typename T>
+template <soa::soaFilteredTable T>
 auto getTableFromFilter(const T& table, soa::SelectionVector&& selection)
 {
-  if constexpr (soa::is_soa_filtered_v<std::decay_t<T>>) {
-    return std::make_unique<o2::soa::Filtered<T>>(std::vector{table}, std::forward<soa::SelectionVector>(selection));
-  } else {
-    return std::make_unique<o2::soa::Filtered<T>>(std::vector{table.asArrowTable()}, std::forward<soa::SelectionVector>(selection));
-  }
+  return std::make_unique<o2::soa::Filtered<T>>(std::vector{table}, std::forward<soa::SelectionVector>(selection));
+}
+
+template <soa::soaTable T> requires (!soa::soaFilteredTable<T>)
+auto getTableFromFilter(const T& table, soa::SelectionVector&& selection)
+{
+  return std::make_unique<o2::soa::Filtered<T>>(std::vector{table.asArrowTable()}, std::forward<soa::SelectionVector>(selection));
+}
+
+template <soa::ngFilteredTable T>
+auto getTableFromFilter(const T& table, soa::SelectionVector&& selection)
+{
+  return std::make_unique<o2::soa::FilteredNG<T>>(std::vector{table}, std::forward<soa::SelectionVector>(selection));
+}
+
+template <soa::ngTable T> requires (!soa::ngFilteredTable<T>)
+auto getTableFromFilter(const T& table, soa::SelectionVector&& selection)
+{
+  return std::make_unique<o2::soa::FilteredNG<T>>(std::vector{table.asArrowTable()}, std::forward<soa::SelectionVector>(selection));
 }
 
 void initializePartitionCaches(std::set<uint32_t> const& hashes, std::shared_ptr<arrow::Schema> const& schema, expressions::Filter const& filter, gandiva::NodePtr& tree, gandiva::FilterPtr& gfilter);

Framework/Core/test/test_AnalysisTask.cxx

@@ -21,6 +21,7 @@ using namespace o2::framework;
 
 namespace o2::aod
 {
+O2HASH("TestA/0");
 namespace test
 {
 DECLARE_SOA_COLUMN(X, x, float);
@@ -404,3 +405,69 @@ TEST_CASE("TestPartitionIteration")
   }
   REQUIRE(i == 1);
 }
+
+TEST_CASE("TestPartitionIterationNG")
+{
+  TableBuilder builderA;
+  auto rowWriterA = builderA.persist<float, float>({"fX", "fY"});
+  rowWriterA(0, 0.0f, 8.0f);
+  rowWriterA(0, 1.0f, 9.0f);
+  rowWriterA(0, 2.0f, 10.0f);
+  rowWriterA(0, 3.0f, 11.0f);
+  rowWriterA(0, 4.0f, 12.0f);
+  rowWriterA(0, 5.0f, 13.0f);
+  rowWriterA(0, 6.0f, 14.0f);
+  rowWriterA(0, 7.0f, 15.0f);
+  auto tableA = builderA.finalize();
+  REQUIRE(tableA->num_rows() == 8);
+
+  using TestA = soa::InPlaceTableNG<"TestA/0"_h, o2::soa::Index<>, aod::test::X, aod::test::Y>;
+  using FilteredTest = o2::soa::FilteredNG<TestA>;
+  using PartitionTest = PartitionNG<TestA>;
+  using PartitionFilteredTest = PartitionNG<o2::soa::FilteredNG<TestA>>;
+  using PartitionNestedFilteredTest = PartitionNG<o2::soa::FilteredNG<o2::soa::FilteredNG<TestA>>>;
+  using namespace o2::framework;
+
+  TestA testA{tableA};
+
+  PartitionTest p1 = aod::test::x < 4.0f;
+  p1.bindTable(testA);
+  REQUIRE(4 == p1.size());
+  REQUIRE(p1.begin() != p1.end());
+  auto i = 0;
+  for (auto& p : p1) {
+    REQUIRE(i == p.x());
+    REQUIRE(i + 8 == p.y());
+    REQUIRE(i == p.index());
+    i++;
+  }
+  REQUIRE(i == 4);
+
+  expressions::Filter f1 = aod::test::x < 4.0f;
+  auto selection = expressions::createSelection(testA.asArrowTable(), f1);
+  FilteredTest filtered{{testA.asArrowTable()}, o2::soa::selectionToVector(selection)};
+  PartitionFilteredTest p2 = aod::test::y > 9.0f;
+  p2.bindTable(filtered);
+
+  REQUIRE(2 == p2.size());
+  i = 0;
+  for (auto& p : p2) {
+    REQUIRE(i + 2 == p.x());
+    REQUIRE(i + 10 == p.y());
+    REQUIRE(i + 2 == p.index());
+    i++;
+  }
+  REQUIRE(i == 2);
+
+  PartitionNestedFilteredTest p3 = aod::test::x < 3.0f;
+  p3.bindTable(*(p2.mFiltered));
+  REQUIRE(1 == p3.size());
+  i = 0;
+  for (auto& p : p3) {
+    REQUIRE(i + 2 == p.x());
+    REQUIRE(i + 10 == p.y());
+    REQUIRE(i + 2 == p.index());
+    i++;
+  }
+  REQUIRE(i == 1);
+}