[df] Preserve schema of empty output TTree in Snapshot

Previously, the behaviour of Snapshot in case no entries pass the selections in
the event loop was two-fold:

* In single thread, Snapshot created an output file with an empty TTree inside,
  without the user-requested branches.

* In multi-threaded mode, Snapshot created an output file, without any TTree
  inside.

This commit aligns the behaviour of the two execution modes. Furthermore, it
ensures the user-requested dataset schema is preserved even when the output
TTree is empty. i.e., all the user-requested branches will be present in the
output TTree.

Author: vepadulano

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

@@ -1503,6 +1503,10 @@ void SetBranchesHelper(TTree *inputTree, TTree &outputTree, const std::string &i
    }
 }
 
+void SetEmptyBranchesHelper(TTree *inputTree, TTree &outputTree, RBranchSet &outputBranches,
+                            const std::string &inputBranchName, const std::string &outputBranchName,
+                            const std::type_info &typeInfo, int basketSize);
+
 /// Ensure that the TTree with the resulting snapshot can be written to the target TFile. This means checking that the
 /// TFile can be opened in the mode specified in `opts`, deleting any existing TTrees in case
 /// `opts.fOverwriteIfExists = true`, or throwing an error otherwise.
@@ -1618,6 +1622,18 @@ public:
       (void)expander; // avoid unused variable warnings for older compilers such as gcc 4.9
    }
 
