Protype TTreeReaderUntypedArray

Author: vepadulano

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

@@ -14,6 +14,7 @@
 
 #include "RColumnReaderBase.hxx"
 #include <ROOT/RVec.hxx>
+#include "ROOT/RDF/Utils.hxx"
 #include <Rtypes.h>  // Long64_t, R__CLING_PTRCHECK
 #include <TTreeReader.h>
 #include <TTreeReaderValue.h>
@@ -22,6 +23,7 @@
 #include <array>
 #include <memory>
 #include <string>
+#include <cstddef>
 
 namespace ROOT {
 namespace Internal {
@@ -59,10 +61,12 @@ public:
 /// TTreeReaderArrays are used whenever the RDF column type is RVec<T>.
 template <typename T>
 class R__CLING_PTRCHECK(off) RTreeColumnReader<RVec<T>> final : public ROOT::Detail::RDF::RColumnReaderBase {
-   std::unique_ptr<TTreeReaderArray<T>> fTreeArray;
+   std::unique_ptr<TTreeReaderUntypedArray> fTreeArray;
+
+   using Value_t = std::byte;
 
    /// We return a reference to this RVec to clients, to guarantee a stable address and contiguous memory layout.
-   RVec<T> fRVec;
+   RVec<Value_t> fRVec;
 
    Long64_t fLastEntry = -1;
 
@@ -86,11 +90,10 @@ class R__CLING_PTRCHECK(off) RTreeColumnReader<RVec<T>> final : public ROOT::Det
             // trigger loading of the contents of the TTreeReaderArray
             // the address of the first element in the reader array is not necessarily equal to
             // the address returned by the GetAddress method
-            auto readerArrayAddr = &readerArray.At(0);
-            RVec<T> rvec(readerArrayAddr, readerArraySize);
+            RVec<Value_t> rvec(readerArray.At(0), readerArraySize);
             swap(fRVec, rvec);
          } else {
-            RVec<T> emptyVec{};
+            RVec<Value_t> emptyVec{};
             swap(fRVec, emptyVec);
          }
       } else {
@@ -107,10 +110,17 @@ class R__CLING_PTRCHECK(off) RTreeColumnReader<RVec<T>> final : public ROOT::Det
          (void)fCopyWarningPrinted;
 #endif
          if (readerArraySize > 0) {
-            RVec<T> rvec(readerArray.begin(), readerArray.end());
-            swap(fRVec, rvec);
+            // Array is not contiguous, make a full copy of it.
+            fRVec = RVec<Value_t>();
+            fRVec.reserve(readerArraySize * sizeof(T));
+            for (std::size_t i{0}; i < readerArraySize; i++)
+            {
+               auto val = readerArray.At(i);
+               std::copy(val, val + sizeof(T), std::back_inserter(fRVec));
+            }
+            fRVec.resize(readerArraySize);
          } else {
-            RVec<T> emptyVec{};
+            RVec<Value_t> emptyVec{};
             swap(fRVec, emptyVec);
          }
       }
@@ -120,7 +130,7 @@ class R__CLING_PTRCHECK(off) RTreeColumnReader<RVec<T>> final : public ROOT::Det
 
 public:
    RTreeColumnReader(TTreeReader &r, const std::string &colName)
-      : fTreeArray(std::make_unique<TTreeReaderArray<T>>(r, colName.c_str()))
+      : fTreeArray(std::make_unique<TTreeReaderUntypedArray>(r, colName, ROOT::Internal::RDF::TypeID2TypeName(typeid(T))))
    {
    }
 };

tree/treeplayer/inc/TTreeReaderArray.h

@@ -15,6 +15,7 @@
 #include "TTreeReaderValue.h"
 #include "TTreeReaderUtils.h"
 #include <type_traits>
