diff --git a/tree/dataframe/inc/ROOT/RNTupleProcessorDS.hxx b/tree/dataframe/inc/ROOT/RNTupleProcessorDS.hxx index 1ca79627a3058..85f4935958fe1 100644 --- a/tree/dataframe/inc/ROOT/RNTupleProcessorDS.hxx +++ b/tree/dataframe/inc/ROOT/RNTupleProcessorDS.hxx @@ -54,14 +54,19 @@ class RNTupleProcessorDS final : public ROOT::RDF::RDataSource { ///////////////////////////////////////////////////////////////////////////// /// \brief Holds useful information about fields added to the RNTupleProcessorDS struct RFieldInfo { - ROOT::DescriptorId_t fFieldId; + ROOT::DescriptorId_t fOnDiskId; std::string fFieldName; std::string fTypeName; std::size_t fNRepetitions; + RNTupleProcessorEntry::FieldIndex_t fEntryFieldIdx; // Enable `std::vector::emplace_back` for this type - RFieldInfo(ROOT::DescriptorId_t fieldId, std::string_view fieldName, std::string_view typeName, - std::size_t nRepetitions) - : fFieldId(fieldId), fFieldName(fieldName), fTypeName(typeName), fNRepetitions(nRepetitions) + RFieldInfo(ROOT::DescriptorId_t onDiskId, std::string_view fieldName, std::string_view typeName, + std::size_t nRepetitions, Internal::RNTupleProcessorEntry::FieldIndex_t entryFieldIdx) + : fOnDiskId(onDiskId), + fFieldName(fieldName), + fTypeName(typeName), + fNRepetitions(nRepetitions), + fEntryFieldIdx(entryFieldIdx) { } }; diff --git a/tree/dataframe/src/RNTupleProcessorDS.cxx b/tree/dataframe/src/RNTupleProcessorDS.cxx index 04bd48828659a..8f73295ac9573 100644 --- a/tree/dataframe/src/RNTupleProcessorDS.cxx +++ b/tree/dataframe/src/RNTupleProcessorDS.cxx @@ -89,7 +89,9 @@ void RNTupleProcessorDS::AddField(const ROOT::RFieldBase &field, std::string_vie // The field is a collection or a fixed-size array. // We open a new collection scope with fieldID being the inner most collection. E.g. for "event.tracks.hits", // fieldInfos would already contain the fieldID of "event.tracks" - fieldInfos.emplace_back(field.GetOnDiskId(), field.GetFieldName(), field.GetTypeName(), nRepetitions); + const std::string name = (procProvenance.Empty() ? "" : procProvenance.Get() + ".") + std::string(colName); + fieldInfos.emplace_back(field.GetOnDiskId(), field.GetFieldName(), field.GetTypeName(), nRepetitions, + *fProcessor->GetEntry().FindFieldIndex(name)); nFieldInfosAdded++; } @@ -127,7 +129,13 @@ void RNTupleProcessorDS::AddField(const ROOT::RFieldBase &field, std::string_vie // The fieldID could be the root field or the class of fieldId might not be loaded. // In these cases, only the inner fields are exposed as RDF columns. - auto fieldOrException = ROOT::RFieldBase::Create(field.GetFieldName(), field.GetTypeName()); + std::string fieldName = field.GetFieldName(); + auto fieldInfoIt = fieldInfos.begin(); + while (fieldName == "_0") { + fieldName = fieldInfoIt->fFieldName; + fieldInfoIt++; + } + auto fieldOrException = ROOT::RFieldBase::Create(fieldName, field.GetTypeName()); if (!fieldOrException) return; auto valueField = fieldOrException.Unwrap(); @@ -160,7 +168,7 @@ void RNTupleProcessorDS::AddField(const ROOT::RFieldBase &field, std::string_vie } else { cardinalityField = std::make_unique>(name); } - cardinalityField->SetOnDiskId(info.fFieldId); + cardinalityField->SetOnDiskId(info.fOnDiskId); } for (auto i = fieldInfos.rbegin(); i != fieldInfos.rend(); ++i) { @@ -183,7 +191,7 @@ void RNTupleProcessorDS::AddField(const ROOT::RFieldBase &field, std::string_vie } } - valueField->SetOnDiskId(fieldInfo.fFieldId); + valueField->SetOnDiskId(fieldInfo.fOnDiskId); // Skip the inner-most collection level to construct the cardinality column // It's taken care of by the `if (!fieldInfos.empty())` scope above @@ -198,29 +206,36 @@ void RNTupleProcessorDS::AddField(const ROOT::RFieldBase &field, std::string_vie cardinalityField = std::make_unique(cardinalityFieldName, cardinalityField->Clone("_0")); } - cardinalityField->SetOnDiskId(fieldInfo.fFieldId); + cardinalityField->SetOnDiskId(fieldInfo.fOnDiskId); } } - if (cardinalityField) { - std::string cardinalityFieldName = - (procProvenance.Empty() ? "R_rdf_sizeof_" : procProvenance.Get() + "_R_rdf_sizeof_") + std::string(colName); - fColumnNames.emplace_back(cardinalityFieldName); - fColumnTypes.emplace_back(cardinalityField->GetTypeName()); - fProcessor->AddFieldToEntry(std::move(cardinalityField), cardinalityFieldName, nullptr, - Internal::RNTupleProcessorProvenance()); - } - - fieldInfos.emplace_back(field.GetOnDiskId(), field.GetFieldName(), field.GetTypeName(), nRepetitions); - nFieldInfosAdded++; std::string canonicalName = (procProvenance.Empty() ? "" : procProvenance.Get() + ".") + std::string(colName); fColumnNames.emplace_back(canonicalName); fColumnTypes.emplace_back(valueField->GetTypeName()); - // Add nested fields explicitly to the entry, because they may be mapped differently. - if (fieldInfos.size() > 1) { - fProcessor->AddFieldToEntry(std::move(valueField), canonicalName, nullptr, - Internal::RNTupleProcessorProvenance()); + RNTupleProcessorEntry::FieldIndex_t fieldIdx; + if (auto optFieldIdx = fProcessor->GetEntry().FindFieldIndex(canonicalName); + optFieldIdx && fProcessor->GetEntry().GetField(*optFieldIdx).GetTypeName() == field.GetTypeName()) { + fieldIdx = *optFieldIdx; + } else { + fieldIdx = fProcessor->AddFieldToEntry(std::move(valueField), canonicalName, nullptr, + Internal::RNTupleProcessorProvenance()); + } + + if (cardinalityField) { + // FIXME(fdegeus): Make cardinality field name work properly with "#colname" + std::string cardinalityFieldName = + (procProvenance.Empty() ? "R_rdf_sizeof_" : procProvenance.Get() + ".R_rdf_sizeof_") + std::string(colName); + fColumnNames.emplace_back(cardinalityFieldName); + fColumnTypes.emplace_back(cardinalityField->GetTypeName()); + auto cardinalityFieldIdx = fProcessor->AddFieldToEntry(std::move(cardinalityField), cardinalityFieldName, nullptr, + Internal::RNTupleProcessorProvenance()); + if (!fieldInfos.empty()) { + const auto &fieldInfo = fieldInfos.back(); + fieldIdx = fieldInfo.fEntryFieldIdx; + } + fProcessor->GetEntry().SetFieldOriginIdx(cardinalityFieldIdx, fieldIdx); } for (std::uint16_t i = 0; i < nFieldInfosAdded; i++) { diff --git a/tree/dataframe/test/datasource_ntuple_processor.cxx b/tree/dataframe/test/datasource_ntuple_processor.cxx index 760178085384f..bf47576a3966f 100644 --- a/tree/dataframe/test/datasource_ntuple_processor.cxx +++ b/tree/dataframe/test/datasource_ntuple_processor.cxx @@ -6,6 +6,8 @@ #include #include +#include + #include #include @@ -290,31 +292,39 @@ TEST_F(RNTupleProcessorDSTest, Chain) ChainTest(fNtplName, fFileName); } -TEST(RNTupleProcessorDS, Join) -{ - FileRaii guardFile1("RNTupleProcessorDS_test_join_1.root"); - FileRaii guardFile2("RNTupleProcessorDS_test_join_2.root"); +class RNTupleProcessorDSAlignedJoinsTest : public ::testing::Test { +protected: + std::vector fFiles = {"RNTupleProcessorDSJoins1.root", "RNTupleProcessorDSJoins2.root"}; + void SetUp() override { - auto model = RNTupleModel::Create(); - auto fldElectron = model->MakeField("electrons"); - auto writer = RNTupleWriter::Recreate(std::move(model), "primary", guardFile1.GetPath()); - for (unsigned i = 0; i < 10; ++i) { - fldElectron->pt = static_cast(i); - writer->Fill(); - } - } - { - auto model = RNTupleModel::Create(); - auto fldJets = model->MakeField>("jets"); - auto writer = RNTupleWriter::Recreate(std::move(model), "auxiliary", guardFile2.GetPath()); - for (unsigned i = 0; i < 10; ++i) { - *fldJets = {static_cast(i) * .1f, static_cast(i) * .2f}; - writer->Fill(); + for (unsigned i = 0; i < fFiles.size(); i++) { + { + auto model = RNTupleModel::Create(); + auto fldElectron = model->MakeField("electrons"); + auto writer = RNTupleWriter::Recreate(std::move(model), "primary", fFiles[i]); + for (unsigned n = 10 * i; n < (10 * i) + 10; ++n) { + fldElectron->pt = static_cast(n); + writer->Fill(); + } + } + { + auto file = std::unique_ptr(TFile::Open(fFiles[i].c_str(), "UPDATE")); + auto model = RNTupleModel::Create(); + auto fldJets = model->MakeField>("jets"); + auto writer = RNTupleWriter::Append(std::move(model), "auxiliary", *file); + for (unsigned n = 10 * i; n < (10 * i) + 10; ++n) { + *fldJets = {static_cast(n) * .1f, static_cast(n) * .2f}; + writer->Fill(); + } + } } } +}; - auto proc = CreateRNTupleJoinProcessor({"primary", guardFile1.GetPath()}, {"auxiliary", guardFile2.GetPath()}, {}); +TEST_F(RNTupleProcessorDSAlignedJoinsTest, SingleJoin) +{ + auto proc = CreateRNTupleJoinProcessor({"primary", fFiles[0]}, {"auxiliary", fFiles[0]}, {}); auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(std::move(proc)); auto count = df.Count(); @@ -330,42 +340,102 @@ TEST(RNTupleProcessorDS, Join) }, [](float a, float b) { return a + b; }, "auxiliary.jets"); EXPECT_FLOAT_EQ(13.5, jetSum.GetValue()); + + auto nJetSum = df.Sum("auxiliary.R_rdf_sizeof_jets"); + EXPECT_FLOAT_EQ(20, nJetSum.GetValue()); } -TEST(RNTupleProcessorDS, JoinUnaligned) +TEST_F(RNTupleProcessorDSAlignedJoinsTest, ChainedJoin) { - FileRaii guardFile1("RNTupleProcessorDS_test_join_unaligned_1.root"); - FileRaii guardFile2("RNTupleProcessorDS_test_join_unaligned_2.root"); + std::vector> joins; + joins.emplace_back(CreateRNTupleJoinProcessor({"primary", fFiles[0]}, {"auxiliary", fFiles[0]}, {})); + joins.emplace_back(CreateRNTupleJoinProcessor({"primary", fFiles[1]}, {"auxiliary", fFiles[1]}, {})); + auto chain = RNTupleProcessor::CreateChain(std::move(joins)); + auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(std::move(chain)); + auto count = df.Count(); + EXPECT_EQ(20, count.GetValue()); + + auto electronFilterCount = df.Filter([](float pt) { return pt > 4.f; }, {"electrons.pt"}).Count(); + EXPECT_EQ(15, electronFilterCount.GetValue()); + + auto jetSum = df.Aggregate( + [](float &acc, const ROOT::RVec &jets) { + for (const auto &jet : jets) + acc += jet; + }, + [](float a, float b) { return a + b; }, "auxiliary.jets"); + EXPECT_FLOAT_EQ(57, jetSum.GetValue()); + + auto nJetSum = df.Sum("auxiliary.R_rdf_sizeof_jets"); + EXPECT_FLOAT_EQ(40, nJetSum.GetValue()); +} + +TEST_F(RNTupleProcessorDSAlignedJoinsTest, JoinedChain) +{ + auto chain_primary = RNTupleProcessor::CreateChain({{"primary", fFiles[0]}, {"primary", fFiles[1]}}); + auto chain_auxiliary = RNTupleProcessor::CreateChain({{"auxiliary", fFiles[0]}, {"auxiliary", fFiles[1]}}); + auto join = CreateRNTupleJoinProcessor(std::move(chain_primary), std::move(chain_auxiliary), {}); + auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(std::move(join)); + + auto count = df.Count(); + EXPECT_EQ(20, count.GetValue()); + + auto electronFilterCount = df.Filter([](float pt) { return pt > 4.f; }, {"electrons.pt"}).Count(); + EXPECT_EQ(15, electronFilterCount.GetValue()); + + auto jetSum = df.Aggregate( + [](float &acc, const ROOT::RVec &jets) { + for (const auto &jet : jets) + acc += jet; + }, + [](float a, float b) { return a + b; }, "auxiliary.jets"); + EXPECT_FLOAT_EQ(57, jetSum.GetValue()); + + auto nJetSum = df.Sum("auxiliary.R_rdf_sizeof_jets"); + EXPECT_FLOAT_EQ(40, nJetSum.GetValue()); +} + +class RNTupleProcessorDSUnalignedJoinsTest : public ::testing::Test { +protected: + std::vector fFiles = {"RNTupleProcessorDSJoins1.root", "RNTupleProcessorDSJoins2.root"}; + + void SetUp() override { - auto model = RNTupleModel::Create(); - auto fldId = model->MakeField("idx"); - auto fldElectron = model->MakeField("electrons"); - auto writer = RNTupleWriter::Recreate(std::move(model), "primary", guardFile1.GetPath()); - for (unsigned i = 0; i < 10; ++i) { - // only fill odd events - if (i % 2 == 0) - continue; - *fldId = i; - fldElectron->pt = static_cast(i); - writer->Fill(); - } - } - { - auto model = RNTupleModel::Create(); - auto fldId = model->MakeField("idx"); - auto fldJets = model->MakeField>("jets"); - auto writer = RNTupleWriter::Recreate(std::move(model), "auxiliary", guardFile2.GetPath()); - // fill in reverse - for (int i = 9; i >= 0; --i) { - *fldId = i; - *fldJets = {static_cast(i) * .1f, static_cast(i) * .2f}; - writer->Fill(); + for (unsigned i = 0; i < fFiles.size(); i++) { + { + auto model = RNTupleModel::Create(); + auto fldIdx = model->MakeField("idx"); + auto fldElectron = model->MakeField("electrons"); + auto writer = RNTupleWriter::Recreate(std::move(model), "primary", fFiles[i]); + for (unsigned n = 10 * i; n < (10 * i) + 10; ++n) { + // only fill odd events + if (n % 2 == 0) + continue; + *fldIdx = n; + fldElectron->pt = static_cast(n); + writer->Fill(); + } + } + { + auto file = std::unique_ptr(TFile::Open(fFiles[i].c_str(), "UPDATE")); + auto model = RNTupleModel::Create(); + auto fldIdx = model->MakeField("idx"); + auto fldJets = model->MakeField>("jets"); + auto writer = RNTupleWriter::Append(std::move(model), "auxiliary", *file); + for (unsigned n = 10 * i; n < (10 * i) + 10; ++n) { + *fldIdx = n; + *fldJets = {static_cast(n) * .1f, static_cast(n) * .2f}; + writer->Fill(); + } + } } } +}; - auto proc = - CreateRNTupleJoinProcessor({"primary", guardFile1.GetPath()}, {"auxiliary", guardFile2.GetPath()}, {"idx"}); +TEST_F(RNTupleProcessorDSUnalignedJoinsTest, SingleJoin) +{ + auto proc = CreateRNTupleJoinProcessor({"primary", fFiles[0]}, {"auxiliary", fFiles[0]}, {"idx"}); auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(std::move(proc)); auto count = df.Count(); @@ -381,6 +451,60 @@ TEST(RNTupleProcessorDS, JoinUnaligned) }, [](float a, float b) { return a + b; }, "auxiliary.jets"); EXPECT_FLOAT_EQ(7.5, jetSum.GetValue()); + + auto nJetSum = df.Sum("auxiliary.R_rdf_sizeof_jets"); + EXPECT_FLOAT_EQ(10, nJetSum.GetValue()); +} + +TEST_F(RNTupleProcessorDSUnalignedJoinsTest, ChainedJoin) +{ + std::vector> joins; + joins.emplace_back(CreateRNTupleJoinProcessor({"primary", fFiles[0]}, {"auxiliary", fFiles[0]}, {"idx"})); + joins.emplace_back(CreateRNTupleJoinProcessor({"primary", fFiles[1]}, {"auxiliary", fFiles[1]}, {"idx"})); + auto chain = RNTupleProcessor::CreateChain(std::move(joins)); + auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(std::move(chain)); + + auto count = df.Count(); + EXPECT_EQ(10, count.GetValue()); + + auto electronFilterCount = df.Filter([](float pt) { return pt > 4.f; }, {"electrons.pt"}).Count(); + EXPECT_EQ(8, electronFilterCount.GetValue()); + + auto jetSum = df.Aggregate( + [](float &acc, const ROOT::RVec &jets) { + for (const auto &jet : jets) + acc += jet; + }, + [](float a, float b) { return a + b; }, "auxiliary.jets"); + EXPECT_FLOAT_EQ(30, jetSum.GetValue()); + + auto nJetSum = df.Sum("auxiliary.R_rdf_sizeof_jets"); + EXPECT_FLOAT_EQ(20, nJetSum.GetValue()); +} + +TEST_F(RNTupleProcessorDSUnalignedJoinsTest, JoinedChain) +{ + auto chain_primary = RNTupleProcessor::CreateChain({{"primary", fFiles[0]}, {"primary", fFiles[1]}}); + auto chain_auxiliary = RNTupleProcessor::CreateChain({{"auxiliary", fFiles[0]}, {"auxiliary", fFiles[1]}}); + auto join = CreateRNTupleJoinProcessor(std::move(chain_primary), std::move(chain_auxiliary), {"idx"}); + auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(std::move(join)); + + auto count = df.Count(); + EXPECT_EQ(10, count.GetValue()); + + auto electronFilterCount = df.Filter([](float pt) { return pt > 4.f; }, {"electrons.pt"}).Count(); + EXPECT_EQ(8, electronFilterCount.GetValue()); + + auto jetSum = df.Aggregate( + [](float &acc, const ROOT::RVec &jets) { + for (const auto &jet : jets) + acc += jet; + }, + [](float a, float b) { return a + b; }, "auxiliary.jets"); + EXPECT_FLOAT_EQ(30, jetSum.GetValue()); + + auto nJetSum = df.Sum("auxiliary.R_rdf_sizeof_jets"); + EXPECT_FLOAT_EQ(20, nJetSum.GetValue()); } #ifdef R__USE_IMT diff --git a/tree/ntuple/inc/ROOT/RNTupleProcessorEntry.hxx b/tree/ntuple/inc/ROOT/RNTupleProcessorEntry.hxx index 98e12e76a09d6..afd994fd397e4 100644 --- a/tree/ntuple/inc/ROOT/RNTupleProcessorEntry.hxx +++ b/tree/ntuple/inc/ROOT/RNTupleProcessorEntry.hxx @@ -103,6 +103,10 @@ private: bool fIsValid; RNTupleProcessorProvenance fProcessorProvenance; + // If the field is artificial (i.e., not actually present on disk), this holds the entry index of the field it + // should get the on-disk ID from. + ROOT::DescriptorId_t fFieldOriginIdx = ROOT::kInvalidDescriptorId; + RProcessorValue(std::unique_ptr field, std::string_view qualifiedFieldName, ROOT::RFieldBase::RValue &&value, bool isValid, RNTupleProcessorProvenance provenance) : fField(std::move(field)), @@ -160,6 +164,32 @@ public: return fProcessorValues[fieldIdx].fIsValid; } + ///////////////////////////////////////////////////////////////////////////// + /// \brief Set a field's origin index + /// + /// \param[in] fieldIdx The index of the field in the entry to set. + /// \param[in] fieldOriginIdx The index of the field in the entry the field identified by `fieldIdx` should + /// point to. + /// + /// Setting the origin index should be used for articifial fields. + void SetFieldOriginIdx(FieldIndex_t fieldIdx, FieldIndex_t fieldOriginIdx) + { + assert(fieldIdx < fProcessorValues.size()); + fProcessorValues[fieldIdx].fFieldOriginIdx = fieldOriginIdx; + } + + ///////////////////////////////////////////////////////////////////////////// + /// \brief Get a field's origin field index. + /// + /// \param[in] fieldIdx The index of the field in the entry to set. + /// + /// \return The index of the field's origin field. + FieldIndex_t GetFieldOriginIdx(FieldIndex_t fieldIdx) + { + assert(fieldIdx < fProcessorValues.size()); + return fProcessorValues[fieldIdx].fFieldOriginIdx; + } + ///////////////////////////////////////////////////////////////////////////// /// \brief Find the name of a field from its field index. /// diff --git a/tree/ntuple/src/RNTupleProcessor.cxx b/tree/ntuple/src/RNTupleProcessor.cxx index 08d157a253bc6..b8fe66b628d48 100644 --- a/tree/ntuple/src/RNTupleProcessor.cxx +++ b/tree/ntuple/src/RNTupleProcessor.cxx @@ -303,6 +303,26 @@ void ROOT::Experimental::RNTupleSingleProcessor::LoadEntryImpl(const REntryMappi fNEntriesProcessed++; } +namespace { +void ResetColumnRepresentative(ROOT::RFieldBase &field) +{ + field.SetColumnRepresentatives(ROOT::RFieldBase::RColumnRepresentations::Selection_t{}); + for (auto &subfield : field.GetMutableSubfields()) { + ResetColumnRepresentative(*subfield); + } +} + +void SetOnDiskId(ROOT::RFieldBase &field, ROOT::DescriptorId_t onDiskId, const ROOT::RNTupleDescriptor &desc) +{ + field.SetOnDiskId(onDiskId); + + for (auto &subfield : field.GetMutableSubfields()) { + onDiskId = desc.FindFieldId(subfield->GetQualifiedFieldName()); + SetOnDiskId(*subfield, onDiskId, desc); + } +} +} // anonymous namespace + void ROOT::Experimental::RNTupleSingleProcessor::Connect( const std::unordered_set &fieldIdxs, const Internal::RNTupleProcessorProvenance & /* provenance */, bool updateFields) @@ -312,9 +332,36 @@ void ROOT::Experimental::RNTupleSingleProcessor::Connect( fFieldIdxs = fieldIdxs; if (updateFields) { + auto descGuard = fPageSource->GetSharedDescriptorGuard(); + const auto &desc = descGuard.GetRef(); + for (const auto &fieldIdx : fFieldIdxs) { const auto &currField = fEntry->GetValue(fieldIdx).GetField(); - auto newField = CreateAndConnectField(fEntry->GetQualifiedFieldName(fieldIdx), currField.GetTypeName()); + + auto onDiskId = desc.FindFieldId(currField.GetQualifiedFieldName()); + bool needsClone = onDiskId != ROOT::kInvalidDescriptorId && + currField.GetTypeName() == desc.GetFieldDescriptor(onDiskId).GetTypeName(); + if (const auto fieldOriginIdx = fEntry->GetFieldOriginIdx(fieldIdx); + fieldOriginIdx != ROOT::kInvalidDescriptorId) { + const auto &originField = fEntry->GetField(fieldOriginIdx); + onDiskId = desc.FindFieldId(originField.GetQualifiedFieldName()); + needsClone = true; + } + + std::unique_ptr newField; + if (needsClone) { + // We have enough information to clone the existing field. + newField = currField.Clone(currField.GetFieldName()); + // The clone's column representatives (+ its subfields') need to be reset before it can be connected to the + // new page source. + ResetColumnRepresentative(*newField); + SetOnDiskId(*newField, onDiskId, desc); + newField = ConnectField(std::move(newField)); + } else { + // We need to create the new field from scratch (e.g. because it is a subfield without a parent field + // attached). + newField = CreateAndConnectField(fEntry->GetQualifiedFieldName(fieldIdx), currField.GetTypeName()); + } fEntry->UpdateField(fieldIdx, std::move(newField)); }