+   template <std::size_t... S>
+   void SetEmptyBranches(TTree *inputTree, TTree &outputTree, std::index_sequence<S...>)
+   {
+      RBranchSet outputBranches{};
+      // We use the expander trick rather than a fold expression to avoid incurring in the bracket depth limit of clang
+      int expander[] = {(SetEmptyBranchesHelper(inputTree, outputTree, outputBranches, fInputBranchNames[S],
+                                                fOutputBranchNames[S], typeid(ColTypes), fOptions.fBasketSize),
+                         0)...,
+                        0};
+      (void)expander;
+   }
+
    void Initialize()
    {
       fOutputFile.reset(
@@ -1648,6 +1664,12 @@ public:
       assert(fOutputTree != nullptr);
       assert(fOutputFile != nullptr);
 
+      // There were no entries to fill the TTree with (either the input TTree was empty or no event passed after
+      // filtering). We have already created an empty TTree, now also create the branches to preserve the schema
+      if (fOutputTree->GetEntries() == 0) {
+         using ind_t = std::index_sequence_for<ColTypes...>;
+         SetEmptyBranches(fInputTree, *fOutputTree, ind_t{});
+      }
       // use AutoSave to flush TTree contents because TTree::Write writes in gDirectory, not in fDirectory
       fOutputTree->AutoSave("flushbaskets");
       // must destroy the TTree first, otherwise TFile will delete it too leading to a double delete
@@ -1715,6 +1737,7 @@ class R__CLING_PTRCHECK(off) SnapshotTTreeHelperMT : public RActionImpl<Snapshot
    std::vector<bool> fIsDefine;
    ROOT::Detail::RDF::RLoopManager *fOutputLoopManager;
    ROOT::Detail::RDF::RLoopManager *fInputLoopManager;
+   TFile *fOutputFile; // Non-owning view on the output file
 
 public:
    using ColumnTypes_t = TypeList<ColTypes...>;
@@ -1847,39 +1870,67 @@ public:
       (void)expander; // avoid unused parameter warnings (gcc 12.1)
    }
 
+   template <std::size_t... S>
+   void SetEmptyBranches(TTree *inputTree, TTree &outputTree, std::index_sequence<S...>)
+   {
+      RBranchSet outputBranches{};
+      // We use the expander trick rather than a fold expression to avoid incurring in the bracket depth limit of clang
+      int expander[] = {(SetEmptyBranchesHelper(inputTree, outputTree, outputBranches, fInputBranchNames[S],
+                                                fOutputBranchNames[S], typeid(ColTypes), fOptions.fBasketSize),
+                         0)...,
+                        0};
+      (void)expander;
+   }
+
    void Initialize()
    {
       const auto cs = ROOT::CompressionSettings(fOptions.fCompressionAlgorithm, fOptions.fCompressionLevel);
-      auto out_file = TFile::Open(fFileName.c_str(), fOptions.fMode.c_str(), /*ftitle=*/fFileName.c_str(), cs);
-      if(!out_file)
+      auto outFile = std::unique_ptr<TFile>{
+         TFile::Open(fFileName.c_str(), fOptions.fMode.c_str(), /*ftitle=*/fFileName.c_str(), cs)};
+      if (!outFile)
          throw std::runtime_error("Snapshot: could not create output file " + fFileName);
-      fMerger = std::make_unique<ROOT::TBufferMerger>(std::unique_ptr<TFile>(out_file));
+      fOutputFile = outFile.get();
+      fMerger = std::make_unique<ROOT::TBufferMerger>(std::move(outFile));
    }
 
    void Finalize()
    {
       assert(std::any_of(fOutputFiles.begin(), fOutputFiles.end(), [](const auto &ptr) { return ptr != nullptr; }));
 
-      auto fileWritten = false;
       for (auto &file : fOutputFiles) {
          if (file) {
             file->Write();
             file->Close();
-            fileWritten = true;
          }
       }
 
-      if (!fileWritten) {
-         Warning("Snapshot",
-                 "No input entries (input TTree was empty or no entry passed the Filters). Output TTree is empty.");
+      // If there were no entries to fill the TTree with (either the input TTree was empty or no event passed after
+      // filtering), create an empty TTree in the output file and create the branches to preserve the schema
+      auto fullTreeName = fDirName.empty() ? fTreeName : fDirName + '/' + fTreeName;
+      assert(fOutputFile && "Missing output file in Snapshot finalization.");
+      if (!fOutputFile->Get(fullTreeName.c_str())) {
+
+         // First find in which directory we need to write the output TTree
+         TDirectory *treeDirectory = fOutputFile;
+         if (!fDirName.empty()) {
+            treeDirectory = fOutputFile->mkdir(fDirName.c_str(), "", true);
+         }
+         ::TDirectory::TContext c{treeDirectory};
+
+         // Create the output TTree and create the user-requested branches
+         auto outTree =
+            std::make_unique<TTree>(fTreeName.c_str(), fTreeName.c_str(), fOptions.fSplitLevel, /*dir=*/treeDirectory);
+         using ind_t = std::index_sequence_for<ColTypes...>;
+         SetEmptyBranches(fInputLoopManager->GetTree(), *outTree, ind_t{});
+
+         fOutputFile->Write();
       }
 
       // flush all buffers to disk by destroying the TBufferMerger
       fOutputFiles.clear();
       fMerger.reset();
 
       // Now connect the data source to the loop manager so it can be used for further processing
-      auto fullTreeName = fDirName.empty() ? fTreeName : fDirName + '/' + fTreeName;
       fOutputLoopManager->SetDataSource(std::make_unique<ROOT::Internal::RDF::RTTreeDS>(fullTreeName, fFileName));
    }
 

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

@@ -1270,14 +1270,13 @@ public:
    /// sub-directory `subdir` of file `f.root` (creating file and sub-directory as needed).
    ///
    /// \attention In multi-thread runs (i.e. when EnableImplicitMT() has been called) threads will loop over clusters of
-   /// entries in an undefined order, so Snapshot will produce outputs in which (clusters of) entries will be shuffled with
-   /// respect to the input TTree. Using such "shuffled" TTrees as friends of the original trees would result in wrong
-   /// associations between entries in the main TTree and entries in the "shuffled" friend. Since v6.22, ROOT will
+   /// entries in an undefined order, so Snapshot will produce outputs in which (clusters of) entries will be shuffled
+   /// with respect to the input TTree. Using such "shuffled" TTrees as friends of the original trees would result in
+   /// wrong associations between entries in the main TTree and entries in the "shuffled" friend. Since v6.22, ROOT will
    /// error out if such a "shuffled" TTree is used in a friendship.
    ///
-   /// \note In case no events are written out (e.g. because no event passes all filters) the behavior of Snapshot in
-   /// single-thread and multi-thread runs is different: in single-thread runs, Snapshot will write out a TTree with
-   /// the specified name and zero entries; in multi-thread runs, no TTree object will be written out to disk.
+   /// \note In case no events are written out (e.g. because no event passes all filters), Snapshot will still write the
+   /// requested output TTree to the file, with all the branches requested to preserve the dataset schema.
    ///
    /// \note Snapshot will refuse to process columns with names of the form `#columnname`. These are special columns
    /// made available by some data sources (e.g. RNTupleDS) that represent the size of column `columnname`, and are

tree/dataframe/src/RDFActionHelpers.cxx

@@ -295,3 +295,100 @@ void EnsureValidSnapshotRNTupleOutput(const RSnapshotOptions &opts, const std::s
 } // end NS RDF
 } // end NS Internal
 } // end NS ROOT
