[ntuple] add support for on-demand loading of page lists

roottest/root/io/hadd/merge_changeComp_check_output.C

@@ -23,6 +23,7 @@ int merge_changeComp_check_output(int expectedCompressionRNT, int expectedCompre
 
    ROOT::Internal::RPageSourceFile source("ntpl", fnameOut, ROOT::RNTupleReadOptions());
    source.Attach();
+   source.LoadAllPageLists();
 
    ROOT::Internal::RClusterPool pool{source};
 

tree/dataframe/src/RNTupleDS.cxx

@@ -552,6 +552,7 @@ void ROOT::RDF::RNTupleDS::PrepareNextRanges()
          source = CreatePageSource(fNTupleName, fFileNames[fNextFileIndex]);
       }
       source->Attach();
+      source->LoadAllPageLists();
       fNextFileIndex++;
 
       auto nEntries = source->GetNEntries();
@@ -565,7 +566,7 @@ void ROOT::RDF::RNTupleDS::PrepareNextRanges()
       const auto rangesByCluster = [&source]() {
          // Take the shared lock of the descriptor just for the time necessary
          auto descGuard = source->GetSharedDescriptorGuard();
-         return ROOT::Internal::GetClusterBoundaries(descGuard.GetRef());
+         return ROOT::Internal::GetClusterBoundaries(descGuard.GetConstRef());
       }();
 
       const unsigned int nRangesByCluster = rangesByCluster.size();
@@ -891,6 +892,8 @@ ROOT::Internal::RDF::GetClustersAndEntries(std::string_view ntupleName, std::str
 {
    auto source = ROOT::Internal::RPageSource::Create(ntupleName, location);
    source->Attach();
+   // TODO(jblomer): improve MT scheduling such that not all page lists are required right away
+   source->LoadAllPageLists();
    auto descGuard = source->GetSharedDescriptorGuard();
-   return std::make_pair(ROOT::Internal::GetClusterBoundaries(descGuard.GetRef()), descGuard->GetNEntries());
+   return std::make_pair(ROOT::Internal::GetClusterBoundaries(descGuard.GetConstRef()), descGuard->GetNEntries());
 }

tree/dataframe/test/dataframe_snapshot_ntuple.cxx

@@ -125,7 +125,9 @@ TEST(RDFSnapshotRNTuple, WriteOpts)
 
    EXPECT_EQ(columns, sdf->GetColumnNames());
 
-   auto reader = RNTupleReader::Open("ntuple", fileGuard.GetPath());
+   ROOT::RNTupleReadOptions ntplOpts;
+   ntplOpts.SetMetadataMode(ROOT::RNTupleReadOptions::EMetadataMode::kEager);
+   auto reader = RNTupleReader::Open("ntuple", fileGuard.GetPath(), ntplOpts);
    EXPECT_EQ(2, reader->GetDescriptor().GetNClusters());
    EXPECT_EQ(2, reader->GetDescriptor().GetClusterDescriptor(0).GetPageRange(0).GetPageInfos().size());
    EXPECT_EQ(2, reader->GetDescriptor().GetClusterDescriptor(1).GetPageRange(0).GetPageInfos().size());
@@ -184,7 +186,9 @@ TEST(RDFSnapshotRNTuple, DefaultCompressionSettings)
 
    EXPECT_EQ(columns, sdf->GetColumnNames());
 
-   auto reader = RNTupleReader::Open("ntuple", fileGuard.GetPath());
+   ROOT::RNTupleReadOptions ntplOpts;
+   ntplOpts.SetMetadataMode(ROOT::RNTupleReadOptions::EMetadataMode::kEager);
+   auto reader = RNTupleReader::Open("ntuple", fileGuard.GetPath(), ntplOpts);
    auto compSettings = *reader->GetDescriptor().GetClusterDescriptor(0).GetColumnRange(0).GetCompressionSettings();
    // The RNTuple default should be 505
    EXPECT_EQ(505, compSettings);
@@ -206,7 +210,9 @@ TEST(RDFSnapshotRNTuple, Compression)
 
    EXPECT_EQ(columns, sdf->GetColumnNames());
 
-   auto reader = RNTupleReader::Open("ntuple", fileGuard.GetPath());
+   ROOT::RNTupleReadOptions ntplOpts;
+   ntplOpts.SetMetadataMode(ROOT::RNTupleReadOptions::EMetadataMode::kEager);
+   auto reader = RNTupleReader::Open("ntuple", fileGuard.GetPath(), ntplOpts);
    auto compSettings = *reader->GetDescriptor().GetClusterDescriptor(0).GetColumnRange(0).GetCompressionSettings();
    EXPECT_EQ(404, compSettings);
 }

