[ntuple] add support for inheritance in SoA fields

jblomer · jblomer · commit 6ca7c23dbaae · 2026-07-06T14:29:26.000+02:00
diff --git a/tree/ntuple/inc/ROOT/RField/RFieldSoA.hxx b/tree/ntuple/inc/ROOT/RField/RFieldSoA.hxx
@@ -21,6 +21,7 @@
 #include <ROOT/RNTupleTypes.hxx>
 
 #include <cstddef>
+#include <functional>
 #include <memory>
 #include <mutex>
 #include <string_view>
@@ -88,8 +89,13 @@ class RSoAField : public RFieldBase {
    RSoAField(std::string_view fieldName, const RSoAField &source); ///< Used by CloneImpl
    RSoAField(std::string_view fieldName, TClass *clSoA);
 
-   /// Called during construction, picks up the (nested) member fields of the underlying record type(s)
+   /// Called during construction, picks up the (nested) member fields of the underlying record type(s) and its
+   /// base classes.
    void CollectRecordMemberFields();
+   /// For a nested SoA struct (either as a member of as a base class), use their fRecordMemberFields in this class,
+   /// i.e. "unroll" the vectors in the nested SoA struct into the SoA base class.
+   void GraftNestedMemberFields(const RSoAField &nestedSoA, std::size_t offsetInParent,
+                                std::function<RFieldBase *(const std::string &)> fnRecordFieldFinder);
 
    void ReconstructSplitFields() const;
 
diff --git a/tree/ntuple/src/RFieldMeta.cxx b/tree/ntuple/src/RFieldMeta.cxx
@@ -686,42 +686,110 @@ ROOT::Experimental::RSoAField::RSoAField(std::string_view fieldName, std::string
 {
 }
 
+void ROOT::Experimental::RSoAField::GraftNestedMemberFields(
+   const RSoAField &nestedSoA, std::size_t offsetInParent,
+   std::function<RFieldBase *(const std::string &)> fnRecordFieldFinder)
+{
+   const std::size_t nNestedRecordMemberFields = nestedSoA.fRecordMemberFields.size();
+
+   // The qualified field name of fields in nestedSoA->fRecordMemberFields will have a "<field name>._0."
+   // prefix because these fields are rooted in a collection named after the nested SoA field.
+   //
+   // E.g., in the following example:
+   //
+   // struct SoA_A {                                   struct Record_A {
+   //   SoA_B fB;                                         Record_B fB;
+   // };                                               };
+   //
+   // struct SoA_B {                                   struct Record_B {
+   //   ROOT::RVec<float> fX;                             float fX;
+   // };                                               };
+   //
+   // The on-disk schema of SoA_A is "collection of Record_A", and the on-disk schema of SoA_B is
+   // "collection of Record_B".
+   // The qualified field name of fX in the subfield hiararchy of SoA_A is <field name>._0.fB.fX.
+   // The qualified field name of fX in the subfield hiararchy of SoA_B is <field name>._0.fX.
+   const auto lenPrefix = nestedSoA.GetFieldName().length() + strlen("._0.");
+
+   for (std::size_t i = 0; i < nNestedRecordMemberFields; ++i) {
+      const auto fieldNameForMatching =
+         nestedSoA.GetFieldName() + "." + nestedSoA.fRecordMemberFields[i]->GetQualifiedFieldName().substr(lenPrefix);
+
+      fRecordMemberFields.emplace_back(fnRecordFieldFinder(fieldNameForMatching));
+      fRecordMemberDeleters.emplace_back(GetDeleterOf(*nestedSoA.fRecordMemberFields[i]));
+      fSoAMemberOffsets.emplace_back(offsetInParent + nestedSoA.fSoAMemberOffsets[i]);
+   }
+}
+
 void ROOT::Experimental::RSoAField::CollectRecordMemberFields()
 {
-   std::vector<RFieldBase *> realRecordMemberFields;
+   // Build a map of all subfields (nested) of the underlying record type. Map the fully qualified name of the
+   // subfields to their field pointer, so that we can later match the subfields of the SoA class to their corresponding
+   // fields in the underlying record type. Note that the members of the SoA class and the underlying record type
+   // can have different ordering. However, the base classes of the SoA class and the underlying record type must match
+   // in order.
+
+   std::vector<RFieldBase *> realRecordMemberFields; // Contains all subfields of the underlying record type
+   // Qualified field name --> index in realRecordMemberFields
+   std::unordered_map<std::string, std::size_t> recordFieldNameToIdx;
+
+   // Count the top-level subfields of the underlying record type for cross-check with the SoA type
    unsigned int nDirectRecordSubfields = 0;
+   unsigned int nDirectRecordBases = 0;
+
    for (auto itr = fSubfields[0]->begin(), iEnd = fSubfields[0]->end(); itr != iEnd; ++itr) {
-      if (itr->GetParent() == fSubfields[0].get())
-         nDirectRecordSubfields++;
+      if (itr->GetParent() == fSubfields[0].get()) {
+         (itr->GetFieldName()[0] == ':') ? nDirectRecordBases++ : nDirectRecordSubfields++;
+      }
+
+      // Build the qualified field name for matching. We root the qualified field name at the underlying record type.
+      auto qualifiedName = itr->GetFieldName();
+      auto parent = itr->GetParent();
+      while (parent != fSubfields[0].get()) {
+         qualifiedName = parent->GetFieldName() + "." + qualifiedName;
+         parent = parent->GetParent();
+      }
+      recordFieldNameToIdx[qualifiedName] = realRecordMemberFields.size();
+
       realRecordMemberFields.emplace_back(&(*itr));
    }
 
-   for (const auto f : realRecordMemberFields) {
-      if (f->GetFieldName()[0] == ':') {
-         throw RException(R__FAIL("SoA fields with inheritance are currently unsupported"));
-      }
+   // Base classes are treated as unrolled nested SoA classes
+   const auto *soaBases = fSoAClass->GetListOfBases();
+   if (soaBases->GetSize() != static_cast<Int_t>(nDirectRecordBases)) {
+      throw RException(R__FAIL(std::string("number of base classes don't match between SoA class ") + GetFieldName() +
+                               " and its underlying record type"));
    }
+   for (unsigned int i = 0; i < static_cast<unsigned int>(soaBases->GetSize()); ++i) {
+      auto base = static_cast<TBaseClass *>(soaBases->At(i));
+      if (base->GetDelta() < 0) {
+         throw RException(R__FAIL(std::string("virtual inheritance is not supported: ") + GetTypeName() +
+                                  " virtually inherits from " + base->GetName()));
+      }
+      TClass *cl = base->GetClassPointer();
 
-   // Build map name --> index in realRecordMemberFields vector. Cut the name prefix up to the collection subfields
-   // so that it matches the data member name
-   std::unordered_map<std::string, std::size_t> recordFieldNameToIdx;
-   {
-      recordFieldNameToIdx.reserve(realRecordMemberFields.size());
-      const auto lenFieldNamePrefix = GetFieldName().length() + strlen("._0.");
-      for (std::size_t i = 0; i < realRecordMemberFields.size(); ++i) {
-         recordFieldNameToIdx[realRecordMemberFields[i]->GetQualifiedFieldName().substr(lenFieldNamePrefix)] = i;
+      const auto baseFieldName = std::string(":_") + std::to_string(i);
+
+      // SoA class `A` inherits from a SoA class `B` whose underlying record type is `X` if and only if the underlying
+      // record type of `A` inherits from a type `X`.
+      const auto underlyingBaseTypeName = ROOT::Internal::GetRNTupleSoARecord(cl);
+      auto recordBaseField = realRecordMemberFields[recordFieldNameToIdx[baseFieldName]];
+      if (underlyingBaseTypeName != recordBaseField->GetTypeName()) {
+         throw RException(R__FAIL(std::string("inheritance of SoA class ") + GetFieldName() +
+                                  " does not match its underlying record type"));
       }
-   }
 
-   const auto *bases = fSoAClass->GetListOfBases();
-   assert(bases);
-   for (auto baseClass : ROOT::Detail::TRangeStaticCast<TBaseClass>(*bases)) {
-      if (baseClass->GetDelta() < 0) {
-         throw RException(R__FAIL(std::string("virtual inheritance is not supported: ") + GetTypeName() +
-                                  " virtually inherits from " + baseClass->GetName()));
+      std::unique_ptr<RSoAField> soaBaseField;
+      try {
+         soaBaseField = std::make_unique<RSoAField>(baseFieldName, cl->GetName());
+      } catch (const RException &e) {
+         throw RException(R__FAIL(std::string("invalid field type in base class: ") + cl->GetName() + " of SoA field " +
+                                  GetFieldName() + " (" + e.what() + ")"));
       }
-      // At a later point, we will support inheritance
-      throw RException(R__FAIL("SoA fields with inheritance are currently unsupported"));
+
+      GraftNestedMemberFields(*soaBaseField, base->GetDelta(), [&](const std::string &name) {
+         return realRecordMemberFields[recordFieldNameToIdx[name]];
+      });
    }
 
    unsigned int nMembers = 0;
@@ -753,17 +821,9 @@ void ROOT::Experimental::RSoAField::CollectRecordMemberFields()
                                      underlyingField->GetTypeName()));
          }
 
-         const auto lenFieldNamePrefix = soaField->GetFieldName().length() + strlen("._0.");
-
-         for (std::size_t i = 0; i < soaField->fRecordMemberFields.size(); ++i) {
-            const auto f = soaField->fRecordMemberFields[i];
-            const auto fieldNameForMatching =
-               dmField->GetFieldName() + "." + f->GetQualifiedFieldName().substr(lenFieldNamePrefix);
-
-            fRecordMemberFields.emplace_back(realRecordMemberFields[recordFieldNameToIdx[fieldNameForMatching]]);
-            fRecordMemberDeleters.emplace_back(GetDeleterOf(*soaField->fRecordMemberFields[i]));
-            fSoAMemberOffsets.emplace_back(dataMember->GetOffset() + soaField->fSoAMemberOffsets[i]);
-         }
+         GraftNestedMemberFields(*soaField, dataMember->GetOffset(), [&](const std::string &name) {
+            return realRecordMemberFields[recordFieldNameToIdx[name]];
+         });
       } else if (auto vecField = dynamic_cast<RRVecField *>(dmField.get())) {
          if (vecField->begin()->GetTypeName() != underlyingField->GetTypeName() ||
              vecField->begin()->GetTypeAlias() != underlyingField->GetTypeAlias()) {
@@ -945,6 +1005,15 @@ void ROOT::Experimental::RSoAField::ReconstructSplitFields() const
    fSplitFields = std::make_unique<std::vector<std::unique_ptr<ROOT::RFieldBase>>>();
    fSplitOffsets = std::make_unique<std::vector<std::size_t>>();
 
+   unsigned int iBase = 0;
+   for (auto base : ROOT::Detail::TRangeStaticCast<TBaseClass>(*fSoAClass->GetListOfBases())) {
+      TClass *cl = base->GetClassPointer();
+      auto baseField = RFieldBase::Create(std::string(":_" + std::to_string(iBase)), cl->GetName()).Unwrap();
+      fSplitFields->emplace_back(std::move(baseField));
+      fSplitOffsets->emplace_back(base->GetDelta());
+      iBase++;
+   }
+
    for (auto dataMember : ROOT::Detail::TRangeStaticCast<TDataMember>(*fSoAClass->GetListOfDataMembers())) {
       if ((dataMember->Property() & kIsStatic) || !dataMember->IsPersistent())
          continue;
diff --git a/tree/ntuple/test/SoAField.hxx b/tree/ntuple/test/SoAField.hxx
@@ -23,26 +23,6 @@ struct SoAVersionMismatch {
    ClassDefNV(SoAVersionMismatch, 4);
 };
 
-struct RecordBase {
-   ClassDefNV(RecordBase, 2);
-};
-
-struct RecordDerived : public RecordBase {
-   ClassDefNV(RecordDerived, 2);
-};
-
-struct SoABase {
-   ClassDefNV(SoABase, 2);
-};
-
-struct SoAOnDerivedRecord {
-   ClassDefNV(SoAOnDerivedRecord, 2);
-};
-
-struct SoADerivedOnBaseRecord : public SoABase {
-   ClassDefNV(SoADerivedOnBaseRecord, 2);
-};
-
 struct RecordSimple {
    float fX;
    float fY;
@@ -152,4 +132,42 @@ struct SoADotBadNestedType {
    ClassDefNV(SoADotBadNestedType, 2);
 };
 
+struct RecordBase {
+   float fBase;
+   ClassDefNV(RecordBase, 2);
+};
+
+struct SoABase {
+   ROOT::RVec<float> fBase;
+   ClassDefNV(SoABase, 2);
+};
+
+struct RecordDerived : public RecordBase {
+   float fDerived;
+   ClassDefNV(RecordDerived, 2);
+};
+
+struct SoADerived : public SoABase {
+   ROOT::RVec<float> fDerived;
+   ClassDefNV(SoADerived, 2);
+};
+
+struct RecordDerivedMulti : public RecordDerived, RecordDot {
+   float fMulti;
+   ClassDefNV(RecordDerivedMulti, 2);
+};
+
+struct SoADerivedMulti : public SoADerived, SoADot {
+   ROOT::RVec<float> fMulti;
+   ClassDefNV(SoADerivedMulti, 2);
+};
+
+struct SoADerivedFail1 : public RecordDerived {
+   ClassDefNV(SoADerivedFail1, 2);
+};
+
+struct SoADerivedFail2 : public SoABase {
+   ClassDefNV(SoADerivedFail2, 2);
+};
+
 #endif // ROOT_RNTuple_Test_SoAField
diff --git a/tree/ntuple/test/SoAFieldLinkDef.h b/tree/ntuple/test/SoAFieldLinkDef.h
@@ -6,12 +6,6 @@
 #pragma link C++ options=rntupleSoARecord(xyz) class SoAUnknownRecord+;
 #pragma link C++ options=rntupleSoARecord(Record) class SoAVersionMismatch+;
 
-#pragma link C++ class RecordBase+;
-#pragma link C++ class RecordDerived+;
-#pragma link C++ options=rntupleSoARecord(RecordBase) class SoABase+;
-#pragma link C++ options=rntupleSoARecord(RecordDerived) class SoAOnDerivedRecord+;
-#pragma link C++ options=rntupleSoARecord(RecordBase) class SoADerivedOnBaseRecord+;
-
 #pragma link C++ class RecordSimple+;
 #pragma link C++ options=rntupleSoARecord(RecordSimple) class SoASimple+;
 #pragma link C++ options=rntupleSoARecord(RecordSimple) class SoASimpleSwapped+;
@@ -31,4 +25,13 @@
 #pragma link C++ options=rntupleSoARecord(RecordDot) class SoADot+;
 #pragma link C++ options=rntupleSoARecord(RecordDot) class SoADotBadNestedType+;
 
+#pragma link C++ class RecordBase+;
+#pragma link C++ class RecordDerived+;
+#pragma link C++ class RecordDerivedMulti+;
+#pragma link C++ options=rntupleSoARecord(RecordBase) class SoABase+;
+#pragma link C++ options=rntupleSoARecord(RecordDerived) class SoADerived+;
+#pragma link C++ options=rntupleSoARecord(RecordDerivedMulti) class SoADerivedMulti+;
+#pragma link C++ options=rntupleSoARecord(RecordDerived) class SoADerivedFail1+;
+#pragma link C++ options=rntupleSoARecord(RecordDerived) class SoADerivedFail2+;
+
 #endif // __CLING__
diff --git a/tree/ntuple/test/ntuple_soa.cxx b/tree/ntuple/test/ntuple_soa.cxx
@@ -70,22 +70,6 @@ TEST(RNTuple, SoACheck)
       EXPECT_THAT(e.what(), ::testing::HasSubstr("version mismatch between SoA type and underlying record type"));
    }
 
-   try {
-      auto f = std::make_unique<RSoAField>("f", "SoAOnDerivedRecord");
-      FAIL() << "creating SoA field on derived record should fail";
-   } catch (const ROOT::RException &e) {
-      EXPECT_THAT(e.what(), ::testing::HasSubstr("SoA fields with inheritance are currently unsupported"));
-   }
-   try {
-      auto f = std::make_unique<RSoAField>("f", "SoADerivedOnBaseRecord");
-      FAIL() << "creating a derived SoA field should fail";
-   } catch (const ROOT::RException &e) {
-      EXPECT_THAT(e.what(), ::testing::HasSubstr("SoA fields with inheritance are currently unsupported"));
-   }
-   {
-      EXPECT_NO_THROW(auto f = std::make_unique<RSoAField>("f", "SoABase"));
-   }
-
    try {
       auto f = std::make_unique<RSoAField>("f", "SoASimpleBadArray");
       FAIL() << "creating SoA field with arrays fail";
@@ -514,3 +498,81 @@ R"({
    // clang-format on
    EXPECT_EQ(expected, os.str());
 }
+
+TEST(RNTuple, SoADerived)
+{
+   ROOT::TestSupport::FileRaii fileGuard("test_rntuple_soa_derived.root");
+
+   {
+      auto model = ROOT::RNTupleModel::Create();
+
+      try {
+         model->AddField(std::make_unique<RSoAField>("x", "SoADerivedFail1"));
+         FAIL() << "SoADerivedFail1 should fail";
+      } catch (const ROOT::RException &e) {
+         EXPECT_THAT(e.what(),
+                     testing::HasSubstr("inheritance of SoA class x does not match its underlying record type"));
+      }
+      try {
+         model->AddField(std::make_unique<RSoAField>("x", "SoADerivedFail2"));
+         FAIL() << "SoADerivedFail2 should fail";
+      } catch (const ROOT::RException &e) {
+         EXPECT_THAT(e.what(), testing::HasSubstr("missing SoA members"));
+      }
+
+      model->AddField(std::make_unique<RSoAField>("derived", "SoADerived"));
+      model->AddField(std::make_unique<RSoAField>("multi", "SoADerivedMulti"));
+      auto writer = ROOT::RNTupleWriter::Recreate(std::move(model), "ntpl", fileGuard.GetPath());
+
+      auto derivedSoA = writer->GetModel().GetDefaultEntry().GetPtr<SoADerived>("derived");
+      derivedSoA->fBase = {1.0, 2.0};
+      derivedSoA->fDerived = {3.0, 4.0};
+
+      auto derivedMultiSoA = writer->GetModel().GetDefaultEntry().GetPtr<SoADerivedMulti>("multi");
+      derivedMultiSoA->fBase = {5.0, 6.0};
+      derivedMultiSoA->fDerived = {7.0, 8.0};
+      derivedMultiSoA->fX = {9.0, 10.0};
+      derivedMultiSoA->fY = {11.0, 12.0};
+      derivedMultiSoA->fProperties.fColor = {13, 14};
+      derivedMultiSoA->fProperties.fSize = {15.0, 16.0};
+      derivedMultiSoA->fMulti = {17.0, 18.0};
+
+      writer->Fill();
+   }
+
+   auto reader = ROOT::RNTupleReader::Open("ntpl", fileGuard.GetPath());
+
+   std::ostringstream os;
+   reader->Show(0, os);
+
+   // clang-format off
+   std::string expected{
+R"({
+  "derived": {
+    ":_0": {
+      "fBase": [1, 2]
+    },
+    "fDerived": [3, 4]
+  },
+  "multi": {
+    ":_0": {
+      ":_0": {
+        "fBase": [5, 6]
+      },
+      "fDerived": [7, 8]
+    },
+    ":_1": {
+      "fX": [9, 10],
+      "fY": [11, 12],
+      "fProperties": {
+        "fColor": [13, 14],
+        "fSize": [15, 16]
+      }
+    },
+    "fMulti": [17, 18]
+  }
+}
+)" };
+   // clang-format on
+   EXPECT_EQ(expected, os.str());
+}
