[df] Prepare for bulk filters

Nodes of the computation graph which may operate a selection on which entries are valid or not are modified to accommodate for multiple entries being evaluated at the same time. For now, the size of the bulk is set to one.

Note that this commit does not touch in any way the reading of data.

Author: vepadulano

roottest/root/dataframe/testIMT.cxx

@@ -55,18 +55,18 @@ void getTracks(unsigned int mu, FourVectors& tracks) {
 // This makes the example stand-alone
 void FillTree(const char* filename, const char* treeName) {
    if (!gSystem->AccessPathName(filename)) return;
-   TFile f(filename,"RECREATE");
-   TTree t(treeName,treeName);
+   auto f = std::make_unique<TFile>(filename, "RECREATE");
+   auto t = std::make_unique<TTree>(treeName, treeName);
    double b1;
    int b2;
    std::vector<FourVector> tracks;
    std::vector<double> dv {-1,2,3,4};
    std::list<int> sl {1,2,3,4};
-   t.Branch("b1", &b1);
-   t.Branch("b2", &b2);
-   t.Branch("tracks", &tracks);
-   t.Branch("dv", &dv);
-   t.Branch("sl", &sl);
+   t->Branch("b1", &b1);
+   t->Branch("b2", &b2);
+   t->Branch("tracks", &tracks);
+   t->Branch("dv", &dv);
+   t->Branch("sl", &sl);
 
    int nevts = 16000;
    for(int i = 0; i < nevts; ++i) {
@@ -77,11 +77,9 @@ void FillTree(const char* filename, const char* treeName) {
 
       dv.emplace_back(i);
       sl.emplace_back(i);
-      t.Fill();
+      t->Fill();
    }
-   t.Write();
-   f.Close();
-   return;
+   f->Write();
 }
 
 auto fileName = "testIMT.root";

tree/dataframe/CMakeLists.txt

@@ -82,6 +82,7 @@ ROOT_STANDARD_LIBRARY_PACKAGE(ROOTDataFrame
     ROOT/RDF/RJittedVariation.hxx
     ROOT/RDF/RLazyDSImpl.hxx
     ROOT/RDF/RLoopManager.hxx
+    ROOT/RDF/RMaskedEntryRange.hxx
     ROOT/RDF/RMergeableValue.hxx
     ROOT/RDF/RMetaData.hxx
     ROOT/RDF/RNodeBase.hxx

tree/dataframe/inc/ROOT/RCsvDS.hxx

@@ -28,7 +28,7 @@ namespace ROOT::Internal::RDF {
 class R__CLING_PTRCHECK(off) RCsvDSColumnReader final : public ROOT::Detail::RDF::RColumnReaderBase {
    void *fValuePtr;
    void *GetImpl(std::size_t) final { return fValuePtr; }
-   void LoadImpl(Long64_t, bool) final {}
+   void LoadImpl(const ROOT::Internal::RDF::RMaskedEntryRange &) final {}
 
 public:
    RCsvDSColumnReader(void *valuePtr) : fValuePtr(valuePtr) {}

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

@@ -122,9 +122,10 @@ public:
    void Run(unsigned int slot, Long64_t entry) final
    {
       const auto mask = fPrevNode.CheckFilters(slot, entry);
-      std::for_each(fValues[slot].begin(), fValues[slot].end(), [entry, mask](auto *v) { v->Load(entry, mask); });
+      std::for_each(fValues[slot].begin(), fValues[slot].end(), [&mask](auto *v) { v->Load(mask); });
 
-      if (mask)
+      // Assume 1-size bulk for now
+      if (mask[0])
          CallExec(slot, /*idx=*/0u, ColumnTypes_t{}, TypeInd_t{});
    }
 

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

@@ -195,28 +195,33 @@ public:
    {
       if constexpr (std::is_same_v<Helper, SnapshotHelperWithVariations>) {
          // check if entry passes all filters
-         std::vector<bool> filterPassed(fPrevNodes.size(), false);
+         std::vector<ROOT::Internal::RDF::RMaskedEntryRange> filterPassed(fPrevNodes.size(), 1ul);
          for (unsigned int variation = 0; variation < fPrevNodes.size(); ++variation) {
             filterPassed[variation] = fPrevNodes[variation]->CheckFilters(slot, entry);
          }
 
          // Currently, every event where any of nominal or variations pass gets written to the output.
          // This logic could be extended for different use cases if the need arises.
-         if (std::any_of(filterPassed.begin(), filterPassed.end(), [](bool val) { return val; })) {
+         // Assume 1-size bulk for now
+         if (std::any_of(filterPassed.begin(), filterPassed.end(),
+                         [](const ROOT::Internal::RDF::RMaskedEntryRange &val) { return val[0]; })) {
             // TODO: Don't allocate
             std::vector<void *> untypedValues;
             auto nReaders = fValues[slot].size();
             untypedValues.reserve(nReaders);
-            std::for_each(fValues[slot].begin(), fValues[slot].end(), [entry](auto *v) { v->Load(entry, true); });
+            std::for_each(fValues[slot].begin(), fValues[slot].end(), [entry](auto *v) {
+               v->Load(ROOT::Internal::RDF::RMaskedEntryRange{1ul, true, static_cast<std::uint64_t>(entry)});
+            });
             for (decltype(nReaders) readerIdx{}; readerIdx < nReaders; readerIdx++)
                untypedValues.push_back(GetValue(slot, readerIdx, /*idx=*/0u));
 
             fHelper.Exec(slot, untypedValues, filterPassed);
          }
       } else {
          const auto mask = fPrevNodes.front()->CheckFilters(slot, entry);
-         std::for_each(fValues[slot].begin(), fValues[slot].end(), [entry, mask](auto *v) { v->Load(entry, mask); });
-         if (mask)
+         std::for_each(fValues[slot].begin(), fValues[slot].end(), [&mask](auto *v) { v->Load(mask); });
+         // Assume 1-size bulk for now
+         if (mask[0])
             CallExec(slot, /*idx=*/0u);
       }
    }

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

@@ -12,6 +12,7 @@
 #define ROOT_INTERNAL_RDF_RCOLUMNREADERBASE
 
 #include <Rtypes.h>
+#include <ROOT/RDF/RMaskedEntryRange.hxx>
 
 namespace ROOT {
 namespace Detail {
@@ -26,22 +27,14 @@ This pure virtual class provides a common base class for the different column re
 RDSColumnReader.
 **/
 class R__CLING_PTRCHECK(off) RColumnReaderBase {
-   Long64_t fLoadedEntry = -1;
 
 public:
    virtual ~RColumnReaderBase() = default;
 
    /// Load the column value for the given entry.
    /// \param entry The entry number to load.
    /// \param mask The entry mask. Values will be loaded only for entries for which the mask equals true.
-   void Load(Long64_t entry, bool mask)
-   {
-      // For now, as `mask` is just a single boolean, as an optimization we can return early here if `mask == false`.
-      if (mask) {
-         fLoadedEntry = entry;
-         this->LoadImpl(entry, mask);
-      }
-   }
+   void Load(const ROOT::Internal::RDF::RMaskedEntryRange &mask) { LoadImpl(mask); }
 
    /// Return the column value for the given entry.
    /// \tparam T The column type
@@ -57,7 +50,7 @@ public:
 
 private:
    virtual void *GetImpl(std::size_t idx) = 0;
-   virtual void LoadImpl(Long64_t /*entry*/, bool /*mask*/) = 0;
+   virtual void LoadImpl(const ROOT::Internal::RDF::RMaskedEntryRange &) = 0;
 };
 
 } // namespace RDF

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

@@ -24,7 +24,7 @@ class R__CLING_PTRCHECK(off) RDSColumnReader final : public ROOT::Detail::RDF::R
    T **fDSValuePtr = nullptr;
 
    void *GetImpl(std::size_t) final { return *fDSValuePtr; }
-   void LoadImpl(Long64_t, bool) final {}
+   void LoadImpl(const ROOT::Internal::RDF::RMaskedEntryRange &) final {}
 
 public:
    RDSColumnReader(void *DSValuePtr) : fDSValuePtr(static_cast<T **>(DSValuePtr)) {}

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

@@ -104,16 +104,20 @@ public:
    }
 
    /// Update the value at the address returned by GetValuePtr with the content corresponding to the given entry
-   void Update(unsigned int slot, Long64_t entry, bool mask) final
+   void Update(unsigned int slot, const ROOT::Internal::RDF::RMaskedEntryRange &mask) final
    {
-      if (entry != fLastCheckedEntry[slot * RDFInternal::CacheLineStep<Long64_t>()]) {
-         // evaluate this define expression, cache the result
-         fValues[slot]->Load(entry, mask);
-         if (mask) {
-            fLastResults[slot * RDFInternal::CacheLineStep<T>()] = GetValueOrDefault(slot, /*idx=*/0u);
-            fLastCheckedEntry[slot * RDFInternal::CacheLineStep<Long64_t>()] = entry;
-         }
+      if (static_cast<Long64_t>(mask.GetFirstEntry()) ==
+          fLastCheckedEntry[slot * RDFInternal::CacheLineStep<Long64_t>()])
+         return;
+
+      // Assume 1-size bulk for now
+      fValues[slot]->Load(mask);
+      const std::size_t bulkSize = 1;
+      for (std::size_t i = 0; i < bulkSize; ++i) {
+         if (mask[i])
+            fLastResults[slot * RDFInternal::CacheLineStep<T>()] = GetValueOrDefault(slot, i);
       }
+      fLastCheckedEntry[slot * RDFInternal::CacheLineStep<Long64_t>()] = mask.GetFirstEntry();
    }
 
    void Update(unsigned int /*slot*/, const ROOT::RDF::RSampleInfo & /*id*/) final {}

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

@@ -83,31 +83,28 @@ class R__CLING_PTRCHECK(off) RDefine final : public RDefineBase {
    }
 
    template <typename... ColTypes, std::size_t... S>
-   void UpdateHelper(unsigned int slot, std::size_t idx, Long64_t /*entry*/, TypeList<ColTypes...>,
+   auto UpdateHelper(unsigned int slot, std::size_t idx, Long64_t /*entry*/, TypeList<ColTypes...>,
                      std::index_sequence<S...>, NoneTag)
    {
-      fLastResults[slot * RDFInternal::CacheLineStep<ret_type>()] =
-         fExpression(GetValueChecked<ColTypes>(slot, S, idx)...);
+      return fExpression(GetValueChecked<ColTypes>(slot, S, idx)...);
       (void)idx; // avoid unused parameter warning (gcc 12.1)
    }
 
    template <typename... ColTypes, std::size_t... S>
-   void UpdateHelper(unsigned int slot, std::size_t idx, Long64_t /*entry*/, TypeList<ColTypes...>,
+   auto UpdateHelper(unsigned int slot, std::size_t idx, Long64_t /*entry*/, TypeList<ColTypes...>,
                      std::index_sequence<S...>, SlotTag)
    {
-      fLastResults[slot * RDFInternal::CacheLineStep<ret_type>()] =
-         fExpression(slot, GetValueChecked<ColTypes>(slot, S, idx)...);
+      return fExpression(slot, GetValueChecked<ColTypes>(slot, S, idx)...);
       (void)idx; // avoid unused parameter warning (gcc 12.1)
    }
 
    template <typename... ColTypes, std::size_t... S>
-   void UpdateHelper(unsigned int slot, std::size_t idx, Long64_t batchFirstEntry, TypeList<ColTypes...>,
+   auto UpdateHelper(unsigned int slot, std::size_t idx, Long64_t entryInBatch, TypeList<ColTypes...>,
                      std::index_sequence<S...>, SlotAndEntryTag)
    {
-      fLastResults[slot * RDFInternal::CacheLineStep<ret_type>()] =
-         fExpression(slot, batchFirstEntry + idx, GetValueChecked<ColTypes>(slot, S, idx)...);
+      return fExpression(slot, entryInBatch, GetValueChecked<ColTypes>(slot, S, idx)...);
       (void)idx;             // avoid unused parameter warning (gcc 12.1)
-      (void)batchFirstEntry; // avoid unused parameter warning (gcc 12.1)
+      (void)entryInBatch;    // avoid unused parameter warning (gcc 12.1)
    }
 
 public:
@@ -138,13 +135,20 @@ public:
    }
 
    /// Update the value at the address returned by GetValuePtr with the content corresponding to the given entry
-   void Update(unsigned int slot, Long64_t entry, bool mask) final
+   void Update(unsigned int slot, const ROOT::Internal::RDF::RMaskedEntryRange &mask) final
    {
-      if (entry != fLastCheckedEntry[slot * RDFInternal::CacheLineStep<Long64_t>()]) {
-         std::for_each(fValues[slot].begin(), fValues[slot].end(), [entry, mask](auto *v) { v->Load(entry, mask); });
-         if (mask) {
-            UpdateHelper(slot, /*idx=*/0u, entry, ColumnTypes_t{}, TypeInd_t{}, ExtraArgsTag{});
-            fLastCheckedEntry[slot * RDFInternal::CacheLineStep<Long64_t>()] = entry;
+      if (static_cast<Long64_t>(mask.GetFirstEntry()) ==
+          fLastCheckedEntry[slot * RDFInternal::CacheLineStep<Long64_t>()])
+         return;
+
+      std::for_each(fValues[slot].begin(), fValues[slot].end(), [&mask](auto *v) { v->Load(mask); });
+      // Assume 1-size bulk for now
+      const std::size_t bulkSize = 1;
+      auto &result = fLastResults[slot * RDFInternal::CacheLineStep<ret_type>()];
+      for (std::size_t i = 0; i < bulkSize; ++i) {
+         if (mask[i]) {
+            result = UpdateHelper(slot, i, mask.GetFirstEntry() + i, ColumnTypes_t{}, TypeInd_t{}, ExtraArgsTag{});
+            fLastCheckedEntry[slot * RDFInternal::CacheLineStep<Long64_t>()] = mask.GetFirstEntry();
          }
       }
    }

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

@@ -16,6 +16,7 @@
 #include "ROOT/RDF/RSampleInfo.hxx"
 #include "ROOT/RDF/Utils.hxx"
 #include "ROOT/RVec.hxx"
+#include <ROOT/RDF/RMaskedEntryRange.hxx>
 
 #include <deque>
 #include <map>
@@ -40,7 +41,9 @@ class RDefineBase {
 protected:
    const std::string fName; ///< The name of the custom column
    const std::string fType; ///< The type of the custom column as a text string
-   std::vector<Long64_t> fLastCheckedEntry;
+   std::vector<Long64_t> fLastCheckedEntry; /// Starting entry of the last bulk processed, per slot
+   /// Which entries in the current bulk are valid, per slot
+   std::vector<ROOT::Internal::RDF::RMaskedEntryRange> fMaskPerSlot;
    RDFInternal::RColumnRegister fColRegister;
    RLoopManager *fLoopManager; // non-owning pointer to the RLoopManager
    const ROOT::RDF::ColumnNames_t fColumnNames;
@@ -63,7 +66,7 @@ public:
    std::string GetName() const;
    std::string GetTypeName() const;
    /// Update the value at the address returned by GetValuePtr with the content corresponding to the given entry
-   virtual void Update(unsigned int slot, Long64_t entry, bool mask) = 0;
+   virtual void Update(unsigned int slot, const ROOT::Internal::RDF::RMaskedEntryRange &mask) = 0;
    /// Update function to be called once per sample, used if the derived type is a RDefinePerSample
    virtual void Update(unsigned int /*slot*/, const ROOT::RDF::RSampleInfo &/*id*/) {}
    /// Clean-up operations to be performed at the end of a task.