[df] Refresh input TTree when changing spec

In-between distributed tasks the input data source of the RDataFrame is
changed (via RLoopManager::ChangeSpec), thus invalidating the input
TTree from the previous task. Ensure the Snapshot action helper is
always aware of the latest, currently valid TTree by retrieving it from
the input RLoopManager directly.

This fixes the failures in the graph caching test. The test was also
extended to check that the output snapshot from the distributed tasks
also contains the right data as it was read from the input TTree.

Author: vepadulano

bindings/experimental/distrdf/test/backend/test_graph_caching.py

@@ -1,11 +1,12 @@
-
 import os
 import unittest
+from array import array
 
+import ROOT
+from DistRDF._graph_cache import _ACTIONS_REGISTER, _RDF_REGISTER
 from DistRDF.Backends import Base
 from DistRDF.DataFrame import RDataFrame
 from DistRDF.HeadNode import get_headnode
-from DistRDF._graph_cache import _ACTIONS_REGISTER, _RDF_REGISTER
 
 
 def clear_caches():
@@ -110,42 +111,74 @@ def test_count_tchain(self):
                 # The RDataFrame should have run as many times as partitions
                 self.assertEqual(cached_rdf.GetNRuns(), npartitions)
 
-    def test_snapshot(self):
-        """The cache is used to Snapshot data."""
-        treename = "myTree"
-        filenames = ["4clusters.root"] * 5
-        nentries = 5000
-        backend = GraphCaching.TestBackend()
-
+    def _test_snapshot_impl(self, backend, treename, filenames, nentries):
         for npartitions in [1, 2, 4, 8, 16]:
             # Start from a fresh cache at each subtest iteration
             clear_caches()
 
             output_basename = "test_graph_caching_test_snapshot"
             output_filenames = [f"{output_basename}_{i}.root" for i in range(npartitions)]
+            try:
+                with self.subTest(npartitions=npartitions):
+                    headnode = get_headnode(backend, npartitions, treename, filenames)
+                    distrdf = RDataFrame(headnode)
+
+                    output_branch = "b1"
+                    snapdf = distrdf.Snapshot(treename, f"{output_basename}.root", (output_branch,))
+
+                    # There should be exactly one cached RDF and set of actions
+                    self.assertEqual(len(_RDF_REGISTER), 1)
+                    self.assertEqual(len(_RDF_REGISTER), len(_ACTIONS_REGISTER))
+                    cached_rdf = tuple(_RDF_REGISTER.values())[0]
+                    # The RDataFrame should have run as many times as partitions
+                    self.assertEqual(cached_rdf.GetNRuns(), npartitions)
+                    # All the correct output files should be present
+                    for output_filename in output_filenames:
+                        self.assertTrue(os.path.exists(output_filename))
+                    # Make sure we have the correct data in the snapshot output
+                    self.assertListEqual(list(snapdf.AsNumpy()[output_branch]), list(range(500)))
+            finally:
+                # Remove output files at each iteration
+                for output_filename in output_filenames:
+                    try:
+                        os.remove(output_filename)
+                    except OSError:
+                        pass
 
-            with self.subTest(npartitions=npartitions):
-                headnode = get_headnode(backend, npartitions, treename, filenames)
-                distrdf = RDataFrame(headnode)
-
-                output_branches = ("b1", )
-                snapdf = distrdf.Snapshot(treename, f"{output_basename}.root", ["b1", ])
+    def test_snapshot(self):
+        """The cache is used to Snapshot data."""
 
