diff --git a/Framework/Core/include/Framework/ASoA.h b/Framework/Core/include/Framework/ASoA.h index 2e478a8ca64a6..4ed8e830fe137 100644 --- a/Framework/Core/include/Framework/ASoA.h +++ b/Framework/Core/include/Framework/ASoA.h @@ -1245,7 +1245,7 @@ struct TableIterator : IP, C... { }; struct ArrowHelpers { - static std::shared_ptr joinTables(std::vector>&& tables); + static std::shared_ptr joinTables(std::vector>&& tables, std::span labels); static std::shared_ptr concatTables(std::vector>&& tables); }; @@ -1683,6 +1683,7 @@ class Table using table_t = self_t; static constexpr const auto originals = computeOriginals(); + static constexpr const auto originalLabels = [] refs, size_t... Is>(std::index_sequence) { return std::array{o2::aod::label()...}; }.template operator()(std::make_index_sequence()); template bindings> requires(ref.origin_hash == "CONC"_h) @@ -1931,7 +1932,7 @@ class Table Table(std::vector>&& tables, uint64_t offset = 0) requires(ref.origin_hash != "CONC"_h) - : Table(ArrowHelpers::joinTables(std::move(tables)), offset) + : Table(ArrowHelpers::joinTables(std::move(tables), std::span{originalLabels}), offset) { } @@ -3213,7 +3214,7 @@ struct JoinFull : Table, D, o2::aod::Hash<"JOIN"_h>, Ts. bindInternalIndicesTo(this); } JoinFull(std::vector>&& tables, uint64_t offset = 0) - : base{ArrowHelpers::joinTables(std::move(tables)), offset} + : base{ArrowHelpers::joinTables(std::move(tables), std::span{base::originalLabels}), offset} { bindInternalIndicesTo(this); } @@ -3223,6 +3224,7 @@ struct JoinFull : Table, D, o2::aod::Hash<"JOIN"_h>, Ts. using self_t = JoinFull; using table_t = base; static constexpr const auto originals = base::originals; + static constexpr const auto originalLabels = base::originalLabels; using columns_t = typename table_t::columns_t; using persistent_columns_t = typename table_t::persistent_columns_t; using iterator = table_t::template iterator_template; @@ -3293,7 +3295,7 @@ using Join = JoinFull, Ts...>; template constexpr auto join(Ts const&... t) { - return Join(ArrowHelpers::joinTables({t.asArrowTable()...})); + return Join(ArrowHelpers::joinTables({t.asArrowTable()...}, std::span{Join::base::originalLabels})); } template diff --git a/Framework/Core/include/Framework/AnalysisManagers.h b/Framework/Core/include/Framework/AnalysisManagers.h index e310f3eef990c..1d894b2b67948 100644 --- a/Framework/Core/include/Framework/AnalysisManagers.h +++ b/Framework/Core/include/Framework/AnalysisManagers.h @@ -282,14 +282,14 @@ template bool prepareOutput(ProcessingContext& context, T& spawns) { using metadata = o2::aod::MetadataTrait>::metadata; - auto originalTable = soa::ArrowHelpers::joinTables(extractOriginals(context)); + auto originalTable = soa::ArrowHelpers::joinTables(extractOriginals(context), std::span{metadata::base_table_t::originalLabels}); if (originalTable->schema()->fields().empty() == true) { using base_table_t = typename T::base_table_t::table_t; originalTable = makeEmptyTable(o2::aod::label()); } spawns.extension = std::make_shared(o2::framework::spawner>(originalTable, o2::aod::label(), spawns.projector)); - spawns.table = std::make_shared(soa::ArrowHelpers::joinTables({spawns.extension->asArrowTable(), originalTable})); + spawns.table = std::make_shared(soa::ArrowHelpers::joinTables({spawns.extension->asArrowTable(), originalTable}, std::span{T::spawnable_t::table_t::originalLabels})); return true; } @@ -304,14 +304,14 @@ template bool prepareOutput(ProcessingContext& context, T& defines) { using metadata = o2::aod::MetadataTrait>::metadata; - auto originalTable = soa::ArrowHelpers::joinTables(extractOriginals(context)); + auto originalTable = soa::ArrowHelpers::joinTables(extractOriginals(context), std::span{metadata::base_table_t::originalLabels}); if (originalTable->schema()->fields().empty() == true) { using base_table_t = typename T::base_table_t::table_t; originalTable = makeEmptyTable(o2::aod::label()); } defines.extension = std::make_shared(o2::framework::spawner>(originalTable, o2::aod::label(), defines.projectors.data(), defines.projector)); - defines.table = std::make_shared(soa::ArrowHelpers::joinTables({defines.extension->asArrowTable(), originalTable})); + defines.table = std::make_shared(soa::ArrowHelpers::joinTables({defines.extension->asArrowTable(), originalTable}, std::span{T::spawnable_t::table_t::originalLabels})); return true; } diff --git a/Framework/Core/include/Framework/AnalysisTask.h b/Framework/Core/include/Framework/AnalysisTask.h index 9bd2e2af173cc..3404cd3eac276 100644 --- a/Framework/Core/include/Framework/AnalysisTask.h +++ b/Framework/Core/include/Framework/AnalysisTask.h @@ -201,9 +201,9 @@ struct AnalysisDataProcessorBuilder { std::shared_ptr table = nullptr; auto joiner = [&record] refs, size_t... Is>(std::index_sequence) { return std::vector{extractTableFromRecord(record)...}; }; if constexpr (soa::is_iterator) { - table = o2::soa::ArrowHelpers::joinTables(joiner.template operator()(std::make_index_sequence())); + table = o2::soa::ArrowHelpers::joinTables(joiner.template operator()(std::make_index_sequence()), std::span{T::parent_t::originalLabels}); } else { - table = o2::soa::ArrowHelpers::joinTables(joiner.template operator()(std::make_index_sequence())); + table = o2::soa::ArrowHelpers::joinTables(joiner.template operator()(std::make_index_sequence()), std::span{T::originalLabels}); } expressions::updateFilterInfo(info, table); if constexpr (!o2::soa::is_smallgroups>) { diff --git a/Framework/Core/include/Framework/TableBuilder.h b/Framework/Core/include/Framework/TableBuilder.h index 32fe78b852eff..5d6a88e7e0886 100644 --- a/Framework/Core/include/Framework/TableBuilder.h +++ b/Framework/Core/include/Framework/TableBuilder.h @@ -863,7 +863,7 @@ template auto spawner(std::vector>&& tables, const char* name, o2::framework::expressions::Projector* projectors, std::shared_ptr& projector) { using placeholders_pack_t = typename o2::aod::MetadataTrait::metadata::placeholders_pack_t; - auto fullTable = soa::ArrowHelpers::joinTables(std::move(tables)); + auto fullTable = soa::ArrowHelpers::joinTables(std::move(tables), std::span{o2::aod::MetadataTrait::metadata::base_table_t::originalLabels}); if (fullTable->num_rows() == 0) { return makeEmptyTable(name, placeholders_pack_t{}); } @@ -892,7 +892,7 @@ template auto spawner(std::vector>&& tables, const char* name, std::shared_ptr& projector) { using expression_pack_t = typename o2::aod::MetadataTrait::metadata::expression_pack_t; - auto fullTable = soa::ArrowHelpers::joinTables(std::move(tables)); + auto fullTable = soa::ArrowHelpers::joinTables(std::move(tables), std::span{o2::aod::MetadataTrait::metadata::base_table_t::originalLabels}); if (fullTable->num_rows() == 0) { return makeEmptyTable(name, expression_pack_t{}); } @@ -929,7 +929,8 @@ auto spawner(std::shared_ptr const& fullTable, const char* name, s template auto spawner(framework::pack columns, std::vector>&& tables, const char* name, std::shared_ptr& projector) { - auto fullTable = soa::ArrowHelpers::joinTables(std::move(tables)); + std::array labels{"original"}; + auto fullTable = soa::ArrowHelpers::joinTables(std::move(tables), std::span{labels}); if (fullTable->num_rows() == 0) { return makeEmptyTable(name, framework::pack{}); } diff --git a/Framework/Core/src/ASoA.cxx b/Framework/Core/src/ASoA.cxx index 5940bc0427225..3a681ee931a2b 100644 --- a/Framework/Core/src/ASoA.cxx +++ b/Framework/Core/src/ASoA.cxx @@ -64,7 +64,7 @@ SelectionVector sliceSelection(gsl::span const& mSelectedRows, in return slicedSelection; } -std::shared_ptr ArrowHelpers::joinTables(std::vector>&& tables) +std::shared_ptr ArrowHelpers::joinTables(std::vector>&& tables, std::span labels) { if (tables.size() == 1) { return tables[0]; @@ -72,10 +72,7 @@ std::shared_ptr ArrowHelpers::joinTables(std::vectornum_rows() != tables[i + 1]->num_rows()) { throw o2::framework::runtime_error_f("Tables %s and %s have different sizes (%d vs %d) and cannot be joined!", - tables[i]->schema()->metadata()->Get("label").ValueOrDie().c_str(), - tables[i + 1]->schema()->metadata()->Get("label").ValueOrDie().c_str(), - tables[i]->num_rows(), - tables[i + 1]->num_rows()); + labels[i], labels[i + 1], tables[i]->num_rows(), tables[i + 1]->num_rows()); } } std::vector> fields; diff --git a/Framework/Core/test/test_ASoA.cxx b/Framework/Core/test/test_ASoA.cxx index 6d9ee16fec84d..80519aebc9ee7 100644 --- a/Framework/Core/test/test_ASoA.cxx +++ b/Framework/Core/test/test_ASoA.cxx @@ -31,6 +31,7 @@ namespace test DECLARE_SOA_COLUMN(X, x, int); DECLARE_SOA_COLUMN(Y, y, int); DECLARE_SOA_COLUMN(Z, z, int); +DECLARE_SOA_COLUMN(W, w, int); DECLARE_SOA_DYNAMIC_COLUMN(Sum, sum, [](int x, int y) { return x + y; }); DECLARE_SOA_EXPRESSION_COLUMN(ESum, esum, int, test::x + test::y); } // namespace test @@ -268,9 +269,17 @@ TEST_CASE("TestJoinedTables") rowWriterZ(0, 8); auto tableZ = builderZ.finalize(); + TableBuilder builderW; + auto rowWriterW = builderW.persist({"fW"}); + rowWriterW(0, 8); + rowWriterW(0, 8); + rowWriterW(0, 8); + auto tableW = builderW.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 TestW = InPlaceTable<"A3"_h, o2::aod::test::W>; using Test = Join; REQUIRE(Test::contains()); @@ -303,6 +312,12 @@ TEST_CASE("TestJoinedTables") for (auto& test : tests4) { REQUIRE(15 == test.x() + test.y() + test.z()); } + + try { + auto testF = join(TestZ{tableZ}, TestW{tableW}); + } catch (RuntimeErrorRef ref) { + REQUIRE(std::string{error_from_ref(ref).what} == "Tables TEST and TEST have different sizes (8 vs 3) and cannot be joined!"); + } } TEST_CASE("TestConcatTables")