From 6c45c41b09d7b4f3040820facabc30780518d55d Mon Sep 17 00:00:00 2001 From: Florine de Geus Date: Mon, 21 Sep 2026 17:52:40 +0200 Subject: [PATCH 1/8] [ntuple] Add `AddAllFieldsFieldsToEntry` --- tree/ntuple/inc/ROOT/RNTupleProcessor.hxx | 33 +++++++++++- tree/ntuple/src/RNTupleProcessor.cxx | 55 +++++++++++++++++++ tree/ntuple/test/ntuple_processor.cxx | 60 +++++++++++++++++++++ tree/ntuple/test/ntuple_processor_chain.cxx | 20 +++++++ tree/ntuple/test/ntuple_processor_join.cxx | 18 +++++++ 5 files changed, 185 insertions(+), 1 deletion(-) diff --git a/tree/ntuple/inc/ROOT/RNTupleProcessor.hxx b/tree/ntuple/inc/ROOT/RNTupleProcessor.hxx index 6403770e62909..edacecd2aa6b0 100644 --- a/tree/ntuple/inc/ROOT/RNTupleProcessor.hxx +++ b/tree/ntuple/inc/ROOT/RNTupleProcessor.hxx @@ -31,9 +31,12 @@ namespace ROOT { namespace Experimental { +class RNTupleProcessor; namespace Internal { -struct RNTupleProcessorEntryLoader; +struct RNTupleProcessorEntryLoader; // for unit tests +const RNTupleProcessorEntry * +LoadFullRNTupleProcessorEntry(ROOT::Experimental::RNTupleProcessor &processor, bool includeSubfields); // for unit tests } // namespace Internal // clang-format off @@ -284,6 +287,9 @@ that is returned by RequestField(). // clang-format on class RNTupleProcessor { friend struct ROOT::Experimental::Internal::RNTupleProcessorEntryLoader; // for unit tests + friend const Internal::RNTupleProcessorEntry * + ROOT::Experimental::Internal::LoadFullRNTupleProcessorEntry(RNTupleProcessor &processor, + bool includeSubfields); // For unit tests friend class RNTupleSingleProcessor; friend class RNTupleChainProcessor; friend class RNTupleJoinProcessor; @@ -449,6 +455,13 @@ protected: AddFieldToEntry(const std::string &fieldName, const std::string &typeName, void *valuePtr, const Internal::RNTupleProcessorProvenance &provenance) = 0; + ///////////////////////////////////////////////////////////////////////////// + /// \brief Add all known fields (from on-disk information) to the entry. + /// + /// \param[in] provenance Provenance of the processor. + /// \param[in] addFieldMode Bitfield indicating how the fields should be added. + virtual void AddAllFieldsToEntry(const Internal::RNTupleProcessorProvenance &provenance, bool includeSubfields) = 0; + ///////////////////////////////////////////////////////////////////////////// /// \brief Add the entry mappings for this processor to the provided join table. /// @@ -764,6 +777,12 @@ private: const std::string &fieldName, const std::string &typeName, void *valuePtr = nullptr, const Internal::RNTupleProcessorProvenance &provenance = Internal::RNTupleProcessorProvenance()) final; + ///////////////////////////////////////////////////////////////////////////// + /// \brief Add all known fields (from on-disk information) to the entry. + /// + /// \sa RNTupleProcessor::AddAllFieldsToEntry() + void AddAllFieldsToEntry(const Internal::RNTupleProcessorProvenance &provenance, bool includeSubfields) final; + ///////////////////////////////////////////////////////////////////////////// /// \brief Add the entry mappings for this processor to the provided join table. /// @@ -870,6 +889,12 @@ private: const std::string &fieldName, const std::string &typeName, void *valuePtr = nullptr, const Internal::RNTupleProcessorProvenance &provenance = Internal::RNTupleProcessorProvenance()) final; + ///////////////////////////////////////////////////////////////////////////// + /// \brief Add all known fields (from on-disk information) to the entry. + /// + /// \sa RNTupleProcessor::AddAllFieldsToEntry() + void AddAllFieldsToEntry(const Internal::RNTupleProcessorProvenance &provenance, bool includeSubfields) final; + ///////////////////////////////////////////////////////////////////////////// /// \brief Add the entry mappings for this processor to the provided join table. /// @@ -980,6 +1005,12 @@ private: const std::string &fieldName, const std::string &typeName, void *valuePtr = nullptr, const Internal::RNTupleProcessorProvenance &provenance = Internal::RNTupleProcessorProvenance()) final; + ///////////////////////////////////////////////////////////////////////////// + /// \brief Add all known fields (from on-disk information) to the entry. + /// + /// \sa RNTupleProcessor::AddAllFieldsToEntry() + void AddAllFieldsToEntry(const Internal::RNTupleProcessorProvenance &provenance, bool includeSubfields) final; + ///////////////////////////////////////////////////////////////////////////// /// \brief Add the entry mappings for this processor to the provided join table. /// diff --git a/tree/ntuple/src/RNTupleProcessor.cxx b/tree/ntuple/src/RNTupleProcessor.cxx index acb8b3e1bd4a7..1ba380a505a4c 100644 --- a/tree/ntuple/src/RNTupleProcessor.cxx +++ b/tree/ntuple/src/RNTupleProcessor.cxx @@ -23,6 +23,15 @@ #include +const ROOT::Experimental::Internal::RNTupleProcessorEntry * +ROOT::Experimental::Internal::LoadFullRNTupleProcessorEntry(RNTupleProcessor &processor, bool includeSubfields) +{ + processor.AddAllFieldsToEntry(RNTupleProcessorProvenance(), includeSubfields); + return processor.fEntry.get(); +} + +//------------------------------------------------------------------------------ + std::unique_ptr ROOT::Experimental::RNTupleOpenSpec::CreatePageSource() const { if (const std::string *storagePath = std::get_if(&fStorage)) @@ -211,6 +220,36 @@ ROOT::Experimental::RNTupleSingleProcessor::AddFieldToEntry(const std::string &f return *fieldIdx; } +void ROOT::Experimental::RNTupleSingleProcessor::AddAllFieldsToEntry( + const Internal::RNTupleProcessorProvenance &provenance, bool includeSubfields) +{ + Initialize(); + auto descGuard = fPageSource->GetSharedDescriptorGuard(); + const auto &desc = descGuard.GetRef(); + auto fnAddSubfields = [this, &desc, &provenance](const ROOT::RFieldDescriptor &field, auto &fn) -> void { + std::string fieldName = desc.GetQualifiedFieldName(field.GetId()); + if (auto prefix = provenance.Get(); !prefix.empty()) + fieldName = prefix + "." + fieldName; + + AddFieldToEntry(fieldName, field.GetTypeName(), nullptr, provenance); + for (const auto &subfield : desc.GetFieldIterable(field.GetId())) { + fn(subfield, fn); + } + }; + + for (const auto &field : desc.GetTopLevelFields()) { + if (includeSubfields) { + fnAddSubfields(field, fnAddSubfields); + } else { + std::string fieldName = desc.GetQualifiedFieldName(field.GetId()); + if (auto prefix = provenance.Get(); !prefix.empty()) + fieldName = prefix + "." + fieldName; + + AddFieldToEntry(fieldName, field.GetTypeName(), nullptr, provenance); + } + } +} + void ROOT::Experimental::RNTupleSingleProcessor::SetEntryMapping(ROOT::NTupleSize_t globalEntryNumber, REntryMapping &mapping) { @@ -370,6 +409,13 @@ ROOT::Experimental::RNTupleChainProcessor::AddFieldToEntry(const std::string &fi return fInnerProcessors[fCurrentProcessorNumber]->AddFieldToEntry(fieldName, typeName, valuePtr, provenance); } +void ROOT::Experimental::RNTupleChainProcessor::AddAllFieldsToEntry( + const Internal::RNTupleProcessorProvenance &provenance, bool includeSubfields) +{ + Initialize(); + fInnerProcessors[0]->AddAllFieldsToEntry(provenance, includeSubfields); +} + void ROOT::Experimental::RNTupleChainProcessor::SetEntryMapping(ROOT::NTupleSize_t globalEntryNumber, REntryMapping &mapping) { @@ -561,6 +607,15 @@ ROOT::Experimental::RNTupleJoinProcessor::AddFieldToEntry(const std::string &fie } } +void ROOT::Experimental::RNTupleJoinProcessor::AddAllFieldsToEntry( + const Internal::RNTupleProcessorProvenance &provenance, bool includeSubfields) +{ + Initialize(); + fPrimaryProcessor->AddAllFieldsToEntry(provenance, includeSubfields); + auto auxProvenance = provenance.Evolve(fAuxiliaryProcessor->GetOptions().GetProcessorName()); + fAuxiliaryProcessor->AddAllFieldsToEntry(auxProvenance, includeSubfields); +} + void ROOT::Experimental::RNTupleJoinProcessor::SetAuxiliaryFieldValidity(bool isValid) { for (const auto &fieldIdx : fAuxiliaryFieldIdxs) { diff --git a/tree/ntuple/test/ntuple_processor.cxx b/tree/ntuple/test/ntuple_processor.cxx index 5172776ccc9b7..c3e788bc3f4aa 100644 --- a/tree/ntuple/test/ntuple_processor.cxx +++ b/tree/ntuple/test/ntuple_processor.cxx @@ -362,6 +362,24 @@ TEST_F(RNTupleProcessorTest, Subfields) } } +TEST_F(RNTupleProcessorTest, AddAllFields) +{ + auto proc = RNTupleProcessor::Create({fNTupleNames[0], fFileNames[0]}); + auto entry = ROOT::Experimental::Internal::LoadFullRNTupleProcessorEntry(*proc, /*includeSubfields=*/false); + auto fieldIdxs = entry->GetFieldIndices(); + + EXPECT_EQ(fieldIdxs.size(), 4); +} + +TEST_F(RNTupleProcessorTest, AddAllFieldsAndSubfields) +{ + auto proc = RNTupleProcessor::Create({fNTupleNames[0], fFileNames[0]}); + auto entry = ROOT::Experimental::Internal::LoadFullRNTupleProcessorEntry(*proc, /*includeSubfields=*/true); + auto fieldIdxs = entry->GetFieldIndices(); + + EXPECT_EQ(fieldIdxs.size(), 13); +} + TEST_F(RNTupleProcessorTest, PrintStructureSingle) { auto proc = RNTupleProcessor::Create({fNTupleNames[0], fFileNames[0]}); @@ -864,6 +882,48 @@ TEST_F(RNTupleProcessorTest, JoinedJoinComposedSameName) } } +TEST_F(RNTupleProcessorTest, AddAllFieldsComposed) +{ + auto primaryProc = RNTupleProcessor::Create({fNTupleNames[0], fFileNames[0]}); + + RNTupleProcessorOptions opts; + opts.SetProcessorName("ntuple_aux2"); + auto auxProcIntermediate = RNTupleProcessor::Create({fNTupleNames[2], fFileNames[2]}, opts); + + auto auxProc = RNTupleProcessor::CreateJoin( + RNTupleProcessor::CreateChain({{fNTupleNames[1], fFileNames[1]}, {fNTupleNames[2], fFileNames[2]}}), + std::move(auxProcIntermediate), {"i"}); + + auto proc = RNTupleProcessor::CreateJoin(std::move(primaryProc), std::move(auxProc), {}); + + auto entry = ROOT::Experimental::Internal::LoadFullRNTupleProcessorEntry(*proc, /*includeSubfields=*/false); + auto fieldIdxs = entry->GetFieldIndices(); + + // 11 fields instead of 10, because the join field is also included. + EXPECT_EQ(fieldIdxs.size(), 11); +} + +TEST_F(RNTupleProcessorTest, AddAllFieldsAndSubfieldsComposed) +{ + auto primaryProc = RNTupleProcessor::Create({fNTupleNames[0], fFileNames[0]}); + + RNTupleProcessorOptions opts; + opts.SetProcessorName("ntuple_aux2"); + auto auxProcIntermediate = RNTupleProcessor::Create({fNTupleNames[2], fFileNames[2]}, opts); + + auto auxProc = RNTupleProcessor::CreateJoin( + RNTupleProcessor::CreateChain({{fNTupleNames[1], fFileNames[1]}, {fNTupleNames[2], fFileNames[2]}}), + std::move(auxProcIntermediate), {"i"}); + + auto proc = RNTupleProcessor::CreateJoin(std::move(primaryProc), std::move(auxProc), {}); + + auto entry = ROOT::Experimental::Internal::LoadFullRNTupleProcessorEntry(*proc, /*includeSubfields=*/true); + auto fieldIdxs = entry->GetFieldIndices(); + + // 36 fields instead of 25, because the join field is also included. + EXPECT_EQ(fieldIdxs.size(), 36); +} + TEST_F(RNTupleProcessorTest, PrintStructureChainedJoin) { std::vector> innerProcs; diff --git a/tree/ntuple/test/ntuple_processor_chain.cxx b/tree/ntuple/test/ntuple_processor_chain.cxx index 60429cc5e3dbc..8219816457359 100644 --- a/tree/ntuple/test/ntuple_processor_chain.cxx +++ b/tree/ntuple/test/ntuple_processor_chain.cxx @@ -144,6 +144,26 @@ TEST_F(RNTupleChainProcessorTest, MissingFields) EXPECT_EQ(15, proc->GetNEntriesProcessed()); } +TEST_F(RNTupleChainProcessorTest, AddAllFields) +{ + auto proc = RNTupleProcessor::CreateChain( + {{fNTupleName, fFileNames[0]}, {fNTupleName, fFileNames[2]}, {fNTupleName, fFileNames[1]}}); + auto entry = ROOT::Experimental::Internal::LoadFullRNTupleProcessorEntry(*proc, /*includeSubfields=*/false); + auto fieldIdxs = entry->GetFieldIndices(); + + EXPECT_EQ(fieldIdxs.size(), 2); +} + +TEST_F(RNTupleChainProcessorTest, AddAllFieldsAndSubfields) +{ + auto proc = RNTupleProcessor::CreateChain( + {{fNTupleName, fFileNames[0]}, {fNTupleName, fFileNames[2]}, {fNTupleName, fFileNames[1]}}); + auto entry = ROOT::Experimental::Internal::LoadFullRNTupleProcessorEntry(*proc, /*includeSubfields=*/true); + auto fieldIdxs = entry->GetFieldIndices(); + + EXPECT_EQ(fieldIdxs.size(), 3); +} + TEST_F(RNTupleChainProcessorTest, EmptyNTuples) { FileRaii fileGuard("test_ntuple_processor_empty_ntuples.root"); diff --git a/tree/ntuple/test/ntuple_processor_join.cxx b/tree/ntuple/test/ntuple_processor_join.cxx index f4a84f04fe73a..ebb7c8eecbf0e 100644 --- a/tree/ntuple/test/ntuple_processor_join.cxx +++ b/tree/ntuple/test/ntuple_processor_join.cxx @@ -165,6 +165,24 @@ TEST(RNTupleJoinProcessor, NameConflict) } } +TEST_F(RNTupleJoinProcessorTest, AddAllFields) +{ + auto proc = RNTupleProcessor::CreateJoin({fNTupleNames[1], fFileNames[1]}, {fNTupleNames[2], fFileNames[2]}, {}); + auto entry = ROOT::Experimental::Internal::LoadFullRNTupleProcessorEntry(*proc, /*includeSubfields=*/false); + auto fieldIdxs = entry->GetFieldIndices(); + + EXPECT_EQ(fieldIdxs.size(), 5); +} + +TEST_F(RNTupleJoinProcessorTest, AddAllFieldsAndSubfields) +{ + auto proc = RNTupleProcessor::CreateJoin({fNTupleNames[1], fFileNames[1]}, {fNTupleNames[2], fFileNames[2]}, {}); + auto entry = ROOT::Experimental::Internal::LoadFullRNTupleProcessorEntry(*proc, /*includeSubfields=*/true); + auto fieldIdxs = entry->GetFieldIndices(); + + EXPECT_EQ(fieldIdxs.size(), 6); +} + TEST_F(RNTupleJoinProcessorTest, UnalignedSingleJoinField) { auto proc = RNTupleProcessor::CreateJoin({fNTupleNames[0], fFileNames[0]}, {fNTupleNames[1], fFileNames[1]}, {"i"}); From 46b852379bd7204f0781f2eec078883628b7b995 Mon Sep 17 00:00:00 2001 From: Florine de Geus Date: Wed, 23 Sep 2026 15:21:30 +0200 Subject: [PATCH 2/8] [ntuple] Enable adding existing fields to a processor entry --- tree/ntuple/inc/ROOT/RNTupleProcessor.hxx | 58 +++++++++++++++- tree/ntuple/src/RNTupleProcessor.cxx | 83 +++++++++++++++++++++-- 2 files changed, 133 insertions(+), 8 deletions(-) diff --git a/tree/ntuple/inc/ROOT/RNTupleProcessor.hxx b/tree/ntuple/inc/ROOT/RNTupleProcessor.hxx index edacecd2aa6b0..0f2b3705d7d51 100644 --- a/tree/ntuple/inc/ROOT/RNTupleProcessor.hxx +++ b/tree/ntuple/inc/ROOT/RNTupleProcessor.hxx @@ -442,7 +442,8 @@ protected: /// \brief Add a field to the entry. /// /// - /// \param[in] fieldName Name of the field to add. + /// \param[in] canonicalFieldName Canonical name of the field to add. This includes the prefix of any parent fields, + /// as well as the auxiliary processor name prefix, if present. /// \param[in] typeName Type of the field to add. /// \param[in] valuePtr Pointer to bind to the field's value in the entry. If this is a `nullptr`, a pointer will be /// created. @@ -452,7 +453,26 @@ protected: /// /// In case the field was already present in the entry, the index of the existing field is returned. virtual Internal::RNTupleProcessorEntry::FieldIndex_t - AddFieldToEntry(const std::string &fieldName, const std::string &typeName, void *valuePtr, + AddFieldToEntry(const std::string &canonicalFieldName, const std::string &typeName, void *valuePtr, + const Internal::RNTupleProcessorProvenance &provenance) = 0; + + ///////////////////////////////////////////////////////////////////////////// + /// \brief Add a field to the entry. + /// + /// + /// \param[in] field The field to add. + /// \param[in] canonicalFieldName Canonical name of the field to add. This includes the prefix of any parent fields, + /// as well as the auxiliary processor name prefix, if present. This is different from the field name stored in + /// `field`. + /// \param[in] valuePtr Pointer to bind to the field's value in the entry. If this is a `nullptr`, a pointer will be + /// created. + /// \param[in] provenance Provenance of the processor. + /// + /// \return The index of the newly added field in the entry. + /// + /// In case the field was already present in the entry, the index of the existing field is returned. + virtual Internal::RNTupleProcessorEntry::FieldIndex_t + AddFieldToEntry(std::unique_ptr field, const std::string &canonicalFieldName, void *valuePtr, const Internal::RNTupleProcessorProvenance &provenance) = 0; ///////////////////////////////////////////////////////////////////////////// @@ -726,6 +746,16 @@ private: std::unique_ptr CreateAndConnectField(const std::string &qualifiedFieldName, const std::string &typeName); + ///////////////////////////////////////////////////////////////////////////// + /// \brief Connect an existing field to the processor's page source. + /// + /// \param[in] field Pointer to the field to add + /// + /// \return The same field, but wrapped in a new unique pointer. This is needed, because in order to add the field to + /// the page source it first needs to be attached to its zero field, in order to correctly handle field + /// substitutions. Afterwards, we release the field again from the zero field. + std::unique_ptr ConnectField(std::unique_ptr field); + ///////////////////////////////////////////////////////////////////////////// /// \brief Initialize the processor by creating an (initially empty) `fEntry`, or setting an existing one. /// @@ -777,6 +807,14 @@ private: const std::string &fieldName, const std::string &typeName, void *valuePtr = nullptr, const Internal::RNTupleProcessorProvenance &provenance = Internal::RNTupleProcessorProvenance()) final; + ///////////////////////////////////////////////////////////////////////////// + /// \brief Add a field to the entry. + /// + /// \sa RNTupleProcessor::AddFieldToEntry() + Internal::RNTupleProcessorEntry::FieldIndex_t + AddFieldToEntry(std::unique_ptr field, const std::string &fieldName, void *valuePtr, + const Internal::RNTupleProcessorProvenance &provenance) final; + ///////////////////////////////////////////////////////////////////////////// /// \brief Add all known fields (from on-disk information) to the entry. /// @@ -889,6 +927,14 @@ private: const std::string &fieldName, const std::string &typeName, void *valuePtr = nullptr, const Internal::RNTupleProcessorProvenance &provenance = Internal::RNTupleProcessorProvenance()) final; + ///////////////////////////////////////////////////////////////////////////// + /// \brief Add a field to the entry. + /// + /// \sa RNTupleProcessor::AddFieldToEntry() + Internal::RNTupleProcessorEntry::FieldIndex_t + AddFieldToEntry(std::unique_ptr field, const std::string &fieldName, void *valuePtr, + const Internal::RNTupleProcessorProvenance &provenance) final; + ///////////////////////////////////////////////////////////////////////////// /// \brief Add all known fields (from on-disk information) to the entry. /// @@ -1005,6 +1051,14 @@ private: const std::string &fieldName, const std::string &typeName, void *valuePtr = nullptr, const Internal::RNTupleProcessorProvenance &provenance = Internal::RNTupleProcessorProvenance()) final; + ///////////////////////////////////////////////////////////////////////////// + /// \brief Add a field to the entry. + /// + /// \sa RNTupleProcessor::AddFieldToEntry() + Internal::RNTupleProcessorEntry::FieldIndex_t + AddFieldToEntry(std::unique_ptr field, const std::string &fieldName, void *valuePtr, + const Internal::RNTupleProcessorProvenance &provenance) final; + ///////////////////////////////////////////////////////////////////////////// /// \brief Add all known fields (from on-disk information) to the entry. /// diff --git a/tree/ntuple/src/RNTupleProcessor.cxx b/tree/ntuple/src/RNTupleProcessor.cxx index 1ba380a505a4c..14c54d6f9fb90 100644 --- a/tree/ntuple/src/RNTupleProcessor.cxx +++ b/tree/ntuple/src/RNTupleProcessor.cxx @@ -164,8 +164,6 @@ ROOT::Experimental::RNTupleSingleProcessor::CreateAndConnectField(const std::str auto descGuard = fPageSource->GetSharedDescriptorGuard(); const auto &desc = descGuard.GetRef(); - ROOT::RFieldZero fieldZero; - ROOT::Internal::SetAllowFieldSubstitutions(fieldZero, true); const auto onDiskFieldId = desc.FindFieldId(onDiskFieldName); @@ -188,20 +186,28 @@ ROOT::Experimental::RNTupleSingleProcessor::CreateAndConnectField(const std::str } field->SetOnDiskId(onDiskFieldId); + return ConnectField(std::move(field)); +} + +std::unique_ptr +ROOT::Experimental::RNTupleSingleProcessor::ConnectField(std::unique_ptr field) +{ + ROOT::RFieldZero fieldZero; + ROOT::Internal::SetAllowFieldSubstitutions(fieldZero, true); fieldZero.Attach(std::move(field)); ROOT::Internal::CallConnectPageSourceOnField(fieldZero, *fPageSource); return std::move(fieldZero.ReleaseSubfields()[0]); } ROOT::Experimental::Internal::RNTupleProcessorEntry::FieldIndex_t -ROOT::Experimental::RNTupleSingleProcessor::AddFieldToEntry(const std::string &fieldName, const std::string &typeName, - void *valuePtr, +ROOT::Experimental::RNTupleSingleProcessor::AddFieldToEntry(const std::string &canonicalFieldName, + const std::string &typeName, void *valuePtr, const Internal::RNTupleProcessorProvenance &provenance) { - auto fieldIdx = fEntry->FindFieldIndex(fieldName, typeName); + auto fieldIdx = fEntry->FindFieldIndex(canonicalFieldName, typeName); if (!fieldIdx) { // Strip the processor name prefix(es), if present. - std::string qualifiedFieldName = fieldName; + std::string qualifiedFieldName = canonicalFieldName; if (provenance.IsPresentInFieldName(qualifiedFieldName)) { qualifiedFieldName = qualifiedFieldName.substr(provenance.Get().size() + 1); } @@ -220,6 +226,33 @@ ROOT::Experimental::RNTupleSingleProcessor::AddFieldToEntry(const std::string &f return *fieldIdx; } +ROOT::Experimental::Internal::RNTupleProcessorEntry::FieldIndex_t +ROOT::Experimental::RNTupleSingleProcessor::AddFieldToEntry(std::unique_ptr field, + const std::string &canonicalFieldName, void *valuePtr, + const Internal::RNTupleProcessorProvenance &provenance) +{ + auto fieldIdx = fEntry->FindFieldIndex(canonicalFieldName, field->GetTypeName()); + if (!fieldIdx) { + // Strip the processor name prefix(es), if present. + std::string qualifiedFieldName = canonicalFieldName; + if (provenance.IsPresentInFieldName(qualifiedFieldName)) { + qualifiedFieldName = qualifiedFieldName.substr(provenance.Get().size() + 1); + } + + field = ConnectField(std::move(field)); + + if (!field) { + throw RException(R__FAIL("cannot register field with name \"" + qualifiedFieldName + + "\" because it is not present in the on-disk information of the RNTuple(s) this " + "processor is created from")); + } + + fieldIdx = fEntry->AddField(qualifiedFieldName, std::move(field), valuePtr, provenance); + } + + return *fieldIdx; +} + void ROOT::Experimental::RNTupleSingleProcessor::AddAllFieldsToEntry( const Internal::RNTupleProcessorProvenance &provenance, bool includeSubfields) { @@ -409,6 +442,14 @@ ROOT::Experimental::RNTupleChainProcessor::AddFieldToEntry(const std::string &fi return fInnerProcessors[fCurrentProcessorNumber]->AddFieldToEntry(fieldName, typeName, valuePtr, provenance); } +ROOT::Experimental::Internal::RNTupleProcessorEntry::FieldIndex_t +ROOT::Experimental::RNTupleChainProcessor::AddFieldToEntry(std::unique_ptr field, + const std::string &fieldName, void *valuePtr, + const Internal::RNTupleProcessorProvenance &provenance) +{ + return fInnerProcessors[fCurrentProcessorNumber]->AddFieldToEntry(std::move(field), fieldName, valuePtr, provenance); +} + void ROOT::Experimental::RNTupleChainProcessor::AddAllFieldsToEntry( const Internal::RNTupleProcessorProvenance &provenance, bool includeSubfields) { @@ -607,6 +648,36 @@ ROOT::Experimental::RNTupleJoinProcessor::AddFieldToEntry(const std::string &fie } } +ROOT::Experimental::Internal::RNTupleProcessorEntry::FieldIndex_t +ROOT::Experimental::RNTupleJoinProcessor::AddFieldToEntry(std::unique_ptr field, + const std::string &fieldName, void *valuePtr, + const Internal::RNTupleProcessorProvenance &provenance) +{ + auto auxProvenance = provenance.Evolve(fAuxiliaryProcessor->GetOptions().GetProcessorName()); + if (auxProvenance.IsPresentInFieldName(fieldName)) { + // If the primaryProcessor has a field with the name of the auxProcessor (either as a "proper" field or because + // the primary processor itself is a join where its auxProcessor bears the same name as the current auxProcessor), + // there will be name conflicts, so error out. + if (fPrimaryProcessor->CanReadFieldFromDisk(fieldName)) { + throw RException(R__FAIL("ambiguous field name: \"" + fieldName + + "\" is present in the primary RNTupleProcessor \"" + + fPrimaryProcessor->GetOptions().GetProcessorName() + + "\", but may also refer to a field in the auxiliary RNTupleProcessor named \"" + + fAuxiliaryProcessor->GetOptions().GetProcessorName() + + "\". To avoid this ambiguity, rename the auxiliary RNTupleProcessor.")); + } + + auto fieldIdx = fAuxiliaryProcessor->AddFieldToEntry(std::move(field), fieldName, valuePtr, auxProvenance); + fAuxiliaryFieldIdxs.insert(fieldIdx); + return fieldIdx; + } else { + auto fieldIdx = fPrimaryProcessor->AddFieldToEntry(std::move(field), fieldName, valuePtr, provenance); + if (fieldIdx) + fFieldIdxs.insert(fieldIdx); + return fieldIdx; + } +} + void ROOT::Experimental::RNTupleJoinProcessor::AddAllFieldsToEntry( const Internal::RNTupleProcessorProvenance &provenance, bool includeSubfields) { From 26b2e3fb86a71be960587fb227373695dd741d88 Mon Sep 17 00:00:00 2001 From: Florine de Geus Date: Wed, 23 Sep 2026 15:42:01 +0200 Subject: [PATCH 3/8] [ntuple] Enable re-requesting fields from a processor --- tree/ntuple/inc/ROOT/RNTupleProcessor.hxx | 25 +++++++++++++++++++++++ tree/ntuple/test/ntuple_processor.cxx | 22 ++++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/tree/ntuple/inc/ROOT/RNTupleProcessor.hxx b/tree/ntuple/inc/ROOT/RNTupleProcessor.hxx index 0f2b3705d7d51..4fdfec2b52bb7 100644 --- a/tree/ntuple/inc/ROOT/RNTupleProcessor.hxx +++ b/tree/ntuple/inc/ROOT/RNTupleProcessor.hxx @@ -175,6 +175,11 @@ public: throw RException(R__FAIL("cannot read \"" + fProcessorEntry->FindFieldName(fFieldIndex) + "\" because it has no value for the current entry")); } + + friend bool operator==(const RNTupleProcessorOptionalPtr &lhs, const RNTupleProcessorOptionalPtr &rhs) + { + return lhs.fProcessorEntry == rhs.fProcessorEntry && lhs.fFieldIndex == rhs.fFieldIndex; + } }; // clang-format off @@ -246,6 +251,11 @@ public: /// invalid data. After binding a pointer to an `RNTupleProcessorOptionalPtr`, we *strongly* recommend only accessing /// its data through this interface, to ensure that only valid data can be read. void Bind(std::shared_ptr valuePtr) { fProcessorEntry->Bind(fFieldIndex, std::move(valuePtr)); } + + friend bool operator==(const RNTupleProcessorOptionalPtr &lhs, const RNTupleProcessorOptionalPtr &rhs) + { + return lhs.fProcessorEntry == rhs.fProcessorEntry && lhs.fFieldIndex == rhs.fFieldIndex; + } }; // clang-format off @@ -543,6 +553,21 @@ public: if constexpr (!std::is_void_v) { typeName = ROOT::Internal::GetRenormalizedTypeName(typeid(T)); } + + // The field already exists, so return the existing one. + if (auto fieldIdx = fEntry->FindFieldIndex(fieldName, typeName)) { + auto value = fEntry->GetValue(*fieldIdx); + // Need to check that the provided pointer is not conflicting with the existing one + + if (valuePtr != nullptr && valuePtr != value.GetPtr().get()) { + throw RException(R__FAIL("attempted to request a field with user-provided value pointer to field \"" + + fieldName + + "\", which already exists in the entry. To change the underlying value pointer, " + "use RNTupleProcessorOptionalPtr::Bind instead.")); + } + return RNTupleProcessorOptionalPtr(fEntry.get(), *fieldIdx); + } + auto fieldIdx = AddFieldToEntry(fieldName, typeName, valuePtr, Internal::RNTupleProcessorProvenance()); return RNTupleProcessorOptionalPtr(fEntry.get(), fieldIdx); } diff --git a/tree/ntuple/test/ntuple_processor.cxx b/tree/ntuple/test/ntuple_processor.cxx index c3e788bc3f4aa..a65919fc7d8fb 100644 --- a/tree/ntuple/test/ntuple_processor.cxx +++ b/tree/ntuple/test/ntuple_processor.cxx @@ -319,6 +319,28 @@ TEST_F(RNTupleProcessorTest, RequestFieldWithTypeString) EXPECT_EQ(5, proc->GetNEntriesProcessed()); } +TEST_F(RNTupleProcessorTest, RequestExistingField) +{ + auto proc = RNTupleProcessor::Create({fNTupleNames[0], fFileNames[0]}); + + auto fldX1 = proc->RequestField("x"); + auto fldX2 = proc->RequestField("x"); + + EXPECT_EQ(fldX1, fldX2); + + try { + float x; + proc->RequestField("x", &x); + FAIL() << "requesting an existing field with a user-provided pointer should throw"; + } catch (const ROOT::RException &err) { + EXPECT_THAT( + err.what(), + testing::HasSubstr( + "attempted to request a field with user-provided value pointer to field \"x\", which already exists in the " + "entry. To change the underlying value pointer, use RNTupleProcessorOptionalPtr::Bind instead.")); + } +} + TEST_F(RNTupleProcessorTest, AlternativeTypes) { auto proc = RNTupleProcessor::Create({fNTupleNames[0], fFileNames[0]}); From 14154074a6f4a4b673abbf20ffd67338564e19d3 Mon Sep 17 00:00:00 2001 From: Florine de Geus Date: Wed, 23 Sep 2026 16:31:03 +0200 Subject: [PATCH 4/8] [ntuple] Add `RNTupleProcessorEntry::GetField` --- tree/ntuple/inc/ROOT/RNTupleProcessorEntry.hxx | 8 ++++++++ tree/ntuple/src/RNTupleProcessorEntry.cxx | 6 ++++++ 2 files changed, 14 insertions(+) diff --git a/tree/ntuple/inc/ROOT/RNTupleProcessorEntry.hxx b/tree/ntuple/inc/ROOT/RNTupleProcessorEntry.hxx index 3bb30f011d84b..558ecc964fa9e 100644 --- a/tree/ntuple/inc/ROOT/RNTupleProcessorEntry.hxx +++ b/tree/ntuple/inc/ROOT/RNTupleProcessorEntry.hxx @@ -195,6 +195,14 @@ public: /// \param[in] field The new field to use in the entry. void UpdateField(FieldIndex_t fieldIdx, std::unique_ptr field); + ///////////////////////////////////////////////////////////////////////////// + /// \brief Get a field by its index. + /// + /// \param[in] fieldIdx Index of the field. + /// + /// \return A constref to the field. + const ROOT::RFieldBase &GetField(FieldIndex_t fieldIdx) const; + ///////////////////////////////////////////////////////////////////////////// /// \brief Bind a new value pointer to a field in the entry. /// diff --git a/tree/ntuple/src/RNTupleProcessorEntry.cxx b/tree/ntuple/src/RNTupleProcessorEntry.cxx index 5b6727d5f4418..1e72e99a5cf46 100644 --- a/tree/ntuple/src/RNTupleProcessorEntry.cxx +++ b/tree/ntuple/src/RNTupleProcessorEntry.cxx @@ -76,6 +76,12 @@ ROOT::Experimental::Internal::RNTupleProcessorEntry::AddField(const std::string return fieldIdx; } +const ROOT::RFieldBase &ROOT::Experimental::Internal::RNTupleProcessorEntry::GetField(FieldIndex_t fieldIdx) const +{ + assert(fieldIdx < fProcessorValues.size()); + return *fProcessorValues[fieldIdx].fField; +} + void ROOT::Experimental::Internal::RNTupleProcessorEntry::UpdateField(FieldIndex_t fieldIdx, std::unique_ptr field) { From f275f1314e29f3ffd5524ec0bf2bdfe23745aafa Mon Sep 17 00:00:00 2001 From: Florine de Geus Date: Wed, 23 Sep 2026 16:32:11 +0200 Subject: [PATCH 5/8] [ntuple] Improve handling multiple present types for `FindFieldIndex` --- tree/ntuple/inc/ROOT/RNTupleProcessorEntry.hxx | 5 ++++- tree/ntuple/src/RNTupleProcessorEntry.cxx | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/tree/ntuple/inc/ROOT/RNTupleProcessorEntry.hxx b/tree/ntuple/inc/ROOT/RNTupleProcessorEntry.hxx index 558ecc964fa9e..e99f3d67816d4 100644 --- a/tree/ntuple/inc/ROOT/RNTupleProcessorEntry.hxx +++ b/tree/ntuple/inc/ROOT/RNTupleProcessorEntry.hxx @@ -171,9 +171,12 @@ public: /// /// \param[in] canonicalFieldName The name of the field in the entry, including its processor name prefixes and /// parent field names, if applicable. + /// \param[in] typeName Type of the field, if relevant. If no type name is provided, the first field corresponding to + /// the provided name is returned. /// /// \return A `std::optional` containing the field index if it was found. - std::optional FindFieldIndex(std::string_view canonicalFieldName, std::string_view typeName) const; + std::optional + FindFieldIndex(std::string_view canonicalFieldName, std::string_view typeName = "") const; ///////////////////////////////////////////////////////////////////////////// /// \brief Add a new field to the entry. diff --git a/tree/ntuple/src/RNTupleProcessorEntry.cxx b/tree/ntuple/src/RNTupleProcessorEntry.cxx index 1e72e99a5cf46..355e8cef185d9 100644 --- a/tree/ntuple/src/RNTupleProcessorEntry.cxx +++ b/tree/ntuple/src/RNTupleProcessorEntry.cxx @@ -43,7 +43,7 @@ ROOT::Experimental::Internal::RNTupleProcessorEntry::FindFieldIndex(std::string_ assert(!fieldIdxs.empty()); for (auto idx : fieldIdxs) { - if (fProcessorValues[idx].fField->GetTypeName() == typeName) { + if (typeName.empty() || fProcessorValues[idx].fField->GetTypeName() == typeName) { return idx; } } From ddca9fa7897e2a82f736b066a6bcbb14af4d6ace Mon Sep 17 00:00:00 2001 From: Florine de Geus Date: Wed, 23 Sep 2026 17:24:50 +0200 Subject: [PATCH 6/8] [df] Move RNTupleDS artificial fields to separate header In preparation for the addition of RNTupleProcessorDS. --- tree/dataframe/src/RNTupleDS.cxx | 147 +-------------------- tree/dataframe/src/RNTupleDSFields.hxx | 170 +++++++++++++++++++++++++ 2 files changed, 173 insertions(+), 144 deletions(-) create mode 100644 tree/dataframe/src/RNTupleDSFields.hxx diff --git a/tree/dataframe/src/RNTupleDS.cxx b/tree/dataframe/src/RNTupleDS.cxx index 5f8a5b45b6400..59ec6f0191ec5 100644 --- a/tree/dataframe/src/RNTupleDS.cxx +++ b/tree/dataframe/src/RNTupleDS.cxx @@ -21,7 +21,8 @@ #include #include #include -#include + +#include "RNTupleDSFields.hxx" #include #include @@ -31,6 +32,7 @@ #include #include #include +#include #include #include #include @@ -51,149 +53,6 @@ // clang-format on namespace ROOT::Internal::RDF { -class RRDFCardinalityFieldBase : public ROOT::RFieldBase { -protected: - // We construct these fields and know that they match the page source - void ReconcileOnDiskField(const RNTupleDescriptor &) final {} - - RRDFCardinalityFieldBase(std::string_view name, std::string_view type) - : ROOT::RFieldBase(name, type, ROOT::ENTupleStructure::kPlain, false /* isSimple */) - { - } - - // Field is only used for reading - void GenerateColumns() final { throw RException(R__FAIL("Cardinality fields must only be used for reading")); } - void GenerateColumns(const ROOT::RNTupleDescriptor &desc) final - { - GenerateColumnsImpl(desc); - } - -public: - RRDFCardinalityFieldBase(const RRDFCardinalityFieldBase &other) = delete; - RRDFCardinalityFieldBase &operator=(const RRDFCardinalityFieldBase &other) = delete; - RRDFCardinalityFieldBase(RRDFCardinalityFieldBase &&other) = default; - RRDFCardinalityFieldBase &operator=(RRDFCardinalityFieldBase &&other) = default; - ~RRDFCardinalityFieldBase() override = default; - - const RColumnRepresentations &GetColumnRepresentations() const final - { - static RColumnRepresentations representations({{ENTupleColumnType::kSplitIndex64}, - {ENTupleColumnType::kIndex64}, - {ENTupleColumnType::kSplitIndex32}, - {ENTupleColumnType::kIndex32}}, - {}); - return representations; - } -}; - -/// An artificial field that transforms an RNTuple column that contains the offset of collections into -/// collection sizes. It is used to provide the "number of" RDF columns for collections, e.g. -/// `R_rdf_sizeof_jets` for a collection named `jets`. -/// -/// This is similar to the RCardinalityField but it presents itself as an integer type. -/// The template argument T must be an integral type. -template -class RRDFCardinalityField final : public RRDFCardinalityFieldBase { - static_assert(std::is_integral_v, "T must be an integral type"); - - inline void CheckSize(ROOT::NTupleSize_t size) const - { - if constexpr (std::is_same_v || std::is_same_v) - return; - if (size > static_cast(std::numeric_limits::max())) { - throw RException(R__FAIL(std::string("integer overflow in field ") + GetFieldName() + - ". Please read the column with a larger-sized integral type.")); - } - } - -protected: - std::unique_ptr CloneImpl(std::string_view newName) const final - { - return std::make_unique(newName); - } - void ConstructValue(void *where) const final { *static_cast(where) = 0; } - -public: - RRDFCardinalityField(std::string_view name) - : RRDFCardinalityFieldBase(name, ROOT::Internal::GetRenormalizedTypeName(typeid(T))) - { - } - RRDFCardinalityField(const RRDFCardinalityField &other) = delete; - RRDFCardinalityField &operator=(const RRDFCardinalityField &other) = delete; - RRDFCardinalityField(RRDFCardinalityField &&other) = default; - RRDFCardinalityField &operator=(RRDFCardinalityField &&other) = default; - ~RRDFCardinalityField() override = default; - - std::size_t GetValueSize() const final { return sizeof(T); } - std::size_t GetAlignment() const final { return alignof(T); } - - /// Get the number of elements of the collection identified by globalIndex - void ReadGlobalImpl(ROOT::NTupleSize_t globalIndex, void *to) final - { - RNTupleLocalIndex collectionStart; - ROOT::NTupleSize_t size; - fPrincipalColumn->GetCollectionInfo(globalIndex, &collectionStart, &size); - CheckSize(size); - *static_cast(to) = size; - } - - /// Get the number of elements of the collection identified by clusterIndex - void ReadInClusterImpl(ROOT::RNTupleLocalIndex localIndex, void *to) final - { - RNTupleLocalIndex collectionStart; - ROOT::NTupleSize_t size; - fPrincipalColumn->GetCollectionInfo(localIndex, &collectionStart, &size); - CheckSize(size); - *static_cast(to) = size; - } -}; - -/** - * @brief An artificial field that provides the size of a fixed-size array - * - * This is the implementation of `R_rdf_sizeof_column` in case `column` contains - * fixed-size arrays on disk. - */ -class RArraySizeField final : public ROOT::RFieldBase { -private: - std::size_t fArrayLength; - - std::unique_ptr CloneImpl(std::string_view newName) const final - { - return std::make_unique(newName, fArrayLength); - } - void GenerateColumns() final { throw RException(R__FAIL("RArraySizeField fields must only be used for reading")); } - void GenerateColumns(const ROOT::RNTupleDescriptor &) final {} - void ReadGlobalImpl(ROOT::NTupleSize_t /*globalIndex*/, void *to) final - { - *static_cast(to) = fArrayLength; - } - void ReadInClusterImpl(RNTupleLocalIndex /*localIndex*/, void *to) final - { - *static_cast(to) = fArrayLength; - } - - // We construct these fields and know that they match the page source - void ReconcileOnDiskField(const RNTupleDescriptor &) final {} - -public: - RArraySizeField(std::string_view name, std::size_t arrayLength) - : ROOT::RFieldBase(name, ROOT::Internal::GetRenormalizedTypeName(typeid(std::size_t)), - ROOT::ENTupleStructure::kPlain, false /* isSimple */), - fArrayLength(arrayLength) - { - } - RArraySizeField(const RArraySizeField &other) = delete; - RArraySizeField &operator=(const RArraySizeField &other) = delete; - RArraySizeField(RArraySizeField &&other) = default; - RArraySizeField &operator=(RArraySizeField &&other) = default; - ~RArraySizeField() final = default; - - void ConstructValue(void *where) const final { *static_cast(where) = 0; } - std::size_t GetValueSize() const final { return sizeof(std::size_t); } - std::size_t GetAlignment() const final { return alignof(std::size_t); } -}; - /// Every RDF column is represented by exactly one RNTuple field class RNTupleColumnReader : public ROOT::Detail::RDF::RColumnReaderBase { using RFieldBase = ROOT::RFieldBase; diff --git a/tree/dataframe/src/RNTupleDSFields.hxx b/tree/dataframe/src/RNTupleDSFields.hxx new file mode 100644 index 0000000000000..06a49b9628e1f --- /dev/null +++ b/tree/dataframe/src/RNTupleDSFields.hxx @@ -0,0 +1,170 @@ +/// \file RNTupleDSFields.cxx +/// \author Jakob Blomer +/// \author Enrico Guiraud +/// \date 2026-09-23 + +#include + +// clang-format off +/** + * \class ROOT::Internal::RDF::RRDFCardinalityFieldBase + * \ingroup dataframe + * \brief Base for an artificial field that transforms an RNTuple column containing the offset of collections into collection sizes. + * + * It is used to provide the "number of" RDF columns for collections, e.g. `R_rdf_sizeof_jets` for a collection named `jets`. + */ +// clang-format on +namespace ROOT::Internal::RDF { +class RRDFCardinalityFieldBase : public ROOT::RFieldBase { +protected: + // We construct these fields and know that they match the page source + void ReconcileOnDiskField(const RNTupleDescriptor &) final {} + + RRDFCardinalityFieldBase(std::string_view name, std::string_view type) + : ROOT::RFieldBase(name, type, ROOT::ENTupleStructure::kPlain, false /* isSimple */) + { + } + + // Field is only used for reading + void GenerateColumns() final { throw RException(R__FAIL("Cardinality fields must only be used for reading")); } + void GenerateColumns(const ROOT::RNTupleDescriptor &desc) final + { + GenerateColumnsImpl(desc); + } + +public: + RRDFCardinalityFieldBase(const RRDFCardinalityFieldBase &other) = delete; + RRDFCardinalityFieldBase &operator=(const RRDFCardinalityFieldBase &other) = delete; + RRDFCardinalityFieldBase(RRDFCardinalityFieldBase &&other) = default; + RRDFCardinalityFieldBase &operator=(RRDFCardinalityFieldBase &&other) = default; + ~RRDFCardinalityFieldBase() override = default; + + const RColumnRepresentations &GetColumnRepresentations() const final + { + static const RColumnRepresentations representations({{ENTupleColumnType::kSplitIndex64}, + {ENTupleColumnType::kIndex64}, + {ENTupleColumnType::kSplitIndex32}, + {ENTupleColumnType::kIndex32}}, + {}); + return representations; + } +}; + +// clang-format off +/** + * \class ROOT::Internal::RDF::RRDFCardinalityField + * \ingroup dataframe + * \brief An artificial field that transforms an RNTuple column containing the offset of collections into collection sizes. + * + * It is used to provide the "number of" RDF columns for collections, e.g. `R_rdf_sizeof_jets` for a collection named `jets`. + * + * This is similar to the RCardinalityField but it presents itself as an integer type. + * The template argument T must be an integral type. + */ +// clang-format on +template +class RRDFCardinalityField final : public RRDFCardinalityFieldBase { + static_assert(std::is_integral_v, "T must be an integral type"); + + inline void CheckSize(ROOT::NTupleSize_t size) const + { + if constexpr (std::is_same_v || std::is_same_v) + return; + if (size > static_cast(std::numeric_limits::max())) { + throw RException(R__FAIL(std::string("integer overflow in field ") + GetFieldName() + + ". Please read the column with a larger-sized integral type.")); + } + } + +protected: + std::unique_ptr CloneImpl(std::string_view newName) const final + { + return std::make_unique(newName); + } + void ConstructValue(void *where) const final { *static_cast(where) = 0; } + +public: + RRDFCardinalityField(std::string_view name) + : RRDFCardinalityFieldBase(name, ROOT::Internal::GetRenormalizedTypeName(typeid(T))) + { + } + RRDFCardinalityField(const RRDFCardinalityField &other) = delete; + RRDFCardinalityField &operator=(const RRDFCardinalityField &other) = delete; + RRDFCardinalityField(RRDFCardinalityField &&other) = default; + RRDFCardinalityField &operator=(RRDFCardinalityField &&other) = default; + ~RRDFCardinalityField() override = default; + + std::size_t GetValueSize() const final { return sizeof(T); } + std::size_t GetAlignment() const final { return alignof(T); } + + /// Get the number of elements of the collection identified by globalIndex + void ReadGlobalImpl(ROOT::NTupleSize_t globalIndex, void *to) final + { + RNTupleLocalIndex collectionStart; + ROOT::NTupleSize_t size; + fPrincipalColumn->GetCollectionInfo(globalIndex, &collectionStart, &size); + CheckSize(size); + *static_cast(to) = size; + } + + /// Get the number of elements of the collection identified by clusterIndex + void ReadInClusterImpl(ROOT::RNTupleLocalIndex localIndex, void *to) final + { + RNTupleLocalIndex collectionStart; + ROOT::NTupleSize_t size; + fPrincipalColumn->GetCollectionInfo(localIndex, &collectionStart, &size); + CheckSize(size); + *static_cast(to) = size; + } +}; + +// clang-format off +/** + * \class ROOT::Internal::RDF::RArraySizeField + * \ingroup dataframe + * \brief An artificial field that provides the size of a fixed-size array + * + * This is the implementation of `R_rdf_sizeof_column` in case `column` contains + * fixed-size arrays on disk. + */ +// clang-format on +class RArraySizeField final : public ROOT::RFieldBase { +private: + std::size_t fArrayLength; + + std::unique_ptr CloneImpl(std::string_view newName) const final + { + return std::make_unique(newName, fArrayLength); + } + void GenerateColumns() final { throw RException(R__FAIL("RArraySizeField fields must only be used for reading")); } + void GenerateColumns(const ROOT::RNTupleDescriptor &) final {} + void ReadGlobalImpl(ROOT::NTupleSize_t /*globalIndex*/, void *to) final + { + *static_cast(to) = fArrayLength; + } + void ReadInClusterImpl(RNTupleLocalIndex /*localIndex*/, void *to) final + { + *static_cast(to) = fArrayLength; + } + + // We construct these fields and know that they match the page source + void ReconcileOnDiskField(const RNTupleDescriptor &) final {} + +public: + RArraySizeField(std::string_view name, std::size_t arrayLength) + : ROOT::RFieldBase(name, ROOT::Internal::GetRenormalizedTypeName(typeid(std::size_t)), + ROOT::ENTupleStructure::kPlain, false /* isSimple */), + fArrayLength(arrayLength) + { + } + RArraySizeField(const RArraySizeField &other) = delete; + RArraySizeField &operator=(const RArraySizeField &other) = delete; + RArraySizeField(RArraySizeField &&other) = default; + RArraySizeField &operator=(RArraySizeField &&other) = default; + ~RArraySizeField() final = default; + + void ConstructValue(void *where) const final { *static_cast(where) = 0; } + std::size_t GetValueSize() const final { return sizeof(std::size_t); } + std::size_t GetAlignment() const final { return alignof(std::size_t); } +}; +} // namespace ROOT::Internal::RDF From 9ae207273b7318192b7c834bbfad0b42a2af5ed2 Mon Sep 17 00:00:00 2001 From: Florine de Geus Date: Wed, 23 Sep 2026 16:34:32 +0200 Subject: [PATCH 7/8] [df] Add initial `RNTupleProcessorDS` In the experimental namespace. For now, it is single-threaded only. In principle, the full type system (as tested for the regular RNTuple DS) works, with the exception of arrays that are part of classes. This will be addressed in a future commit. --- tree/dataframe/CMakeLists.txt | 2 + .../dataframe/inc/ROOT/RNTupleProcessorDS.hxx | 138 +++ tree/dataframe/src/RNTupleProcessorDS.cxx | 433 +++++++++ tree/dataframe/test/CMakeLists.txt | 2 + .../test/datasource_ntuple_processor.cxx | 855 ++++++++++++++++++ tree/ntuple/inc/ROOT/RNTupleProcessor.hxx | 14 + .../ntuple/inc/ROOT/RNTupleProcessorEntry.hxx | 4 +- tree/ntuple/src/RNTupleProcessor.cxx | 1 + 8 files changed, 1448 insertions(+), 1 deletion(-) create mode 100644 tree/dataframe/inc/ROOT/RNTupleProcessorDS.hxx create mode 100644 tree/dataframe/src/RNTupleProcessorDS.cxx create mode 100644 tree/dataframe/test/datasource_ntuple_processor.cxx diff --git a/tree/dataframe/CMakeLists.txt b/tree/dataframe/CMakeLists.txt index b16e21c7506b0..6792071dcd5c8 100644 --- a/tree/dataframe/CMakeLists.txt +++ b/tree/dataframe/CMakeLists.txt @@ -102,6 +102,7 @@ ROOT_STANDARD_LIBRARY_PACKAGE(ROOTDataFrame ROOT/RDF/PyROOTHelpers.hxx ROOT/RDF/RDFDescription.hxx ROOT/RNTupleDS.hxx + ROOT/RNTupleProcessorDS.hxx ${RDATAFRAME_EXTRA_HEADERS} SOURCES src/RActionBase.cxx @@ -144,6 +145,7 @@ ROOT_STANDARD_LIBRARY_PACKAGE(ROOTDataFrame src/RTrivialDS.cxx src/RDFDescription.cxx src/RNTupleDS.cxx + src/RNTupleProcessorDS.cxx DICTIONARY_OPTIONS -writeEmptyRootPCM ${RDATAFRAME_EXTRA_INCLUDES} diff --git a/tree/dataframe/inc/ROOT/RNTupleProcessorDS.hxx b/tree/dataframe/inc/ROOT/RNTupleProcessorDS.hxx new file mode 100644 index 0000000000000..1ca79627a3058 --- /dev/null +++ b/tree/dataframe/inc/ROOT/RNTupleProcessorDS.hxx @@ -0,0 +1,138 @@ +/// \file RNTupleProcessorDS.hxx +/// \author Florine de Geus +/// \author Jakob Blomer +/// \author Enrico Guiraud +/// \date 2026-06-18 +/// \warning This is part of the ROOT 7 prototype! It will change without notice. It might trigger earthquakes. Feedback +/// is welcome! + +#ifndef ROOT_RNTupleProcessorDS +#define ROOT_RNTupleProcessorDS + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ROOT { +class RDataFrame; +} // namespace ROOT + +// clang-format off +/** +* \class ROOT::Experimental::RDF::RNTupleProcessorDS +* \ingroup dataframe +* \brief The RDataSource implementation for the RNTupleProcessor. It lets RDataFrame read combined (chained and joined) RNTuple data. +* +* An RDataFrame that reads combined (chained and joined) RNTuple data, through the RNTupleProcessor. +* It can be constructed using FromRNTupleProcessor(). +* +* For each column containing an array or a collection, a corresponding column `#colname` is available to access +* `colname.size()` without reading and deserializing the collection values. +* +**/ +// clang-format on +namespace ROOT::Experimental::Internal::RDF { +class RNTupleProcessorDS final : public ROOT::RDF::RDataSource { + // TODO: turn into vector of processors and entry mappings for MT + std::unique_ptr fProcessor; + RNTupleProcessor::REntryMapping fEntryMapping; + + std::vector fColumnNames; + std::vector fColumnTypes; + + std::vector fColumnReaders; + + std::vector> fCurrentRanges; + + ///////////////////////////////////////////////////////////////////////////// + /// \brief Holds useful information about fields added to the RNTupleProcessorDS + struct RFieldInfo { + ROOT::DescriptorId_t fFieldId; + std::string fFieldName; + std::string fTypeName; + std::size_t fNRepetitions; + // 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) + { + } + }; + + ///////////////////////////////////////////////////////////////////////////// + /// \brief Provides the RDF column "colName" given the field. + /// + /// \param[in] field Field to add. + /// \param[in] colName RDF column name corresponding to the field. + /// \param[in] procProvenance Processor provenance of the field to add. + /// \param[in] fieldInfos Additional information about the field's outer collections. + /// \param[in] convertToRVec Whether to add collections (`std::vector` or `std::array`) as ROOT::VecOps::RVec + /// (enabled by default). + /// + /// For records and collections, AddField recurses into the subfields. For instance, if the provided field refers to + /// an `std::vector`, with + /// ~~~{.cpp} + /// struct Jet { + /// float pt; + /// float eta; + /// }; + /// ~~~ + /// AddField will recurse into `Jet.pt` and `Jet.eta` and provide the two inner fields as `ROOT::VecOps::RVec` + /// each. + /// + /// When `convertToRVec` is true, the field's corresponding column is added as a `ROOT::VecOps::RVec`. Otherwise, the + /// collection field's on-disk type is used. Note, however, that inner record members of such collections will still + /// be added as `ROOT::VecOps::RVec` (e.g., `std::set will be added as a `std::set`, but `Jet.[pt|eta] will be + /// added as `ROOT::VecOps::RVec). + void AddField(const ROOT::RFieldBase &field, std::string_view colName, + Internal::RNTupleProcessorProvenance procProvenance, std::deque &fieldInfos, + bool convertToRVec = true); + +public: + RNTupleProcessorDS(std::unique_ptr processor); + RNTupleProcessorDS(const RNTupleProcessorDS &) = delete; + RNTupleProcessorDS &operator=(const RNTupleProcessorDS &) = delete; + RNTupleProcessorDS(RNTupleProcessorDS &&) = delete; + RNTupleProcessorDS &operator=(RNTupleProcessorDS &&) = delete; + ~RNTupleProcessorDS() final; + + void SetNSlots(unsigned int nSlots) final; + // FIXME(fdegeus) get correct number of files (needs to be added in RNTupleProcessor) + std::size_t GetNFiles() const final { return 1; } + const std::vector &GetColumnNames() const final { return fColumnNames; } + bool HasColumn(std::string_view colName) const final; + std::string GetTypeName(std::string_view colName) const final; + std::vector> GetEntryRanges() final; + std::string GetLabel() final { return "RNTupleProcessorDS"; } + + void Initialize() final; + void InitSlot(unsigned int slot, ULong64_t firstEntry) final; + void FinalizeSlot(unsigned int slot) final; + void Finalize() final; + + std::unique_ptr + GetColumnReaders(unsigned int slot, std::string_view name, const std::type_info &) final; + + ROOT::RDF::RSampleInfo + CreateSampleInfo(unsigned int, + const std::unordered_map &) const final; + + bool SetEntry(unsigned int, ULong64_t) final; + +protected: + Record_t GetColumnReadersImpl(std::string_view name, const std::type_info &) final; +}; +} // namespace ROOT::Experimental::Internal::RDF + +namespace ROOT::Experimental::RDF { +RDataFrame FromRNTupleProcessor(std::unique_ptr processor); +} // namespace ROOT::Experimental::RDF + +#endif // ROOT_RNTupleProcessorDS diff --git a/tree/dataframe/src/RNTupleProcessorDS.cxx b/tree/dataframe/src/RNTupleProcessorDS.cxx new file mode 100644 index 0000000000000..04bd48828659a --- /dev/null +++ b/tree/dataframe/src/RNTupleProcessorDS.cxx @@ -0,0 +1,433 @@ +/// \file RNTupleProcessorDS.cxx +/// \author Florine de Geus +/// \author Jakob Blomer +/// \author Enrico Guiraud +/// \date 2026-06-18 +/// \warning This is part of the ROOT 7 prototype! It will change without notice. It might trigger earthquakes. Feedback +/// is welcome! + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "RNTupleDSFields.hxx" + +#include +#include + +#include +#include +#include +#include + +namespace ROOT::Experimental::Internal::RDF { +/// Every RDF column is represented by exactly one RNTuple field +class RNTupleProcessorColumnReader : public ROOT::Detail::RDF::RColumnReaderBase { + RNTupleProcessorOptionalPtr fValuePtr; ///< Container for the value of the column + Long64_t fLastLoadedEntry = ROOT::kInvalidNTupleIndex; + +public: + RNTupleProcessorColumnReader(RNTupleProcessorOptionalPtr &valuePtr) : fValuePtr(std::move(valuePtr)) {} + ~RNTupleProcessorColumnReader() override = default; + + void *GetImpl(Long64_t entry) final + { + assert(entry == fLastLoadedEntry); + return fValuePtr.GetPtr().get(); + } + + void SetEntryLoaded(Long64_t entry) { fLastLoadedEntry = entry; } +}; + +RNTupleProcessorDS::~RNTupleProcessorDS() = default; + +// TODO: see if common parts of this and RNTupleDS::AddField can be factored out. +void RNTupleProcessorDS::AddField(const ROOT::RFieldBase &field, std::string_view colName, + Internal::RNTupleProcessorProvenance procProvenance, + std::deque &fieldInfos, bool convertToRVec) +{ + // As an example for the mapping of RNTuple fields to RDF columns, let's consider an RNTuple + // using the following types and with a top-level field named "event" of type Event: + // + // struct Event { + // int id; + // std::vector tracks; + // }; + // struct Track { + // std::vector hits; + // }; + // struct Hit { + // float x; + // float y; + // }; + // + // AddField() will be called from the constructor with the RNTuple root field (ENTupleStructure::kRecord). + // From there, we recurse into the "event" sub field (also ENTupleStructure::kRecord) and further down the + // tree of sub fields and expose the following RDF columns: + // + // "event" [Event] + // "event.id" [int] + // "event.tracks" [RVec] + // "R_rdf_sizeof_event.tracks" [unsigned int] + // "event.tracks.hits" [RVec>] + // "R_rdf_sizeof_event.tracks.hits" [RVec] + // "event.tracks.hits.x" [RVec>] + // "R_rdf_sizeof_event.tracks.hits.x" [RVec] + // "event.tracks.hits.y" [RVec>] + // "R_rdf_sizeof_event.tracks.hits.y" [RVec] + + std::uint16_t nFieldInfosAdded = 0; + + const auto &nRepetitions = field.GetNRepetitions(); + if ((field.GetStructure() == ROOT::ENTupleStructure::kCollection) || (nRepetitions > 0)) { + // 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); + nFieldInfosAdded++; + } + + if (field.GetStructure() == ROOT::ENTupleStructure::kCollection) { + // Inner fields of collections are provided as projected collections of only that inner field, + // E.g. we provide a projected collection RVec> for "event.tracks.hits.x" in the example + // above. + bool representableAsRVec = + convertToRVec && (field.GetTypeName().substr(0, 19) == "ROOT::VecOps::RVec<" || + field.GetTypeName().substr(0, 12) == "std::vector<" || field.GetTypeName() == ""); + const auto *f = field.GetConstSubfields()[0]; + AddField(*f, colName, procProvenance, fieldInfos, representableAsRVec); + + // Note that at the end of the recursion, we handled the inner sub collections as well as the + // collection as whole, so we are done. + return; + + } else if (nRepetitions > 0) { + // Fixed-size array, same logic as ROOT::RVec. + const auto *f = field.GetConstSubfields()[0]; + AddField(*f, colName, procProvenance, fieldInfos); + return; + } else if (field.GetStructure() == ROOT::ENTupleStructure::kRecord) { + // Inner fields of records are provided as individual RDF columns, e.g. "event.id" + for (const auto &f : field.GetConstSubfields()) { + auto innerName = colName.empty() ? f->GetFieldName() : (std::string(colName) + "." + f->GetFieldName()); + // Inner fields of collections of records are always exposed as ROOT::RVec + AddField(*f, innerName, procProvenance, fieldInfos); + } + + // Do not add untyped record fields + if (field.GetTypeName() == "") + return; + } + + // 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()); + if (!fieldOrException) + return; + auto valueField = fieldOrException.Unwrap(); + if (const auto cardinalityField = dynamic_cast(valueField.get())) { + // Cardinality fields in RDataFrame are presented as integers + if (cardinalityField->As32Bit()) { + valueField = std::make_unique>(field.GetFieldName()); + } else if (cardinalityField->As64Bit()) { + valueField = std::make_unique>(field.GetFieldName()); + } else { + R__ASSERT(false && "cardinality field stored with an unexpected integer type"); + } + } + + valueField->SetOnDiskId(field.GetOnDiskId()); + auto valueSubfields = valueField->GetMutableSubfields(); + const auto fieldSubfields = field.GetConstSubfields(); + for (unsigned i = 0; i < valueSubfields.size(); ++i) { + valueSubfields[i]->SetOnDiskId(fieldSubfields[i]->GetOnDiskId()); + } + + std::unique_ptr cardinalityField; + // Collections get the additional "number of" RDF column (e.g. "R_rdf_sizeof_tracks") + if (!fieldInfos.empty()) { + const auto &info = fieldInfos.back(); + const std::string name = + (procProvenance.Empty() ? "R_rdf_sizeof_" : procProvenance.Get() + "_R_rdf_sizeof_") + info.fFieldName; + if (info.fNRepetitions > 0) { + cardinalityField = std::make_unique(name, info.fNRepetitions); + } else { + cardinalityField = std::make_unique>(name); + } + cardinalityField->SetOnDiskId(info.fFieldId); + } + + for (auto i = fieldInfos.rbegin(); i != fieldInfos.rend(); ++i) { + const auto &fieldInfo = *i; + + const auto valueFieldName = valueField->GetFieldName(); + + if (fieldInfo.fNRepetitions > 0) { + // Fixed-size array, read it as ROOT::RVec in memory + valueField = + std::make_unique(valueFieldName, valueField->Clone("_0"), fieldInfo.fNRepetitions); + } else { + // Actual collection. A std::vector or ROOT::RVec gets added as a ROOT::RVec. All other collection types keep + // their original type. + if (convertToRVec) { + valueField = std::make_unique(valueFieldName, valueField->Clone("_0")); + } else { + auto outerFieldType = fieldInfo.fTypeName; + valueField = ROOT::RFieldBase::Create(valueFieldName, outerFieldType).Unwrap(); + } + } + + valueField->SetOnDiskId(fieldInfo.fFieldId); + + // Skip the inner-most collection level to construct the cardinality column + // It's taken care of by the `if (!fieldInfos.empty())` scope above + if (i != fieldInfos.rbegin()) { + const auto cardinalityFieldName = cardinalityField->GetFieldName(); + if (fieldInfo.fNRepetitions > 0) { + // This collection level refers to a fixed-size array + cardinalityField = std::make_unique( + cardinalityFieldName, cardinalityField->Clone("_0"), fieldInfo.fNRepetitions); + } else { + // This collection level refers to an RVec + cardinalityField = std::make_unique(cardinalityFieldName, cardinalityField->Clone("_0")); + } + + cardinalityField->SetOnDiskId(fieldInfo.fFieldId); + } + } + + 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()); + } + + for (std::uint16_t i = 0; i < nFieldInfosAdded; i++) { + fieldInfos.pop_back(); + } +} + +RNTupleProcessorDS::RNTupleProcessorDS(std::unique_ptr processor) + : fProcessor(std::move(processor)) +{ + // Do not add the subfields now, this is handled by AddField. + fProcessor->AddAllFieldsToEntry(Internal::RNTupleProcessorProvenance(), /*includeSubfields=*/false); + const auto &entry = fProcessor->GetEntry(); + for (auto fieldIdx : entry.GetFieldIndices()) { + const auto &field = entry.GetValue(fieldIdx).GetField(); + std::deque fieldInfos{}; + AddField(field, entry.GetQualifiedFieldName(fieldIdx), entry.GetFieldProvenance(fieldIdx), fieldInfos); + } +} + +ROOT::RDF::RDataSource::Record_t +RNTupleProcessorDS::GetColumnReadersImpl(std::string_view /* name */, const std::type_info & /* ti */) +{ + // This datasource uses the newer GetColumnReaders() API + return {}; +} + +std::unique_ptr +RNTupleProcessorDS::GetColumnReaders(unsigned int /* slot */, std::string_view name, const std::type_info &tid) +{ + // At this point we can assume that `name` will be found in fColumnNames + const auto requestedType = ROOT::Internal::GetRenormalizedTypeName(ROOT::Internal::RDF::TypeID2TypeName(tid)); + + // First check if a field with the requested type already exists. If that is the case, we can immediately create a + // column reader for it. + if (fProcessor->GetEntry().FindFieldIndex(name, requestedType)) { + auto valuePtr = fProcessor->RequestField(std::string(name), requestedType); + auto reader = std::make_unique(valuePtr); + fColumnReaders.emplace_back(reader.get()); + + return reader; + } + + // Secondly, check whether a field with the same name, but a different on-disk type exists. If this is not the case, + // first try to add the field. If the field cannot be added, throw an exception. + // Otherwise, we create a new field with the requested type and add it to the processor entry + // before creating a column reader for it. + auto fieldIdx = fProcessor->GetEntry().FindFieldIndex(name); + if (!fieldIdx) { + std::unique_ptr reader; + try { + auto valuePtr = fProcessor->RequestField(std::string(name), requestedType); + reader = std::make_unique(valuePtr); + fColumnReaders.emplace_back(reader.get()); + } catch (const ROOT::RException &) { + throw std::runtime_error("RNTupleProcessorDS: Column \"" + std::string(name) + "\" does not exist"); + } + return reader; + } + + const auto &field = fProcessor->GetEntry().GetField(*fieldIdx); + const std::string strName = std::string(name); + std::unique_ptr newField; + if (dynamic_cast(&field)) { + if (requestedType == "bool") { + newField = std::make_unique>(strName); + } else if (requestedType == "char") { + newField = std::make_unique>(strName); + } else if (requestedType == "std::int8_t") { + newField = std::make_unique>(strName); + } else if (requestedType == "std::uint8_t") { + newField = std::make_unique>(strName); + } else if (requestedType == "std::int16_t") { + newField = std::make_unique>(strName); + } else if (requestedType == "std::uint16_t") { + newField = std::make_unique>(strName); + } else if (requestedType == "std::int32_t") { + newField = std::make_unique>(strName); + } else if (requestedType == "std::uint32_t") { + newField = std::make_unique>(strName); + } else if (requestedType == "std::int64_t") { + newField = std::make_unique>(strName); + } else if (requestedType == "std::uint64_t") { + newField = std::make_unique>(strName); + } else { + throw std::runtime_error("RNTupleProcessorDS: Could not create field with type \"" + requestedType + + "\" for column \"" + strName + "\""); + } + } else { + auto baseName = ROOT::Split(strName, ".").back(); + auto newAltProtoFieldOrException = ROOT::RFieldBase::Create(baseName, requestedType); + if (!newAltProtoFieldOrException) { + throw std::runtime_error("RNTupleProcessorDS: Could not create field with type \"" + requestedType + + "\" for column \"" + strName + "\""); + } + newField = newAltProtoFieldOrException.Unwrap(); + } + newField->SetOnDiskId(field.GetOnDiskId()); + + try { + fProcessor->AddFieldToEntry(std::move(newField), strName, nullptr, Internal::RNTupleProcessorProvenance()); + } catch (const ROOT::RException &) { + std::string msg = "RNTupleProcessorDS: invalid type \"" + requestedType + "\" for column \"" + strName + + "\" with on-disk type \"" + field.GetTypeName() + "\""; + throw std::runtime_error(msg); + } + + auto valuePtr = fProcessor->RequestField(std::string(name), requestedType); + auto reader = std::make_unique(valuePtr); + fColumnReaders.emplace_back(reader.get()); + + return reader; +} + +std::vector> RNTupleProcessorDS::GetEntryRanges() +{ + if (fCurrentRanges.empty()) { + fCurrentRanges.emplace_back(0, fProcessor->GetNEntries()); + return fCurrentRanges; + } + return {}; +} + +void RNTupleProcessorDS::InitSlot(unsigned int /* slot */, ULong64_t /* firstEntry */) +{ + assert(fNSlots == 1); +} + +void RNTupleProcessorDS::FinalizeSlot(unsigned int /* slot */) +{ + assert(fNSlots == 1); +} + +std::string RNTupleProcessorDS::GetTypeName(std::string_view colName) const +{ + auto colNamePos = std::find(fColumnNames.begin(), fColumnNames.end(), colName); + + if (colNamePos == fColumnNames.end()) { + auto msg = std::string("RNTupleProcessorDS: There is no column with name \"") + std::string(colName) + "\""; + throw std::runtime_error(msg); + } + + const auto index = std::distance(fColumnNames.begin(), colNamePos); + return fColumnTypes[index]; +} + +bool RNTupleProcessorDS::HasColumn(std::string_view colName) const +{ + return std::find(fColumnNames.begin(), fColumnNames.end(), colName) != fColumnNames.end(); +} + +void RNTupleProcessorDS::Initialize() +{ + fCurrentRanges.clear(); + fProcessor->Connect(fProcessor->GetEntry().GetFieldIndices(), Internal::RNTupleProcessorProvenance(), + /*updateFields=*/false); +} + +void RNTupleProcessorDS::Finalize() +{ + // noop +} + +void RNTupleProcessorDS::SetNSlots(unsigned int nSlots) +{ + assert(fNSlots == 0); + // FIXME(fdegeus): Set to actual nSlots once MT support lands + if (nSlots > 1) + Warning("RNTupleProcessorDS", "MT is not yet supported. Using single-threaded processing instead."); + fNSlots = 1; +} + +bool RNTupleProcessorDS::SetEntry(unsigned int /* slot */, ULong64_t entry) +{ + fProcessor->SetEntryMapping(entry, fEntryMapping); + if (fEntryMapping.GetEntryNumber() == ROOT::kInvalidNTupleIndex) + return false; + + fProcessor->LoadEntry(fEntryMapping); + + for (auto colReader : fColumnReaders) { + colReader->SetEntryLoaded(entry); + } + + return true; +} + +ROOT::RDF::RSampleInfo RNTupleProcessorDS::CreateSampleInfo( + unsigned int /* slot */, const std::unordered_map &sampleMap) const +{ + const auto &ntupleID = fProcessor->GetOptions().GetProcessorName(); + + // TODO: There is no support for RNTuple in RDatasetSpec, thus the sample map + // is always empty at the moment. + if (sampleMap.empty()) + return ROOT::RDF::RSampleInfo(ntupleID, std::make_pair(0, fProcessor->GetNEntries())); + + if (sampleMap.find(ntupleID) == sampleMap.end()) + throw std::runtime_error("Full sample identifier '" + ntupleID + "' cannot be found in the available samples."); + + return ROOT::RDF::RSampleInfo(ntupleID, std::make_pair(0, fProcessor->GetNEntries()), sampleMap.at(ntupleID)); +} +} // namespace ROOT::Experimental::Internal::RDF + +ROOT::RDataFrame +ROOT::Experimental::RDF::FromRNTupleProcessor(std::unique_ptr processor) +{ + return ROOT::RDataFrame( + std::make_unique(std::move(processor))); +} diff --git a/tree/dataframe/test/CMakeLists.txt b/tree/dataframe/test/CMakeLists.txt index d6085886c7110..3e4ed9be6a86b 100644 --- a/tree/dataframe/test/CMakeLists.txt +++ b/tree/dataframe/test/CMakeLists.txt @@ -113,6 +113,8 @@ ROOT_GENERATE_DICTIONARY(ClassWithNestedSameNameDict DEPENDENCIES RIO) ROOT_ADD_GTEST(datasource_ntuple datasource_ntuple.cxx LIBRARIES ROOTDataFrame) +ROOT_ADD_GTEST(datasource_ntuple_processor datasource_ntuple_processor.cxx LIBRARIES ROOTDataFrame) + if(MSVC) set_source_files_properties(datasource_ntuple.cxx PROPERTIES COMPILE_FLAGS /bigobj) endif() diff --git a/tree/dataframe/test/datasource_ntuple_processor.cxx b/tree/dataframe/test/datasource_ntuple_processor.cxx new file mode 100644 index 0000000000000..c83b574b8c1a2 --- /dev/null +++ b/tree/dataframe/test/datasource_ntuple_processor.cxx @@ -0,0 +1,855 @@ +#include +#include +#include + +#include +#include +#include + +#include + +#include +#include + +#include + +#include + +using ROOT::TestSupport::FileRaii; + +using ROOT::RNTupleModel; +using ROOT::RNTupleWriter; +using ROOT::Experimental::RNTupleProcessor; +using ROOT::Experimental::Internal::RDF::RNTupleProcessorDS; + +/* + * DISCLAIMER: The majority of these tests are directly copied from datasource_ntuple.cxx. + * + * Test specifically for joins have been added. + * + * TODO: Create set of common tests for both data sources, with specific tests for either data source where relevant. + */ + +template +void EXPECT_VEC_EQ(const V1 &v1, const V2 &v2) +{ + ASSERT_EQ(v1.size(), v2.size()); + for (std::size_t i = 0ul; i < v1.size(); ++i) { + EXPECT_EQ(v1[i], v2[i]); + } +} + +class RNTupleProcessorDSTest : public ::testing::Test { +protected: + std::string fFileName = "RNTupleProcessorDS_test.root"; + std::string fNtplName = "ntuple"; + + void SetUp() override + { + auto model = RNTupleModel::Create(); + *model->MakeField("pt") = 42; + *model->MakeField("energy") = 7; + *model->MakeField("tag") = "xyz"; + *model->MakeField>("jets") = std::vector{1.f, 2.f}; + auto fldNnlo = model->MakeField>>("nnlo"); + fldNnlo->push_back(std::vector()); + fldNnlo->push_back(std::vector{1.0}); + fldNnlo->push_back(std::vector{1.0, 2.0, 4.0, 8.0}); + *model->MakeField("rvec") = ROOT::RVecI{1, 2, 3}; + auto fldElectron = model->MakeField("electron"); + fldElectron->pt = 137.0; + auto fldVecElectron = model->MakeField>("VecElectron"); + for (int i = 0; i < 128; ++i) + fldVecElectron->push_back(*fldElectron); + auto fldNElectron = std::make_unique>>("nElectron"); + model->AddProjectedField(std::move(fldNElectron), [](const std::string &) { return "VecElectron"; }); + { + auto ntuple = RNTupleWriter::Recreate(std::move(model), fNtplName, fFileName); + ntuple->Fill(); + } + } + + void TearDown() override { std::remove(fFileName.c_str()); } +}; + +TEST_F(RNTupleProcessorDSTest, ColTypeNames) +{ + RNTupleProcessorDS ds(RNTupleProcessor::Create({fNtplName, fFileName})); + + auto colNames = ds.GetColumnNames(); + ASSERT_EQ(16, colNames.size()); + + EXPECT_TRUE(ds.HasColumn("pt")); + EXPECT_TRUE(ds.HasColumn("energy")); + EXPECT_TRUE(ds.HasColumn("rvec")); + EXPECT_TRUE(ds.HasColumn("R_rdf_sizeof_nnlo")); + EXPECT_TRUE(ds.HasColumn("electron")); + EXPECT_TRUE(ds.HasColumn("electron.pt")); + EXPECT_TRUE(ds.HasColumn("VecElectron")); + EXPECT_TRUE(ds.HasColumn("R_rdf_sizeof_VecElectron")); + EXPECT_TRUE(ds.HasColumn("VecElectron.pt")); + EXPECT_TRUE(ds.HasColumn("R_rdf_sizeof_VecElectron.pt")); + EXPECT_TRUE(ds.HasColumn("nElectron")); + EXPECT_FALSE(ds.HasColumn("Address")); + + EXPECT_STREQ("std::string", ds.GetTypeName("tag").c_str()); + EXPECT_STREQ("float", ds.GetTypeName("energy").c_str()); + EXPECT_EQ(ROOT::Internal::GetRenormalizedTypeName(typeid(std::size_t)), ds.GetTypeName("R_rdf_sizeof_jets")); + EXPECT_STREQ("ROOT::VecOps::RVec", ds.GetTypeName("rvec").c_str()); + EXPECT_STREQ("std::uint64_t", ds.GetTypeName("nElectron").c_str()); + + try { + ds.GetTypeName("Address"); + FAIL() << "should not be able to get a type for a non-existent column"; + } catch (const std::runtime_error &err) { + EXPECT_THAT(err.what(), testing::HasSubstr("RNTupleProcessorDS: There is no column with name \"Address\"")); + } +} + +TEST_F(RNTupleProcessorDSTest, NFiles) +{ + RNTupleProcessorDS ds(RNTupleProcessor::Create({fNtplName, fFileName})); + + EXPECT_EQ(1, ds.GetNFiles()); +} + +TEST_F(RNTupleProcessorDSTest, Basic) +{ + auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(RNTupleProcessor::Create({fNtplName, fFileName})); + auto pt_max = df.Max("pt"); + auto pt_sum = df.Sum("pt"); + EXPECT_EQ(42.f, *pt_max); + EXPECT_EQ(42.f, *pt_sum); + EXPECT_EQ(*pt_sum, *pt_max); +} + +TEST_F(RNTupleProcessorDSTest, CardinalityColumn) +{ + auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(RNTupleProcessor::Create({fNtplName, fFileName})); + + // Check that the special column # works without jitting... + auto identity = [](std::size_t sz) { return sz; }; + auto max_njets = df.Define("njets", identity, {"R_rdf_sizeof_jets"}).Max("njets"); + auto max_njets2 = df.Max("#jets"); + auto max_rvec = df.Max("#rvec"); + EXPECT_EQ(*max_njets, *max_njets2); + EXPECT_EQ(*max_njets, 2); + EXPECT_EQ(*max_rvec, 3); + + // ...and with jitting + auto max_njets_jitted = df.Define("njets", "R_rdf_sizeof_jets").Max("njets"); + auto max_njets_jitted2 = df.Define("njets", "#jets").Max("njets"); + auto max_njets_jitted3 = df.Max("#jets"); + auto max_rvec_jitted = df.Max("#rvec"); + EXPECT_EQ(*max_njets_jitted, *max_njets_jitted2); + EXPECT_EQ(*max_njets_jitted3, *max_njets_jitted2); + EXPECT_EQ(2, *max_njets_jitted); + EXPECT_EQ(3, *max_rvec_jitted); +} + +TEST_F(RNTupleProcessorDSTest, ProjectedCardinalityColumn) +{ + auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(RNTupleProcessor::Create({fNtplName, fFileName})); + + EXPECT_EQ(128u, *df.Filter("nElectron == 128").Max("nElectron")); + + EXPECT_EQ(128u, *df.Filter([](std::uint64_t x) { return x == 128; }, {"nElectron"}).Max("nElectron")); + EXPECT_EQ(128u, *df.Filter([](std::int32_t x) { return x == 128; }, {"nElectron"}).Max("nElectron")); + EXPECT_EQ(128u, *df.Filter([](std::uint32_t x) { return x == 128; }, {"nElectron"}).Max("nElectron")); + EXPECT_EQ(128u, *df.Filter([](std::int16_t x) { return x == 128; }, {"nElectron"}).Max("nElectron")); + EXPECT_EQ(128u, *df.Filter([](std::uint16_t x) { return x == 128; }, {"nElectron"}).Max("nElectron")); + EXPECT_EQ(128u, *df.Filter([](std::uint8_t x) { return x == 128; }, {"nElectron"}).Max("nElectron")); + EXPECT_EQ(128u, *df.Filter([](bool x) { return x; }, {"nElectron"}).Max("nElectron")); + try { + *df.Filter([](std::int8_t x) { return x == 0; }, {"nElectron"}).Count(); + FAIL() << "integer overflow should fail"; + } catch (const ROOT::RException &e) { + EXPECT_THAT(e.what(), ::testing::HasSubstr("integer overflow")); + } +} + +static void ReadTest(const std::string &name, const std::string &fname) +{ + auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(RNTupleProcessor::Create({name, fname})); + + auto count = df.Count(); + auto sumpt = df.Sum("pt"); + auto tag = df.Take("tag"); + auto njets = df.Take("R_rdf_sizeof_jets"); + auto sumjets = df.Sum>("jets"); + auto sumnnlosize = df.Sum>("R_rdf_sizeof_nnlo"); + auto sumvec = [](float red, const ROOT::RVec> &nnlo) { + auto sum = 0.f; + for (auto &v : nnlo) + for (auto e : v) + sum += e; + return red + sum; + }; + auto sumnnlo = df.Aggregate(sumvec, std::plus{}, "nnlo", 0.f); + auto rvec = df.Take("rvec"); + auto vectorasrvec = df.Take("jets"); + auto sumElectronPt = df.Aggregate([](float &acc, const Electron &e) { acc += e.pt; }, + [](float a, float b) { return a + b; }, "electron"); + auto sumVecElectronPt1 = df.Aggregate( + [](float &acc, const ROOT::RVec &ve) { + for (const auto &e : ve) + acc += e.pt; + }, + [](float a, float b) { return a + b; }, "VecElectron"); + auto sumVecElectronPt2 = df.Aggregate( + [](float &acc, const ROOT::RVec &pts) { + for (auto pt : pts) + acc += pt; + }, + [](float a, float b) { return a + b; }, "VecElectron.pt"); + + EXPECT_EQ(1ull, count.GetValue()); + EXPECT_DOUBLE_EQ(42.f, sumpt.GetValue()); + EXPECT_EQ(1ull, tag.GetValue().size()); + EXPECT_EQ(std::string("xyz"), tag.GetValue()[0]); + EXPECT_EQ(1ull, njets.GetValue().size()); + EXPECT_EQ(2u, njets.GetValue()[0]); + EXPECT_EQ(3.f, sumjets.GetValue()); + EXPECT_EQ(16.f, sumnnlo.GetValue()); + EXPECT_EQ(5u, sumnnlosize.GetValue()); + EXPECT_TRUE(All(rvec->at(0) == ROOT::RVecI{1, 2, 3})); + EXPECT_TRUE(All(vectorasrvec->at(0) == ROOT::RVecF{1.f, 2.f})); + EXPECT_FLOAT_EQ(137.0, sumElectronPt.GetValue()); + EXPECT_FLOAT_EQ(128. * 137.0, sumVecElectronPt1.GetValue()); + EXPECT_FLOAT_EQ(128. * 137.0, sumVecElectronPt2.GetValue()); +} + +static void ChainTest(const std::string &name, const std::string &fname) +{ + auto proc1 = RNTupleProcessor::Create({name, fname}); + auto df1 = ROOT::Experimental::RDF::FromRNTupleProcessor(std::move(proc1)); + EXPECT_DOUBLE_EQ(42.0, df1.Sum("pt").GetValue()); + + auto proc2 = RNTupleProcessor::CreateChain({{name, fname}, {name, fname}}); + auto df2 = ROOT::Experimental::RDF::FromRNTupleProcessor(std::move(proc2)); + EXPECT_DOUBLE_EQ(84.0, df2.Sum("pt").GetValue()); + + std::vector openSpecs(1000, {name, fname}); + auto proc1000 = RNTupleProcessor::CreateChain(openSpecs); + auto df1000 = ROOT::Experimental::RDF::FromRNTupleProcessor(std::move(proc1000)); + EXPECT_DOUBLE_EQ(42000.0, df1000.Sum("pt").GetValue()); + + FileRaii guardFile1("RNTupleProcessorDS_test_chain_1.root"); + FileRaii guardFile2("RNTupleProcessorDS_test_chain_2.root"); + FileRaii guardFile3("RNTupleProcessorDS_test_chain_3.root"); + + { + auto model = RNTupleModel::Create(); + auto fldElectron = model->MakeField("e"); + auto writer = RNTupleWriter::Recreate(std::move(model), "chain", guardFile1.GetPath()); + fldElectron->pt = 1.0; + writer->Fill(); + } + { + auto model = RNTupleModel::Create(); + // Add dummy field to ensure that the Electron fields in the files have different field IDs + model->MakeField("dummy1"); + auto fldElectron = model->MakeField("e"); + auto writer = RNTupleWriter::Recreate(std::move(model), "chain", guardFile2.GetPath()); + // empty file in chain, no entries + } + { + auto model = RNTupleModel::Create(); + model->MakeField("dummy1"); + model->MakeField("dummy2"); + auto fldElectron = model->MakeField("e"); + auto writer = RNTupleWriter::Recreate(std::move(model), "chain", guardFile3.GetPath()); + fldElectron->pt = 2.0; + writer->Fill(); + fldElectron->pt = 3.0; + writer->Fill(); + } + + auto proc3 = RNTupleProcessor::CreateChain( + {{"chain", guardFile1.GetPath()}, {"chain", guardFile2.GetPath()}, {"chain", guardFile3.GetPath()}}); + auto df3 = ROOT::Experimental::RDF::FromRNTupleProcessor(std::move(proc3)); + auto sumElectronPt = + df3.Aggregate([](float &acc, const Electron &e) { acc += e.pt; }, [](float a, float b) { return a + b; }, "e"); + EXPECT_FLOAT_EQ(6.0, sumElectronPt.GetValue()); + // Trigger the event loop again + EXPECT_FLOAT_EQ(6.0, sumElectronPt.GetValue()); +} + +TEST_F(RNTupleProcessorDSTest, Read) +{ + ReadTest(fNtplName, fFileName); +} + +TEST_F(RNTupleProcessorDSTest, Chain) +{ + ChainTest(fNtplName, fFileName); +} + +TEST(RNTupleProcessorDS, Join) +{ + FileRaii guardFile1("RNTupleProcessorDS_test_join_1.root"); + FileRaii guardFile2("RNTupleProcessorDS_test_join_2.root"); + + { + 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(); + } + } + + auto proc = RNTupleProcessor::CreateJoin({"primary", guardFile1.GetPath()}, {"auxiliary", guardFile2.GetPath()}, {}); + auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(std::move(proc)); + + auto count = df.Count(); + EXPECT_EQ(10, count.GetValue()); + + auto electronFilterCount = df.Filter([](float pt) { return pt > 4.f; }, {"electrons.pt"}).Count(); + EXPECT_EQ(5, 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(13.5, jetSum.GetValue()); +} + +TEST(RNTupleProcessorDS, JoinUnaligned) +{ + FileRaii guardFile1("RNTupleProcessorDS_test_join_unaligned_1.root"); + FileRaii guardFile2("RNTupleProcessorDS_test_join_unaligned_2.root"); + + { + 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(); + } + } + + auto proc = + RNTupleProcessor::CreateJoin({"primary", guardFile1.GetPath()}, {"auxiliary", guardFile2.GetPath()}, {"idx"}); + auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(std::move(proc)); + + auto count = df.Count(); + EXPECT_EQ(5, count.GetValue()); + + auto electronFilterCount = df.Filter([](float pt) { return pt > 4.f; }, {"electrons.pt"}).Count(); + EXPECT_EQ(3, 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(7.5, jetSum.GetValue()); +} + +#ifdef R__USE_IMT +struct IMTRAII { + IMTRAII() { ROOT::EnableImplicitMT(); } + ~IMTRAII() { ROOT::DisableImplicitMT(); } +}; + +// TODO(fwdg): implement MT support and remove expected warnings +TEST_F(RNTupleProcessorDSTest, ReadMT) +{ + IMTRAII _; + + ROOT_EXPECT_WARNING(ReadTest(fNtplName, fFileName), "RNTupleProcessorDS", + "MT is not yet supported. Using single-threaded processing instead."); +} + +TEST_F(RNTupleProcessorDSTest, ChainMT) +{ + IMTRAII _; + + ROOT_EXPECT_WARNING(ChainTest(fNtplName, fFileName), "RNTupleProcessorDS", + "MT is not yet supported. Using single-threaded processing instead."); +} +#endif + +TEST_F(RNTupleProcessorDSTest, ModifyColumnValues) +{ + auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(RNTupleProcessor::Create({fNtplName, fFileName})); + auto dfCorrected = + df.Define("jetsCorrected", + [](ROOT::RVec &jets) { + for (auto &jet : jets) { + jet *= 1.1; + } + return jets; + }, + {"jets"}) + .Define("jetsFiltered", [](const ROOT::RVec &jets) { return jets[jets <= 2.f]; }, {"jetsCorrected"}); + + ROOT::RVec jetsExpected{1.f, 2.f}; + ROOT::RVec jetsCorrectedExpected{1.1f, 2.2f}; + ROOT::RVec jetsFilteredExpected{1.1f}; + + // In the same action, we expect "jets" and "jetsCorrected" to be equal, with the modified values after the first + // "Define", because "jets" is modified in-place here. + dfCorrected.Foreach( + [&jetsCorrectedExpected, &jetsFilteredExpected](const ROOT::RVec &jets, const ROOT::RVec &jetsC, + const ROOT::RVec &jetsF) { + EXPECT_VEC_EQ(jetsCorrectedExpected, jets); + EXPECT_VEC_EQ(jets, jetsC); + EXPECT_VEC_EQ(jetsFilteredExpected, jetsF); + }, + {"jets", "jetsCorrected", "jetsFiltered"}); + + // Even though "jetsCorrected" is not used, "jets" should still be modified because "jetsFiltered" is defined from + // "jetsCorrected". + dfCorrected.Foreach( + [&jetsCorrectedExpected, &jetsFilteredExpected](const ROOT::RVec &jets, const ROOT::RVec &jetsF) { + EXPECT_VEC_EQ(jetsCorrectedExpected, jets); + EXPECT_VEC_EQ(jetsFilteredExpected, jetsF); + }, + {"jets", "jetsFiltered"}); + + // In separate actions, we expect "jets" to have its original on-disk value, but "jetsCorrected" (and by + // extension, "jetsFiltered") to have the modified values. + auto jets = dfCorrected.Take>("jets").GetValue(); + auto jetsCorrected = dfCorrected.Take>("jetsCorrected").GetValue(); + auto jetsFiltered = dfCorrected.Take>("jetsFiltered").GetValue(); + + ASSERT_EQ(1ull, jets.size()); + ASSERT_EQ(1ull, jetsCorrected.size()); + ASSERT_EQ(1ull, jetsFiltered.size()); + EXPECT_VEC_EQ(jetsExpected, jets[0]); + EXPECT_VEC_EQ(jetsCorrectedExpected, jetsCorrected[0]); + EXPECT_VEC_EQ(jetsFilteredExpected, jetsFiltered[0]); +} + +TEST(RNTupleProcessorDS, CollectionFieldTypes) +{ + // NB: The other tests already cover std::vector and std::array (and nestings thereof). + FileRaii fileGuard{"RNTupleProcessorDS_test_collection_field_types.root"}; + { + auto model = RNTupleModel::Create(); + *model->MakeField>("intSet") = std::set{3, 1, 2}; + *model->MakeField>("electronSet") = + std::set{Electron{1.f}, Electron{2.f}, Electron{3.f}}; + *model->MakeField>>("electronSetVec") = + std::set>{{Electron{1.f}, Electron{2.f}}, {Electron{3.f}}}; + *model->MakeField>>("electronSetSet") = + std::set>{{Electron{1.f}, Electron{2.f}}, {Electron{3.f}}}; + + // Untyped collection + auto fldJetPt = ROOT::RVectorField::CreateUntyped("jet_pt", std::make_unique>("_0")); + model->AddField(std::move(fldJetPt)); + + // Untyped collection with an untyped record, with a projection + std::vector> muon; + muon.emplace_back(std::make_unique>("muon_pt")); + auto fldMuonRecord = std::make_unique("_0", std::move(muon)); + auto fldMuons = ROOT::RVectorField::CreateUntyped("muon", std::move(fldMuonRecord)); + model->AddField(std::move(fldMuons)); + auto muonPtField = ROOT::RFieldBase::Create("muon_pt", "ROOT::VecOps::RVec").Unwrap(); + model->AddProjectedField(std::move(muonPtField), [](const std::string &fieldName) { + if (fieldName == "muon_pt") + return "muon"; + else + return "muon._0.muon_pt"; + }); + + auto ntuple = RNTupleWriter::Recreate(std::move(model), "ntuple", fileGuard.GetPath()); + ntuple->Fill(); + } + + RNTupleProcessorDS ds(RNTupleProcessor::Create({"ntuple", fileGuard.GetPath()})); + + auto colNames = ds.GetColumnNames(); + + ASSERT_EQ(20, colNames.size()); + + EXPECT_TRUE(ds.HasColumn("intSet")); + EXPECT_TRUE(ds.HasColumn("R_rdf_sizeof_intSet")); + EXPECT_TRUE(ds.HasColumn("electronSet")); + EXPECT_TRUE(ds.HasColumn("R_rdf_sizeof_electronSet")); + EXPECT_TRUE(ds.HasColumn("electronSet.pt")); + EXPECT_TRUE(ds.HasColumn("R_rdf_sizeof_electronSet.pt")); + EXPECT_TRUE(ds.HasColumn("jet_pt")); + EXPECT_TRUE(ds.HasColumn("R_rdf_sizeof_jet_pt")); + EXPECT_TRUE(ds.HasColumn("muon_pt")); + EXPECT_TRUE(ds.HasColumn("R_rdf_sizeof_muon_pt")); + EXPECT_TRUE(ds.HasColumn("muon.muon_pt")); + EXPECT_TRUE(ds.HasColumn("R_rdf_sizeof_muon.muon_pt")); + + EXPECT_STREQ("std::set", ds.GetTypeName("intSet").c_str()); + EXPECT_STREQ("std::set", ds.GetTypeName("electronSet").c_str()); + EXPECT_STREQ("ROOT::VecOps::RVec", ds.GetTypeName("electronSet.pt").c_str()); + EXPECT_STREQ("std::set>", ds.GetTypeName("electronSetSet").c_str()); + EXPECT_STREQ("ROOT::VecOps::RVec>", ds.GetTypeName("electronSetSet.pt").c_str()); + + // TODO(fdegeus) figure out how to (cleanly) still add inner vectors etc. as RVecs. + EXPECT_STREQ("std::set>", ds.GetTypeName("electronSetVec").c_str()); + EXPECT_STREQ("ROOT::VecOps::RVec", ds.GetTypeName("jet_pt").c_str()); + EXPECT_STREQ("ROOT::VecOps::RVec", ds.GetTypeName("muon_pt").c_str()); + EXPECT_STREQ(ds.GetTypeName("muon.muon_pt").c_str(), ds.GetTypeName("muon_pt").c_str()); +} + +TEST_F(RNTupleProcessorDSTest, AlternativeColumnTypes) +{ + auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(RNTupleProcessor::Create({fNtplName, fFileName})); + + // Alternative inner type + auto usingDouble = df.Define("nJets", [](const ROOT::RVec &jets) { return jets.size(); }, {"jets"}) + .Take>("nJets") + .GetValue(); + EXPECT_EQ(2ull, ROOT::VecOps::Sum(usingDouble)); + + // Alternative outer type (original on-disk type) + auto asStdVec = df.Define("nJets", [](const std::vector &jets) { return jets.size(); }, {"jets"}) + .Take>("nJets") + .GetValue(); + EXPECT_EQ(2ull, ROOT::VecOps::Sum(asStdVec)); + + // Original datasource protofield type + auto asRVec = df.Define("nJets", [](const ROOT::RVec &jets) { return jets.size(); }, {"jets"}) + .Take>("nJets") + .GetValue(); + EXPECT_EQ(2ull, ROOT::VecOps::Sum(asRVec)); + + auto multipleAlternativeTypes = + df.Define("nJets", [](const std::vector &jets) { return jets.size(); }, {"jets"}) + .Define("smallestJet", [](const std::multiset &jets) { return *(jets.begin()); }, {"jets"}) + .Min("smallestJet") + .GetValue(); + EXPECT_FLOAT_EQ(1.f, multipleAlternativeTypes); + + auto jitted = df.Define("jetsType", "ROOT::Internal::RDF::TypeID2TypeName(typeid(jets))") + .Take("jetsType") + .GetValue(); + EXPECT_EQ("ROOT::VecOps::RVec", jitted[0]); + + // Original protofield type of ROOT::RVec>, test with different ROOT::RVec/std::vector combinations + auto nestedStdVecStdVec = + df.Define("nNnlo", [](const std::vector> &nnlo) { return nnlo.size(); }, {"nnlo"}) + .Take>("nNnlo") + .GetValue(); + EXPECT_EQ(3ull, ROOT::VecOps::Sum(nestedStdVecStdVec)); + + auto nestedStdVecRVec = + df.Define("nNnlo", [](const std::vector> &nnlo) { return nnlo.size(); }, {"nnlo"}) + .Take>("nNnlo") + .GetValue(); + EXPECT_EQ(3ull, ROOT::VecOps::Sum(nestedStdVecRVec)); + + auto nestedRVecStdVec = + df.Define("nNnlo", [](const ROOT::RVec> &nnlo) { return nnlo.size(); }, {"nnlo"}) + .Take>("nNnlo") + .GetValue(); + EXPECT_EQ(3ull, ROOT::VecOps::Sum(nestedRVecStdVec)); + + // Check that the ROOT RtypesCore typedefs are handled properly + auto usingTypeAlias1 = df.Define("nJets", [](const std::vector &jets) { return jets.size(); }, {"jets"}) + .Take>("nJets") + .GetValue(); + EXPECT_EQ(2ull, ROOT::VecOps::Sum(usingTypeAlias1)); + + auto usingTypeAlias2 = + df.Define("vecSum", [](const ROOT::RVec &rvec) { return ROOT::VecOps::Sum(rvec); }, {"rvec"}) + .Take("vecSum") + .GetValue(); + EXPECT_EQ(6, usingTypeAlias2[0]); + + try { + // Invalid outer field type + auto dfInvalid = ROOT::Experimental::RDF::FromRNTupleProcessor(RNTupleProcessor::Create({fNtplName, fFileName})); + dfInvalid.Define("firstJet", [](const std::pair &jets) { return jets.first; }, {"jets"}) + .Take>("firstJet") + .GetValue(); + FAIL() << "specifying templated actions with incompatible column types should throw"; + } catch (const std::runtime_error &err) { + EXPECT_THAT(err.what(), + testing::HasSubstr("RNTupleProcessorDS: invalid type \"std::pair\" for column " + "\"jets\" with on-disk type \"std::vector\"")); + } + + try { + // Invalid inner field types + auto dfInvalid = ROOT::Experimental::RDF::FromRNTupleProcessor(RNTupleProcessor::Create({fNtplName, fFileName})); + dfInvalid.Define("nJets", [](const std::vector &jets) { return jets.size(); }, {"jets"}) + .Take>("nJets") + .GetValue(); + FAIL() << "specifying templated actions with incompatible column types should throw"; + } catch (const std::runtime_error &err) { + EXPECT_THAT(err.what(), + testing::HasSubstr("RNTupleProcessorDS: invalid type \"std::vector\" for column " + "\"jets\" with on-disk type \"std::vector\"")); + } +} + +const static std::array>, 3> arraysDatasetCol4El{ + ROOT::RVec>{ + {ROOT::RVecI{1, 2}, ROOT::RVecI{4, 5, 6}, ROOT::RVecI{42, 43, 44, 45}}, + }, + ROOT::RVec>{ + {ROOT::RVecI{55, 66}, ROOT::RVecI{-18, 33, std::numeric_limits::max()}, + ROOT::RVecI{42, std::numeric_limits::min(), 44, 1888}}, + {ROOT::RVecI{10, 11}, ROOT::RVecI{-32, -33, -34}, ROOT::RVecI{2953, -20, 343212}}, + }, + ROOT::RVec>{ + {ROOT::RVecI{-32, -33, -34}, ROOT::RVecI{42, -43, 44, -45}, ROOT::RVecI{1}}, + {ROOT::RVecI{0, 0, std::numeric_limits::min(), 42}, ROOT::RVecI{-32, -33, -34}, + ROOT::RVecI{30000, 40000, 50000}}, + {ROOT::RVecI{0, 0, std::numeric_limits::min(), 42}, ROOT::RVecI{-32, -33, -34}, + ROOT::RVecI{std::numeric_limits::min(), std::numeric_limits::max(), + std::numeric_limits::min()}}, + }}; + +class RNTupleProcessorDSArraysDataset : public ::testing::Test { +protected: + std::string fFileName = "rntupleProcessorDS_arrays_dataset.root"; + std::string fNtplName = "ntuple"; + + void SetUp() override + { + auto model = RNTupleModel::Create(); + auto col1_arr = model->MakeField>("col1_arr"); + auto col2_arr_rvec = model->MakeField>("col2_arr_rvec"); + auto col3_rvec_arr = model->MakeField>>("col3_rvec_arr"); + auto col4_arr_rvec_arr_rvec = + model->MakeField>, 3>>("col4_arr_rvec_arr_rvec"); + // TODO(fdegeus): test classes with arrays + auto ntuple = RNTupleWriter::Recreate(std::move(model), fNtplName, fFileName); + for (int i = 1; i <= 5; i++) { + + *col1_arr = {1 * i, 2 * i, 3 * i}; + + *col2_arr_rvec = {ROOT::RVecI{i}, ROOT::RVecI{1 * i, 2 * i, 3 * i}, + ROOT::RVecI{-1 * i, -2 * i, -3 * i, -4 * i, -5 * i}}; + + ROOT::RVec> rvecVal; + rvecVal.reserve(i); + for (auto j = 0; j < i; j++) { + rvecVal.emplace_back(std::array{1 * i, 2 * i, 3 * i}); + } + *col3_rvec_arr = rvecVal; + + *col4_arr_rvec_arr_rvec = arraysDatasetCol4El; + + std::array col5_member_1{42.f, std::numeric_limits::max(), 0.f}; + std::array col5_member_2{ROOT::RVecF{1.f}, ROOT::RVecF{2.f, 3.f}, ROOT::RVecF{4.f, 5.f, 6.f}}; + ROOT::RVec> col5_member_3{col5_member_1, col5_member_1, col5_member_1, col5_member_1}; + + ntuple->Fill(); + } + } + + void TearDown() override { std::remove(fFileName.c_str()); } +}; + +void ReadArraysTest(const std::string &name, const std::string &fname) +{ + // These tests use the columns that contain std::array data on disk as RVecs + // reading them into RVecs and checking their values. + auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(RNTupleProcessor::Create({name, fname})); + + auto count = df.Count(); + + auto col1Take = df.Take("col1_arr"); + auto col2Take = df.Take>("col2_arr_rvec"); + auto col3Take = df.Take>("col3_rvec_arr"); + auto col4Take = df.Take>>>("col4_arr_rvec_arr_rvec"); + + EXPECT_EQ(5ull, count.GetValue()); + + auto col1Val = col1Take.GetValue(); + EXPECT_EQ(col1Val.size(), 5); + for (int i = 1; i <= 5; i++) { + EXPECT_VEC_EQ(col1Val[i - 1], ROOT::RVecI{1 * i, 2 * i, 3 * i}); + } + + auto col2Val = col2Take.GetValue(); + EXPECT_EQ(col2Val.size(), 5); + for (int i = 1; i <= 5; i++) { + EXPECT_VEC_EQ(col2Val[i - 1][0], ROOT::RVecI{i}); + EXPECT_VEC_EQ(col2Val[i - 1][1], ROOT::RVecI{1 * i, 2 * i, 3 * i}); + EXPECT_VEC_EQ(col2Val[i - 1][2], ROOT::RVecI{-1 * i, -2 * i, -3 * i, -4 * i, -5 * i}); + } + + auto col3Val = col3Take.GetValue(); + EXPECT_EQ(col3Val.size(), 5); + for (int i = 1; i <= 5; i++) { + for (auto j = 0; j < i; j++) { + EXPECT_VEC_EQ(col3Val[i - 1][j], std::array{1 * i, 2 * i, 3 * i}); + } + } + + auto col4Val = col4Take.GetValue(); + EXPECT_EQ(col4Val.size(), 5); + + // Each event contains the values of arraysDatasetCol4El in column 4 + for (unsigned i = 0; i < 5; i++) { + // Unpack the three top-level std::arrays + const auto &arr1 = col4Val[i][0]; + const auto &arr2 = col4Val[i][1]; + const auto &arr3 = col4Val[i][2]; + + // The first contains an RVec with only one element + for (unsigned k = 0; k < arr1.size(); k++) { + EXPECT_VEC_EQ(arr1[0][k], arraysDatasetCol4El[0][0][k]); + } + + // The others contain more than one element in their sub-level RVecs + for (unsigned j = 0; j < arr2.size(); j++) { + for (unsigned k = 0; k < arr2[j].size(); k++) { + EXPECT_VEC_EQ(arr2[j][k], arraysDatasetCol4El[1][j][k]); + } + } + + for (unsigned j = 0; j < arr3.size(); j++) { + for (unsigned k = 0; k < arr3[j].size(); k++) { + EXPECT_VEC_EQ(arr3[j][k], arraysDatasetCol4El[2][j][k]); + } + } + } +} + +TEST_F(RNTupleProcessorDSArraysDataset, Read) +{ + ReadArraysTest(fNtplName, fFileName); +} + +#ifdef R__USE_IMT + +TEST_F(RNTupleProcessorDSArraysDataset, ReadMT) +{ + IMTRAII _; + + ROOT_EXPECT_WARNING(ReadArraysTest(fNtplName, fFileName), "RNTupleProcessorDS", + "MT is not yet supported. Using single-threaded processing instead."); +} +#endif + +void UseArraysAsRVec(const std::string &name, const std::string &fname) +{ + // These tests use the columns that contain std::array data on disk as RVecs + // passing them as arguments to functions that expect RVecs. + auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(RNTupleProcessor::Create({name, fname})); + + auto df1 = df.Define("col1_short", [](const ROOT::RVecI &arr) { return ROOT::VecOps::Take(arr, 2); }, {"col1_arr"}); + auto take1 = df1.Take("col1_short"); + + // Exercise jitting + auto df2 = df1.Define("sum_col2", + "int sum = 0; for (const auto &v: col3_rvec_arr) sum += ROOT::VecOps::Sum(v); return sum;"); + auto take2 = df2.Take("sum_col2"); + + auto take1Val = take1.GetValue(); + for (int i = 0; i < 5; i++) { + EXPECT_VEC_EQ(take1Val[i], ROOT::RVecI{(i + 1), (i + 1) * 2}); + } + + std::array expectedVals{6, 24, 54, 96, 150}; + EXPECT_VEC_EQ(take2.GetValue(), expectedVals); +} + +TEST_F(RNTupleProcessorDSArraysDataset, UseArrays) +{ + UseArraysAsRVec(fNtplName, fFileName); +} + +#ifdef R__USE_IMT + +TEST_F(RNTupleProcessorDSArraysDataset, UseArraysMT) +{ + IMTRAII _; + + ROOT_EXPECT_WARNING(UseArraysAsRVec(fNtplName, fFileName), "RNTupleProcessorDS", + "MT is not yet supported. Using single-threaded processing instead."); +} +#endif + +void UseArraySizeColumn(const std::string &name, const std::string &fname) +{ + // These tests use the columns that contain std::array data on disk as RVecs + // checking the size of the collection with the R_rdf_sizeof_* columns + auto df = ROOT::Experimental::RDF::FromRNTupleProcessor(RNTupleProcessor::Create({name, fname})); + + auto sizeOfCol1 = df.Take("R_rdf_sizeof_col1_arr"); + // Use # here to exercise that too + auto sizeOfCol2 = df.Take>("#col2_arr_rvec"); + auto sizeOfCol4 = df.Take>>>("R_rdf_sizeof_col4_arr_rvec_arr_rvec"); + + EXPECT_VEC_EQ(sizeOfCol1.GetValue(), std::array{3, 3, 3, 3, 3}); + + for (const auto &sizeVec : sizeOfCol2) { + EXPECT_VEC_EQ(sizeVec, ROOT::RVecULL{1, 3, 5}); + } + + // The sizes of elements in col4 are + // { { { 2, 3, 4 } }, { { 2, 3, 4 }, { 2, 3, 3 } }, { { 3, 4, 1 }, { 4, 3, 3 }, { 4, 3, 3 } } } + auto sizeOfCol4Val = sizeOfCol4.GetValue(); + for (unsigned i = 0; i < 5; i++) { + // Unpack the three top-level vectors + const auto &rvec1 = sizeOfCol4Val[i][0]; + const auto &rvec2 = sizeOfCol4Val[i][1]; + const auto &rvec3 = sizeOfCol4Val[i][2]; + + // The first contains an RVec with only one element + EXPECT_VEC_EQ(rvec1[0], ROOT::RVecULL{2, 3, 4}); + + // The others contain more than one element in their sub-level RVecs + ROOT::RVec expectedRVecs2{{2, 3, 4}, {2, 3, 3}}; + for (unsigned j = 0; j < rvec2.size(); j++) { + EXPECT_VEC_EQ(rvec2[j], expectedRVecs2[j]); + } + + ROOT::RVec expectedRVecs3{{3, 4, 1}, {4, 3, 3}, {4, 3, 3}}; + for (unsigned j = 0; j < rvec3.size(); j++) { + EXPECT_VEC_EQ(rvec3[j], expectedRVecs3[j]); + } + } +} + +TEST_F(RNTupleProcessorDSArraysDataset, UseArraySizeColumn) +{ + UseArraySizeColumn(fNtplName, fFileName); +} + +#ifdef R__USE_IMT + +TEST_F(RNTupleProcessorDSArraysDataset, UseArraySizeColumnMT) +{ + IMTRAII _; + + ROOT_EXPECT_WARNING(UseArraySizeColumn(fNtplName, fFileName), "RNTupleProcessorDS", + "MT is not yet supported. Using single-threaded processing instead."); +} +#endif diff --git a/tree/ntuple/inc/ROOT/RNTupleProcessor.hxx b/tree/ntuple/inc/ROOT/RNTupleProcessor.hxx index 4fdfec2b52bb7..0dfdd88eb5999 100644 --- a/tree/ntuple/inc/ROOT/RNTupleProcessor.hxx +++ b/tree/ntuple/inc/ROOT/RNTupleProcessor.hxx @@ -37,6 +37,11 @@ namespace Internal { struct RNTupleProcessorEntryLoader; // for unit tests const RNTupleProcessorEntry * LoadFullRNTupleProcessorEntry(ROOT::Experimental::RNTupleProcessor &processor, bool includeSubfields); // for unit tests + +namespace RDF { +class RNTupleProcessorColumnReader; +class RNTupleProcessorDS; +} // namespace RDF } // namespace Internal // clang-format off @@ -303,6 +308,7 @@ class RNTupleProcessor { friend class RNTupleSingleProcessor; friend class RNTupleChainProcessor; friend class RNTupleJoinProcessor; + friend class Internal::RDF::RNTupleProcessorDS; protected: ///////////////////////////////////////////////////////////////////////////// @@ -438,6 +444,14 @@ protected: /// \param[in] mapping Mapping to the entry to load virtual void LoadEntryImpl(const REntryMapping &mapping) = 0; + ///////////////////////////////////////////////////////////////////////////// + /// \brief Get a reference to the processor's internal entry. + Internal::RNTupleProcessorEntry &GetEntry() { return *fEntry; } + + ///////////////////////////////////////////////////////////////////////////// + /// \brief Get a reference to the processor's internal entry. + const Internal::RNTupleProcessorEntry &GetEntry() const { return *fEntry; } + ///////////////////////////////////////////////////////////////////////////// /// \brief Get the total number of entries in this processor virtual ROOT::NTupleSize_t GetNEntries() = 0; diff --git a/tree/ntuple/inc/ROOT/RNTupleProcessorEntry.hxx b/tree/ntuple/inc/ROOT/RNTupleProcessorEntry.hxx index e99f3d67816d4..98e12e76a09d6 100644 --- a/tree/ntuple/inc/ROOT/RNTupleProcessorEntry.hxx +++ b/tree/ntuple/inc/ROOT/RNTupleProcessorEntry.hxx @@ -76,8 +76,10 @@ public: /// \param[in] fieldName Field name to check. bool IsPresentInFieldName(std::string_view fieldName) const { - return !fProvenance.empty() && fieldName.find(fProvenance + ".") == 0; + return !fProvenance.empty() && fieldName.find(fProvenance) == 0; } + + bool Empty() const { return fProvenance.empty(); } }; // clang-format off diff --git a/tree/ntuple/src/RNTupleProcessor.cxx b/tree/ntuple/src/RNTupleProcessor.cxx index 14c54d6f9fb90..354438aaef9cf 100644 --- a/tree/ntuple/src/RNTupleProcessor.cxx +++ b/tree/ntuple/src/RNTupleProcessor.cxx @@ -395,6 +395,7 @@ ROOT::NTupleSize_t ROOT::Experimental::RNTupleChainProcessor::GetNEntries() for (unsigned i = 0; i < fInnerProcessors.size(); ++i) { if (fInnerNEntries[i] == kInvalidNTupleIndex) { + fInnerProcessors[i]->Initialize(fEntry); fInnerNEntries[i] = fInnerProcessors[i]->GetNEntries(); } From 4c09dc95573f5b97b714861279bfbd88529dfdfb Mon Sep 17 00:00:00 2001 From: Florine de Geus Date: Tue, 13 Oct 2026 16:16:54 +0200 Subject: [PATCH 8/8] [df] Disable large chain test for Windows --- tree/dataframe/test/datasource_ntuple_processor.cxx | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tree/dataframe/test/datasource_ntuple_processor.cxx b/tree/dataframe/test/datasource_ntuple_processor.cxx index c83b574b8c1a2..854da14a3e048 100644 --- a/tree/dataframe/test/datasource_ntuple_processor.cxx +++ b/tree/dataframe/test/datasource_ntuple_processor.cxx @@ -229,10 +229,14 @@ static void ChainTest(const std::string &name, const std::string &fname) auto df2 = ROOT::Experimental::RDF::FromRNTupleProcessor(std::move(proc2)); EXPECT_DOUBLE_EQ(84.0, df2.Sum("pt").GetValue()); + // Causes "Too many open files" errors on Windows + // TODO(fdegeus) Investigate root cause + fix +#if !defined(_MSC_VER) || defined(R__ENABLE_BROKEN_WIN_TESTS) std::vector openSpecs(1000, {name, fname}); auto proc1000 = RNTupleProcessor::CreateChain(openSpecs); auto df1000 = ROOT::Experimental::RDF::FromRNTupleProcessor(std::move(proc1000)); EXPECT_DOUBLE_EQ(42000.0, df1000.Sum("pt").GetValue()); +#endif FileRaii guardFile1("RNTupleProcessorDS_test_chain_1.root"); FileRaii guardFile2("RNTupleProcessorDS_test_chain_2.root");