[df] Fix interaction between RVecDS and Snapshot

This commit updates the implementation of the column reader retrieval in
RVecDS to use GetColumnReaders with type_info instead of the older
version. This in turn allows Snapshot to request the correct column
reader when reading a Numpy-based dataset.

Author: vepadulano

bindings/pyroot/pythonizations/test/rdataframe_misc.py

@@ -5,6 +5,8 @@
 
 import ROOT
 
+import numpy
+
 
 class DatasetContext:
     """A helper class to create the dataset for the tutorial below."""
@@ -129,7 +131,30 @@ def test_ttree_ownership(self):
 
             self.assertEqual(rdf.Count().GetValue(), 9)
 
-
+    def test_regression_gh_20291(self):
+        """
+        Regression test for https://github.com/root-project/root/issues/20291
+        """
+        # Issues on Windows with contention on file deletion
+        if platform.system() == "Windows":
+            return
+        out_path = "dataframe_misc_regression_gh20291.root"
+        try:
+            x, y = numpy.array([1, 2, 3]), numpy.array([4, 5, 6])
+            df = ROOT.RDF.FromNumpy({"x": x, "y": y})
+
+            df.Snapshot("tree", out_path)
+
+            df_out = ROOT.RDataFrame("tree", out_path)
+            count = df_out.Count()
+            take_x = df_out.Take["Long64_t"]("x")
+            take_y = df_out.Take["Long64_t"]("y")
+
+            self.assertEqual(count.GetValue(), 3)
+            self.assertSequenceEqual(take_x.GetValue(), [1, 2, 3])
+            self.assertSequenceEqual(take_y.GetValue(), [4, 5, 6])
+        finally:
+            os.remove(out_path)
 
 
 if __name__ == '__main__':

tree/dataframe/inc/ROOT/RVecDS.hxx

@@ -32,6 +32,14 @@ namespace Internal {
 
 namespace RDF {
 
+class R__CLING_PTRCHECK(off) RVecDSColumnReader final : public ROOT::Detail::RDF::RColumnReaderBase {
+   TPointerHolder *fPtrHolder;
+   void *GetImpl(Long64_t) final { return fPtrHolder->GetPointer(); }
+
+public:
+   RVecDSColumnReader(TPointerHolder *ptrHolder) : fPtrHolder(ptrHolder) {}
+};
+
 ////////////////////////////////////////////////////////////////////////////////////////////////
 /// \brief A RDataSource implementation which takes a collection of RVecs, which
 /// are able to adopt data from Numpy arrays
@@ -46,46 +54,18 @@ class RVecDS final : public ROOT::RDF::RDataSource {
    using PointerHolderPtrs_t = std::vector<ROOT::Internal::RDF::TPointerHolder *>;
 
    std::tuple<ROOT::RVec<ColumnTypes>...> fColumns;
-   const std::vector<std::string> fColNames;
-   const std::map<std::string, std::string> fColTypesMap;
+   std::vector<std::string> fColNames;
+   std::unordered_map<std::string, std::string> fColTypesMap;
    // The role of the fPointerHoldersModels is to be initialised with the pack
    // of arguments in the constrcutor signature at construction time
    // Once the number of slots is known, the fPointerHolders are initialised
    // according to the models.
-   const PointerHolderPtrs_t fPointerHoldersModels;
+   PointerHolderPtrs_t fPointerHoldersModels;
    std::vector<PointerHolderPtrs_t> fPointerHolders;
    std::vector<std::pair<ULong64_t, ULong64_t>> fEntryRanges{};
    std::function<void()> fDeleteRVecs;
 
-   Record_t GetColumnReadersImpl(std::string_view colName, const std::type_info &id)
-   {
-      auto colNameStr = std::string(colName);
-      // This could be optimised and done statically
-      const auto idName = ROOT::Internal::RDF::TypeID2TypeName(id);
-      auto it = fColTypesMap.find(colNameStr);
-      if (fColTypesMap.end() == it) {
-         std::string err = "The specified column name, \"" + colNameStr + "\" is not known to the data source.";
-         throw std::runtime_error(err);
-      }
-
-      const auto colIdName = it->second;
-      if (colIdName != idName) {
-         std::string err = "Column " + colNameStr + " has type " + colIdName +
-                           " while the id specified is associated to type " + idName;
-         throw std::runtime_error(err);
-      }
-
-      const auto colBegin = fColNames.begin();
-      const auto colEnd = fColNames.end();
-      const auto namesIt = std::find(colBegin, colEnd, colName);
-      const auto index = std::distance(colBegin, namesIt);
-
-      Record_t ret(fNSlots);
-      for (auto slot : ROOT::TSeqU(fNSlots)) {
-         ret[slot] = fPointerHolders[index][slot]->GetPointerAddr();
-      }
-      return ret;
-   }
+   Record_t GetColumnReadersImpl(std::string_view, const std::type_info &) { return {}; }
 
    size_t GetEntriesNumber() { return std::get<0>(fColumns).size(); }
    template <std::size_t... S>
@@ -146,6 +126,33 @@ public:
       fDeleteRVecs();
    }
 
+   std::unique_ptr<ROOT::Detail::RDF::RColumnReaderBase>
+   GetColumnReaders(unsigned int slot, std::string_view colName, const std::type_info &id) final
+   {
+      auto colNameStr = std::string(colName);
+
+      auto it = fColTypesMap.find(colNameStr);
+      if (fColTypesMap.end() == it) {
+         std::string err = "The specified column name, \"" + colNameStr + "\" is not known to the data source.";
+         throw std::runtime_error(err);
+      }
+
+      const auto &colIdName = it->second;
+      const auto idName = ROOT::Internal::RDF::TypeID2TypeName(id);
+      if (colIdName != idName) {
+         std::string err = "Column " + colNameStr + " has type " + colIdName +
+                           " while the id specified is associated to type " + idName;
+         throw std::runtime_error(err);
+      }
+
+      if (auto colNameIt = std::find(fColNames.begin(), fColNames.end(), colNameStr); colNameIt != fColNames.end()) {
+         const auto index = std::distance(fColNames.begin(), colNameIt);
+         return std::make_unique<ROOT::Internal::RDF::RVecDSColumnReader>(fPointerHolders[index][slot]);
+      }
+
+      throw std::runtime_error("Could not find column name \"" + colNameStr + "\" in available column names.");
+   }
+
    const std::vector<std::string> &GetColumnNames() const { return fColNames; }
 
    std::vector<std::pair<ULong64_t, ULong64_t>> GetEntryRanges()