From 69baf3d2f10aec02a45e1c19ffa912811f012996 Mon Sep 17 00:00:00 2001
From: Jonas Hahnfeld <jonas.hahnfeld@cern.ch>
Date: Fri, 18 Oct 2024 12:00:16 +0200
Subject: [PATCH] [ntuple] Add non-const RNTupleModel::GetField

---
 tree/ntuple/v7/inc/ROOT/RNTupleModel.hxx | 1 +
 tree/ntuple/v7/src/RNTupleModel.cxx      | 9 +++++++++
 tree/ntuple/v7/test/ntuple_model.cxx     | 3 +++
 3 files changed, 13 insertions(+)

diff --git a/tree/ntuple/v7/inc/ROOT/RNTupleModel.hxx b/tree/ntuple/v7/inc/ROOT/RNTupleModel.hxx
index 72e974444d6d3..5864fd720d8c2 100644
--- a/tree/ntuple/v7/inc/ROOT/RNTupleModel.hxx
+++ b/tree/ntuple/v7/inc/ROOT/RNTupleModel.hxx
@@ -351,6 +351,7 @@ public:
 
    RFieldZero &GetFieldZero() { return *fFieldZero; }
    const RFieldZero &GetFieldZero() const { return *fFieldZero; }
+   RFieldBase &GetField(std::string_view fieldName);
    const RFieldBase &GetField(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
index 727e6e40fbf22..bb89028021737 100644
--- a/tree/ntuple/v7/src/RNTupleModel.cxx
+++ b/tree/ntuple/v7/src/RNTupleModel.cxx
@@ -336,6 +336,15 @@ ROOT::Experimental::RNTupleModel::AddProjectedField(std::unique_ptr<RFieldBase>
    return RResult<void>::Success();
 }
 
+ROOT::Experimental::RFieldBase &ROOT::Experimental::RNTupleModel::GetField(std::string_view fieldName)
+{
+   auto f = FindField(fieldName);
+   if (!f)
+      throw RException(R__FAIL("invalid field: " + std::string(fieldName)));
+
+   return *f;
+}
+
 const ROOT::Experimental::RFieldBase &ROOT::Experimental::RNTupleModel::GetField(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
index fa5cd0adaa8d3..965c4d4296072 100644
--- a/tree/ntuple/v7/test/ntuple_model.cxx
+++ b/tree/ntuple/v7/test/ntuple_model.cxx
@@ -4,14 +4,17 @@ TEST(RNTupleModel, GetField)
 {
    auto model = RNTupleModel::CreateBare();
    model->MakeField<std::vector<float>>("v1");
+   model->MakeField<std::vector<float>>("v2");
 
    for (auto &subfield : model->GetFieldZero()) {
       if (subfield.GetTypeName() == "float") {
          subfield.SetColumnRepresentatives({{EColumnType::kReal32}});
       }
    }
+   model->GetField("v2._0").SetColumnRepresentatives({{EColumnType::kReal16}});
 
    EXPECT_EQ(EColumnType::kReal32, model->GetField("v1._0").GetColumnRepresentatives()[0][0]);
+   EXPECT_EQ(EColumnType::kReal16, model->GetField("v2._0").GetColumnRepresentatives()[0][0]);
 }
 
 TEST(RNTupleModel, EstimateWriteMemoryUsage)