tree/ntuple/inc/ROOT/RNTupleReadOptions.hxx

@@ -66,6 +66,16 @@ This is useful, e.g., in the context of RDataFrame where the threads are fully m
 </td>
 </tr>
 
+<tr>
+<td>`EMetadataMode`</td>
+<td>EMetadataMode</td>
+<td>EMetadataMode::kDefault</td>
+<td>
+Set whether cluster metadata (page locations) are loaded entirely on open or on demand during entry iteration.
+By default, the RNTuple reader loads metadata on demand.
+</td>
+</tr>
+
 <tr>
 <td>`EnableMetrics`</td>
 <td>`bool`</td>
@@ -93,12 +103,19 @@ public:
       kDefault,
    };
 
+   enum class EMetadataMode {
+      kEager,
+      kOnDemand,
+      kDefault = kOnDemand,
+   };
+
 private:
    EClusterCache fClusterCache = EClusterCache::kDefault;
    /// The number of cluster to be prefetched in a single batch; this option is transitional and will be replaced
    /// by an option that allows to control the amount of memory that the prefetcher uses.
    unsigned int fClusterBunchSize = 1;
    EImplicitMT fUseImplicitMT = EImplicitMT::kDefault;
+   EMetadataMode fMetadataMode = EMetadataMode::kDefault;
    bool fEnableMetrics = false;
    /// Largest header or footer envelope a page source will accept from an anchor it does not fully trust,
    /// such as the S3 anchor, which is a separately stored object. Guards the allocation made from those sizes.
@@ -111,6 +128,9 @@ public:
    EImplicitMT GetUseImplicitMT() const { return fUseImplicitMT; }
    void SetUseImplicitMT(EImplicitMT val) { fUseImplicitMT = val; }
 
+   EMetadataMode GetMetadataMode() const { return fMetadataMode; }
+   void SetMetadataMode(EMetadataMode val) { fMetadataMode = val; }
+
    bool GetEnableMetrics() const { return fEnableMetrics; }
    void SetEnableMetrics(bool val) { fEnableMetrics = val; }
 }; // class RNTupleReadOptions

tree/ntuple/inc/ROOT/RNTupleView.hxx

