[ntuple] add GetAttributes methods to RNTupleAttrSetReader

Author: silverweed

tree/ntuple/inc/ROOT/RNTupleAttrReading.hxx

@@ -25,6 +25,8 @@ class RNTupleModel;
 
 namespace Experimental {
 
+class RNTupleAttrEntryIterable;
+
 // clang-format off
 /**
 \class ROOT::Experimental::RNTupleAttrRange
@@ -111,6 +113,7 @@ for (auto idx : attrSet->GetAttributes(10)) {
 // clang-format on
 class RNTupleAttrSetReader final {
    friend class ROOT::RNTupleReader;
+   friend class RNTupleAttrEntryIterable;
 
    /// List containing pairs { entryRange, entryIndex }, used to quickly find out which entries in the Attribute
    /// RNTuple contain entries that overlap a given range. The list is sorted by range start, i.e.
@@ -121,8 +124,13 @@ class RNTupleAttrSetReader final {
    /// The reconstructed user model
    std::unique_ptr<ROOT::RNTupleModel> fUserModel;
 
+   static bool EntryRangesAreSorted(const decltype(fEntryRanges) &ranges);
+
    explicit RNTupleAttrSetReader(std::unique_ptr<RNTupleReader> reader);
 
+   std::vector<ROOT::NTupleSize_t>
+   GetAttributesRangeInternal(NTupleSize_t startEntry, NTupleSize_t endEntry, bool rangeIsContained);
+
 public:
    RNTupleAttrSetReader(const RNTupleAttrSetReader &) = delete;
    RNTupleAttrSetReader &operator=(const RNTupleAttrSetReader &) = delete;
@@ -148,6 +156,92 @@ public:
 
    /// Returns the number of all attribute entries in this attribute set.
    std::size_t GetNEntries() const { return fEntryRanges.size(); }
+
+   /// Returns all the attributes in this Set. The returned attributes are sorted by entry range start.
+   RNTupleAttrEntryIterable GetAttributes();
+   /// Returns all the attributes whose range contains index `entryIndex`.
+   RNTupleAttrEntryIterable GetAttributes(NTupleSize_t entryIndex);
+   /// Returns all the attributes whose range fully contains `[startEntry, endEntry)`
+   RNTupleAttrEntryIterable GetAttributesContainingRange(NTupleSize_t startEntry, NTupleSize_t endEntry);
+   /// Returns all the attributes whose range is fully contained in `[startEntry, endEntry)`
+   RNTupleAttrEntryIterable GetAttributesInRange(NTupleSize_t startEntry, NTupleSize_t endEntry);
+};
+
+// clang-format off
+/**
+\class ROOT::Experimental::RNTupleAttrEntryIterable
+\ingroup NTuple
+\brief Iterable class used to loop over attribute entries.
+
+This class allows to perform range-for iteration on some set of attributes, typically returned by the
+RNTupleAttrSetReader::GetAttributes family of methods.
+
+See the documentation of RNTupleAttrSetReader for example usage.
+*/
+// clang-format on
+class RNTupleAttrEntryIterable final {
+public:
+   struct RFilter {
+      RNTupleAttrRange fRange;
+      bool fIsContained;
+   };
+
+private:
+   RNTupleAttrSetReader &fReader;
+   std::optional<RFilter> fFilter;
+
+public:
+   class RIterator final {
+   private:
+      using Iter_t = decltype(std::declval<RNTupleAttrSetReader>().fEntryRanges.begin());
+      Iter_t fCur, fEnd;
+      std::optional<RFilter> fFilter;
+
+      Iter_t Next() const;
+      bool FullyContained(RNTupleAttrRange range) const;
+
+   public:
+      using iterator_category = std::forward_iterator_tag;
+      using iterator = RIterator;
+      using value_type = NTupleSize_t;
+      using difference_type = std::ptrdiff_t;
+      using pointer = const value_type *;
+      using reference = const value_type &;
+
+      RIterator(Iter_t iter, Iter_t end, std::optional<RFilter> filter) : fCur(iter), fEnd(end), fFilter(filter)
+      {
+         if (fFilter) {
+            if (fFilter->fRange.GetLength() == 0)
+               fCur = end;
+            else
+               fCur = Next();
+         }
+      }
+      iterator operator++()
+      {
+         ++fCur;
+         fCur = Next();
+         return *this;
+      }
+      iterator operator++(int)
+      {
+         iterator it = *this;
+         ++fCur;
+         fCur = Next();
+         return it;
+      }
+      reference operator*() { return fCur->second; }
+      bool operator!=(const iterator &rh) const { return !operator==(rh); }
+      bool operator==(const iterator &rh) const { return fCur == rh.fCur; }
+   };
+
+   explicit RNTupleAttrEntryIterable(RNTupleAttrSetReader &reader, std::optional<RFilter> filter = {})
+      : fReader(reader), fFilter(filter)
+   {
+   }
+
+   RIterator begin() { return RIterator{fReader.fEntryRanges.begin(), fReader.fEntryRanges.end(), fFilter}; }
+   RIterator end() { return RIterator{fReader.fEntryRanges.end(), fReader.fEntryRanges.end(), fFilter}; }
 };
 
 } // namespace Experimental