-                # There should be exactly one cached RDF and set of actions
-                self.assertEqual(len(_RDF_REGISTER), 1)
-                self.assertEqual(len(_RDF_REGISTER), len(_ACTIONS_REGISTER))
-                cached_rdf = tuple(_RDF_REGISTER.values())[0]
-                # The RDataFrame should have run as many times as partitions
-                self.assertEqual(cached_rdf.GetNRuns(), npartitions)
-                # All the correct output files should be present
-                for output_filename in output_filenames:
-                    self.assertTrue(os.path.exists(output_filename))
-                # The snapshotted dataframe should be usable
-                self.assertEqual(snapdf.Count().GetValue(), nentries)
+        def write_data(treenames, filenames):
+            # Create a dataset with 100 entries per file with sequential values
+            b1 = array("i", [0])
+            for i in range(1, len(filenames) + 1):
+                with ROOT.TFile.Open(filenames[i - 1], "recreate") as f:
+                    t = ROOT.TTree(treenames[i - 1], treenames[i - 1])
+                    t.Branch("b1", b1, "b1/I")
+                    nentries = 0
+                    for val in range((i - 1) * 100, i * 100):
+                        b1[0] = val
+                        t.Fill()
+                        nentries += 1
+                        # 5 clusters per file
+                        if nentries % 20 == 0:
+                            t.FlushBaskets()
+                    f.Write()
+
+        treename = "Events"
+        filenames = [f"test_graph_caching_test_snapshot_input_{i}.root" for i in range(1, 6)]
+        nentries = 100 * len(filenames)
+        backend = GraphCaching.TestBackend()
 
-            # Remove output files at each iteration
-            for output_filename in output_filenames:
-                os.remove(output_filename)
+        try:
+            write_data([treename] * len(filenames), filenames)
+            self._test_snapshot_impl(backend, treename, filenames, nentries)
+        finally:
+            for fn in filenames:
+                try:
+                    os.remove(fn)
+                except OSError:
+                    pass
 
     def test_multiple_graphs(self):
         """The caches are used with multiple executions."""

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

@@ -1528,14 +1528,14 @@ class R__CLING_PTRCHECK(off) SnapshotTTreeHelper : public RActionImpl<SnapshotTT
    RBranchSet fOutputBranches;
    std::vector<bool> fIsDefine;
    ROOT::Detail::RDF::RLoopManager *fOutputLoopManager;
-   ROOT::RDF::RDataSource *fInputDataSource;
+   ROOT::Detail::RDF::RLoopManager *fInputLoopManager;
 
 public:
    using ColumnTypes_t = TypeList<ColTypes...>;
    SnapshotTTreeHelper(std::string_view filename, std::string_view dirname, std::string_view treename,
                        const ColumnNames_t &vbnames, const ColumnNames_t &bnames, const RSnapshotOptions &options,
                        std::vector<bool> &&isDefine, ROOT::Detail::RDF::RLoopManager *loopManager,
-                       ROOT::RDF::RDataSource *inputDataSource)
+                       ROOT::Detail::RDF::RLoopManager *inputLM)
       : fFileName(filename),
         fDirName(dirname),
         fTreeName(treename),
@@ -1546,7 +1546,7 @@ public:
         fBranchAddresses(vbnames.size(), nullptr),
         fIsDefine(std::move(isDefine)),
         fOutputLoopManager(loopManager),
-        fInputDataSource(inputDataSource)
+        fInputLoopManager(inputLM)
    {
       EnsureValidSnapshotTTreeOutput(fOptions, fTreeName, fFileName);
    }
@@ -1569,12 +1569,11 @@ public:
       }
    }
 
-   void InitTask(TTreeReader *r, unsigned int /* slot */)
+   void InitTask(TTreeReader * /*treeReader*/, unsigned int /* slot */)
    {
-      if (r)
-         fInputTree = r->GetTree();
-      else if (auto treeDS = dynamic_cast<ROOT::Internal::RDF::RTTreeDS *>(fInputDataSource))
-         fInputTree = treeDS->GetTree();
+      // We ask the input RLoopManager if it has a TTree. We cannot rely on getting this information when constructing
+      // this action helper, since the TTree might change e.g. when ChangeSpec is called in-between distributed tasks.
+      fInputTree = fInputLoopManager->GetTree();
       fBranchAddressesNeedReset = true;
    }
 
@@ -1689,7 +1688,7 @@ public:
                                  fOptions,
                                  std::vector<bool>(fIsDefine),
                                  fOutputLoopManager,
-                                 fInputDataSource};
+                                 fInputLoopManager};
    }
 };
 
