diff --git a/Framework/Core/include/Framework/ASoA.h b/Framework/Core/include/Framework/ASoA.h index 6ff77f8facaec..73b70b1eab779 100644 --- a/Framework/Core/include/Framework/ASoA.h +++ b/Framework/Core/include/Framework/ASoA.h @@ -2308,6 +2308,28 @@ class Table template using InPlaceTable = Table, o2::aod::Hash, o2::aod::Hash<"TEST"_h>, C...>; +template +std::shared_ptr projectColumns(std::shared_ptr const& src) +{ + auto indices = [](framework::pack, std::shared_ptr const& sch) { + return std::vector{sch->GetFieldIndex(C::columnLabel())...}; + }(typename Narrow::persistent_columns_t{}, src->schema()); + return src->SelectColumns(indices).ValueOrDie(); +} + +template + requires(!is_filtered_table) +Narrow as_table(T const& wide) +{ + return Narrow{projectColumns(wide.asArrowTable()), wide.offset()}; +} + +template +soa::Filtered as_table(T const& wide) +{ + return soa::Filtered({projectColumns(wide.asArrowTable())}, wide.getSelectedRows(), wide.offset()); +} + void getterNotFound(const char* targetColumnLabel); void emptyColumnLabel(); diff --git a/Framework/Core/test/test_ASoA.cxx b/Framework/Core/test/test_ASoA.cxx index 117dddff4c548..9c90e2c07c6f2 100644 --- a/Framework/Core/test/test_ASoA.cxx +++ b/Framework/Core/test/test_ASoA.cxx @@ -320,6 +320,70 @@ TEST_CASE("TestJoinedTables") } } +TEST_CASE("TestAsTableProjection") +{ + TableBuilder builderX; + auto rowWriterX = builderX.persist({"fX"}); + for (int i = 0; i < 8; ++i) { + rowWriterX(0, i); + } + auto tableX = builderX.finalize(); + + TableBuilder builderY; + auto rowWriterY = builderY.persist({"fY"}); + for (int i = 0; i < 8; ++i) { + rowWriterY(0, 7 - i); + } + auto tableY = builderY.finalize(); + + TableBuilder builderZ; + auto rowWriterZ = builderZ.persist({"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; + + 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(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>(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 wideF{{tableX, tableY, tableZ}, SelectionVector{1, 3, 5}, 0}; + auto narrowF = as_table(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;