[ntuple] add support for leaf count arrays in classes

tree/ntuple/inc/ROOT/RField/RFieldSequenceContainer.hxx

@@ -417,6 +417,51 @@ public:
    void AcceptVisitor(ROOT::Detail::RFieldVisitor &visitor) const final;
 };
 
+/**
+\class ROOT::RLeafCountArrayField
+\brief A field for an RVec field as a class member on disk that is represented as leaf count array` in memory.
+\ingroup NTuple
+*/
+class RLeafCountArrayField final : public RFieldBase {
+   friend class RClassField; // only class fields are able to construct this
+
+private:
+   std::size_t fItemSize{0};                  ///< The size of a child field's item
+   ROOT::Internal::RColumnIndex fNWritten{0}; ///< Number of items written in the current cluster
+   std::ptrdiff_t fCountLeafDelta{0};         ///< Offset delta of the count leaf int member relative to the array
+   bool fHasPersistentCountLeaf{true};        ///< The count leaf may be only transient in RNTuple
+
+   RLeafCountArrayField(std::string_view fieldName, std::unique_ptr<RFieldBase> itemField,
+                        std::ptrdiff_t countLeafDelta, bool hasPersistentCountLeaf);
+
+protected:
+   std::unique_ptr<RFieldBase> CloneImpl(std::string_view newName) const final;
+
+   const RColumnRepresentations &GetColumnRepresentations() const final;
+   void GenerateColumns() final;
+   void GenerateColumns(const ROOT::RNTupleDescriptor &desc) final;
+
+   void ConstructValue(void *where) const final { *reinterpret_cast<void **>(where) = nullptr; }
+
+   std::size_t AppendImpl(const void *from) final;
+   void ReadGlobalImpl(ROOT::NTupleSize_t globalIndex, void *to) final;
+
+   void ReconcileOnDiskField(const RNTupleDescriptor &desc) final;
+
+   void CommitClusterImpl() final { fNWritten = 0; }
+
+public:
+   RLeafCountArrayField(RLeafCountArrayField &&other) = default;
+   RLeafCountArrayField &operator=(RLeafCountArrayField &&other) = default;
+   ~RLeafCountArrayField() final = default;
+
+   std::vector<RValue> SplitValue(const RValue &value) const final;
+
+   size_t GetValueSize() const final { return sizeof(void *); }
+   size_t GetAlignment() const final { return std::alignment_of<void *>(); }
+   void AcceptVisitor(ROOT::Detail::RFieldVisitor &visitor) const final;
+};
+
 } // namespace ROOT
 
 #endif

tree/ntuple/inc/ROOT/RFieldVisitor.hxx

@@ -91,6 +91,7 @@ public:
    virtual void VisitVectorField(const ROOT::RVectorField &field) { VisitField(field); }
    virtual void VisitVectorBoolField(const ROOT::RField<std::vector<bool>> &field) { VisitField(field); }
    virtual void VisitRVecField(const ROOT::RRVecField &field) { VisitField(field); }
+   virtual void VisitLeafCountArrayField(const RLeafCountArrayField &field) { VisitField(field); }
 }; // class RFieldVisitor
 
 } // namespace Detail
@@ -245,6 +246,7 @@ public:
    void VisitVectorField(const ROOT::RVectorField &field) final;
    void VisitVectorBoolField(const ROOT::RField<std::vector<bool>> &field) final;
    void VisitRVecField(const ROOT::RRVecField &field) final;
+   void VisitLeafCountArrayField(const RLeafCountArrayField &field) final;
    void VisitBitsetField(const ROOT::RBitsetField &field) final;
    void VisitNullableField(const ROOT::RNullableField &field) final;
    void VisitEnumField(const ROOT::REnumField &field) final;

tree/ntuple/src/RFieldSequenceContainer.cxx