@@ -93,7 +93,7 @@ protected:
       std::unique_ptr<ROOT::RFieldBase> field;
       {
          auto descGuard = pageSource.GetSharedDescriptorGuard();
-         const auto &desc = descGuard.GetRef();
+         const auto &desc = descGuard.GetConstRef();
          const auto &fieldDesc = desc.GetFieldDescriptor(fieldId);
          if constexpr (std::is_void_v<T>) {
             if (typeName.empty())
@@ -275,7 +275,7 @@ protected:
    static ROOT::RField<T> CreateField(ROOT::DescriptorId_t fieldId, ROOT::Internal::RPageSource &pageSource)
    {
       auto descGuard = pageSource.GetSharedDescriptorGuard();
-      const auto &desc = descGuard.GetRef();
+      const auto &desc = descGuard.GetConstRef();
       const auto &fieldDesc = desc.GetFieldDescriptor(fieldId);
       if (!Internal::IsMatchingFieldType<T>(fieldDesc.GetTypeName())) {
          throw RException(R__FAIL("type mismatch for field " + fieldDesc.GetFieldName() + ": " +
@@ -340,7 +340,7 @@ private:
       std::string fieldName;
       {
          auto descGuard = source->GetSharedDescriptorGuard();
-         const auto &desc = descGuard.GetRef();
+         const auto &desc = descGuard.GetConstRef();
          const auto &fieldDesc = desc.GetFieldDescriptor(fieldId);
          if (fieldDesc.GetStructure() != ROOT::ENTupleStructure::kCollection) {
             throw RException(

tree/ntuple/inc/ROOT/RPageStorage.hxx

@@ -40,6 +40,7 @@
 #include <shared_mutex>
 #include <unordered_map>
 #include <unordered_set>
+#include <variant>
 #include <vector>
 
 namespace ROOT {
@@ -619,42 +620,120 @@ public:
       std::shared_mutex *fLock = nullptr;
 
    public:
+      RSharedDescriptorGuard() = default;
       RSharedDescriptorGuard(const ROOT::RNTupleDescriptor &desc, std::shared_mutex &lock)
          : fDescriptor(&desc), fLock(&lock)
       {
          fLock->lock_shared();
       }
       RSharedDescriptorGuard(const RSharedDescriptorGuard &) = delete;
       RSharedDescriptorGuard &operator=(const RSharedDescriptorGuard &) = delete;
-      RSharedDescriptorGuard(RSharedDescriptorGuard &&) = default;
-      RSharedDescriptorGuard &operator=(RSharedDescriptorGuard &&) = default;
-      ~RSharedDescriptorGuard() { fLock->unlock_shared(); }
+      RSharedDescriptorGuard(RSharedDescriptorGuard &&other)
+      {
+         std::swap(fDescriptor, other.fDescriptor);
+         std::swap(fLock, other.fLock);
+      }
+      RSharedDescriptorGuard &operator=(RSharedDescriptorGuard &&other)
+      {
+         if (this == &other)
+            return *this;
+         std::swap(fDescriptor, other.fDescriptor);
+         std::swap(fLock, other.fLock);
+         return *this;
+      }
+      ~RSharedDescriptorGuard()
+      {
+         if (IsValid())
+            Release();
+      }
       const ROOT::RNTupleDescriptor *operator->() const { return fDescriptor; }
-      const ROOT::RNTupleDescriptor &GetRef() const { return *fDescriptor; }
+      const ROOT::RNTupleDescriptor &GetConstRef() const { return *fDescriptor; }
+      bool IsValid() const { return fLock; }
+      void Release()
+      {
+         fLock->unlock_shared();
+         fLock = nullptr;
+         fDescriptor = nullptr;
+      }
    };
 
    /// An RAII wrapper used for the writable access to `RPageSource::fDescriptor`. See `GetSharedDescriptorGuard()`.
    class RExclDescriptorGuard {
-      ROOT::RNTupleDescriptor &fDescriptor;
-      std::shared_mutex &fLock;
+      ROOT::RNTupleDescriptor *fDescriptor = nullptr;
+      std::shared_mutex *fLock = nullptr;
 
    public:
-      RExclDescriptorGuard(ROOT::RNTupleDescriptor &desc, std::shared_mutex &lock) : fDescriptor(desc), fLock(lock)
+      RExclDescriptorGuard() = default;
+      RExclDescriptorGuard(ROOT::RNTupleDescriptor &desc, std::shared_mutex &lock) : fDescriptor(&desc), fLock(&lock)
       {
-         fLock.lock();
+         fLock->lock();
       }
       RExclDescriptorGuard(const RExclDescriptorGuard &) = delete;
       RExclDescriptorGuard &operator=(const RExclDescriptorGuard &) = delete;
-      RExclDescriptorGuard(RExclDescriptorGuard &&) = delete;
-      RExclDescriptorGuard &operator=(RExclDescriptorGuard &&) = delete;
+      RExclDescriptorGuard(RExclDescriptorGuard &&other)
+      {
+         std::swap(fDescriptor, other.fDescriptor);
+         std::swap(fLock, other.fLock);
+      }
+      RExclDescriptorGuard &operator=(RExclDescriptorGuard &&other)
+      {
+         if (this == &other)
+            return *this;
+         std::swap(fDescriptor, other.fDescriptor);
+         std::swap(fLock, other.fLock);
+         return *this;
+      }
       ~RExclDescriptorGuard()
       {
-         fDescriptor.IncGeneration();
-         fLock.unlock();
+         if (IsValid())
+            Release();
+      }
+      ROOT::RNTupleDescriptor &operator*() const { return *fDescriptor; }
+      ROOT::RNTupleDescriptor *operator->() const { return fDescriptor; }
+      void MoveIn(ROOT::RNTupleDescriptor desc) { *fDescriptor = std::move(desc); }
+      bool IsValid() const { return fLock; }
+      void Release()
+      {
+         fDescriptor->IncGeneration();
+         fLock->unlock();
+         fLock = nullptr;
+         fDescriptor = nullptr;
+      }
+   };
+
+   /// Either a shared or an exclusive lock guard
+   class RAnyDescriptorGuard {
+   private:
+      std::variant<RSharedDescriptorGuard, RExclDescriptorGuard> fAnyGuard;
+
+   public:
+      RAnyDescriptorGuard(RSharedDescriptorGuard sharedGuard) : fAnyGuard(std::move(sharedGuard)) {}
+      RAnyDescriptorGuard(RExclDescriptorGuard exclGuard) : fAnyGuard(std::move(exclGuard)) {}
+      RAnyDescriptorGuard(const RAnyDescriptorGuard &) = delete;
+      RAnyDescriptorGuard &operator=(const RAnyDescriptorGuard &) = delete;
+      RAnyDescriptorGuard(RAnyDescriptorGuard &&) = default;
+      RAnyDescriptorGuard &operator=(RAnyDescriptorGuard &&) = default;
+      ~RAnyDescriptorGuard() = default;
+
+      const ROOT::RNTupleDescriptor &GetConstRef() const
+      {
+         return (fAnyGuard.index() == 0) ? std::get<0>(fAnyGuard).GetConstRef() : std::get<1>(fAnyGuard).operator*();
+      }
+      const ROOT::RNTupleDescriptor *operator->() const
+      {
+         return (fAnyGuard.index() == 0) ? std::get<0>(fAnyGuard).operator->() : std::get<1>(fAnyGuard).operator->();
+      }
+      bool IsValid() const
+      {
+         return (fAnyGuard.index() == 0) ? std::get<0>(fAnyGuard).IsValid() : std::get<1>(fAnyGuard).IsValid();
+      }
+      void Release()
+      {
+         if (fAnyGuard.index() == 0)
+            std::get<0>(fAnyGuard).Release();
+         else
+            std::get<1>(fAnyGuard).Release();
       }
-      ROOT::RNTupleDescriptor &operator*() const { return fDescriptor; }
-      ROOT::RNTupleDescriptor *operator->() const { return &fDescriptor; }
-      void MoveIn(ROOT::RNTupleDescriptor desc) { fDescriptor = std::move(desc); }
    };
 
 private:
@@ -696,6 +775,17 @@ private:
    bool fIsAttached = false;                 ///< Set to true once `Attach()` is called
    bool fHasStreamerInfosRegistered = false; ///< Set to true when RegisterStreamerInfos() is called.
 
+   /// The interpretation of page lists, set in Attach()
+   Internal::RNTupleSerializer::EDescriptorDeserializeMode fDeserializationMode =
+      Internal::RNTupleSerializer::EDescriptorDeserializeMode::kRaw;
+
+   /// This vector is aligned with the cluster groups in the descriptor. It stores the cumulative number of clusters
+   /// in the cluster groups. Given a cluster ID, we can thus quickly determine the cluster group that the cluster
+   /// comes from. While descriptor IDs in a descriptor in general are arbitrary, for the descriptor in the page source,
+   /// that was created from a serialized on-disk representation, we know that cluster and cluster group IDs are
+   /// issued consecutively.
+   std::vector<NTupleSize_t> fCumulativeClusterCounts;
+
    /// The active columns are implicitly defined by the model fields or views
    RActivePhysicalColumns fActivePhysicalColumns;
 
@@ -719,11 +809,21 @@ private:
    /// Pages of pinned clusters won't be evicted from the page pool.
    std::unordered_set<ROOT::DescriptorId_t> fPinnedClusters;
 
+   /// Ensures that fDescriptor has cluster details loaded for the given cluster group ID. This method is expected
+   /// to be called while holding the lock passed by descGuard. The returned descriptor guard is either the
+   /// passed one or a new, exclusive guard if cluster details needed to be loaded.
+   RAnyDescriptorGuard EnsureClusterDetails(DescriptorId_t cgId, RAnyDescriptorGuard descGuard);
+   /// Uses binary search in fCumulativeClusterCounts to determine the cluster group that the cluster ID belongs in.
+   DescriptorId_t FindClusterGroupId(DescriptorId_t clusterId) const;
+
    /// Does nothing if fLastUsedCluster == clusterId. Otherwise, updated fLastUsedCluster
    /// and evict unused paged from the page pool of all previous clusters.
    /// Must not be called when the descriptor guard is taken.
    void UpdateLastUsedCluster(ROOT::DescriptorId_t clusterId);
 
+   // Populate the cluster details of the given cluster group. Must hold an exclusive descriptor lock when calling.
+   // No-op if the cluster details are already present.
+   void LoadPageList(DescriptorId_t clusterGroupId, const RExclDescriptorGuard &exclGuard);
    // Common treatment of zero pages in LoadPageFromSummary()
    ROOT::Internal::RPageRef LoadZeroPage(ColumnHandle_t columnHandle, const RPageSummary &pageSummary);
    // Once the page is found to be missing in the page cache and all information about the page is collected,
@@ -857,18 +957,20 @@ public:
    /// Open the physical storage container and deserialize header and footer
    void Attach(ROOT::Internal::RNTupleSerializer::EDescriptorDeserializeMode mode =
                   ROOT::Internal::RNTupleSerializer::EDescriptorDeserializeMode::kForReading);
+   /// Load all remaining cluster details
+   void LoadAllPageLists();
 
    ROOT::NTupleSize_t GetNEntries();
    ROOT::NTupleSize_t GetNElements(ROOT::DescriptorId_t physicalColumnId);
    /// Returns a shared descriptor guard to ensure that the returned cluster id is useable, i.e. that the
    /// corresponding cluster was not meanwhile evicted from the set of active clusters.
-   RSharedDescriptorGuard
+   RAnyDescriptorGuard
    FindClusterId(ROOT::DescriptorId_t physicalColumnId, ROOT::NTupleSize_t index, ROOT::DescriptorId_t &cid);
    /// An overload of FindClusterId that searches using a certain column element index.
-   RSharedDescriptorGuard FindClusterId(ROOT::NTupleSize_t entryIdx, ROOT::DescriptorId_t &cid);
+   RAnyDescriptorGuard FindClusterId(ROOT::NTupleSize_t entryIdx, ROOT::DescriptorId_t &cid);
    /// Uses FindClusterId to search for the cluster with the entry index following the last entry index of the
    /// given cluster.
-   RSharedDescriptorGuard FindNextClusterId(ROOT::DescriptorId_t clusterId, ROOT::DescriptorId_t &nextId);
+   RAnyDescriptorGuard FindNextClusterId(ROOT::DescriptorId_t clusterId, ROOT::DescriptorId_t &nextId);
 
    /// Promise to only read from the given entry range. If set, prevents the cluster pool from reading-ahead beyond
    /// the given range. The range needs to be within `[0, GetNEntries())`.

tree/ntuple/inc/ROOT/RPageStorageDaos.hxx

@@ -26,6 +26,7 @@
 #include <atomic>
 #include <cstdio>
 #include <memory>
+#include <mutex>
 #include <string>
 #include <optional>
 
@@ -151,6 +152,8 @@ private:
 
    /// A container that stores object data (header/footer, pages, etc.)
    std::unique_ptr<RDaosContainer> fDaosContainer;
+   /// Protects the DAOS container
+   std::mutex fContainerLock;
    /// A URI to a DAOS pool of the form 'daos://pool-label/container-label'
    std::string fURI;
 

tree/ntuple/inc/ROOT/RPageStorageFile.hxx

@@ -24,6 +24,7 @@
 #include <array>
 #include <cstdio>
 #include <memory>
+#include <mutex>
 #include <optional>
 #include <string>
 #include <utility>
@@ -139,6 +140,8 @@ private:
    RNTupleDescriptorBuilder fDescriptorBuilder;
    /// Tracks the last read offset for seek distance calculation
    std::uint64_t fLastOffset = 0;
+   /// Protects the raw file, the reader, and fLastOffset
+   std::mutex fReaderLock;
 
    /// File-specific I/O performance counters
    struct RFileCounters {
@@ -160,6 +163,9 @@ private:
    std::unique_ptr<ROOT::Internal::RCluster>
    PrepareSingleCluster(const ROOT::Internal::RCluster::RKey &clusterKey, std::vector<RRawFile::RIOVec> &readRequests);
 
+   // Helper to calculate the distance between fLastOffset and offset and to update the skip counter accordingly
+   void UpdateSkipCounter(std::uint64_t offset);
+
 protected:
    void LoadStructureImpl() final;
    ROOT::RNTupleDescriptor AttachImpl() final;