@@ -1715,15 +1714,15 @@ class R__CLING_PTRCHECK(off) SnapshotTTreeHelperMT : public RActionImpl<Snapshot
    std::vector<RBranchSet> fOutputBranches;
    std::vector<bool> fIsDefine;
    ROOT::Detail::RDF::RLoopManager *fOutputLoopManager;
-   ROOT::RDF::RDataSource *fInputDataSource;
+   ROOT::Detail::RDF::RLoopManager *fInputLoopManager;
 
 public:
    using ColumnTypes_t = TypeList<ColTypes...>;
 
    SnapshotTTreeHelperMT(const unsigned int nSlots, std::string_view filename, std::string_view dirname,
                          std::string_view treename, const ColumnNames_t &vbnames, const ColumnNames_t &bnames,
                          const RSnapshotOptions &options, std::vector<bool> &&isDefine,
-                         ROOT::Detail::RDF::RLoopManager *loopManager, ROOT::RDF::RDataSource *inputDataSource)
+                         ROOT::Detail::RDF::RLoopManager *loopManager, ROOT::Detail::RDF::RLoopManager *inputLM)
       : fNSlots(nSlots),
         fOutputFiles(fNSlots),
         fOutputTrees(fNSlots),
@@ -1740,7 +1739,7 @@ public:
         fOutputBranches(fNSlots),
         fIsDefine(std::move(isDefine)),
         fOutputLoopManager(loopManager),
-        fInputDataSource(inputDataSource)
+        fInputLoopManager(inputLM)
    {
       EnsureValidSnapshotTTreeOutput(fOptions, fTreeName, fFileName);
    }
@@ -1785,11 +1784,11 @@ public:
       if (fOptions.fAutoFlush)
          fOutputTrees[slot]->SetAutoFlush(fOptions.fAutoFlush);
       if (r) {
-         // not an empty-source RDF
+         // We could be getting a task-local TTreeReader from the TTreeProcessorMT.
          fInputTrees[slot] = r->GetTree();
-      } else if (auto treeDS = dynamic_cast<ROOT::Internal::RDF::RTTreeDS *>(fInputDataSource))
-         fInputTrees[slot] = treeDS->GetTree();
-
+      } else {
+         fInputTrees[slot] = fInputLoopManager->GetTree();
+      }
       fBranchAddressesNeedReset[slot] = 1; // reset first event flag for this slot
    }
 
@@ -1914,7 +1913,7 @@ public:
                                    fOptions,
                                    std::vector<bool>(fIsDefine),
                                    fOutputLoopManager,
-                                   fInputDataSource};
+                                   fInputLoopManager};
    }
 };
 

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

@@ -250,8 +250,8 @@ struct SnapshotHelperArgs {
    std::string fTreeName;
    std::vector<std::string> fOutputColNames;
    ROOT::RDF::RSnapshotOptions fOptions;
-   ROOT::Detail::RDF::RLoopManager *fLoopManager;
-   ROOT::RDF::RDataSource *fDataSource;
+   ROOT::Detail::RDF::RLoopManager *fOutputLoopManager;
+   ROOT::Detail::RDF::RLoopManager *fInputLoopManager;
    bool fToNTuple;
 };
 
@@ -267,8 +267,8 @@ 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->fLoopManager;
-   const auto &dataSource = snapHelperArgs->fDataSource;
+   const auto &lmPtr = snapHelperArgs->fOutputLoopManager;
+   const auto &inputLM = snapHelperArgs->fInputLoopManager;
 
    auto sz = sizeof...(ColTypes);
    std::vector<bool> isDefine(sz);
