[df] Support MT snapshotting to RNTuple

Switch the existing code to use the RNTupleParallelWriter with one
RNTupleFillContext per slot. This should be (almost) as efficient
as the RNTupleWriter (one additional cloned RNTupleModel for the
only fill context), but save quite a bit of code duplication and
in testing effort.

Author: hahnjo

tree/dataframe/inc/ROOT/RDF/InterfaceUtils.hxx

@@ -293,19 +293,13 @@ BuildAction(const ColumnNames_t &colNames, const std::shared_ptr<SnapshotHelperA
 
    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>;
+      // We use the same helper for single- and multi-thread snapshot.
+      using Helper_t = UntypedSnapshotRNTupleHelper;
+      using Action_t = RActionSnapshot<Helper_t, PrevNodeType>;
 
-         actionPtr.reset(new Action_t(
-            Helper_t(filename, dirname, treename, colNames, outputColNames, options, inputLM, outputLM, 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.");
-      }
+      actionPtr.reset(new Action_t(Helper_t(nSlots, filename, dirname, treename, colNames, outputColNames, options,
+                                            inputLM, outputLM, colTypeIDs),
+                                   colNames, colTypeIDs, prevNode, colRegister));
    } else {
       if (!ROOT::IsImplicitMTEnabled()) {
          // single-thread snapshot

tree/dataframe/inc/ROOT/RDF/SnapshotHelpers.hxx

@@ -30,7 +30,10 @@ class TBranch;
 class TFile;
 
 namespace ROOT {
-class RNTupleWriter;
+namespace Experimental {
+class RNTupleFillContext;
+class RNTupleParallelWriter;
+} // namespace Experimental
 class REntry;
 class RFieldToken;
 class TBufferMerger;
@@ -64,16 +67,18 @@ class R__CLING_PTRCHECK(off) UntypedSnapshotRNTupleHelper final : public RAction
    ROOT::Detail::RDF::RLoopManager *fOutputLoopManager;
    ColumnNames_t fInputFieldNames; // This contains the resolved aliases
    ColumnNames_t fOutputFieldNames;
-   std::unique_ptr<ROOT::RNTupleWriter> fWriter;
+   std::unique_ptr<ROOT::Experimental::RNTupleParallelWriter> fWriter;
    std::vector<ROOT::RFieldToken> fFieldTokens;
 
-   std::unique_ptr<ROOT::REntry> fOutputEntry;
+   unsigned int fNSlots;
+   std::vector<std::shared_ptr<ROOT::Experimental::RNTupleFillContext>> fFillContexts;
+   std::vector<std::unique_ptr<ROOT::REntry>> fEntries;
 
    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,
+   UntypedSnapshotRNTupleHelper(unsigned int nSlots, 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,
                                 const std::vector<const std::type_info *> &colTypeIDs);
@@ -86,9 +91,11 @@ public:
 
    void Initialize();
 
-   void InitTask(TTreeReader *, unsigned int /* slot */) {}
+   void Exec(unsigned int slot, const std::vector<void *> &values);
+
+   void InitTask(TTreeReader *, unsigned int slot);
 
-   void Exec(unsigned int /* slot */, const std::vector<void *> &values);
+   void FinalizeTask(unsigned int slot);
 
    void Finalize();
 

tree/dataframe/src/RDFSnapshotHelpers.cxx

@@ -23,7 +23,8 @@
 #include <ROOT/RFieldToken.hxx>
 #include <ROOT/RNTuple.hxx>
 #include <ROOT/RNTupleDS.hxx>
-#include <ROOT/RNTupleWriter.hxx>
+#include <ROOT/RNTupleFillContext.hxx>
+#include <ROOT/RNTupleParallelWriter.hxx>
 #include <ROOT/RTTreeDS.hxx>
 #include <ROOT/TBufferMerger.hxx>
 
@@ -800,9 +801,10 @@ ROOT::Internal::RDF::UntypedSnapshotTTreeHelperMT::MakeNew(void *newName, std::s
 }
 
 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, const std::vector<const std::type_info *> &colTypeIDs)
+   unsigned int nSlots, 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,
+   const std::vector<const std::type_info *> &colTypeIDs)
    : fFileName(filename),
      fDirName(dirname),
      fNTupleName(ntuplename),
@@ -811,6 +813,9 @@ ROOT::Internal::RDF::UntypedSnapshotRNTupleHelper::UntypedSnapshotRNTupleHelper(
      fOutputLoopManager(outputLM),
      fInputFieldNames(vfnames),
      fOutputFieldNames(ReplaceDotWithUnderscore(fnames)),
+     fNSlots(nSlots),
+     fFillContexts(nSlots),
+     fEntries(nSlots),
      fInputColumnTypeIDs(colTypeIDs)
 {
    EnsureValidSnapshotRNTupleOutput(fOptions, fNTupleName, fFileName);
@@ -845,7 +850,6 @@ void ROOT::Internal::RDF::UntypedSnapshotRNTupleHelper::Initialize()
       fFieldTokens[i] = model->GetToken(fOutputFieldNames[i]);
    }
    model->Freeze();
-   fOutputEntry = model->CreateBareEntry();
 
    ROOT::RNTupleWriteOptions writeOptions;
    writeOptions.SetCompression(fOptions.fCompressionAlgorithm, fOptions.fCompressionLevel);
@@ -864,20 +868,41 @@ void ROOT::Internal::RDF::UntypedSnapshotRNTupleHelper::Initialize()
          outputDir = fOutputFile->mkdir(fDirName.c_str());
    }
 
-   fWriter = ROOT::RNTupleWriter::Append(std::move(model), fNTupleName, *outputDir, writeOptions);
+   fWriter = ROOT::Experimental::RNTupleParallelWriter::Append(std::move(model), fNTupleName, *outputDir, writeOptions);
+}
+
+void ROOT::Internal::RDF::UntypedSnapshotRNTupleHelper::InitTask(TTreeReader *, unsigned int slot)
+{
+   if (!fFillContexts[slot]) {
+      fFillContexts[slot] = fWriter->CreateFillContext();
+      fEntries[slot] = fFillContexts[slot]->GetModel().CreateBareEntry();
+   }
 }
 
-void ROOT::Internal::RDF::UntypedSnapshotRNTupleHelper::Exec(unsigned int /* slot */, const std::vector<void *> &values)
+void ROOT::Internal::RDF::UntypedSnapshotRNTupleHelper::Exec(unsigned int slot, const std::vector<void *> &values)
 {
+   auto &fillContext = fFillContexts[slot];
+   auto &outputEntry = fEntries[slot];
    assert(values.size() == fFieldTokens.size());
    for (decltype(values.size()) i = 0; i < values.size(); i++) {
-      fOutputEntry->BindRawPtr(fFieldTokens[i], values[i]);
+      outputEntry->BindRawPtr(fFieldTokens[i], values[i]);
    }
-   fWriter->Fill(*fOutputEntry);
+   fillContext->Fill(*outputEntry);
+}
+
+void ROOT::Internal::RDF::UntypedSnapshotRNTupleHelper::FinalizeTask(unsigned int slot)
+{
+   // In principle we would not need to flush a cluster here, but we want to benefit from parallelism for compression.
+   // NB: RNTupleFillContext::FlushCluster() is a nop if there is no new entry since the last flush.
+   fFillContexts[slot]->FlushCluster();
 }
 
 void ROOT::Internal::RDF::UntypedSnapshotRNTupleHelper::Finalize()
 {
+   // First clear and destroy all entries, which were created from the RNTupleFillContexts.
+   fEntries.clear();
+   fFillContexts.clear();
+   // Then destroy the RNTupleParallelWriter and write the metadata.
    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));
@@ -898,7 +923,7 @@ 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, fInputColumnTypeIDs};
+   return UntypedSnapshotRNTupleHelper{
+      fNSlots,           finalName, fDirName,          fNTupleName,        fInputFieldNames,
+      fOutputFieldNames, fOptions,  fInputLoopManager, fOutputLoopManager, fInputColumnTypeIDs};
 }