+
+namespace {
+void CreateCStyleArrayBranch(TTree *inputTree, TTree &outputTree, ROOT::Internal::RDF::RBranchSet &outputBranches,
+                             const std::string &inputBranchName, const std::string &outputBranchName, int basketSize)
+{
+   TBranch *inputBranch = nullptr;
+   if (inputTree) {
+      inputBranch = inputTree->GetBranch(inputBranchName.c_str());
+      if (!inputBranch) // try harder
+         inputBranch = inputTree->FindBranch(inputBranchName.c_str());
+   }
+   if (!inputBranch)
+      return;
+   const auto STLKind = TClassEdit::IsSTLCont(inputBranch->GetClassName());
+   if (STLKind == ROOT::ESTLType::kSTLvector || STLKind == ROOT::ESTLType::kROOTRVec)
+      return;
+   // must construct the leaflist for the output branch and create the branch in the output tree
+   const auto *leaf = static_cast<TLeaf *>(inputBranch->GetListOfLeaves()->UncheckedAt(0));
+   if (!leaf)
+      return;
+   const auto bname = leaf->GetName();
+   auto *sizeLeaf = leaf->GetLeafCount();
+   const auto sizeLeafName = sizeLeaf ? std::string(sizeLeaf->GetName()) : std::to_string(leaf->GetLenStatic());
+
+   // We proceed only if branch is a fixed-or-variable-sized array
+   if (sizeLeaf || leaf->GetLenStatic() > 1) {
+      if (sizeLeaf && !outputBranches.Get(sizeLeafName)) {
+         // The output array branch `bname` has dynamic size stored in leaf `sizeLeafName`, but that leaf has not been
+         // added to the output tree yet. However, the size leaf has to be available for the creation of the array
+         // branch to be successful. So we create the size leaf here.
+         const auto sizeTypeStr = ROOT::Internal::RDF::TypeName2ROOTTypeName(sizeLeaf->GetTypeName());
+         // Use Original basket size for Existing Branches otherwise use Custom basket Size.
+         const auto bufSize = (basketSize > 0) ? basketSize : sizeLeaf->GetBranch()->GetBasketSize();
+         auto *outputBranch = outputTree.Branch(sizeLeafName.c_str(), static_cast<void *>(nullptr),
+                                                (sizeLeafName + '/' + sizeTypeStr).c_str(), bufSize);
+         outputBranches.Insert(sizeLeafName, outputBranch);
+      }
+
+      const auto btype = leaf->GetTypeName();
+      const auto rootbtype = ROOT::Internal::RDF::TypeName2ROOTTypeName(btype);
+      if (rootbtype == ' ') {
+         Warning("Snapshot",
+                 "RDataFrame::Snapshot: could not correctly construct a leaflist for C-style array in column %s. The "
+                 "leaf is of type '%s'. This column will not be written out.",
+                 bname, btype);
+         return;
+      }
+
+      const auto leaflist = std::string(bname) + "[" + sizeLeafName + "]/" + rootbtype;
+      // Use original basket size for existing branches and new basket size for new branches
+      const auto bufSize = (basketSize > 0) ? basketSize : inputBranch->GetBasketSize();
+      auto *outputBranch =
+         outputTree.Branch(outputBranchName.c_str(), static_cast<void *>(nullptr), leaflist.c_str(), bufSize);
+      outputBranch->SetTitle(inputBranch->GetTitle());
+      outputBranches.Insert(outputBranchName, outputBranch);
+   }
+}
+} // namespace
+
+void ROOT::Internal::RDF::SetEmptyBranchesHelper(TTree *inputTree, TTree &outputTree,
+                                                 ROOT::Internal::RDF::RBranchSet &outputBranches,
+                                                 const std::string &inputBranchName,
+                                                 const std::string &outputBranchName, const std::type_info &typeInfo,
+                                                 int basketSize)
+{
+   const auto bufSize = (basketSize > 0) ? basketSize : 32000;
+   auto *classPtr = TClass::GetClass(typeInfo);
+   if (!classPtr) {
+      // Case of a leaflist of fundamental type, logic taken from
+      // TTree::BranchImpRef(const char* branchname, TClass* ptrClass, EDataType datatype, void* addobj, Int_t bufsize,
+      // Int_t splitlevel)
+      auto typeName = ROOT::Internal::RDF::TypeID2TypeName(typeInfo);
+      auto rootTypeChar = ROOT::Internal::RDF::TypeName2ROOTTypeName(typeName);
+      if (rootTypeChar == ' ') {
+         Warning(
+            "Snapshot",
+            "RDataFrame::Snapshot: could not correctly construct a leaflist for fundamental type in column %s. This "
+            "column will not be written out.",
+            outputBranchName.c_str());
+         return;
+      }
+      std::string leafList{outputBranchName + '/' + rootTypeChar};
+      auto *outputBranch =
+         outputTree.Branch(outputBranchName.c_str(), static_cast<void *>(nullptr), leafList.c_str(), bufSize);
+      outputBranches.Insert(outputBranchName, outputBranch);
+      return;
+   }
+
+   // Find if there is an input branch, check for cases where we need a leaflist (e.g. C-style arrays)
+   CreateCStyleArrayBranch(inputTree, outputTree, outputBranches, inputBranchName, outputBranchName, basketSize);
+
+   // General case
+   if (!outputBranches.Get(outputBranchName)) {
+      auto *outputBranch = outputTree.Branch(outputBranchName.c_str(), classPtr->GetName(), nullptr, bufSize);
+      outputBranches.Insert(outputBranchName, outputBranch);
+   }
+}

