[ntuple] add GetAttributes methods to RNTupleAttrSetReader

silverweed · silverweed · commit 0c52e4e22712 · 2026-04-14T11:24:12.000+02:00
diff --git a/tree/ntuple/inc/ROOT/RNTupleAttrReading.hxx b/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,6 +124,8 @@ 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);
 
 public:
@@ -148,6 +153,91 @@ 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 = nullptr;
+   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 SkipFiltered() 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 = SkipFiltered();
+         }
+      }
+      iterator operator++()
+      {
+         ++fCur;
+         fCur = SkipFiltered();
+         return *this;
+      }
+      iterator operator++(int)
+      {
+         iterator it = *this;
+         operator++();
+         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
diff --git a/tree/ntuple/src/RNTupleAttrReading.cxx b/tree/ntuple/src/RNTupleAttrReading.cxx
@@ -102,3 +102,104 @@ 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;
+};
+
+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::SkipFiltered() const
+{
+   // If we have no filter, every entry is valid.
+   if (!fFilter)
+      return fCur;
+
+   // 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;
+      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;
+}
diff --git a/tree/ntuple/test/ntuple_attributes.cxx b/tree/ntuple/test/ntuple_attributes.cxx
@@ -102,11 +102,61 @@ TEST(RNTupleAttributes, BasicReadingWriting)
       EXPECT_EQ(ntuple, nullptr);
    }
 
+   /// Reading
    auto reader = RNTupleReader::Open("ntuple", fileGuard.GetPath());
    EXPECT_EQ(reader->GetDescriptor().GetNAttributeSets(), 1);
    for (const auto &attrSetIt : reader->GetDescriptor().GetAttrSetIterable()) {
       EXPECT_EQ(attrSetIt.GetName(), "AttrSet1");
    }
+
+   auto attrSetReader = reader->OpenAttributeSet("AttrSet1");
+   EXPECT_EQ(attrSetReader->GetNEntries(), 1);
+   auto pAttr = attrSetReader->GetModel().GetDefaultEntry().GetPtr<std::string>("attr");
+   {
+      int nAttrs = 0;
+      // iterate all attributes
+      for (auto idx : attrSetReader->GetAttributes()) {
+         attrSetReader->LoadEntry(idx);
+         EXPECT_EQ(*pAttr, "My Attribute");
+         nAttrs += 1;
+      }
+      EXPECT_EQ(nAttrs, 1);
+   }
+   {
+      int nAttrs = 0;
+      // attributes containing entry 99
+      for (auto idx : attrSetReader->GetAttributes(99)) {
+         attrSetReader->LoadEntry(idx);
+         EXPECT_EQ(*pAttr, "My Attribute");
+         nAttrs += 1;
+      }
+      EXPECT_EQ(nAttrs, 1);
+   }
+   {
+      // attributes containing entry 100 (no entry)
+      auto iter = attrSetReader->GetAttributes(100);
+      EXPECT_EQ(iter.begin(), iter.end());
+   }
+   {
+      // attributes contained in entry range 50-200 (no entry)
+      auto iter = attrSetReader->GetAttributesInRange(50, 200);
+      EXPECT_EQ(iter.begin(), iter.end());
+   }
+   {
+      int nAttrs = 0;
+      // attributes contained in entry range 0-1000
+      for (auto idx : attrSetReader->GetAttributesInRange(0, 1000)) {
+         attrSetReader->LoadEntry(idx);
+         EXPECT_EQ(*pAttr, "My Attribute");
+         nAttrs += 1;
+      }
+      EXPECT_EQ(nAttrs, 1);
+   }
+   {
+      // attributes containing entry range 200-300 (no entry)
+      auto iter = attrSetReader->GetAttributesContainingRange(200, 300);
+      EXPECT_EQ(iter.begin(), iter.end());
+   }
 }
 
 TEST(RNTupleAttributes, BasicWritingWithExplicitEntry)
@@ -201,11 +251,11 @@ TEST(RNTupleAttributes, MultipleSets)
 
       auto attrModel1 = RNTupleModel::Create();
       auto pInt1 = attrModel1->MakeField<int>("int");
-      auto attrSet1 = writer->CreateAttributeSet(attrModel1->Clone(), "MyAttrSet1");
+      auto attrSet1 = writer->CreateAttributeSet(std::move(attrModel1), "MyAttrSet1");
 
       auto attrModel2 = RNTupleModel::Create();
       auto pString2 = attrModel2->MakeField<std::string>("string");
-      auto attrSet2 = writer->CreateAttributeSet(attrModel2->Clone(), "MyAttrSet2");
+      auto attrSet2 = writer->CreateAttributeSet(std::move(attrModel2), "MyAttrSet2");
 
       auto attrRange2 = attrSet2->BeginRange();
       for (int i = 0; i < 100; ++i) {
@@ -231,6 +281,55 @@ TEST(RNTupleAttributes, MultipleSets)
    EXPECT_EQ(attrSetReader1->GetNEntries(), 100);
    auto attrSetReader2 = reader->OpenAttributeSet("MyAttrSet2");
    EXPECT_EQ(attrSetReader2->GetNEntries(), 1);
+
+   auto attrEntry1 = attrSetReader1->CreateEntry();
+   auto pAttrInt = attrEntry1->GetPtr<int>("int");
+   auto attrEntry2 = attrSetReader2->CreateEntry();
+   auto pAttrString = attrEntry2->GetPtr<std::string>("string");
+   {
+      int nAttrs = 0;
+      for (auto idx : attrSetReader1->GetAttributesInRange(0, 1000)) {
+         auto range = attrSetReader1->LoadEntry(idx, *attrEntry1);
+         EXPECT_EQ(*pAttrInt, idx);
+         EXPECT_EQ(range.GetStart(), idx);
+         EXPECT_EQ(range.GetLength(), 1);
+         nAttrs += 1;
+      }
+      EXPECT_EQ(nAttrs, 100);
+   }
+   {
+      int nAttrs = 0;
+      for (auto idx : attrSetReader1->GetAttributes(42)) {
+         auto range = attrSetReader1->LoadEntry(idx, *attrEntry1);
+         EXPECT_EQ(*pAttrInt, 42);
+         EXPECT_EQ(range.GetStart(), 42);
+         EXPECT_EQ(range.GetLength(), 1);
+         nAttrs += 1;
+      }
+      EXPECT_EQ(nAttrs, 1);
+   }
+   {
+      int nAttrs = 0;
+      for (auto idx : attrSetReader2->GetAttributes()) {
+         auto range = attrSetReader2->LoadEntry(idx, *attrEntry2);
+         EXPECT_EQ(*pAttrString, "Run 1");
+         EXPECT_EQ(range.GetStart(), 0);
+         EXPECT_EQ(range.GetLength(), 100);
+         nAttrs += 1;
+      }
+      EXPECT_EQ(nAttrs, 1);
+   }
+   {
+      for (auto idx : attrSetReader2->GetAttributes()) {
+         // Reading into the wrong entry
+         try {
+            attrSetReader2->LoadEntry(idx, *attrEntry1);
+            FAIL() << "reading into an unrelated entry should fail";
+         } catch (const ROOT::RException &ex) {
+            EXPECT_THAT(ex.what(), testing::HasSubstr("mismatch between entry and model"));
+         }
+      }
+   }
 }
 
 TEST(RNTupleAttributes, AttributeInvalidModel)