tree/ntuple/src/RNTupleAttrReading.cxx

@@ -98,3 +98,145 @@ std::unique_ptr<ROOT::REntry> ROOT::Experimental::RNTupleAttrSetReader::CreateEn
    return fUserModel->CreateEntry();
 }
 
+// Entry ranges should be sorted with respect to GetStart by construction.
+bool ROOT::Experimental::RNTupleAttrSetReader::EntryRangesAreSorted(const decltype(fEntryRanges) &ranges)
+{
+   ROOT::NTupleSize_t prevStart = 0;
+   for (const auto &[range, _] : ranges) {
+      if (range.GetStart() < prevStart)
+         return false;
+      prevStart = range.GetStart();
+   }
+   return true;
+};
+
+std::vector<ROOT::NTupleSize_t>
+ROOT::Experimental::RNTupleAttrSetReader::GetAttributesRangeInternal(NTupleSize_t startEntry, NTupleSize_t endEntry,
+                                                                     bool rangeIsContained)
+{
+   std::vector<ROOT::NTupleSize_t> result;
+
+   if (endEntry < startEntry) {
+      R__LOG_WARNING(ROOT::Internal::NTupleLog())
+         << "end < start when getting attributes from Attribute Set '" << GetDescriptor().GetName()
+         << "' (range given: [" << startEntry << ", " << endEntry << "].";
+      return result;
+   }
+
+   assert(EntryRangesAreSorted(fEntryRanges));
+
+   const auto FullyContained = [rangeIsContained](auto startInner, auto endInner, auto startOuter, auto endOuter) {
+      if (rangeIsContained) {
+         std::swap(startOuter, startInner);
+         std::swap(endOuter, endInner);
+      }
+      return startOuter <= startInner && endInner <= endOuter;
+   };
+
+   // TODO: consider using binary search, since fEntryRanges is sorted
+   // (maybe it should be done only if the size of the list is bigger than a threshold).
+   for (const auto &[range, index] : fEntryRanges) {
+      const auto &firstLast = range.GetFirstLast();
+      if (!firstLast)
+         continue;
+
+      const auto &[first, last] = *firstLast;
+      if (first >= endEntry)
+         break; // We can break here because fEntryRanges is sorted.
+
+      if (FullyContained(startEntry, endEntry, first, last + 1)) {
+         result.push_back(index);
+      }
+   }
+
+   return result;
+}
+
+ROOT::Experimental::RNTupleAttrEntryIterable
+ROOT::Experimental::RNTupleAttrSetReader::GetAttributesContainingRange(NTupleSize_t startEntry, NTupleSize_t endEntry)
+{
+   RNTupleAttrRange range;
+   if (endEntry <= startEntry) {
+      R__LOG_WARNING(ROOT::Internal::NTupleLog())
+         << "empty range given when getting attributes from Attribute Set '" << GetDescriptor().GetName()
+         << "' (range given: [" << startEntry << ", " << endEntry << ")).";
+      // Make sure we find 0 entries
+      range = RNTupleAttrRange::FromStartLength(startEntry, 0);
+   } else {
+      range = RNTupleAttrRange::FromStartEnd(startEntry, endEntry);
+   }
+   RNTupleAttrEntryIterable::RFilter filter{range, false};
+   return RNTupleAttrEntryIterable{*this, filter};
+}
+
+ROOT::Experimental::RNTupleAttrEntryIterable
+ROOT::Experimental::RNTupleAttrSetReader::GetAttributesInRange(NTupleSize_t startEntry, NTupleSize_t endEntry)
+{
+   RNTupleAttrRange range;
+   if (endEntry <= startEntry) {
+      R__LOG_WARNING(ROOT::Internal::NTupleLog())
+         << "empty range given when getting attributes from Attribute Set '" << GetDescriptor().GetName()
+         << "' (range given: [" << startEntry << ", " << endEntry << ")).";
+      // Make sure we find 0 entries
+      range = RNTupleAttrRange::FromStartLength(startEntry, 0);
+   } else {
+      range = RNTupleAttrRange::FromStartEnd(startEntry, endEntry);
+   }
+   RNTupleAttrEntryIterable::RFilter filter{range, true};
+   return RNTupleAttrEntryIterable{*this, filter};
+}
+
+ROOT::Experimental::RNTupleAttrEntryIterable
+ROOT::Experimental::RNTupleAttrSetReader::GetAttributes(NTupleSize_t entryIndex)
+{
+   RNTupleAttrEntryIterable::RFilter filter{RNTupleAttrRange::FromStartEnd(entryIndex, entryIndex + 1), false};
+   return RNTupleAttrEntryIterable{*this, filter};
+}
+
+ROOT::Experimental::RNTupleAttrEntryIterable ROOT::Experimental::RNTupleAttrSetReader::GetAttributes()
+{
+   return RNTupleAttrEntryIterable{*this};
+}
+
+//
+//  RNTupleAttrEntryIterable
+//
+bool ROOT::Experimental::RNTupleAttrEntryIterable::RIterator::FullyContained(RNTupleAttrRange range) const
+{
+   assert(fFilter);
+   if (fFilter->fIsContained) {
+      return fFilter->fRange.GetStart() <= range.GetStart() && range.GetEnd() <= fFilter->fRange.GetEnd();
+   } else {
+      return range.GetStart() <= fFilter->fRange.GetStart() && fFilter->fRange.GetEnd() <= range.GetEnd();
+   }
+}
+
+ROOT::Experimental::RNTupleAttrEntryIterable::RIterator::Iter_t
+ROOT::Experimental::RNTupleAttrEntryIterable::RIterator::Next() const
+{
+   // TODO: consider using binary search, since fEntryRanges is sorted
+   // (maybe it should be done only if the size of the list is bigger than a threshold).
+   for (auto it = fCur; it != fEnd; ++it) {
+      const auto &[range, index] = *it;
+      // If we have no filter, every entry is valid.
+      if (!fFilter)
+         return it;
+
+      const auto &firstLast = range.GetFirstLast();
+      // If this is nullopt it means this is a zero-length entry: we always skip those except
+      // for the "catch-all" GetAttributes() (which is when fFilter is also nullopt).
+      if (!firstLast)
+         continue;
+
+      const auto &[first, last] = *firstLast;
+      if (first >= fFilter->fRange.GetEnd()) {
+         // Since fEntryRanges is sorted we know we are at the end of the iteration
+         // TODO: tweak fEnd to directly pass the last entry?
+         return fEnd;
+      }
+
+      if (FullyContained(RNTupleAttrRange::FromStartEnd(first, last + 1)))
+         return it;
+   }
+   return fEnd;
+}