RNTupleProcessorDS.cxx

Go to the documentation of this file.

/// \file RNTupleProcessorDS.cxx
/// \author Florine de Geus <florine.de.geus@cern.ch>
/// \author Jakob Blomer <jblomer@cern.ch>
/// \author Enrico Guiraud <enrico.guiraud@cern.ch>
/// \date 2026-06-18
/// \warning This is part of the ROOT 7 prototype! It will change without notice. It might trigger earthquakes. Feedback
/// is welcome!
 
#include <ROOT/RDF/RColumnReaderBase.hxx>
#include <ROOT/RDataFrame.hxx>
#include <ROOT/RDF/Utils.hxx>
#include <ROOT/RField.hxx>
#include <ROOT/RFieldUtils.hxx>
#include <ROOT/RPageStorageFile.hxx>
#include <ROOT/RNTupleDescriptor.hxx>
#include <ROOT/RNTupleProcessorDS.hxx>
#include <ROOT/RPageStorage.hxx>
 
#include "RNTupleDSFields.hxx"
 
#include <TError.h>
#include <TSystem.h>
 
#include <cassert>
#include <memory>
#include <string>
#include <string_view>
 
namespace ROOT::Experimental::Internal::RDF {
/// Every RDF column is represented by exactly one RNTuple field
class RNTupleProcessorColumnReader : public ROOT::Detail::RDF::RColumnReaderBase {
   RNTupleProcessorValuePtr<void> fValuePtr; ///< Container for the value of the column
   Long64_t fLastLoadedEntry = ROOT::kInvalidNTupleIndex;
 
public:
   RNTupleProcessorColumnReader(RNTupleProcessorValuePtr<void> &valuePtr) : fValuePtr(std::move(valuePtr)) {}
   ~RNTupleProcessorColumnReader() override = default;
 
   void *GetImpl(Long64_t entry [[maybe_unused]]) final
   {
      assert(entry == fLastLoadedEntry);
      return fValuePtr.GetPtr().get();
   }
 