@@ -1038,3 +1038,129 @@ void ROOT::RArrayAsVectorField::AcceptVisitor(ROOT::Detail::RFieldVisitor &visit
 {
    visitor.VisitArrayAsVectorField(*this);
 }
+
+//------------------------------------------------------------------------------
+
+ROOT::RLeafCountArrayField::RLeafCountArrayField(std::string_view fieldName, std::unique_ptr<RFieldBase> itemField,
+                                                 std::ptrdiff_t countLeafDelta, bool hasPersistentCountLeaf)
+   : ROOT::RFieldBase(fieldName, "ROOT::VecOps::RVec<" + itemField->GetTypeName() + ">",
+                      ROOT::ENTupleStructure::kCollection, false /* isSimple */),
+     fItemSize(itemField->GetValueSize()),
+     fCountLeafDelta(countLeafDelta),
+     fHasPersistentCountLeaf(hasPersistentCountLeaf)
+{
+   static_assert(sizeof(int) == sizeof(std::int32_t)); // implicit assumption in this field
+   assert(fCountLeafDelta < 0);                        // the count leaf must be defined before the array
+
+   if (!(itemField->GetTraits() & kTraitTriviallyDestructible)) {
+      throw RException(R__FAIL("RLeafCountArrayField only supports trivially destructible item types"));
+   }
+   // Note that we expect the array pointer to be deleted by the user.
+   fTraits |= kTraitTriviallyDestructible | (itemField->GetTraits() & kTraitTriviallyConstructible);
+   Attach(std::move(itemField));
+}
+
+std::unique_ptr<ROOT::RFieldBase> ROOT::RLeafCountArrayField::CloneImpl(std::string_view newName) const
+{
+   auto newItemField = fSubfields[0]->Clone(fSubfields[0]->GetFieldName());
+   return std::unique_ptr<RLeafCountArrayField>(
+      new RLeafCountArrayField(newName, std::move(newItemField), fCountLeafDelta, fHasPersistentCountLeaf));
+}
+
+const ROOT::RFieldBase::RColumnRepresentations &ROOT::RLeafCountArrayField::GetColumnRepresentations() const
+{
+   static RColumnRepresentations representations({{ENTupleColumnType::kSplitIndex64},
+                                                  {ENTupleColumnType::kIndex64},
+                                                  {ENTupleColumnType::kSplitIndex32},
+                                                  {ENTupleColumnType::kIndex32}},
+                                                 {});
+   return representations;
+}
+
+void ROOT::RLeafCountArrayField::GenerateColumns()
+{
+   GenerateColumnsImpl<ROOT::Internal::RColumnIndex>();
+}
+
+void ROOT::RLeafCountArrayField::GenerateColumns(const ROOT::RNTupleDescriptor &desc)
+{
+   GenerateColumnsImpl<ROOT::Internal::RColumnIndex>(desc);
+}
+
+std::size_t ROOT::RLeafCountArrayField::AppendImpl(const void *from)
+{
+   auto arrPtr = *reinterpret_cast<const unsigned char *const *>(from);
+   auto count = *reinterpret_cast<const std::uint32_t *>(static_cast<const unsigned char *>(from) + fCountLeafDelta);
+   std::size_t nbytes = 0;
+
+   if (fSubfields[0]->IsSimple() && count) {
+      GetPrincipalColumnOf(*fSubfields[0])->AppendV(arrPtr, count);
+      nbytes += count * GetPrincipalColumnOf(*fSubfields[0])->GetElement()->GetPackedSize();
+   } else {
+      for (unsigned i = 0; i < count; ++i) {
+         nbytes += CallAppendOn(*fSubfields[0], arrPtr + (i * fItemSize));
+      }
+   }
+
+   fNWritten += count;
+   fPrincipalColumn->Append(&fNWritten);
+   return nbytes + fPrincipalColumn->GetElement()->GetPackedSize();
+}
+
+void ROOT::RLeafCountArrayField::ReadGlobalImpl(ROOT::NTupleSize_t globalIndex, void *to)
+{
+   auto typedValue = reinterpret_cast<unsigned char **>(to);
+   auto countPtr = reinterpret_cast<std::uint32_t *>(static_cast<unsigned char *>(to) + fCountLeafDelta);
+
+   ROOT::NTupleSize_t nItems;
+   RNTupleLocalIndex collectionStart;
+   fPrincipalColumn->GetCollectionInfo(globalIndex, &collectionStart, &nItems);
+   if (nItems > std::numeric_limits<std::uint32_t>::max()) {
+      throw RException(R__FAIL("count leaf overflow"));
+   }
+
+   if (fHasPersistentCountLeaf) {
+      if (nItems != *countPtr) {
+         throw RException(R__FAIL("count leaf value different from collection size on disk"));
+      }
+   } else {
+      *countPtr = nItems;
+   }
+
+   operator delete(*typedValue);
+   *typedValue = static_cast<unsigned char *>(operator new(nItems * fItemSize));
+
+   if (fSubfields[0]->IsSimple()) {
+      if (nItems)
+         GetPrincipalColumnOf(*fSubfields[0])->ReadV(collectionStart, nItems, *typedValue);
+   } else {
+      for (std::size_t i = 0; i < nItems; ++i) {
+         CallReadOn(*fSubfields[0], collectionStart + i, *typedValue + (i * fItemSize));
+      }
+   }
+}
+
+void ROOT::RLeafCountArrayField::ReconcileOnDiskField(const RNTupleDescriptor &desc)
+{
+   EnsureMatchingOnDiskField(desc, kDiffTypeName).ThrowOnError();
+}
+
+std::vector<ROOT::RFieldBase::RValue> ROOT::RLeafCountArrayField::SplitValue(const RValue &value) const
+{
+   auto arrPtr = *reinterpret_cast<unsigned char **>(value.GetPtr<void>().get());
+   auto count =
+      *reinterpret_cast<std::uint32_t *>(static_cast<unsigned char *>(value.GetPtr<void>().get()) + fCountLeafDelta);
+
+   std::vector<RValue> result;
+   result.reserve(count);
+   for (std::uint32_t i = 0; i < count; ++i) {
+      result.emplace_back(
+         fSubfields[0]->BindValue(std::shared_ptr<void>(value.GetPtr<void>(), arrPtr + i * fItemSize)));
+   }
+   return result;
+}
+
+void ROOT::RLeafCountArrayField::AcceptVisitor(ROOT::Detail::RFieldVisitor &visitor) const
+{
+   visitor.VisitLeafCountArrayField(*this);
+}

tree/ntuple/src/RFieldVisitor.cxx

@@ -412,6 +412,11 @@ void ROOT::Internal::RPrintValueVisitor::VisitRVecField(const ROOT::RRVecField &
    PrintCollection(field);
 }
 
+void ROOT::Internal::RPrintValueVisitor::VisitLeafCountArrayField(const ROOT::RLeafCountArrayField &field)
+{
+   PrintCollection(field);
+}
+
 //---------------------------- RNTupleFormatter --------------------------------
 
 std::string ROOT::Internal::RNTupleFormatter::FitString(const std::string &str, int availableSpace)