@@ -304,14 +304,14 @@ BuildAction(const ColumnNames_t &colNames, const std::shared_ptr<SnapshotHelperA
          using Helper_t = SnapshotTTreeHelper<ColTypes...>;
          using Action_t = RAction<Helper_t, PrevNodeType>;
          actionPtr.reset(new Action_t(Helper_t(filename, dirname, treename, colNames, outputColNames, options,
-                                               std::move(isDefine), lmPtr, dataSource),
+                                               std::move(isDefine), lmPtr, inputLM),
                                       colNames, prevNode, colRegister));
       } else {
          // multi-thread snapshot
          using Helper_t = SnapshotTTreeHelperMT<ColTypes...>;
          using Action_t = RAction<Helper_t, PrevNodeType>;
          actionPtr.reset(new Action_t(Helper_t(nSlots, filename, dirname, treename, colNames, outputColNames, options,
-                                               std::move(isDefine), lmPtr, dataSource),
+                                               std::move(isDefine), lmPtr, inputLM),
                                       colNames, prevNode, colRegister));
       }
    }

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

@@ -1359,7 +1359,7 @@ public:
 
          auto snapHelperArgs = std::make_shared<RDFInternal::SnapshotHelperArgs>(RDFInternal::SnapshotHelperArgs{
             std::string(filename), std::string(dirname), std::string(treename), colListWithAliasesAndSizeBranches,
-            options, newRDF->GetLoopManager(), GetDataSource(), true /* fToNTuple */});
+            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
@@ -1388,7 +1388,7 @@ public:
 
          auto snapHelperArgs = std::make_shared<RDFInternal::SnapshotHelperArgs>(RDFInternal::SnapshotHelperArgs{
             std::string(filename), std::string(dirname), std::string(treename), colListWithAliasesAndSizeBranches,
-            options, newRDF->GetLoopManager(), GetDataSource(), false /* fToRNTuple */});
+            options, newRDF->GetLoopManager(), GetLoopManager(), false /* fToRNTuple */});
 
          resPtr = CreateAction<RDFInternal::ActionTags::Snapshot, RDFDetail::RInferredType>(
             colListNoAliasesWithSizeBranches, newRDF, snapHelperArgs, fProxiedPtr,
@@ -3247,7 +3247,7 @@ private:
 
          auto snapHelperArgs = std::make_shared<RDFInternal::SnapshotHelperArgs>(RDFInternal::SnapshotHelperArgs{
             std::string(filename), std::string(dirname), std::string(treename), columnListWithoutSizeColumns, options,
-            newRDF->GetLoopManager(), GetDataSource(), true /* fToRNTuple */});
+            newRDF->GetLoopManager(), GetLoopManager(), true /* fToRNTuple */});
 
          // The Snapshot helper will use validCols (with aliases resolved) as input columns, and
          // columnListWithoutSizeColumns (still with aliases in it, passed through snapHelperArgs) as output column
@@ -3275,7 +3275,7 @@ private:
 
          auto snapHelperArgs = std::make_shared<RDFInternal::SnapshotHelperArgs>(RDFInternal::SnapshotHelperArgs{
             std::string(filename), std::string(dirname), std::string(treename), columnListWithoutSizeColumns, options,
-            newRDF->GetLoopManager(), GetDataSource(), false /* fToRNTuple */});
+            newRDF->GetLoopManager(), GetLoopManager(), false /* fToRNTuple */});
 
          // The Snapshot helper will use validCols (with aliases resolved) as input columns, and
          // columnListWithoutSizeColumns (still with aliases in it, passed through snapHelperArgs) as output column

tree/dataframe/src/RLoopManager.cxx

@@ -1070,7 +1070,13 @@ const ColumnNames_t &RLoopManager::GetDefaultColumnNames() const
 
 TTree *RLoopManager::GetTree() const
 {
-   return fTree.get();
+   // This is currently called in SnapshotTTreeHelper[MT] to retrieve the task-local input TTree. It is not guaranteed
+   // that the same RLoopManager will always have the same input TTree for its entire lifetime, notably it could be
+   // changed by ChangeSpec when moving to a different entry range.
+   if (auto *treeDS = dynamic_cast<ROOT::Internal::RDF::RTTreeDS *>(fDataSource.get())) {
+      return treeDS->GetTree();
+   }
+   return nullptr;
 }
 
 void RLoopManager::Register(RDFInternal::RActionBase *actionPtr)