tree/dataframe/test/CMakeLists.txt

@@ -48,7 +48,14 @@ if(MSVC)
 endif()
 
 if(NOT (MSVC OR (APPLE AND CMAKE_SYSTEM_PROCESSOR MATCHES arm64)) OR win_broken_tests OR M1_BROKEN_TESTS)
-  ROOT_ADD_GTEST(dataframe_snapshot dataframe_snapshot.cxx LIBRARIES ROOTDataFrame)
+  ROOT_ADD_GTEST(dataframe_snapshot dataframe_snapshot.cxx LIBRARIES ROOTDataFrame GenVector)
+endif()
+
+if(NOT MSVC OR win_broken_tests)
+ROOT_ADD_GTEST(dataframe_snapshot_emptyoutput dataframe_snapshot_emptyoutput.cxx LIBRARIES ROOTDataFrame GenVector)
+ROOT_GENERATE_DICTIONARY(DummyDict ${CMAKE_CURRENT_SOURCE_DIR}/DummyHeader.hxx
+                         MODULE dataframe_snapshot_emptyoutput LINKDEF DummyHeaderLinkDef.hxx OPTIONS -inlineInputHeader
+                         DEPENDENCIES ROOTVecOps GenVector)
 endif()
 
 ROOT_ADD_GTEST(dataframe_datasetspec dataframe_datasetspec.cxx LIBRARIES ROOTDataFrame)

tree/dataframe/test/DummyHeader.hxx

@@ -0,0 +1,7 @@
+#ifndef RDF_TEST_DUMMYHEADER
+#define RDF_TEST_DUMMYHEADER
+
+#include <ROOT/RVec.hxx>
+#include <Math/Vector4D.h>
+
+#endif // RDF_TEST_DUMMYHEADER

tree/dataframe/test/DummyHeaderLinkDef.hxx

@@ -0,0 +1,5 @@
+#ifdef __CLING__
+
+#pragma link C++ class ROOT::RVec < ROOT::Math::PtEtaPhiMVector> + ;
+
+#endif

tree/dataframe/test/dataframe_snapshot.cxx

@@ -15,6 +15,8 @@
 #include <TTreeReader.h>
 #include <TTreeReaderArray.h>
 
+#include "DummyHeader.hxx"
+
 using namespace ROOT;              // RDataFrame
 using namespace ROOT::RDF;         // RInterface
 using namespace ROOT::VecOps;      // RVec