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/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 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..854da14a3e048 --- /dev/null +++ b/tree/dataframe/test/datasource_ntuple_processor.cxx @@ -0,0 +1,859 @@ +#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()); + + // 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"); + 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 6403770e62909..0dfdd88eb5999 100644 --- a/tree/ntuple/inc/ROOT/RNTupleProcessor.hxx +++ b/tree/ntuple/inc/ROOT/RNTupleProcessor.hxx @@ -31,9 +31,17 @@ 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 RDF { +class RNTupleProcessorColumnReader; +class RNTupleProcessorDS; +} // namespace RDF } // namespace Internal // clang-format off @@ -172,6 +180,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 @@ -243,6 +256,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 @@ -284,9 +302,13 @@ 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; + friend class Internal::RDF::RNTupleProcessorDS; protected: ///////////////////////////////////////////////////////////////////////////// @@ -422,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; @@ -436,7 +466,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. @@ -446,9 +477,35 @@ 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; + + ///////////////////////////////////////////////////////////////////////////// + /// \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. /// @@ -510,6 +567,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); } @@ -713,6 +785,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. /// @@ -764,6 +846,20 @@ 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. + /// + /// \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 +966,20 @@ 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. + /// + /// \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 +1090,20 @@ 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. + /// + /// \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/inc/ROOT/RNTupleProcessorEntry.hxx b/tree/ntuple/inc/ROOT/RNTupleProcessorEntry.hxx index 3bb30f011d84b..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 @@ -171,9 +173,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. @@ -195,6 +200,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/RNTupleProcessor.cxx b/tree/ntuple/src/RNTupleProcessor.cxx index acb8b3e1bd4a7..354438aaef9cf 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)) @@ -155,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); @@ -179,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); } @@ -211,6 +226,63 @@ 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) +{ + 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) { @@ -323,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(); } @@ -370,6 +443,21 @@ 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) +{ + Initialize(); + fInnerProcessors[0]->AddAllFieldsToEntry(provenance, includeSubfields); +} + void ROOT::Experimental::RNTupleChainProcessor::SetEntryMapping(ROOT::NTupleSize_t globalEntryNumber, REntryMapping &mapping) { @@ -561,6 +649,45 @@ 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) +{ + 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/src/RNTupleProcessorEntry.cxx b/tree/ntuple/src/RNTupleProcessorEntry.cxx index 5b6727d5f4418..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; } } @@ -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) { diff --git a/tree/ntuple/test/ntuple_processor.cxx b/tree/ntuple/test/ntuple_processor.cxx index 5172776ccc9b7..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]}); @@ -362,6 +384,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 +904,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"});