Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[ntuple] Make RNTupleChainProcessor composable #17393

Open
wants to merge 6 commits into
base: master
Choose a base branch
from
+370 −192
Open

[ntuple] Make RNTupleChainProcessor composable #17393

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
bc62961
[ntuple] Lazily initialize `RNTupleSingleProcessor`
enirolf Dec 18, 2024
84e35ab
[ntuple] Make `LoadEntry` random access
enirolf Dec 18, 2024
5b31c5d
[ntuple] Store model in base processor class
enirolf Jan 9, 2025
8494e63
[ntuple] Make `RNTupleChainProcessor` composable
enirolf Jan 9, 2025
fdceaaf
[ntuple] Add composed chain tests
enirolf Jan 9, 2025
b538fcc
[ntuple] Update chain processing tutorial
enirolf Jan 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
163 changes: 107 additions & 56 deletions tree/ntuple/v7/inc/ROOT/RNTupleProcessor.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,10 @@ entry. Additional bookkeeping information can be obtained through the RNTuplePro
*/
// clang-format on
class RNTupleProcessor {
friend class RNTupleSingleProcessor;
friend class RNTupleChainProcessor;
friend class RNTupleJoinProcessor;

protected:
// clang-format off
/**
Expand Down Expand Up @@ -120,37 +124,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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Perhaps add a comment that this is costly for a chain processor (all underlying files need to be opened). I guess this expensive operation is only necessary if a chain is built from other chains..?

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)
silverweed marked this conversation as resolved.
Show resolved Hide resolved
: fNTuples(ntuples), fModel(std::move(model))
{
}

Expand All @@ -163,25 +174,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
Expand All @@ -208,17 +218,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) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, here it is. Perhaps move to the previous commit.

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;
}

Expand All @@ -229,7 +237,7 @@ public:
return obj;
}

reference operator*() { return *fProcessor.fEntry; }
reference operator*() { return fProcessor.GetEntry(); }

friend bool operator!=(const iterator &lh, const iterator &rh)
{
Expand All @@ -244,8 +252,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.
Expand All @@ -258,6 +266,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.
///
Expand Down Expand Up @@ -290,17 +309,37 @@ class RNTupleSingleProcessor : public RNTupleProcessor {
friend class RNTupleProcessor;

private:
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
Expand All @@ -314,19 +353,23 @@ 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.
///
/// \return The number of entries in the newly-connected RNTuple.
/// \brief Load the entry identified by the provided (global) entry number (i.e., considering all RNTuples 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);
/// \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.
NTupleSize_t GetNEntries() final;

/////////////////////////////////////////////////////////////////////////////
/// \brief Constructs a new RNTupleChainProcessor.
Expand All @@ -337,7 +380,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
Expand All @@ -351,19 +395,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.
Expand Down
Loading
Loading