[ntuple] Add RNTupleModelChangeset::AddColumnRepr

This is for now an internal-only functionality that will be used by the
Merger to extend column representations as needed during merging.

Note that we already support reading back extended column reprs, but
there is currently no way to write them.

Author: silverweed

tree/ntuple/inc/ROOT/RNTupleModel.hxx

@@ -402,11 +402,18 @@ You will not normally use this directly; see `RNTupleModel::RUpdater` instead.
 */
 // clang-format on
 struct RNTupleModelChangeset {
+   struct RAddedColumnRepr {
+      ROOT::RFieldBase *fField = nullptr;
+      std::uint16_t fNumNewColumns = 0;
+   };
+
    RNTupleModel &fModel;
    /// Points to the fields in fModel that were added as part of an updater transaction
    std::vector<ROOT::RFieldBase *> fAddedFields;
    /// Points to the projected fields in fModel that were added as part of an updater transaction
    std::vector<ROOT::RFieldBase *> fAddedProjectedFields;
+   /// Points to fields in fModel that had new column representations appended to them.
+   std::vector<RAddedColumnRepr> fAddedColumnReprs;
 
    RNTupleModelChangeset(RNTupleModel &model) : fModel(model) {}
    bool IsEmpty() const { return fAddedFields.empty() && fAddedProjectedFields.empty(); }
@@ -420,6 +427,9 @@ struct RNTupleModelChangeset {
    /// \see RNTupleModel::AddProjectedField()
    ROOT::RResult<void>
    AddProjectedField(std::unique_ptr<ROOT::RFieldBase> field, RNTupleModel::FieldMappingFunc_t mapping);
+
+   ROOT::RResult<void>
+   AddColumnRepr(ROOT::RFieldBase *field, const ROOT::RFieldBase::RColumnRepresentations::Selection_t &reprs);
 };
 
 } // namespace Internal

tree/ntuple/src/RFieldBase.cxx

@@ -830,7 +830,15 @@ ROOT::RFieldBase::RColumnRepresentations::Selection_t ROOT::RFieldBase::GetColum
 
 void ROOT::RFieldBase::SetColumnRepresentatives(const RColumnRepresentations::Selection_t &representatives)
 {
-   if (fState != EState::kUnconnected)
+   const auto isNewSelectionASuperset = [&]() {
+      if (representatives.size() < fColumnRepresentatives.size())
+         return false;
+      for (auto i = 0u; i < fColumnRepresentatives.size(); ++i)
+         if (representatives[i] != fColumnRepresentatives[i].get())
+            return false;
+      return true;
+   };
+   if (fState != EState::kUnconnected && !(fState == EState::kConnectedToSink && isNewSelectionASuperset()))
       throw RException(R__FAIL("cannot set column representative once field is connected"));
    const auto &validTypes = GetColumnRepresentations().GetSerializationTypes();
    fColumnRepresentatives.clear();

tree/ntuple/src/RNTupleModel.cxx

@@ -287,6 +287,27 @@ ROOT::RNTupleModel::RUpdater::AddProjectedField(std::unique_ptr<ROOT::RFieldBase
    return R__FORWARD_RESULT(fOpenChangeset.AddProjectedField(std::move(field), std::move(mapping)));
 }
 
+ROOT::RResult<void> ROOT::Internal::RNTupleModelChangeset::AddColumnRepr(
+   ROOT::RFieldBase *field, const ROOT::RFieldBase::RColumnRepresentations::Selection_t &newReprs)
+{
+   auto reprs = field->GetColumnRepresentatives();
+   const auto nPrev = reprs.size();
+   for (auto newRepr : newReprs)
+      // NOTE: we don't need to check for duplicates because SetColumnRepresentatives will do that for us.
+      reprs.push_back(newRepr);
+
+   field->SetColumnRepresentatives(reprs);
+   const auto nNew = field->GetColumnRepresentatives().size();
+   assert(reprs.size() > 0 && reprs[0].size() <= 2);
+   const auto cardinality = static_cast<std::uint16_t>(reprs[0].size());
+   assert(nNew >= nPrev);
+   std::uint16_t nAdded = (nNew - nPrev) * cardinality;
+   if (nAdded > 0)
+      fAddedColumnReprs.push_back(RAddedColumnRepr{field, nAdded});
+
+   return RResult<void>::Success();
+}
+
 void ROOT::RNTupleModel::EnsureValidFieldName(std::string_view fieldName)
 {
    RResult<void> nameValid = ROOT::Internal::EnsureValidNameForRNTuple(fieldName, "Field");

tree/ntuple/src/RPageSinkBuf.cxx

@@ -125,12 +125,20 @@ void ROOT::Internal::RPageSinkBuf::UpdateSchema(const ROOT::Internal::RNTupleMod
       GetProjectedFieldsOfModel(*fInnerModel).Add(std::move(cloned), fieldMap);
       return p;
    };
+   auto cloneAddColumnRepr = [&](const RNTupleModelChangeset::RAddedColumnRepr &repr) {
+      auto &innerField = fInnerModel->GetMutableField(repr.fField->GetFieldName());
+      innerField.SetColumnRepresentatives(repr.fField->GetColumnRepresentatives());
+      ROOT::Internal::CallConnectExtendedColumnsToPageSinkOnField(innerField, *this, firstEntry);
+      return repr;
+   };
    RNTupleModelChangeset innerChangeset{*fInnerModel};
    fInnerModel->Unfreeze();
    std::transform(changeset.fAddedFields.cbegin(), changeset.fAddedFields.cend(),
                   std::back_inserter(innerChangeset.fAddedFields), cloneAddField);
    std::transform(changeset.fAddedProjectedFields.cbegin(), changeset.fAddedProjectedFields.cend(),
                   std::back_inserter(innerChangeset.fAddedProjectedFields), cloneAddProjectedField);
+   std::transform(changeset.fAddedColumnReprs.cbegin(), changeset.fAddedColumnReprs.cend(),
+                  std::back_inserter(innerChangeset.fAddedColumnReprs), cloneAddColumnRepr);
    fInnerModel->Freeze();
    fInnerSink->UpdateSchema(innerChangeset, firstEntry);
 }

tree/ntuple/src/RPageStorage.cxx

@@ -868,6 +868,9 @@ void ROOT::Internal::RPagePersistentSink::UpdateSchema(const ROOT::Internal::RNT
          for (const auto &descendant : *f)
             nNewPhysicalColumns += getNColumns(descendant);
       }
+      for (auto added : changeset.fAddedColumnReprs) {
+         nNewPhysicalColumns += added.fNumNewColumns;
+      }
       fDescriptorBuilder.ShiftAliasColumns(nNewPhysicalColumns);
    }
 
@@ -913,6 +916,9 @@ void ROOT::Internal::RPagePersistentSink::UpdateSchema(const ROOT::Internal::RNT
       for (auto &descendant : *f)
          addProjectedField(descendant);
    }
+   for (const auto& [f, _] : changeset.fAddedColumnReprs) {
+      ROOT::Internal::CallConnectExtendedColumnsToPageSinkOnField(*f, *this, firstEntry);
+   }
 
    const auto nColumns = descriptor.GetNPhysicalColumns();
    fOpenColumnRanges.reserve(fOpenColumnRanges.size() + (nColumns - nColumnsBeforeUpdate));