[ntuple] Add de/serialization of RNTupleAttributes

Author: silverweed

tree/ntuple/inc/ROOT/RNTupleDescriptor.hxx

@@ -67,6 +67,53 @@ struct RNTupleClusterBoundaries {
 std::vector<ROOT::Internal::RNTupleClusterBoundaries> GetClusterBoundaries(const RNTupleDescriptor &desc);
 } // namespace Internal
 
+namespace Experimental {
+namespace Internal {
+class RNTupleAttrSetDescriptorBuilder;
+}
+
+// clang-format off
+/**
+\class ROOT::Experimental::RNTupleAttrSetDescriptor
+\ingroup NTuple
+\brief Metadata stored for every Attribute Set linked to an RNTuple.
+*/
+// clang-format on
+class RNTupleAttrSetDescriptor final {
+   friend class Experimental::Internal::RNTupleAttrSetDescriptorBuilder;
+
+   std::uint16_t fSchemaVersionMajor = 0;
+   std::uint16_t fSchemaVersionMinor = 0;
+   std::uint32_t fAnchorLength = 0; ///< uncompressed size of the linked anchor
+   // The locator of the AttributeSet anchor.
+   // In case of kTypeFile, it points to the beginning of the Anchor's payload.
+   // NOTE: Only kTypeFile is supported at the moment.
+   RNTupleLocator fAnchorLocator;
+   std::string fName;
+
+public:
+   RNTupleAttrSetDescriptor() = default;
+   RNTupleAttrSetDescriptor(const RNTupleAttrSetDescriptor &other) = delete;
+   RNTupleAttrSetDescriptor &operator=(const RNTupleAttrSetDescriptor &other) = delete;
+   RNTupleAttrSetDescriptor(RNTupleAttrSetDescriptor &&other) = default;
+   RNTupleAttrSetDescriptor &operator=(RNTupleAttrSetDescriptor &&other) = default;
+
+   bool operator==(const RNTupleAttrSetDescriptor &other) const;
+   bool operator!=(const RNTupleAttrSetDescriptor &other) const { return !(*this == other); }
+
+   const std::string &GetName() const { return fName; }
+   std::uint16_t GetSchemaVersionMajor() const { return fSchemaVersionMajor; }
+   std::uint16_t GetSchemaVersionMinor() const { return fSchemaVersionMinor; }
+   std::uint32_t GetAnchorLength() const { return fAnchorLength; }
+   const RNTupleLocator &GetAnchorLocator() const { return fAnchorLocator; }
+
+   RNTupleAttrSetDescriptor Clone() const;
+};
+
+class RNTupleAttrSetDescriptorIterable;
+
+} // namespace Experimental
+
 // clang-format off
 /**
 \class ROOT::RFieldDescriptor
@@ -701,6 +748,8 @@ private:
    std::vector<ROOT::DescriptorId_t> fSortedClusterGroupIds;
    /// Potentially a subset of all the available clusters
    std::unordered_map<ROOT::DescriptorId_t, RClusterDescriptor> fClusterDescriptors;
+   /// List of AttributeSets linked to this RNTuple
+   std::vector<Experimental::RNTupleAttrSetDescriptor> fAttributeSets;
 
    // We don't expose this publicly because when we add sharded clusters, this interface does not make sense anymore
    ROOT::DescriptorId_t FindClusterId(ROOT::NTupleSize_t entryIdx) const;
@@ -718,6 +767,7 @@ public:
    class RClusterGroupDescriptorIterable;
    class RClusterDescriptorIterable;
    class RExtraTypeInfoDescriptorIterable;
+   friend class Experimental::RNTupleAttrSetDescriptorIterable;
 
    /// Modifiers passed to CreateModel()
    struct RCreateModelOptions {
@@ -806,6 +856,8 @@ public:
 
    RExtraTypeInfoDescriptorIterable GetExtraTypeInfoIterable() const;
 
+   ROOT::Experimental::RNTupleAttrSetDescriptorIterable GetAttrSetIterable() const;
+
    const std::string &GetName() const { return fName; }
    const std::string &GetDescription() const { return fDescription; }
 
@@ -816,6 +868,7 @@ public:
    std::size_t GetNClusters() const { return fNClusters; }
    std::size_t GetNActiveClusters() const { return fClusterDescriptors.size(); }
    std::size_t GetNExtraTypeInfos() const { return fExtraTypeInfoDescriptors.size(); }
+   std::size_t GetNAttributeSets() const { return fAttributeSets.size(); }
 
    /// We know the number of entries from adding the cluster summaries
    ROOT::NTupleSize_t GetNEntries() const { return fNEntries; }
@@ -1145,6 +1198,59 @@ public:
    RIterator end() { return RIterator(fNTuple.fExtraTypeInfoDescriptors.cend()); }
 };
 
+namespace Experimental {
+// clang-format off
+/**
+\class ROOT::Experimental::RNTupleAttrSetDescriptorIterable
+\ingroup NTuple
+\brief Used to loop over all the Attribute Sets linked to an RNTuple
+*/
+// clang-format on
+// TODO: move this to RNTupleDescriptor::RNTupleAttrSetDescriptorIterable when it moves out of Experimental.
+class RNTupleAttrSetDescriptorIterable final {
+private:
+   /// The associated RNTuple for this range.
+   const RNTupleDescriptor &fNTuple;
+
+public:
+   class RIterator final {
+   private:
+      using Iter_t = std::vector<RNTupleAttrSetDescriptor>::const_iterator;
+      /// The wrapped vector iterator
+      Iter_t fIter;
+
+   public:
+      using iterator_category = std::forward_iterator_tag;
+      using iterator = RIterator;
+      using value_type = RNTupleAttrSetDescriptor;
+      using difference_type = std::ptrdiff_t;
+      using pointer = const value_type *;
+      using reference = const value_type &;
+
+      RIterator(Iter_t iter) : fIter(iter) {}
+      iterator &operator++() /* prefix */
+      {
+         ++fIter;
+         return *this;
+      }
+      iterator operator++(int) /* postfix */
+      {
+         auto old = *this;
+         operator++();
+         return old;
+      }
+      reference operator*() const { return *fIter; }
+      pointer operator->() const { return &*fIter; }
+      bool operator!=(const iterator &rh) const { return fIter != rh.fIter; }
+      bool operator==(const iterator &rh) const { return fIter == rh.fIter; }
+   };
+
+   RNTupleAttrSetDescriptorIterable(const RNTupleDescriptor &ntuple) : fNTuple(ntuple) {}
+   RIterator begin() { return RIterator(fNTuple.fAttributeSets.cbegin()); }
+   RIterator end() { return RIterator(fNTuple.fAttributeSets.cend()); }
+};
+} // namespace Experimental
+
 // clang-format off
 /**
 \class ROOT::RNTupleDescriptor::RHeaderExtension
@@ -1218,6 +1324,39 @@ public:
    }
 };
 
+namespace Experimental::Internal {
+class RNTupleAttrSetDescriptorBuilder final {
+   ROOT::Experimental::RNTupleAttrSetDescriptor fDesc;
+
+public:
+   RNTupleAttrSetDescriptorBuilder &Name(std::string_view name)
+   {
+      fDesc.fName = name;
+      return *this;
+   }
+   RNTupleAttrSetDescriptorBuilder &SchemaVersion(std::uint16_t major, std::uint16_t minor)
+   {
+      fDesc.fSchemaVersionMajor = major;
+      fDesc.fSchemaVersionMinor = minor;
+      return *this;
+   }
+   RNTupleAttrSetDescriptorBuilder &AnchorLocator(const RNTupleLocator &loc)
+   {
+      fDesc.fAnchorLocator = loc;
+      return *this;
+   }
+   RNTupleAttrSetDescriptorBuilder &AnchorLength(std::uint32_t length)
+   {
+      fDesc.fAnchorLength = length;
+      return *this;
+   }
+
+   /// Attempt to make an AttributeSet descriptor. This may fail if the builder
+   /// was not given enough information to make a proper descriptor.
+   RResult<ROOT::Experimental::RNTupleAttrSetDescriptor> MoveDescriptor();
+};
+} // namespace Experimental::Internal
+
 namespace Internal {
 
 // clang-format off
@@ -1601,6 +1740,8 @@ public:
    RResult<void> AddExtraTypeInfo(RExtraTypeInfoDescriptor &&extraTypeInfoDesc);
    void ReplaceExtraTypeInfo(RExtraTypeInfoDescriptor &&extraTypeInfoDesc);
 
+   RResult<void> AddAttributeSet(Experimental::RNTupleAttrSetDescriptor &&attrSetDesc);
+
    /// Mark the beginning of the header extension; any fields and columns added after a call to this function are
    /// annotated as begin part of the header extension.
    void BeginHeaderExtension();
@@ -1634,6 +1775,7 @@ inline RNTupleDescriptor CloneDescriptorSchema(const RNTupleDescriptor &desc)
 }
 
 } // namespace Internal
+
 } // namespace ROOT
 
 #endif // ROOT_RNTupleDescriptor

tree/ntuple/inc/ROOT/RNTupleSerialize.hxx

@@ -36,6 +36,13 @@ class RNTupleDescriptor;
 class RClusterDescriptor;
 enum class EExtraTypeInfoIds;
 
+namespace Experimental {
+class RNTupleAttrSetDescriptor;
+namespace Internal {
+class RNTupleAttrSetDescriptorBuilder;
+} // namespace Internal
+} // namespace Experimental
+
 namespace Internal {
 
 class RClusterDescriptorBuilder;
@@ -271,6 +278,12 @@ public:
    static RResult<std::uint32_t> DeserializeSchemaDescription(const void *buffer, std::uint64_t bufSize,
                                                               ROOT::Internal::RNTupleDescriptorBuilder &descBuilder);
 
+   static RResult<std::uint32_t>
+   SerializeAttributeSet(const Experimental::RNTupleAttrSetDescriptor &attrSetDesc, void *buffer);
+   static RResult<std::uint32_t>
+   DeserializeAttributeSet(const void *buffer, std::uint64_t bufSize,
+                           Experimental::Internal::RNTupleAttrSetDescriptorBuilder &attrSetDescBld);
+
    static RResult<RContext> SerializeHeader(void *buffer, const RNTupleDescriptor &desc);
    static RResult<std::uint32_t> SerializePageList(void *buffer, const RNTupleDescriptor &desc,
                                                    std::span<ROOT::DescriptorId_t> physClusterIDs,

tree/ntuple/src/RNTupleDescriptor.cxx

@@ -803,6 +803,8 @@ ROOT::RNTupleDescriptor ROOT::RNTupleDescriptor::Clone() const
    clone.fSortedClusterGroupIds = fSortedClusterGroupIds;
    for (const auto &d : fClusterDescriptors)
       clone.fClusterDescriptors.emplace(d.first, d.second.Clone());
+   for (const auto &d : fAttributeSets)
+      clone.fAttributeSets.emplace_back(d.Clone());
    return clone;
 }
 
@@ -1123,6 +1125,19 @@ void ROOT::Internal::RNTupleDescriptorBuilder::SetFeature(unsigned int flag)
    fDescriptor.fFeatureFlags.insert(flag);
 }
 
+ROOT::RResult<ROOT::Experimental::RNTupleAttrSetDescriptor>
+ROOT::Experimental::Internal::RNTupleAttrSetDescriptorBuilder::MoveDescriptor()
+{
+   if (fDesc.fName.empty())
+      return R__FAIL("attribute set name cannot be empty");
+   if (fDesc.fAnchorLength == 0)
+      return R__FAIL("invalid anchor length");
+   if (fDesc.fAnchorLocator.GetType() == RNTupleLocator::kTypeUnknown)
+      return R__FAIL("invalid locator type");
+
+   return std::move(fDesc);
+}
+
 ROOT::RResult<ROOT::RColumnDescriptor> ROOT::Internal::RColumnDescriptorBuilder::MakeDescriptor() const
 {
    if (fColumn.GetLogicalId() == ROOT::kInvalidDescriptorId)
@@ -1377,6 +1392,19 @@ void ROOT::Internal::RNTupleDescriptorBuilder::ReplaceExtraTypeInfo(RExtraTypeIn
       fDescriptor.fExtraTypeInfoDescriptors.emplace_back(std::move(extraTypeInfoDesc));
 }
 
+ROOT::RResult<void>
+ROOT::Internal::RNTupleDescriptorBuilder::AddAttributeSet(Experimental::RNTupleAttrSetDescriptor &&attrSetDesc)
+{
+   auto &attrSets = fDescriptor.fAttributeSets;
+   if (std::find_if(attrSets.begin(), attrSets.end(), [&name = attrSetDesc.GetName()](const auto &desc) {
+          return desc.GetName() == name;
+       }) != attrSets.end()) {
+      return R__FAIL("attribute sets with duplicate names");
+   }
+   attrSets.push_back(std::move(attrSetDesc));
+   return RResult<void>::Success();
+}
+
 RNTupleSerializer::StreamerInfoMap_t ROOT::Internal::RNTupleDescriptorBuilder::BuildStreamerInfos() const
 {
    RNTupleSerializer::StreamerInfoMap_t streamerInfoMap;
@@ -1492,3 +1520,26 @@ ROOT::RNTupleDescriptor::RExtraTypeInfoDescriptorIterable ROOT::RNTupleDescripto
 {
    return RExtraTypeInfoDescriptorIterable(*this);
 }
+
+ROOT::Experimental::RNTupleAttrSetDescriptorIterable ROOT::RNTupleDescriptor::GetAttrSetIterable() const
+{
+   return Experimental::RNTupleAttrSetDescriptorIterable(*this);
+}
+
+bool ROOT::Experimental::RNTupleAttrSetDescriptor::operator==(const RNTupleAttrSetDescriptor &other) const
+{
+   return fAnchorLength == other.fAnchorLength && fSchemaVersionMajor == other.fSchemaVersionMajor &&
+          fSchemaVersionMinor == other.fSchemaVersionMinor && fAnchorLocator == other.fAnchorLocator &&
+          fName == other.fName;
+};
+
+ROOT::Experimental::RNTupleAttrSetDescriptor ROOT::Experimental::RNTupleAttrSetDescriptor::Clone() const
+{
+   RNTupleAttrSetDescriptor desc;
+   desc.fAnchorLength = fAnchorLength;
+   desc.fSchemaVersionMajor = fSchemaVersionMajor;
+   desc.fSchemaVersionMinor = fSchemaVersionMinor;
+   desc.fAnchorLocator = fAnchorLocator;
+   desc.fName = fName;
+   return desc;
+}