[ntuple] Add RNTupleModel::GetMutableField

Similar to GetMutableFieldZero(), it must only be called on unfrozen models.
root-project · Oct 20, 2024 · df519dd · df519dd
1 parent c0d6df6
commit df519dd
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 0 deletions.
diff --git a/tree/ntuple/v7/inc/ROOT/RNTupleModel.hxx b/tree/ntuple/v7/inc/ROOT/RNTupleModel.hxx
@@ -363,6 +363,7 @@ public:
    /// Mutable access to the root field is used to make adjustments to the fields.
    RFieldZero &GetMutableFieldZero();
    const RFieldZero &GetConstFieldZero() const { return *fFieldZero; }
+   RFieldBase &GetMutableField(std::string_view fieldName);
    const RFieldBase &GetConstField(std::string_view fieldName) const;
 
    const std::string &GetDescription() const { return fDescription; }

diff --git a/tree/ntuple/v7/src/RNTupleModel.cxx b/tree/ntuple/v7/src/RNTupleModel.cxx
@@ -354,6 +354,17 @@ ROOT::Experimental::RFieldZero &ROOT::Experimental::RNTupleModel::GetMutableFiel
    return *fFieldZero;
 }
 
+ROOT::Experimental::RFieldBase &ROOT::Experimental::RNTupleModel::GetMutableField(std::string_view fieldName)
+{
+   if (IsFrozen())
+      throw RException(R__FAIL("invalid attempt to get mutable field of frozen model"));
+   auto f = FindField(fieldName);
+   if (!f)
+      throw RException(R__FAIL("invalid field: " + std::string(fieldName)));
+
+   return *f;
+}
+
 const ROOT::Experimental::RFieldBase &ROOT::Experimental::RNTupleModel::GetConstField(std::string_view fieldName) const
 {
    auto f = FindField(fieldName);

diff --git a/tree/ntuple/v7/test/ntuple_model.cxx b/tree/ntuple/v7/test/ntuple_model.cxx
@@ -67,11 +67,13 @@ TEST(RNTupleModel, GetField)
    auto m = RNTupleModel::Create();
    m->MakeField<int>("x");
    m->MakeField<CustomStruct>("cs");
+   m->GetMutableField("cs.v1._0").SetColumnRepresentatives({{EColumnType::kReal32}});
    m->Freeze();
    EXPECT_EQ(m->GetConstField("x").GetFieldName(), "x");
    EXPECT_EQ(m->GetConstField("x").GetTypeName(), "std::int32_t");
    EXPECT_EQ(m->GetConstField("cs.v1").GetFieldName(), "v1");
    EXPECT_EQ(m->GetConstField("cs.v1").GetTypeName(), "std::vector<float>");
+   EXPECT_EQ(m->GetConstField("cs.v1._0").GetColumnRepresentatives()[0][0], EColumnType::kReal32);
    try {
       m->GetConstField("nonexistent");
       FAIL() << "invalid field name should throw";
@@ -90,6 +92,12 @@ TEST(RNTupleModel, GetField)
    } catch (const RException &err) {
       EXPECT_THAT(err.what(), testing::HasSubstr("frozen model"));
    }
+   try {
+      m->GetMutableField("x");
+      FAIL() << "GetMutableField should throw";
+   } catch (const RException &err) {
+      EXPECT_THAT(err.what(), testing::HasSubstr("frozen model"));
+   }
    EXPECT_EQ("", m->GetConstFieldZero().GetQualifiedFieldName());
    EXPECT_EQ("x", m->GetConstField("x").GetQualifiedFieldName());
    EXPECT_EQ("cs", m->GetConstField("cs").GetQualifiedFieldName());