   void SetEntryLoaded(Long64_t entry) { fLastLoadedEntry = entry; }
};
 
RNTupleProcessorDS::~RNTupleProcessorDS() = default;
 
// TODO: see if common parts of this and RNTupleDS::AddField can be factored out.
void RNTupleProcessorDS::AddField(const ROOT::RFieldBase &field, std::string_view colName,
                                  ROOT::Internal::RNTupleProcessorProvenance procProvenance,
                                  std::deque<RNTupleProcessorDS::RFieldInfo> &fieldInfos, bool convertToRVec)
{
   // As an example for the mapping of RNTuple fields to RDF columns, let's consider an RNTuple
   // using the following types and with a top-level field named "event" of type Event:
   //
   // struct Event {
   //    int id;
   //    std::vector<Track> tracks;
   // };
   // struct Track {
   //    std::vector<Hit> hits;
   // };
   // struct Hit {
   //    float x;
   //    float y;
   // };
   //
   // AddField() will be called from the constructor with the RNTuple root field (ENTupleStructure::kRecord).
   // From there, we recurse into the "event" sub field (also ENTupleStructure::kRecord) and further down the
   // tree of sub fields and expose the following RDF columns:
   //
   // "event"                             [Event]
   // "event.id"                          [int]
   // "event.tracks"                      [RVec<Track>]
   // "R_rdf_sizeof_event.tracks"         [unsigned int]
   // "event.tracks.hits"                 [RVec<RVec<Hit>>]
   // "R_rdf_sizeof_event.tracks.hits"    [RVec<unsigned int>]
   // "event.tracks.hits.x"               [RVec<RVec<float>>]
   // "R_rdf_sizeof_event.tracks.hits.x"  [RVec<unsigned int>]
   // "event.tracks.hits.y"               [RVec<RVec<float>>]
   // "R_rdf_sizeof_event.tracks.hits.y"  [RVec<unsigned int>]
 
   std::uint16_t nFieldInfosAdded = 0;
 
   const auto &nRepetitions = field.GetNRepetitions();
   if ((field.GetStructure() == ROOT::ENTupleStructure::kCollection) || (nRepetitions > 0)) {
      // The field is a collection or a fixed-size array.
      // We open a new collection scope with fieldID being the inner most collection. E.g. for "event.tracks.hits",
      // fieldInfos would already contain the fieldID of "event.tracks"
      const std::string name = (procProvenance.Empty() ? "" : procProvenance.Get() + ".") + std::string(colName);
      fieldInfos.emplace_back(field.GetOnDiskId(), field.GetFieldName(), field.GetTypeName(), nRepetitions,
                              *fProcessor->GetEntry().FindFieldIndex(name));
      nFieldInfosAdded++;
   }
 
   if (field.GetStructure() == ROOT::ENTupleStructure::kCollection) {
      // Inner fields of collections are provided as projected collections of only that inner field,
      // E.g. we provide a projected collection RVec<RVec<float>> for "event.tracks.hits.x" in the example
      // above.
      bool representableAsRVec =
         convertToRVec && (field.GetTypeName().substr(0, 19) == "ROOT::VecOps::RVec<" ||
                           field.GetTypeName().substr(0, 12) == "std::vector<" || field.GetTypeName() == "");
      const auto *f = field.GetConstSubfields()[0];
      AddField(*f, colName, procProvenance, fieldInfos, representableAsRVec);
 
      // Note that at the end of the recursion, we handled the inner sub collections as well as the
      // collection as whole, so we are done.
      return;
 
   } else if (nRepetitions > 0) {
      // Fixed-size array, same logic as ROOT::RVec.
      const auto *f = field.GetConstSubfields()[0];
      AddField(*f, colName, procProvenance, fieldInfos);
      return;
   } else if (field.GetStructure() == ROOT::ENTupleStructure::kRecord) {
      // Inner fields of records are provided as individual RDF columns, e.g. "event.id"
      for (const auto &f : field.GetConstSubfields()) {
         auto innerName = colName.empty() ? f->GetFieldName() : (std::string(colName) + "." + f->GetFieldName());
         // Inner fields of collections of records are always exposed as ROOT::RVec
         AddField(*f, innerName, procProvenance, fieldInfos);
      }
 
      // Do not add untyped record fields
      if (field.GetTypeName() == "")
         return;
   }
 
   // The fieldID could be the root field or the class of fieldId might not be loaded.
   // In these cases, only the inner fields are exposed as RDF columns.
   std::string fieldName = field.GetFieldName();
   auto fieldInfoIt = fieldInfos.begin();
   while (fieldName == "_0") {
      assert(fieldInfoIt != fieldInfos.end());
      fieldName = fieldInfoIt->fFieldName;
      fieldInfoIt++;
   }
   auto fieldOrException = ROOT::RFieldBase::Create(fieldName, field.GetTypeName());
   if (!fieldOrException)
      return;
   auto valueField = fieldOrException.Unwrap();
   if (const auto cardinalityField = dynamic_cast<const ROOT::RCardinalityField *>(valueField.get())) {
      // Cardinality fields in RDataFrame are presented as integers
      if (cardinalityField->As32Bit()) {
         valueField = std::make_unique<ROOT::Internal::RDF::RRDFCardinalityField<std::uint32_t>>(field.GetFieldName());
      } else if (cardinalityField->As64Bit()) {
         valueField = std::make_unique<ROOT::Internal::RDF::RRDFCardinalityField<std::uint64_t>>(field.GetFieldName());
      } else {
         R__ASSERT(false && "cardinality field stored with an unexpected integer type");
      }
   }
 
   valueField->SetOnDiskId(field.GetOnDiskId());
   auto valueSubfields = valueField->GetMutableSubfields();
   const auto fieldSubfields = field.GetConstSubfields();
   for (unsigned i = 0; i < valueSubfields.size(); ++i) {
      valueSubfields[i]->SetOnDiskId(fieldSubfields[i]->GetOnDiskId());
   }
 
   std::unique_ptr<ROOT::RFieldBase> cardinalityField;
   // Collections get the additional "number of" RDF column (e.g. "R_rdf_sizeof_tracks")
   if (!fieldInfos.empty()) {
      const auto &info = fieldInfos.back();
      const std::string name =
         (procProvenance.Empty() ? "R_rdf_sizeof_" : procProvenance.Get() + "_R_rdf_sizeof_") + info.fFieldName;
      if (info.fNRepetitions > 0) {
         cardinalityField = std::make_unique<ROOT::Internal::RDF::RArraySizeField>(name, info.fNRepetitions);
      } else {
         cardinalityField = std::make_unique<ROOT::Internal::RDF::RRDFCardinalityField<std::size_t>>(name);
      }
      cardinalityField->SetOnDiskId(info.fOnDiskId);
   }
 
   for (auto i = fieldInfos.rbegin(); i != fieldInfos.rend(); ++i) {
      const auto &fieldInfo = *i;
 
      const auto valueFieldName = valueField->GetFieldName();
 
      if (fieldInfo.fNRepetitions > 0) {
         // Fixed-size array, read it as ROOT::RVec in memory
         valueField =
            std::make_unique<ROOT::RArrayAsRVecField>(valueFieldName, valueField->Clone("_0"), fieldInfo.fNRepetitions);
      } else {
         // Actual collection. A std::vector or ROOT::RVec gets added as a ROOT::RVec. All other collection types keep
         // their original type.
         if (convertToRVec) {
            valueField = std::make_unique<ROOT::RRVecField>(valueFieldName, valueField->Clone("_0"));
         } else {
            auto outerFieldType = fieldInfo.fTypeName;
            valueField = ROOT::RFieldBase::Create(valueFieldName, outerFieldType).Unwrap();
         }
      }
 
      valueField->SetOnDiskId(fieldInfo.fOnDiskId);
 
      // Skip the inner-most collection level to construct the cardinality column
      // It's taken care of by the `if (!fieldInfos.empty())` scope above
      if (i != fieldInfos.rbegin()) {
         const auto cardinalityFieldName = cardinalityField->GetFieldName();
         if (fieldInfo.fNRepetitions > 0) {
            // This collection level refers to a fixed-size array
            cardinalityField = std::make_unique<ROOT::RArrayAsRVecField>(
               cardinalityFieldName, cardinalityField->Clone("_0"), fieldInfo.fNRepetitions);
         } else {
            // This collection level refers to an RVec
            cardinalityField = std::make_unique<ROOT::RRVecField>(cardinalityFieldName, cardinalityField->Clone("_0"));
         }
 
         cardinalityField->SetOnDiskId(fieldInfo.fOnDiskId);
      }
   }
 
   std::string canonicalName = (procProvenance.Empty() ? "" : procProvenance.Get() + ".") + std::string(colName);
   fColumnNames.emplace_back(canonicalName);
   fColumnTypes.emplace_back(valueField->GetTypeName());
 
   ROOT::Internal::RNTupleProcessorEntry::FieldIndex_t fieldIdx;
   if (auto optFieldIdx = fProcessor->GetEntry().FindFieldIndex(canonicalName);
       optFieldIdx && fProcessor->GetEntry().GetField(*optFieldIdx).GetTypeName() == field.GetTypeName()) {
      fieldIdx = *optFieldIdx;
   } else {
      fieldIdx = fProcessor->AddFieldToEntry(std::move(valueField), canonicalName, nullptr,
                                             ROOT::Internal::RNTupleProcessorProvenance());
   }
 
   if (cardinalityField) {
      // FIXME(fdegeus): Make cardinality field name work properly with "#colname"
      std::string cardinalityFieldName =
         (procProvenance.Empty() ? "R_rdf_sizeof_" : procProvenance.Get() + ".R_rdf_sizeof_") + std::string(colName);
      fColumnNames.emplace_back(cardinalityFieldName);
      fColumnTypes.emplace_back(cardinalityField->GetTypeName());
      auto cardinalityFieldIdx = fProcessor->AddFieldToEntry(std::move(cardinalityField), cardinalityFieldName, nullptr,
                                                             ROOT::Internal::RNTupleProcessorProvenance());
      if (!fieldInfos.empty()) {
         const auto &fieldInfo = fieldInfos.back();
         fieldIdx = fieldInfo.fEntryFieldIdx;
      }
      fProcessor->GetEntry().SetFieldOriginIdx(cardinalityFieldIdx, fieldIdx);
   }
 
   for (std::uint16_t i = 0; i < nFieldInfosAdded; i++) {
      fieldInfos.pop_back();
   }
}
 
RNTupleProcessorDS::RNTupleProcessorDS(std::unique_ptr<ROOT::RNTupleProcessor> processor)
   : fProcessor(std::move(processor))
{
   // Do not add the subfields now, this is handled by AddField.
   fProcessor->AddAllFieldsToEntry(ROOT::Internal::RNTupleProcessorProvenance(), /*includeSubfields=*/false);
   const auto &entry = fProcessor->GetEntry();
   for (auto fieldIdx : entry.GetFieldIndices()) {
      const auto &field = entry.GetValue(fieldIdx).GetField();
      std::deque<RNTupleProcessorDS::RFieldInfo> fieldInfos{};
      AddField(field, entry.GetQualifiedFieldName(fieldIdx), entry.GetFieldProvenance(fieldIdx), fieldInfos);
   }
}
 
ROOT::RDF::RDataSource::Record_t
RNTupleProcessorDS::GetColumnReadersImpl(std::string_view /* name */, const std::type_info & /* ti */)
{
   // This datasource uses the newer GetColumnReaders() API
   return {};
}
 
std::unique_ptr<ROOT::Detail::RDF::RColumnReaderBase>
RNTupleProcessorDS::GetColumnReaders(unsigned int /* slot */, std::string_view name, const std::type_info &tid)
{
   // At this point we can assume that `name` will be found in fColumnNames
   const auto requestedType = ROOT::Internal::GetRenormalizedTypeName(ROOT::Internal::RDF::TypeID2TypeName(tid));
 
   // First check if a field with the requested type already exists. If that is the case, we can immediately create a
   // column reader for it.
   if (fProcessor->GetEntry().FindFieldIndex(name, requestedType)) {
      auto valuePtr = fProcessor->RequestField(std::string(name), requestedType);
      auto reader = std::make_unique<Internal::RDF::RNTupleProcessorColumnReader>(valuePtr);
      fColumnReaders.emplace_back(reader.get());
 
      return reader;
   }
 
   // Secondly, check whether a field with the same name, but a different on-disk type exists. If this is not the case,
   // first try to add the field. If the field cannot be added, throw an exception.
   // Otherwise, we create a new field with the requested type and add it to the processor entry
   // before creating a column reader for it.
   auto fieldIdx = fProcessor->GetEntry().FindFieldIndex(name);
   if (!fieldIdx) {
      std::unique_ptr<Internal::RDF::RNTupleProcessorColumnReader> reader;
      try {
         auto valuePtr = fProcessor->RequestField(std::string(name), requestedType);
         reader = std::make_unique<Internal::RDF::RNTupleProcessorColumnReader>(valuePtr);
         fColumnReaders.emplace_back(reader.get());
      } catch (const ROOT::RException &) {
         throw std::runtime_error("RNTupleProcessorDS: Column \"" + std::string(name) + "\" does not exist");
      }
      return reader;
   }
 
   const auto &field = fProcessor->GetEntry().GetField(*fieldIdx);
   const std::string strName = std::string(name);
   std::unique_ptr<ROOT::RFieldBase> newField;
   if (dynamic_cast<const ROOT::RCardinalityField *>(&field)) {
      if (requestedType == "bool") {
         newField = std::make_unique<ROOT::Internal::RDF::RRDFCardinalityField<bool>>(strName);
      } else if (requestedType == "char") {
         newField = std::make_unique<ROOT::Internal::RDF::RRDFCardinalityField<char>>(strName);
      } else if (requestedType == "std::int8_t") {
         newField = std::make_unique<ROOT::Internal::RDF::RRDFCardinalityField<std::int8_t>>(strName);
      } else if (requestedType == "std::uint8_t") {
         newField = std::make_unique<ROOT::Internal::RDF::RRDFCardinalityField<std::uint8_t>>(strName);
      } else if (requestedType == "std::int16_t") {
         newField = std::make_unique<ROOT::Internal::RDF::RRDFCardinalityField<std::int16_t>>(strName);
      } else if (requestedType == "std::uint16_t") {
         newField = std::make_unique<ROOT::Internal::RDF::RRDFCardinalityField<std::uint16_t>>(strName);
      } else if (requestedType == "std::int32_t") {
         newField = std::make_unique<ROOT::Internal::RDF::RRDFCardinalityField<std::int32_t>>(strName);
      } else if (requestedType == "std::uint32_t") {
         newField = std::make_unique<ROOT::Internal::RDF::RRDFCardinalityField<std::uint32_t>>(strName);
      } else if (requestedType == "std::int64_t") {
         newField = std::make_unique<ROOT::Internal::RDF::RRDFCardinalityField<std::int64_t>>(strName);
      } else if (requestedType == "std::uint64_t") {
         newField = std::make_unique<ROOT::Internal::RDF::RRDFCardinalityField<std::uint64_t>>(strName);
      } else {
         throw std::runtime_error("RNTupleProcessorDS: Could not create field with type \"" + requestedType +
                                  "\" for column \"" + strName + "\"");
      }
   } else {
      auto baseName = ROOT::Split(strName, ".").back();
      auto newAltProtoFieldOrException = ROOT::RFieldBase::Create(baseName, requestedType);
      if (!newAltProtoFieldOrException) {
         throw std::runtime_error("RNTupleProcessorDS: Could not create field with type \"" + requestedType +
                                  "\" for column \"" + strName + "\"");
      }
      newField = newAltProtoFieldOrException.Unwrap();
   }
   newField->SetOnDiskId(field.GetOnDiskId());
 
   try {
      fProcessor->AddFieldToEntry(std::move(newField), strName, nullptr, ROOT::Internal::RNTupleProcessorProvenance());
   } catch (const ROOT::RException &) {
      std::string msg = "RNTupleProcessorDS: invalid type \"" + requestedType + "\" for column \"" + strName +
                        "\" with on-disk type \"" + field.GetTypeName() + "\"";
      throw std::runtime_error(msg);
   }
 
   auto valuePtr = fProcessor->RequestField(std::string(name), requestedType);
   auto reader = std::make_unique<Internal::RDF::RNTupleProcessorColumnReader>(valuePtr);
   fColumnReaders.emplace_back(reader.get());
 
   return reader;
}
 
std::vector<std::pair<ULong64_t, ULong64_t>> RNTupleProcessorDS::GetEntryRanges()
{
   if (fCurrentRanges.empty()) {
      fCurrentRanges.emplace_back(0, fProcessor->GetNEntries());
      return fCurrentRanges;
   }
   return {};
}
 
void RNTupleProcessorDS::InitSlot(unsigned int /* slot */, ULong64_t /* firstEntry */)
{
   assert(fNSlots == 1);
}
 
void RNTupleProcessorDS::FinalizeSlot(unsigned int /* slot */)
{
   assert(fNSlots == 1);
}
 
std::string RNTupleProcessorDS::GetTypeName(std::string_view colName) const
{
   auto colNamePos = std::find(fColumnNames.begin(), fColumnNames.end(), colName);
 
   if (colNamePos == fColumnNames.end()) {
      auto msg = std::string("RNTupleProcessorDS: There is no column with name \"") + std::string(colName) + "\"";
      throw std::runtime_error(msg);
   }
 
   const auto index = std::distance(fColumnNames.begin(), colNamePos);
   return fColumnTypes[index];
}
 
bool RNTupleProcessorDS::HasColumn(std::string_view colName) const
{
   return std::find(fColumnNames.begin(), fColumnNames.end(), colName) != fColumnNames.end();
}
 
void RNTupleProcessorDS::Initialize()
{
   fCurrentRanges.clear();
   fProcessor->Connect(fProcessor->GetEntry().GetFieldIndices(), ROOT::Internal::RNTupleProcessorProvenance(),
                       /*updateFields=*/false);
}
 
void RNTupleProcessorDS::Finalize()
{
   // noop
}
 
void RNTupleProcessorDS::SetNSlots(unsigned int nSlots)
{
   assert(fNSlots == 0);
   // FIXME(fdegeus): Set to actual nSlots once MT support lands
   if (nSlots > 1)
      Warning("RNTupleProcessorDS", "MT is not yet supported. Using single-threaded processing instead.");
   fNSlots = 1;
}
 
bool RNTupleProcessorDS::SetEntry(unsigned int /* slot */, ULong64_t entry)
{
   fProcessor->SetEntryMapping(entry, fEntryMapping);
   if (fEntryMapping.GetEntryNumber() == ROOT::kInvalidNTupleIndex)
      return false;
 
   fProcessor->LoadEntry(fEntryMapping);
 
   for (auto colReader : fColumnReaders) {
      colReader->SetEntryLoaded(entry);
   }
 
   return true;
}
 
ROOT::RDF::RSampleInfo RNTupleProcessorDS::CreateSampleInfo(
   unsigned int /* slot */, const std::unordered_map<std::string, ROOT::RDF::Experimental::RSample *> &sampleMap) const
{
   const auto &ntupleID = fProcessor->GetOptions().GetProcessorName();
 
   // TODO: There is no support for RNTuple in RDatasetSpec, thus the sample map
   // is always empty at the moment.
   if (sampleMap.empty())
      return ROOT::RDF::RSampleInfo(ntupleID, std::make_pair(0, fProcessor->GetNEntries()));
 
   if (sampleMap.find(ntupleID) == sampleMap.end())
      throw std::runtime_error("Full sample identifier '" + ntupleID + "' cannot be found in the available samples.");
 
   return ROOT::RDF::RSampleInfo(ntupleID, std::make_pair(0, fProcessor->GetNEntries()), sampleMap.at(ntupleID));
}
} // namespace ROOT::Experimental::Internal::RDF
 
ROOT::RDataFrame ROOT::Experimental::RDF::FromRNTupleProcessor(std::unique_ptr<ROOT::RNTupleProcessor> processor)
{
   return ROOT::RDataFrame(
      std::make_unique<ROOT::Experimental::Internal::RDF::RNTupleProcessorDS>(std::move(processor)));
}