tree/dataframe/test/dataframe_snapshot_ntuple.cxx

@@ -614,23 +614,22 @@ struct TIMTEnabler {
    ~TIMTEnabler() { ROOT::DisableImplicitMT(); }
 };
 
-TEST(RDFSnapshotRNTuple, ThrowIfMT)
+TEST(RDFSnapshotRNTuple, WithMT)
 {
    TIMTEnabler _(4);
 
-   FileRAII fileGuard{"RDFSnapshotRNTuple_throw_if_mt.root"};
+   FileRAII fileGuard{"RDFSnapshotRNTuple_mt.root"};
 
    auto df = ROOT::RDataFrame(25ull).Define("x", [] { return 10; });
 
    RSnapshotOptions opts;
    opts.fOutputFormat = ROOT::RDF::ESnapshotOutputFormat::kRNTuple;
 
-   try {
-      auto sdf = df.Snapshot("ntuple", fileGuard.GetPath(), {"x"}, opts);
-      *sdf;
-      FAIL() << "MT snapshotting to RNTuple is not supported yet";
-   } catch (const std::runtime_error &err) {
-      EXPECT_STREQ(err.what(), "Snapshot: Snapshotting to RNTuple with IMT enabled is not supported yet.");
-   }
+   auto sdf = df.Snapshot("ntuple", fileGuard.GetPath(), {"x"}, opts);
+   *sdf;
+
+   auto reader = RNTupleReader::Open("ntuple", fileGuard.GetPath());
+   EXPECT_EQ(25, reader->GetNEntries());
+   // There should be more than one cluster, but this is not guaranteed because of scheduling...
 }
 #endif // R__USE_IMT