[df] Add initial RNTupleProcessor data source

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}

tree/dataframe/inc/ROOT/RNTupleProcessorDS.hxx

@@ -0,0 +1,138 @@
+/// \file RNTupleProcessorDS.hxx
+/// \author Florine de Geus <florine.de.geus@cern.ch>
+/// \author Jakob Blomer <jblomer@cern.ch>
+/// \author Enrico Guiraud <enrico.guiraud@cern.ch>
+/// \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 <ROOT/RDataSource.hxx>
+#include <ROOT/RNTupleProcessor.hxx>
+#include <ROOT/RNTupleUtils.hxx>
+#include <ROOT/RNTupleDescriptor.hxx>
+
+#include <memory>
+#include <string>
+#include <string_view>
+#include <vector>
+#include <unordered_map>
+
+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<RNTupleProcessor> fProcessor;
+   RNTupleProcessor::REntryMapping fEntryMapping;
+
+   std::vector<std::string> fColumnNames;
+   std::vector<std::string> fColumnTypes;
+
+   std::vector<ROOT::Experimental::Internal::RDF::RNTupleProcessorColumnReader *> fColumnReaders;
+
+   std::vector<std::pair<ULong64_t, ULong64_t>> 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<Jet>`, 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<float>`
+   /// 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<Jet> will be added as a `std::set`, but `Jet.[pt|eta] will be
+   /// added as `ROOT::VecOps::RVec<float>).
+   void AddField(const ROOT::RFieldBase &field, std::string_view colName,
+                 Internal::RNTupleProcessorProvenance procProvenance, std::deque<RFieldInfo> &fieldInfos,
+                 bool convertToRVec = true);
+
+public:
+   RNTupleProcessorDS(std::unique_ptr<ROOT::Experimental::RNTupleProcessor> 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<std::string> &GetColumnNames() const final { return fColumnNames; }
+   bool HasColumn(std::string_view colName) const final;
+   std::string GetTypeName(std::string_view colName) const final;
+   std::vector<std::pair<ULong64_t, ULong64_t>> 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<ROOT::Detail::RDF::RColumnReaderBase>
+   GetColumnReaders(unsigned int slot, std::string_view name, const std::type_info &) final;
+
+   ROOT::RDF::RSampleInfo
+   CreateSampleInfo(unsigned int,
+                    const std::unordered_map<std::string, ROOT::RDF::Experimental::RSample *> &) 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<ROOT::Experimental::RNTupleProcessor> processor);
+} // namespace ROOT::Experimental::RDF
+
+#endif // ROOT_RNTupleProcessorDS

tree/dataframe/src/RNTupleDS.cxx

@@ -21,7 +21,8 @@
 #include <ROOT/RNTupleDS.hxx>
 #include <ROOT/RNTupleTypes.hxx>
 #include <ROOT/RPageStorage.hxx>
-#include <string_view>
+
+#include "RNTupleDSFields.hxx"
 
 #include <TError.h>
 #include <TSystem.h>
@@ -31,6 +32,7 @@
 #include <memory>
 #include <mutex>
 #include <string>
+#include <string_view>
 #include <vector>
 #include <typeinfo>
 #include <type_traits>
@@ -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<ROOT::Internal::RColumnIndex>(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 <typename T>
-class RRDFCardinalityField final : public RRDFCardinalityFieldBase {
-   static_assert(std::is_integral_v<T>, "T must be an integral type");
-
-   inline void CheckSize(ROOT::NTupleSize_t size) const
-   {
-      if constexpr (std::is_same_v<T, bool> || std::is_same_v<T, std::uint64_t>)
-         return;
-      if (size > static_cast<ROOT::NTupleSize_t>(std::numeric_limits<T>::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<ROOT::RFieldBase> CloneImpl(std::string_view newName) const final
-   {
-      return std::make_unique<RRDFCardinalityField>(newName);
-   }
-   void ConstructValue(void *where) const final { *static_cast<T *>(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<T *>(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<T *>(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<ROOT::RFieldBase> CloneImpl(std::string_view newName) const final
-   {
-      return std::make_unique<RArraySizeField>(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<std::size_t *>(to) = fArrayLength;
-   }
-   void ReadInClusterImpl(RNTupleLocalIndex /*localIndex*/, void *to) final
-   {
-      *static_cast<std::size_t *>(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<std::size_t *>(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;