[ntuple] allow creating an AttrSetReader from RNTupleReader

Author: silverweed

tree/ntuple/inc/ROOT/RNTupleReader.hxx

@@ -527,6 +527,8 @@ public:
    /// ~~~
    void EnableMetrics() { fMetrics.Enable(); }
    const Experimental::Detail::RNTupleMetrics &GetMetrics() const { return fMetrics; }
+
+   std::unique_ptr<Experimental::RNTupleAttrSetReader> OpenAttributeSet(std::string_view attrSetName);
 }; // class RNTupleReader
 
 } // namespace ROOT

tree/ntuple/inc/ROOT/RPageStorage.hxx

@@ -854,6 +854,10 @@ public:
    /// Forces the loading of ROOT StreamerInfo from the underlying file. This currently only has an effect for
    /// TFile-backed sources.
    virtual void LoadStreamerInfo() = 0;
+
+   /// Reads an attribute set at the given location and returns a page source for it.
+   virtual std::unique_ptr<ROOT::Internal::RPageSource>
+   ReadAttributeSet(ROOT::RNTupleLocator anchorLocator, std::uint64_t anchorUncompLen);
 }; // class RPageSource
 
 } // namespace Internal

tree/ntuple/inc/ROOT/RPageStorageFile.hxx

@@ -200,6 +200,9 @@ public:
    LoadClusters(std::span<ROOT::Internal::RCluster::RKey> clusterKeys) final;
 
    void LoadStreamerInfo() final;
+
+   std::unique_ptr<ROOT::Internal::RPageSource>
+   ReadAttributeSet(ROOT::RNTupleLocator anchorLocator, std::uint64_t anchorUncompLen) final;
 }; // class RPageSourceFile
 
 } // namespace Internal

tree/ntuple/src/RNTupleReader.cxx

@@ -15,6 +15,7 @@
 
 #include <ROOT/RField.hxx>
 #include <ROOT/RFieldVisitor.hxx>
+#include <ROOT/RNTupleAttrReading.hxx>
 #include <ROOT/RNTupleImtTaskScheduler.hxx>
 #include <ROOT/RNTuple.hxx>
 #include <ROOT/RNTupleModel.hxx>
@@ -366,3 +367,20 @@ ROOT::DescriptorId_t ROOT::RNTupleReader::RetrieveFieldId(std::string_view field
    }
    return fieldId;
 }
+
+std::unique_ptr<ROOT::Experimental::RNTupleAttrSetReader>
+ROOT::RNTupleReader::OpenAttributeSet(std::string_view attrSetName)
+{
+   auto attrSets = GetDescriptor().GetAttrSetIterable();
+   const auto it =
+      std::find_if(attrSets.begin(), attrSets.end(), [&](const auto &d) { return d.GetName() == attrSetName; });
+   if (it == attrSets.end())
+      throw ROOT::RException(R__FAIL(std::string("No such attribute set: ") + std::string(attrSetName)));
+
+   auto attrSource = fSource->ReadAttributeSet(it->GetAnchorLocator(), it->GetAnchorLength());
+   auto newReader = std::unique_ptr<RNTupleReader>(new RNTupleReader(std::move(attrSource), RNTupleReadOptions{}));
+   R__ASSERT(newReader);
+   auto attrSetReader = std::unique_ptr<ROOT::Experimental::RNTupleAttrSetReader>(
+      new ROOT::Experimental::RNTupleAttrSetReader(std::move(newReader)));
+   return attrSetReader;
+}

tree/ntuple/src/RPageStorage.cxx

@@ -1331,3 +1331,9 @@ void ROOT::Internal::RPagePersistentSink::EnableDefaultMetrics(const std::string
       *fMetrics.MakeCounter<RNTupleTickCounter<RNTupleAtomicCounter> *>("timeCpuZip", "ns",
                                                                         "CPU time spent compressing")});
 }
+
+std::unique_ptr<ROOT::Internal::RPageSource>
+ROOT::Internal::RPageSource::ReadAttributeSet(ROOT::RNTupleLocator /*locator*/, std::uint64_t /*uncompLen*/)
+{
+   throw ROOT::RException(R__FAIL("reading attribute sets is not implemented for this Page Source"));
+}

tree/ntuple/src/RPageStorageFile.cxx

@@ -744,3 +744,15 @@ void ROOT::Internal::RPageSourceFile::LoadStreamerInfo()
 {
    fReader.LoadStreamerInfo();
 }
+
+std::unique_ptr<ROOT::Internal::RPageSource>
+ROOT::Internal::RPageSourceFile::ReadAttributeSet(ROOT::RNTupleLocator anchorLocator, std::uint64_t anchorUncompLen)
+{
+   assert(anchorLocator.GetType() == RNTupleLocator::kTypeFile);
+
+   const auto anchorPos = anchorLocator.GetPosition<std::uint64_t>();
+   auto attrAnchor =
+      fReader.GetNTupleProperAtOffset(anchorPos, anchorLocator.GetNBytesOnStorage(), anchorUncompLen).Unwrap();
+   auto attrSource = OpenWithDifferentAnchor(attrAnchor);
+   return attrSource;
+}