+#include <cstddef>
 
 namespace ROOT {
 namespace Internal {
@@ -37,6 +38,9 @@ class TTreeReaderArrayBase : public TTreeReaderValueBase {
 
    bool IsContiguous() const { return fImpl->IsContiguous(GetProxy()); }
 
+   // Returns the `sizeof` of the collection element type.
+   UInt_t GetValueSize() const { return fImpl ? fImpl->GetValueSize(GetProxy()): 0; }
+
 protected:
    void *UntypedAt(std::size_t idx) const { return fImpl->At(GetProxy(), idx); }
    void CreateProxy() override;
@@ -51,6 +55,22 @@ class TTreeReaderArrayBase : public TTreeReaderValueBase {
    // ClassDefOverride(TTreeReaderArrayBase, 0);//Accessor to member of an object stored in a collection
 };
 
+class R__CLING_PTRCHECK(off) TTreeReaderUntypedArray final : public TTreeReaderArrayBase {
+   std::string fArrayElementTypeName;
+
+public:
+   TTreeReaderUntypedArray(TTreeReader &tr, std::string_view branchName, std::string_view innerTypeName)
+      : TTreeReaderArrayBase(&tr, branchName.data(), TDictionary::GetDictionary(innerTypeName.data())),
+        fArrayElementTypeName(innerTypeName)
+   {
+   }
+
+   std::byte *At(std::size_t idx) const { return reinterpret_cast<std::byte *>(UntypedAt(idx)); }
+
+protected:
+   const char *GetDerivedTypeName() const final { return fArrayElementTypeName.c_str(); }
+};
+
 } // namespace Internal
 } // namespace ROOT
 

tree/treeplayer/inc/TTreeReaderUtils.h

@@ -89,6 +89,7 @@ namespace Internal {
       virtual size_t GetSize(Detail::TBranchProxy*) = 0;
       virtual void* At(Detail::TBranchProxy*, size_t /*idx*/) = 0;
       virtual bool IsContiguous(Detail::TBranchProxy *) = 0;
+      virtual UInt_t GetValueSize(Detail::TBranchProxy *) = 0;
    };
 
 }

tree/treeplayer/src/TTreeReaderArray.cxx

@@ -32,6 +32,7 @@
 
 #include <memory>
 #include <optional>
+#include <iostream>
 
 // pin vtable
 ROOT::Internal::TVirtualCollectionReader::~TVirtualCollectionReader() {}
@@ -72,6 +73,8 @@ class TClonesReader : public TVirtualCollectionReader {
    }
 
    bool IsContiguous(ROOT::Detail::TBranchProxy *) override { return false; }
+
+   UInt_t GetValueSize(ROOT::Detail::TBranchProxy *) override { return sizeof(TObject *); }
 };
 
 bool IsCPContiguous(const TVirtualCollectionProxy &cp)
@@ -86,6 +89,11 @@ bool IsCPContiguous(const TVirtualCollectionProxy &cp)
    }
 }
 
+UInt_t GetCPValueSize(const TVirtualCollectionProxy &cp)
+{
+   return cp.Sizeof();
+}
+
 // Reader interface for STL
 class TSTLReader final : public TVirtualCollectionReader {
 public:
@@ -131,6 +139,12 @@ class TSTLReader final : public TVirtualCollectionReader {
       auto cp = GetCP(proxy);
       return IsCPContiguous(*cp);
    }
+
+   UInt_t GetValueSize(ROOT::Detail::TBranchProxy *proxy) override
+   {
+      auto cp = GetCP(proxy);
+      return GetCPValueSize(*cp);
+   }
 };
 
 class TCollectionLessSTLReader final : public TVirtualCollectionReader {
@@ -190,6 +204,12 @@ class TCollectionLessSTLReader final : public TVirtualCollectionReader {
       auto cp = GetCP(proxy);
       return IsCPContiguous(*cp);
    }
+
+   UInt_t GetValueSize(ROOT::Detail::TBranchProxy *proxy) override
+   {
+      auto cp = GetCP(proxy);
+      return GetCPValueSize(*cp);
+   }
 };
 
 // Reader interface for leaf list
@@ -243,6 +263,12 @@ class TObjectArrayReader : public TVirtualCollectionReader {
    void SetBasicTypeSize(Int_t size) { fBasicTypeSize = size; }
 
    bool IsContiguous(ROOT::Detail::TBranchProxy *) override { return true; }
+
+   UInt_t GetValueSize(ROOT::Detail::TBranchProxy *proxy) override
+   {
+      auto cp = GetCP(proxy);
+      return GetCPValueSize(*cp);
+   }
 };
 
 template <class BASE>
@@ -387,6 +413,12 @@ class TBasicTypeArrayReader final : public TVirtualCollectionReader {
    }
 
    bool IsContiguous(ROOT::Detail::TBranchProxy *) override { return false; }
+
+   UInt_t GetValueSize(ROOT::Detail::TBranchProxy *proxy) override
+   {
+      auto cp = GetCP(proxy);
+      return GetCPValueSize(*cp);
+   }
 };
 
 class TBasicTypeClonesReader final : public TClonesReader {
@@ -434,6 +466,12 @@ class TLeafReader : public TVirtualCollectionReader {
 
    bool IsContiguous(ROOT::Detail::TBranchProxy *) override { return true; }
 
+   UInt_t GetValueSize(ROOT::Detail::TBranchProxy *) override
+   {
+      auto *leaf = fValueReader->GetLeaf();
+      return leaf ? leaf->GetLenType(): 0;
+   }
+
 protected:
    void ProxyRead() { fValueReader->ProxyRead(); }
 };