DPL: allow narrowing of large joins

This allows narrowing of large joins when they are not strictly necessary, allowing to simplify iteration and reducing the scope of helper functions and alikes.

Author: ktf

Framework/Core/include/Framework/ASoA.h

@@ -2308,6 +2308,28 @@ class Table
 template <uint32_t D, soa::is_column... C>
 using InPlaceTable = Table<o2::aod::Hash<"TEST"_h>, o2::aod::Hash<D>, o2::aod::Hash<"TEST"_h>, C...>;
 
+template <typename Narrow>
+std::shared_ptr<arrow::Table> projectColumns(std::shared_ptr<arrow::Table> const& src)
+{
+  auto indices = []<typename... C>(framework::pack<C...>, std::shared_ptr<arrow::Schema> const& sch) {
+    return std::vector<int>{sch->GetFieldIndex(C::columnLabel())...};
+  }(typename Narrow::persistent_columns_t{}, src->schema());
+  return src->SelectColumns(indices).ValueOrDie();
+}
+
+template <typename Narrow, is_table T>
+  requires(!is_filtered_table<T>)
+Narrow as_table(T const& wide)
+{
+  return Narrow{projectColumns<Narrow>(wide.asArrowTable()), wide.offset()};
+}
+
+template <typename Narrow, is_filtered_table T>
+soa::Filtered<Narrow> as_table(T const& wide)
+{
+  return soa::Filtered<Narrow>({projectColumns<Narrow>(wide.asArrowTable())}, wide.getSelectedRows(), wide.offset());
+}
+
 void getterNotFound(const char* targetColumnLabel);
 void emptyColumnLabel();
 

Framework/Core/test/test_ASoA.cxx

@@ -320,6 +320,70 @@ TEST_CASE("TestJoinedTables")
   }
 }
 
+TEST_CASE("TestAsTableProjection")
+{
+  TableBuilder builderX;
+  auto rowWriterX = builderX.persist<int32_t>({"fX"});
+  for (int i = 0; i < 8; ++i) {
+    rowWriterX(0, i);
+  }
+  auto tableX = builderX.finalize();
+
+  TableBuilder builderY;
+  auto rowWriterY = builderY.persist<int32_t>({"fY"});
+  for (int i = 0; i < 8; ++i) {
+    rowWriterY(0, 7 - i);
+  }
+  auto tableY = builderY.finalize();
+
+  TableBuilder builderZ;
+  auto rowWriterZ = builderZ.persist<int32_t>({"fZ"});
+  for (int i = 0; i < 8; ++i) {
+    rowWriterZ(0, 8);
+  }
+  auto tableZ = builderZ.finalize();
+
+  using TestX = InPlaceTable<"A0"_h, o2::aod::test::X>;
+  using TestY = InPlaceTable<"A1"_h, o2::aod::test::Y>;
+  using TestZ = InPlaceTable<"A2"_h, o2::aod::test::Z>;
+  using Wide = Join<TestX, TestY, TestZ>;
+
+  Wide wide{{tableX, tableY, tableZ}, 0};
+  REQUIRE(wide.asArrowTable()->num_columns() == 3);
+
+  // Project the wide join onto just the X table: the resulting arrow table
+  // shares the data but exposes a single column, and iteration sees only x.
+  auto narrow = as_table<TestX>(wide);
+  REQUIRE(narrow.asArrowTable()->num_columns() == 1);
+  REQUIRE(narrow.size() == 8);
+  int sum = 0;
+  for (auto& r : narrow) {
+    sum += r.x();
+  }
+  REQUIRE(sum == (0 + 1 + 2 + 3 + 4 + 5 + 6 + 7));
+
+  // Projecting onto two of the three tables keeps two columns.
+  auto narrow2 = as_table<Join<TestX, TestY>>(wide);
+  REQUIRE(narrow2.asArrowTable()->num_columns() == 2);
+  for (auto& r : narrow2) {
+    REQUIRE(7 == r.x() + r.y());
+  }
+
+  // Filtered projection: the selection carries over to the narrow table, so the
+  // narrow view iterates exactly the selected rows of the projected column.
+  Filtered<Wide> wideF{{tableX, tableY, tableZ}, SelectionVector{1, 3, 5}, 0};
+  auto narrowF = as_table<TestX>(wideF);
+  REQUIRE(narrowF.asArrowTable()->num_columns() == 1);
+  int sumF = 0;
+  int nF = 0;
+  for (auto& r : narrowF) {
+    sumF += r.x();
+    ++nF;
+  }
+  REQUIRE(nF == 3);
+  REQUIRE(sumF == (1 + 3 + 5));
+}
+
 TEST_CASE("TestConcatTables")
 {
   TableBuilder builderA;