[df] Lift template requirement for RNTuple Snapshot

vepadulano · vepadulano · commit 8c5860e5ec23 · 2025-06-25T16:02:40.000+02:00
A new action helper is introduced to allow calling Snapshot without template
arguments and without the need for JITting to write RNTuple data to disk.

Unlike TTree, RNTuple is actually able to write certain types to disk even
without explicit dictionary information or RTTI. Notably, all the class of
composed types such as `std::vector&lt;std::vector&lt;...&gt;&gt;` and similar, for which
RField is able to create recursively the inner fields if all the components are
known.

This required a slight modification in the logic. A new tag type
`UseNativeDataType` is used in contexts that require `std::type_info` such as
`GetColumnReaders` to signal that we let the RNTuple data source figure out the
type. The tag is also used in the new action helper to retrieve the correct type
name to crete the RField with.
diff --git a/tree/dataframe/inc/ROOT/RDF/ActionHelpers.hxx b/tree/dataframe/inc/ROOT/RDF/ActionHelpers.hxx
@@ -2247,6 +2247,57 @@ public:
    }
 };
 
+class R__CLING_PTRCHECK(off) UntypedSnapshotRNTupleHelper final : public RActionImpl<UntypedSnapshotRNTupleHelper> {
+   std::string fFileName;
+   std::string fDirName;
+   std::string fNTupleName;
+
+   std::unique_ptr<TFile> fOutputFile{nullptr};
+
+   RSnapshotOptions fOptions;
+   ROOT::Detail::RDF::RLoopManager *fInputLoopManager;
+   ROOT::Detail::RDF::RLoopManager *fOutputLoopManager;
+   ColumnNames_t fInputFieldNames; // This contains the resolved aliases
+   ColumnNames_t fOutputFieldNames;
+   std::unique_ptr<ROOT::RNTupleWriter> fWriter{nullptr};
+
+   ROOT::REntry *fOutputEntry;
+
+   std::vector<bool> fIsDefine;
+
+   std::vector<const std::type_info *> fInputColumnTypeIDs; // Types for the input columns
+
+public:
+   UntypedSnapshotRNTupleHelper(std::string_view filename, std::string_view dirname, std::string_view ntuplename,
+                                const ColumnNames_t &vfnames, const ColumnNames_t &fnames,
+                                const RSnapshotOptions &options, ROOT::Detail::RDF::RLoopManager *inputLM,
+                                ROOT::Detail::RDF::RLoopManager *outputLM, std::vector<bool> &&isDefine,
+                                const std::vector<const std::type_info *> &colTypeIDs);
+
+   UntypedSnapshotRNTupleHelper(const UntypedSnapshotRNTupleHelper &) = delete;
+   UntypedSnapshotRNTupleHelper &operator=(const UntypedSnapshotRNTupleHelper &) = delete;
+   UntypedSnapshotRNTupleHelper(UntypedSnapshotRNTupleHelper &&) = default;
+   UntypedSnapshotRNTupleHelper &operator=(UntypedSnapshotRNTupleHelper &&) = default;
+   ~UntypedSnapshotRNTupleHelper() final;
+
+   void InitTask(TTreeReader *, unsigned int /* slot */) {}
+
+   void Exec(unsigned int /* slot */, const std::vector<void *> &values);
+
+   void Initialize();
+
+   void Finalize();
+
+   std::string GetActionName() { return "Snapshot"; }
+
+   ROOT::RDF::SampleCallback_t GetSampleCallback() final
+   {
+      return [](unsigned int, const RSampleInfo &) mutable {};
+   }
+
+   UntypedSnapshotRNTupleHelper MakeNew(void *newName);
+};
+
 class R__CLING_PTRCHECK(off) UntypedSnapshotTTreeHelper final : public RActionImpl<UntypedSnapshotTTreeHelper> {
    std::string fFileName;
    std::string fDirName;
diff --git a/tree/dataframe/inc/ROOT/RDF/InterfaceUtils.hxx b/tree/dataframe/inc/ROOT/RDF/InterfaceUtils.hxx
@@ -345,7 +345,7 @@ BuildAction(const ColumnNames_t &colNames, const std::shared_ptr<SnapshotHelperA
    const auto &treename = snapHelperArgs->fTreeName;
    const auto &outputColNames = snapHelperArgs->fOutputColNames;
    const auto &options = snapHelperArgs->fOptions;
-   const auto &lmPtr = snapHelperArgs->fOutputLoopManager;
+   const auto &outputLM = snapHelperArgs->fOutputLoopManager;
    const auto &inputLM = snapHelperArgs->fInputLoopManager;
 
    auto sz = colNames.size();
@@ -354,21 +354,36 @@ BuildAction(const ColumnNames_t &colNames, const std::shared_ptr<SnapshotHelperA
       isDefine[i] = colRegister.IsDefineOrAlias(colNames[i]);
 
    std::unique_ptr<RActionBase> actionPtr;
+   if (snapHelperArgs->fToNTuple) {
+      if (!ROOT::IsImplicitMTEnabled()) {
+         // single-thread snapshot
+         using Helper_t = UntypedSnapshotRNTupleHelper;
+         using Action_t = RActionSnapshot<Helper_t, PrevNodeType>;
 
-   if (!ROOT::IsImplicitMTEnabled()) {
-      // single-thread snapshot
-      using Helper_t = UntypedSnapshotTTreeHelper;
-      using Action_t = RActionSnapshot<Helper_t, PrevNodeType>;
-      actionPtr.reset(new Action_t(Helper_t(filename, dirname, treename, colNames, outputColNames, options,
-                                            std::move(isDefine), lmPtr, inputLM, colTypeIDs),
-                                   colNames, colTypeIDs, prevNode, colRegister));
+         actionPtr.reset(new Action_t(Helper_t(filename, dirname, treename, colNames, outputColNames, options, inputLM,
+                                               outputLM, std::move(isDefine), colTypeIDs),
+                                      colNames, colTypeIDs, prevNode, colRegister));
+      } else {
+         // multi-thread snapshot to RNTuple is not yet supported
+         // TODO(fdegeus) Add MT snapshotting
+         throw std::runtime_error("Snapshot: Snapshotting to RNTuple with IMT enabled is not supported yet.");
+      }
    } else {
-      // multi-thread snapshot
-      using Helper_t = UntypedSnapshotTTreeHelperMT;
-      using Action_t = RActionSnapshot<Helper_t, PrevNodeType>;
-      actionPtr.reset(new Action_t(Helper_t(nSlots, filename, dirname, treename, colNames, outputColNames, options,
-                                            std::move(isDefine), lmPtr, inputLM, colTypeIDs),
-                                   colNames, colTypeIDs, prevNode, colRegister));
+      if (!ROOT::IsImplicitMTEnabled()) {
+         // single-thread snapshot
+         using Helper_t = UntypedSnapshotTTreeHelper;
+         using Action_t = RActionSnapshot<Helper_t, PrevNodeType>;
+         actionPtr.reset(new Action_t(Helper_t(filename, dirname, treename, colNames, outputColNames, options,
+                                               std::move(isDefine), outputLM, inputLM, colTypeIDs),
+                                      colNames, colTypeIDs, prevNode, colRegister));
+      } else {
+         // multi-thread snapshot
+         using Helper_t = UntypedSnapshotTTreeHelperMT;
+         using Action_t = RActionSnapshot<Helper_t, PrevNodeType>;
+         actionPtr.reset(new Action_t(Helper_t(nSlots, filename, dirname, treename, colNames, outputColNames, options,
+                                               std::move(isDefine), outputLM, inputLM, colTypeIDs),
+                                      colNames, colTypeIDs, prevNode, colRegister));
+      }
    }
 
    return actionPtr;
diff --git a/tree/dataframe/inc/ROOT/RDF/RInterface.hxx b/tree/dataframe/inc/ROOT/RDF/RInterface.hxx
@@ -1344,10 +1344,14 @@ public:
 
       RResultPtr<RInterface<RLoopManager>> resPtr;
 
-      auto retrieveTypeID = [](const std::string &colName, const std::string &colTypeName) -> const std::type_info * {
+      auto retrieveTypeID = [](const std::string &colName, const std::string &colTypeName,
+                               bool isRNTuple = false) -> const std::type_info * {
          try {
             return &ROOT::Internal::RDF::TypeName2TypeID(colTypeName);
          } catch (const std::runtime_error &err) {
+            if (isRNTuple)
+               return &typeid(ROOT::Internal::RDF::UseNativeDataType);
+
             if (std::string(err.what()).find("Cannot extract type_info of type") != std::string::npos) {
                // We could not find RTTI for this column, thus we cannot write it out at the moment.
                std::string trueTypeName{colTypeName};
@@ -1380,12 +1384,26 @@ public:
             std::string(filename), std::string(dirname), std::string(treename), colListWithAliasesAndSizeBranches,
             options, newRDF->GetLoopManager(), GetLoopManager(), true /* fToNTuple */});
 
-         // The Snapshot helper will use colListNoAliasesWithSizeBranches (with aliases resolved) as input columns, and
-         // colListWithAliasesAndSizeBranches (still with aliases in it, passed through snapHelperArgs) as output column
-         // names.
-         resPtr = CreateAction<RDFInternal::ActionTags::Snapshot, RDFDetail::RInferredType>(
-            colListNoAliasesWithSizeBranches, newRDF, snapHelperArgs, fProxiedPtr,
-            colListNoAliasesWithSizeBranches.size());
+         auto &&nColumns = colListNoAliasesWithSizeBranches.size();
+         const auto validColumnNames = GetValidatedColumnNames(nColumns, colListNoAliasesWithSizeBranches);
+
+         const auto nSlots = fLoopManager->GetNSlots();
+         std::vector<const std::type_info *> colTypeIDs;
+         colTypeIDs.reserve(nColumns);
+         for (decltype(nColumns) i{}; i < nColumns; i++) {
+            const auto &colName = validColumnNames[i];
+            const auto colTypeName = ROOT::Internal::RDF::ColumnName2ColumnTypeName(
+               colName, /*tree*/ nullptr, GetDataSource(), fColRegister.GetDefine(colName), options.fVector2RVec);
+            const std::type_info *colTypeID = retrieveTypeID(colName, colTypeName, /*isRNTuple*/ true);
+            colTypeIDs.push_back(colTypeID);
+         }
+         // Crucial e.g. if the column names do not correspond to already-available column readers created by the data
+         // source
+         CheckAndFillDSColumns(validColumnNames, colTypeIDs);
+
+         auto action =
+            RDFInternal::BuildAction(validColumnNames, snapHelperArgs, nSlots, fProxiedPtr, fColRegister, colTypeIDs);
+         resPtr = MakeResultPtr(newRDF, *GetLoopManager(), std::move(action));
       } else {
          if (RDFInternal::GetDataSourceLabel(*this) == "RNTupleDS" &&
              options.fOutputFormat == ESnapshotOutputFormat::kDefault) {
diff --git a/tree/dataframe/inc/ROOT/RDF/Utils.hxx b/tree/dataframe/inc/ROOT/RDF/Utils.hxx
@@ -336,6 +336,13 @@ auto MakeAliasedSharedPtr(T *rawPtr)
  */
 ROOT::RDF::Experimental::RDatasetSpec RetrieveSpecFromJson(const std::string &jsonFile);
 
+/**
+ * Tag to let data sources use the native data type when creating a column reader.
+ *
+ * See usage of this in RNTupleDS
+ */
+struct UseNativeDataType {};
+
 } // end NS RDF
 } // end NS Internal
 } // end NS ROOT
diff --git a/tree/dataframe/inc/ROOT/RNTupleDS.hxx b/tree/dataframe/inc/ROOT/RNTupleDS.hxx
@@ -170,6 +170,8 @@ class RNTupleDS final : public ROOT::RDF::RDataSource {
 
    explicit RNTupleDS(std::unique_ptr<ROOT::Internal::RPageSource> pageSource);
 
+   ROOT::RFieldBase *GetFieldWithTypeChecks(std::string_view fieldName, const std::type_info &tid);
+
 public:
    RNTupleDS(std::string_view ntupleName, std::string_view fileName);
    RNTupleDS(std::string_view ntupleName, const std::vector<std::string> &fileNames);
diff --git a/tree/dataframe/src/RDFActionHelpers.cxx b/tree/dataframe/src/RDFActionHelpers.cxx
@@ -892,3 +892,101 @@ ROOT::Internal::RDF::UntypedSnapshotTTreeHelperMT::MakeNew(void *newName, std::s
                                                             fInputLoopManager,
                                                             fInputColumnTypeIDs};
 }
+
+ROOT::Internal::RDF::UntypedSnapshotRNTupleHelper::UntypedSnapshotRNTupleHelper(
+   std::string_view filename, std::string_view dirname, std::string_view ntuplename, const ColumnNames_t &vfnames,
+   const ColumnNames_t &fnames, const RSnapshotOptions &options, ROOT::Detail::RDF::RLoopManager *inputLM,
+   ROOT::Detail::RDF::RLoopManager *outputLM, std::vector<bool> &&isDefine,
+   const std::vector<const std::type_info *> &colTypeIDs)
+   : fFileName(filename),
+     fDirName(dirname),
+     fNTupleName(ntuplename),
+     fOptions(options),
+     fInputLoopManager(inputLM),
+     fOutputLoopManager(outputLM),
+     fInputFieldNames(vfnames),
+     fOutputFieldNames(ReplaceDotWithUnderscore(fnames)),
+     fIsDefine(std::move(isDefine)),
+     fInputColumnTypeIDs(colTypeIDs)
+{
+   EnsureValidSnapshotRNTupleOutput(fOptions, fNTupleName, fFileName);
+}
+
+ROOT::Internal::RDF::UntypedSnapshotRNTupleHelper::~UntypedSnapshotRNTupleHelper()
+{
+   if (!fNTupleName.empty() && !fOutputLoopManager->GetDataSource() && fOptions.fLazy)
+      Warning("Snapshot", "A lazy Snapshot action was booked but never triggered.");
+}
+
+void ROOT::Internal::RDF::UntypedSnapshotRNTupleHelper::Exec(unsigned int /* slot */, const std::vector<void *> &values)
+{
+   assert(values.size() == fOutputFieldNames.size());
+   for (decltype(values.size()) i = 0; i < values.size(); i++) {
+      fOutputEntry->BindRawPtr(fOutputFieldNames[i], values[i]);
+   }
+   fWriter->Fill();
+}
+
+void ROOT::Internal::RDF::UntypedSnapshotRNTupleHelper::Initialize()
+{
+   auto model = ROOT::RNTupleModel::Create();
+   auto nFields = fOutputFieldNames.size();
+   for (decltype(nFields) i = 0; i < nFields; i++) {
+      // Need to retrieve the type of every field to create as a string
+      // If the input type for a field does not have RTTI, internally we store it as the tag UseNativeDataType. When
+      // that is detected, we need to ask the data source which is the type name based on the on-disk information.
+      const auto typeName = *fInputColumnTypeIDs[i] == typeid(ROOT::Internal::RDF::UseNativeDataType)
+                               ? ROOT::Internal::RDF::GetTypeNameWithOpts(*fInputLoopManager->GetDataSource(),
+                                                                          fInputFieldNames[i], fOptions.fVector2RVec)
+                               : ROOT::Internal::RDF::TypeID2TypeName(*fInputColumnTypeIDs[i]);
+      model->AddField(ROOT::RFieldBase::Create(fOutputFieldNames[i], typeName).Unwrap());
+   }
+   fOutputEntry = &model->GetDefaultEntry();
+
+   ROOT::RNTupleWriteOptions writeOptions;
+   writeOptions.SetCompression(fOptions.fCompressionAlgorithm, fOptions.fCompressionLevel);
+
+   fOutputFile.reset(TFile::Open(fFileName.c_str(), fOptions.fMode.c_str()));
+   if (!fOutputFile)
+      throw std::runtime_error("Snapshot: could not create output file " + fFileName);
+
+   TDirectory *outputDir = fOutputFile.get();
+   if (!fDirName.empty()) {
+      TString checkupdate = fOptions.fMode;
+      checkupdate.ToLower();
+      if (checkupdate == "update")
+         outputDir = fOutputFile->mkdir(fDirName.c_str(), "", true); // do not overwrite existing directory
+      else
+         outputDir = fOutputFile->mkdir(fDirName.c_str());
+   }
+
+   fWriter = ROOT::RNTupleWriter::Append(std::move(model), fNTupleName, *outputDir, writeOptions);
+}
+
+void ROOT::Internal::RDF::UntypedSnapshotRNTupleHelper::Finalize()
+{
+   fWriter.reset();
+   // We can now set the data source of the loop manager for the RDataFrame that is returned by the Snapshot call.
+   fOutputLoopManager->SetDataSource(std::make_unique<ROOT::RDF::RNTupleDS>(fDirName + "/" + fNTupleName, fFileName));
+}
+
+/**
+ * Create a new UntypedSnapshotRNTupleHelper with a different output file name.
+ *
+ * \param[in] newName A type-erased string with the output file name
+ * \return UntypedSnapshotRNTupleHelper
+ *
+ * This MakeNew implementation is tied to the cloning feature of actions
+ * of the computation graph. In particular, cloning a Snapshot node usually
+ * also involves changing the name of the output file, otherwise the cloned
+ * Snapshot would overwrite the same file.
+ */
+ROOT::Internal::RDF::UntypedSnapshotRNTupleHelper
+ROOT::Internal::RDF::UntypedSnapshotRNTupleHelper::MakeNew(void *newName)
+{
+   const std::string finalName = *reinterpret_cast<const std::string *>(newName);
+   return UntypedSnapshotRNTupleHelper{finalName,          fDirName,           fNTupleName,
+                                       fInputFieldNames,   fOutputFieldNames,  fOptions,
+                                       fInputLoopManager,  fOutputLoopManager, std::vector<bool>(fIsDefine),
+                                       fInputColumnTypeIDs};
+}
diff --git a/tree/dataframe/src/RNTupleDS.cxx b/tree/dataframe/src/RNTupleDS.cxx
@@ -430,40 +430,59 @@ ROOT::RDF::RNTupleDS::GetColumnReadersImpl(std::string_view /* name */, const st
    return {};
 }
 
-std::unique_ptr<ROOT::Detail::RDF::RColumnReaderBase>
-ROOT::RDF::RNTupleDS::GetColumnReaders(unsigned int slot, std::string_view name, const std::type_info &tid)
+ROOT::RFieldBase *ROOT::RDF::RNTupleDS::GetFieldWithTypeChecks(std::string_view fieldName, const std::type_info &tid)
 {
    // At this point we can assume that `name` will be found in fColumnNames
-   const auto index = std::distance(fColumnNames.begin(), std::find(fColumnNames.begin(), fColumnNames.end(), name));
+   const auto index =
+      std::distance(fColumnNames.begin(), std::find(fColumnNames.begin(), fColumnNames.end(), fieldName));
+
+   // A reader was requested but we don't have RTTI for it, this is encoded with the tag UseNativeDataType. We can just
+   // return the available protofield
+   if (tid == typeid(ROOT::Internal::RDF::UseNativeDataType)) {
+      return fProtoFields[index].get();
+   }
+
+   // The user explicitly requested a type
    const auto requestedType = ROOT::Internal::GetRenormalizedTypeName(ROOT::Internal::RDF::TypeID2TypeName(tid));
 
-   ROOT::RFieldBase *field;
    // If the field corresponding to the provided name is not a cardinality column and the requested type is different
    // from the proto field that was created when the data source was constructed, we first have to create an
    // alternative proto field for the column reader. Otherwise, we can directly use the existing proto field.
-   if (name.substr(0, 13) != "R_rdf_sizeof_" && requestedType != fColumnTypes[index]) {
+   if (fieldName.substr(0, 13) != "R_rdf_sizeof_" && requestedType != fColumnTypes[index]) {
       auto &altProtoFields = fAlternativeProtoFields[index];
-      auto altProtoField = std::find_if(altProtoFields.begin(), altProtoFields.end(),
-                                        [&requestedType](const std::unique_ptr<ROOT::RFieldBase> &fld) {
-                                           return fld->GetTypeName() == requestedType;
-                                        });
-      if (altProtoField != altProtoFields.end()) {
-         field = altProtoField->get();
-      } else {
-         auto newAltProtoFieldOrException = ROOT::RFieldBase::Create(std::string(name), requestedType);
-         if (!newAltProtoFieldOrException) {
-            throw std::runtime_error("RNTupleDS: Could not create field with type \"" + requestedType +
-                                     "\" for column \"" + std::string(name));
-         }
-         auto newAltProtoField = newAltProtoFieldOrException.Unwrap();
-         newAltProtoField->SetOnDiskId(fProtoFields[index]->GetOnDiskId());
-         field = newAltProtoField.get();
-         altProtoFields.emplace_back(std::move(newAltProtoField));
+
+      // If we can find the requested type in the registered alternative protofields, return the corresponding field
+      if (auto altProtoField = std::find_if(altProtoFields.begin(), altProtoFields.end(),
+                                            [&requestedType](const std::unique_ptr<ROOT::RFieldBase> &fld) {
+                                               return fld->GetTypeName() == requestedType;
+                                            });
+          altProtoField != altProtoFields.end()) {
+         return altProtoField->get();
       }
-   } else {
-      field = fProtoFields[index].get();
+
+      // Otherwise, create a new protofield and register it in the alternatives before returning
+      auto newAltProtoFieldOrException = ROOT::RFieldBase::Create(std::string(fieldName), requestedType);
+      if (!newAltProtoFieldOrException) {
+         throw std::runtime_error("RNTupleDS: Could not create field with type \"" + requestedType +
+                                  "\" for column \"" + std::string(fieldName));
+      }
+      auto newAltProtoField = newAltProtoFieldOrException.Unwrap();
+      newAltProtoField->SetOnDiskId(fProtoFields[index]->GetOnDiskId());
+      auto *newField = newAltProtoField.get();
+      altProtoFields.emplace_back(std::move(newAltProtoField));
+      return newField;
    }
 
+   // General case: there was a correspondence between the user-requested type and the corresponding column type
+   return fProtoFields[index].get();
+}
+
+std::unique_ptr<ROOT::Detail::RDF::RColumnReaderBase>
+ROOT::RDF::RNTupleDS::GetColumnReaders(unsigned int slot, std::string_view name, const std::type_info &tid)
+{
+   ROOT::RFieldBase *field = GetFieldWithTypeChecks(name, tid);
+   assert(field != nullptr);
+
    // Map the field's and subfields' IDs to qualified names so that we can later connect the fields to
    // other page sources from the chain
    fFieldId2QualifiedName[field->GetOnDiskId()] = fPrincipalDescriptor.GetQualifiedFieldName(field->GetOnDiskId());
diff --git a/tree/dataframe/test/CMakeLists.txt b/tree/dataframe/test/CMakeLists.txt
@@ -59,7 +59,7 @@ ROOT_ADD_GTEST(dataframe_snapshot_emptyoutput dataframe_snapshot_emptyoutput.cxx
 ROOT_GENERATE_DICTIONARY(DummyDict ${CMAKE_CURRENT_SOURCE_DIR}/DummyHeader.hxx
                          MODULE dataframe_snapshot_emptyoutput LINKDEF DummyHeaderLinkDef.hxx OPTIONS -inlineInputHeader
                          DEPENDENCIES ROOTVecOps GenVector)
-ROOT_ADD_GTEST(dataframe_snapshot_ntuple dataframe_snapshot_ntuple.cxx LIBRARIES ROOTDataFrame ROOTNTuple)
+ROOT_ADD_GTEST(dataframe_snapshot_ntuple dataframe_snapshot_ntuple.cxx LIBRARIES ROOTDataFrame ROOTNTuple NTupleStruct)
 endif()
 
 ROOT_ADD_GTEST(dataframe_datasetspec dataframe_datasetspec.cxx LIBRARIES ROOTDataFrame)
