[ntuple] Make RNTupleChainProcessor composable

tree/ntuple/v7/inc/ROOT/RNTupleProcessor.hxx

@@ -32,11 +32,14 @@
 namespace ROOT {
 namespace Experimental {
 
+namespace Internal {
+struct RNTupleProcessorEntryLoader;
+} // namespace Internal
+
 /// Used to specify the underlying RNTuples in RNTupleProcessor
 struct RNTupleOpenSpec {
    std::string fNTupleName;
    std::string fStorage;
-   RNTupleReadOptions fOptions;
 
    RNTupleOpenSpec(std::string_view n, std::string_view s) : fNTupleName(n), fStorage(s) {}
 };
@@ -74,6 +77,11 @@ entry. Additional bookkeeping information can be obtained through the RNTuplePro
 */
 // clang-format on
 class RNTupleProcessor {
+   friend struct ROOT::Experimental::Internal::RNTupleProcessorEntryLoader; // for unit tests
+   friend class RNTupleSingleProcessor;
+   friend class RNTupleChainProcessor;
+   friend class RNTupleJoinProcessor;
+
 protected:
    // clang-format off
    /**
@@ -120,37 +128,44 @@ protected:
    // Maps the (qualified) field name to its corresponding field context.
    std::unordered_map<std::string, RFieldContext> fFieldContexts;
 
-   NTupleSize_t fNEntriesProcessed;  //< Total number of entries processed so far
-   std::size_t fCurrentNTupleNumber; //< Index of the currently open RNTuple
-   NTupleSize_t fLocalEntryNumber;   //< Entry number within the current ntuple
+   std::unique_ptr<RNTupleModel> fModel;
+
+   NTupleSize_t fNEntriesProcessed = 0;     //< Total number of entries processed so far
+   NTupleSize_t fCurrentEntryNumber = 0;    //< Current processor entry number
+   std::size_t fCurrentProcessorNumber = 0; //< Number of the currently open inner processor
 
    /////////////////////////////////////////////////////////////////////////////
    /// \brief Creates and connects a concrete field to the current page source, based on its proto field.
    void ConnectField(RFieldContext &fieldContext, Internal::RPageSource &pageSource, REntry &entry);
 
-   //////////////////////////////////////////////////////////////////////////
-   /// \brief Advance the processor to the next available entry.
+   /////////////////////////////////////////////////////////////////////////////
+   /// \brief Load the entry identified by the provided entry number.
    ///
-   /// \return The number of the entry loaded after advancing, or kInvalidNTupleIndex if there was no entry to advance
-   /// to.
+   /// \param[in] entryNumber Entry number to load
    ///
-   /// Checks if the end of the currently connected RNTuple is reached. If this is the case, either the next RNTuple
-   /// is connected or the iterator has reached the end.
-   virtual NTupleSize_t Advance() = 0;
+   /// \return `entryNumber` if the entry was successfully loaded, `kInvalidNTupleIndex` otherwise.
+   virtual NTupleSize_t LoadEntry(NTupleSize_t entryNumber) = 0;
 
    /////////////////////////////////////////////////////////////////////////////
-   /// \brief Fill the entry with values belonging to the current entry number.
-   virtual void LoadEntry() = 0;
+   /// \brief Points processor entry's field values to those in the provided entry.
+   ///
+   /// \param[in] entry The entry whose field values to use.
+   virtual void SetEntryPointers(const REntry &entry) = 0;
 
    /////////////////////////////////////////////////////////////////////////////
-   /// \brief Set the local (i.e. relative to the page source currently openend) entry number. Used by
-   /// `RNTupleProcessor::RIterator`.
-   ///
-   /// \param[in] entryNumber
-   void SetLocalEntryNumber(NTupleSize_t entryNumber) { fLocalEntryNumber = entryNumber; }
+   /// \brief Get the total number of entries in this processor
+   virtual NTupleSize_t GetNEntries() = 0;
 
-   RNTupleProcessor(const std::vector<RNTupleOpenSpec> &ntuples)
-      : fNTuples(ntuples), fNEntriesProcessed(0), fCurrentNTupleNumber(0), fLocalEntryNumber(0)
+   /////////////////////////////////////////////////////////////////////////////
+   /// \brief Create a new base RNTupleProcessor.
+   ///
+   /// \param[in] ntuples The input RNTuples for processing
+   /// \param[in] model The RNTupleModel representing the entries returned by the processor.
+   ///
+   /// \note Before processing, a model *must* exist. However, this is handled downstream by the RNTupleProcessor's
+   /// factory functions (CreateSingle, CreateChain and CreateJoin) and constructors.
+   RNTupleProcessor(const std::vector<RNTupleOpenSpec> &ntuples, std::unique_ptr<RNTupleModel> model)
+      : fNTuples(ntuples), fModel(std::move(model))
    {
    }
 
@@ -163,25 +178,24 @@ public:
 
    /////////////////////////////////////////////////////////////////////////////
    /// \brief Get the total number of entries processed so far.
-   ///
-   /// When only one RNTuple is present in the processor chain, the return value is equal to GetLocalEntryNumber.
    NTupleSize_t GetNEntriesProcessed() const { return fNEntriesProcessed; }
 
    /////////////////////////////////////////////////////////////////////////////
-   /// \brief Get the index to the RNTuple currently being processed, according to the sources specified upon creation.
-   std::size_t GetCurrentNTupleNumber() const { return fCurrentNTupleNumber; }
+   /// \brief Get the entry number that is currently being processed.
+   NTupleSize_t GetCurrentEntryNumber() const { return fCurrentEntryNumber; }
 
    /////////////////////////////////////////////////////////////////////////////
-   /// \brief Get the entry number local to the RNTuple that is currently being processed.
+   /// \brief Get the number of the inner processor currently being read.
    ///
-   /// When only one RNTuple is present in the processor chain, the return value is equal to GetGlobalEntryNumber.
-   NTupleSize_t GetLocalEntryNumber() const { return fLocalEntryNumber; }
+   /// This method is only relevant for the RNTupleChainProcessor. For the other processors, 0 is always returned.
+   std::size_t GetCurrentProcessorNumber() const { return fCurrentProcessorNumber; }
+
+   const RNTupleModel &GetModel() const { return *fModel; }
 
    /////////////////////////////////////////////////////////////////////////////
    /// \brief Returns a reference to the entry used by the processor.
    ///
    /// \return A reference to the entry used by the processor.
-   ///
    const REntry &GetEntry() const { return *fEntry; }
 
    // clang-format off
@@ -208,17 +222,15 @@ public:
          : fProcessor(processor), fCurrentEntryNumber(entryNumber)
       {
          // This constructor is called with kInvalidNTupleIndex for RNTupleProcessor::end(). In that case, we already
-         // know there is nothing to advance to.
+         // know there is nothing to load.
          if (fCurrentEntryNumber != kInvalidNTupleIndex) {
-            fProcessor.SetLocalEntryNumber(fCurrentEntryNumber);
-            fCurrentEntryNumber = fProcessor.Advance();
+            fCurrentEntryNumber = fProcessor.LoadEntry(fCurrentEntryNumber);
          }
       }
 
       iterator operator++()
       {
-         fProcessor.SetLocalEntryNumber(fCurrentEntryNumber + 1);
-         fCurrentEntryNumber = fProcessor.Advance();
+         fCurrentEntryNumber = fProcessor.LoadEntry(fCurrentEntryNumber + 1);
          return *this;
       }
 
@@ -229,7 +241,7 @@ public:
          return obj;
       }
 
-      reference operator*() { return *fProcessor.fEntry; }
+      reference operator*() { return fProcessor.GetEntry(); }
 
       friend bool operator!=(const iterator &lh, const iterator &rh)
       {
@@ -244,8 +256,8 @@ public:
    RIterator begin() { return RIterator(*this, 0); }
    RIterator end() { return RIterator(*this, kInvalidNTupleIndex); }
 
-   static std::unique_ptr<RNTupleProcessor> Create(const RNTupleOpenSpec &ntuple);
-   static std::unique_ptr<RNTupleProcessor> Create(const RNTupleOpenSpec &ntuple, RNTupleModel &model);
+   static std::unique_ptr<RNTupleProcessor>
+   Create(const RNTupleOpenSpec &ntuple, std::unique_ptr<RNTupleModel> model = nullptr);
 
    /////////////////////////////////////////////////////////////////////////////
    /// \brief Create a new RNTuple processor chain for vertical concatenation of RNTuples.
@@ -258,6 +270,17 @@ public:
    static std::unique_ptr<RNTupleProcessor>
    CreateChain(const std::vector<RNTupleOpenSpec> &ntuples, std::unique_ptr<RNTupleModel> model = nullptr);
 
+   /////////////////////////////////////////////////////////////////////////////
+   /// \brief Create a new RNTuple processor chain for vertical concatenation of previously created processors.
+   ///
+   /// \param[in] innerProcessors A list with the processors to chain.
+   /// \param[in] model An RNTupleModel specifying which fields can be read by the processor. If no model is provided,
+   /// one will be created based on the descriptor of the first ntuple specified.
+   ///
+   /// \return A pointer to the newly created RNTupleProcessor.
+   static std::unique_ptr<RNTupleProcessor> CreateChain(std::vector<std::unique_ptr<RNTupleProcessor>> innerProcessors,
+                                                        std::unique_ptr<RNTupleModel> model = nullptr);
+
    /////////////////////////////////////////////////////////////////////////////
    /// \brief Create a new RNTuple processor for horizontallly concatenated RNTuples.
    ///
@@ -290,17 +313,38 @@ class RNTupleSingleProcessor : public RNTupleProcessor {
    friend class RNTupleProcessor;
 
 private:
+   RNTupleOpenSpec fNTupleSpec;
+   bool fIsConnected = false;
+
+   /////////////////////////////////////////////////////////////////////////////
+   /// \brief Connects the page source of the underlying RNTuple.
+   void Connect();
+
+   /////////////////////////////////////////////////////////////////////////////
+   /// \brief Load the entry identified by the provided (global) entry number (i.e., considering all RNTuples in this
+   /// processor).
+   ///
+   /// \sa ROOT::Experimental::RNTupleProcessor::LoadEntry
+   NTupleSize_t LoadEntry(NTupleSize_t entryNumber) final;
+
+   /////////////////////////////////////////////////////////////////////////////
+   /// \sa ROOT::Experimental::RNTupleProcessor::SetEntryPointers.
+   void SetEntryPointers(const REntry &entry) final;
+
+   /////////////////////////////////////////////////////////////////////////////
+   /// \brief Get the total number of entries in this processor.
+   NTupleSize_t GetNEntries() final
+   {
+      Connect();
+      return fPageSource->GetNEntries();
+   }
+
    /////////////////////////////////////////////////////////////////////////////
    /// \brief Constructs a new RNTupleProcessor for processing a single RNTuple.
    ///
    /// \param[in] ntuple The source specification (name and storage location) for the RNTuple to process.
    /// \param[in] model The model that specifies which fields should be read by the processor.
-   RNTupleSingleProcessor(const RNTupleOpenSpec &ntuple, RNTupleModel &model);
-
-   NTupleSize_t Advance() final;
-
-public:
-   void LoadEntry() { fEntry->Read(fLocalEntryNumber); }
+   RNTupleSingleProcessor(const RNTupleOpenSpec &ntuple, std::unique_ptr<RNTupleModel> model);
 };
 
 // clang-format off
@@ -314,19 +358,25 @@ class RNTupleChainProcessor : public RNTupleProcessor {
    friend class RNTupleProcessor;
 
 private:
-   NTupleSize_t Advance() final;
-   void LoadEntry() final { fEntry->Read(fLocalEntryNumber); }
+   std::vector<std::unique_ptr<RNTupleProcessor>> fInnerProcessors;
+   std::vector<NTupleSize_t> fInnerNEntries;
 
    /////////////////////////////////////////////////////////////////////////////
-   /// \brief Connect an RNTuple for processing.
-   ///
-   /// \param[in] ntuple The RNTupleOpenSpec describing the RNTuple to connect.
+   /// \brief Load the entry identified by the provided (global) entry number (i.e., considering all RNTuples in this
+   /// processor).
    ///
-   /// \return The number of entries in the newly-connected RNTuple.
+   /// \sa ROOT::Experimental::RNTupleProcessor::LoadEntry
+   NTupleSize_t LoadEntry(NTupleSize_t entryNumber) final;
+
+   /////////////////////////////////////////////////////////////////////////////
+   /// \sa ROOT::Experimental::RNTupleProcessor::SetEntryPointers.
+   void SetEntryPointers(const REntry &) final;
+
+   /////////////////////////////////////////////////////////////////////////////
+   /// \brief Get the total number of entries in this processor.
    ///
-   /// Creates and attaches new page source for the specified RNTuple, and connects the fields that are known by
-   /// the processor to it.
-   NTupleSize_t ConnectNTuple(const RNTupleOpenSpec &ntuple);
+   /// \note This requires opening all underlying RNTuples being processed in the chain, and could become costly!
+   NTupleSize_t GetNEntries() final;
 
    /////////////////////////////////////////////////////////////////////////////
    /// \brief Constructs a new RNTupleChainProcessor.
@@ -337,7 +387,8 @@ private:
    /// specified, it is created from the descriptor of the first RNTuple specified in `ntuples`.
    ///
    /// RNTuples are processed in the order in which they are specified.
-   RNTupleChainProcessor(const std::vector<RNTupleOpenSpec> &ntuples, std::unique_ptr<RNTupleModel> model = nullptr);
+   RNTupleChainProcessor(std::vector<std::unique_ptr<RNTupleProcessor>> processors,
+                         std::unique_ptr<RNTupleModel> model);
 };
 
 // clang-format off
@@ -351,19 +402,26 @@ class RNTupleJoinProcessor : public RNTupleProcessor {
    friend class RNTupleProcessor;
 
 private:
-   std::unique_ptr<RNTupleModel> fJoinModel;
    std::vector<std::unique_ptr<Internal::RPageSource>> fAuxiliaryPageSources;
    /// Tokens representing the join fields present in the main RNTuple
    std::vector<REntry::RFieldToken> fJoinFieldTokens;
    std::vector<std::unique_ptr<Internal::RNTupleIndex>> fJoinIndices;
 
    bool IsUsingIndex() const { return fJoinIndices.size() > 0; }
 
-   NTupleSize_t Advance() final;
+   /////////////////////////////////////////////////////////////////////////////
+   /// \brief Load the entry identified by the provided entry number of the primary RNTuple.
+   ///
+   /// \sa ROOT::Experimental::RNTupleProcessor::LoadEntry
+   NTupleSize_t LoadEntry(NTupleSize_t entryNumber) final;
+
+   /////////////////////////////////////////////////////////////////////////////
+   /// \sa ROOT::Experimental::RNTupleProcessor::SetEntryPointers.
+   void SetEntryPointers(const REntry &) final;
 
    /////////////////////////////////////////////////////////////////////////////
-   /// \brief Fill the entry with values belonging to the current entry number of the primary RNTuple.
-   void LoadEntry() final;
+   /// \brief Get the total number of entries in this processor.
+   NTupleSize_t GetNEntries() final { return fPageSource->GetNEntries(); }
 
    /////////////////////////////////////////////////////////////////////////////
    /// \brief Constructs a new RNTupleJoinProcessor.