[RF] Implement batch data access for RooCompositeDataStore

Implement getBatches() and getCategoryBatches() for
RooCompositeDataStore, which so far threw an exception. Like the
existing getWeightBatch(), the implementation lazily concatenates the
columns of the component datasets into internal buffers by loading the
composite rows one by one. The index category, which is not stored in
any of the component datasets, is synthesized from the row lookup.

This is needed so that combined datasets backed by a composite storage
(e.g. the output of generating from a RooSimultaneous with AllBinned())
can be loaded by the generic RooFit::Evaluator data path, which so far
only worked when the dataset was split into its channel components
first.

Author: guitargeek

roofit/roofitcore/inc/RooCompositeDataStore.h

@@ -92,26 +92,27 @@ class RooCompositeDataStore : public RooAbsDataStore {
   void loadValues(const RooAbsDataStore *tds, const RooFormulaVar* select=nullptr, const char* rangeName=nullptr,
       std::size_t nStart=0, std::size_t nStop = std::numeric_limits<std::size_t>::max()) override;
 
-  RooAbsData::RealSpans getBatches(std::size_t first, std::size_t len) const override {
-    //TODO
-    std::cerr << "This functionality is not yet implemented for composite data stores." << std::endl;
-    throw std::logic_error("getBatches() not implemented for RooCompositeDataStore.");
-    (void)first; (void)len;
-    return {};
-  }
+  RooAbsData::RealSpans getBatches(std::size_t first, std::size_t len) const override;
+  RooAbsData::CategorySpans getCategoryBatches(std::size_t first, std::size_t len) const override;
   std::span<const double> getWeightBatch(std::size_t first, std::size_t len) const override;
 
 
  protected:
-
-  std::map<Int_t,RooAbsDataStore*> _dataMap ;
-  RooCategory* _indexCat = nullptr;
-  mutable RooAbsDataStore* _curStore = nullptr; ///<! Datastore associated with current event
-  mutable Int_t _curIndex = 0; ///<! Index associated with current event
-  mutable std::unique_ptr<std::vector<double>> _weightBuffer; ///<! Buffer for weights in case a batch of values is requested.
-  bool _ownComps = false; ///<!
-
-  ClassDefOverride(RooCompositeDataStore,1) // Composite Data Storage class
+    void fillBatchBuffers() const;
+
+    std::map<Int_t, RooAbsDataStore *> _dataMap;
+    RooCategory *_indexCat = nullptr;
+    mutable RooAbsDataStore *_curStore = nullptr; ///<! Datastore associated with current event
+    mutable Int_t _curIndex = 0;                  ///<! Index associated with current event
+    mutable std::unique_ptr<std::vector<double>>
+       _weightBuffer; ///<! Buffer for weights in case a batch of values is requested.
+    mutable std::map<RooAbsArg const *, std::vector<double>>
+       _realBatchBuffers; ///<! Buffers for real-valued columns in case batches of values are requested.
+    mutable std::map<RooAbsArg const *, std::vector<RooAbsCategory::value_type>>
+       _catBatchBuffers;    ///<! Buffers for category columns in case batches of values are requested.
+    bool _ownComps = false; ///<!
+
+    ClassDefOverride(RooCompositeDataStore, 1) // Composite Data Storage class
 };
 
 

roofit/roofitcore/src/RooCompositeDataStore.cxx

@@ -35,6 +35,7 @@ When iterated from start to finish, datasets will be traversed in the order of t
 #include "RooRealVar.h"
 #include "RooCategory.h"
 
+#include <algorithm>
 #include <iomanip>
 #include <iostream>
 
@@ -408,3 +409,86 @@ std::span<const double> RooCompositeDataStore::getWeightBatch(std::size_t first,
 
   return {_weightBuffer->data() + first, len};
 }
+
+////////////////////////////////////////////////////////////////////////////////
+/// Fill the internal per-column buffers for batch access by loading the
+/// composite rows one by one, like getWeightBatch() does for the weights.
+/// This also synthesizes a column for the index category, which is not
+/// stored in any of the component datasets.
+void RooCompositeDataStore::fillBatchBuffers() const
+{
+   const auto n = static_cast<std::size_t>(numEntries());
+
+   if (!_realBatchBuffers.empty() || !_catBatchBuffers.empty()) {
+      // Refill from scratch if entries were added or removed since the last
+      // fill. Value mutations that keep the number of entries are not
+      // detected, like for the weight buffer above.
+      const std::size_t nFilled = !_realBatchBuffers.empty() ? _realBatchBuffers.begin()->second.size()
+                                                             : _catBatchBuffers.begin()->second.size();
+      if (nFilled == n)
+         return;
+      _realBatchBuffers.clear();
+      _catBatchBuffers.clear();
+   }
+
+   std::vector<std::pair<RooAbsReal const *, std::vector<double> *>> realCols;
+   std::vector<std::pair<RooAbsCategory const *, std::vector<RooAbsCategory::value_type> *>> catCols;
+
+   for (RooAbsArg const *arg : _vars) {
+      if (auto cat = dynamic_cast<RooAbsCategory const *>(arg)) {
+         auto &buf = _catBatchBuffers[arg];
+         buf.reserve(n);
+         catCols.emplace_back(cat, &buf);
+      } else if (auto real = dynamic_cast<RooAbsReal const *>(arg)) {
+         auto &buf = _realBatchBuffers[arg];
+         buf.reserve(n);
+         realCols.emplace_back(real, &buf);
+      }
+   }
+
+   for (std::size_t i = 0; i < n; ++i) {
+      get(i);
+      for (auto &col : realCols)
+         col.second->push_back(col.first->getVal());
+      for (auto &col : catCols)
+         col.second->push_back(col.first->getCurrentIndex());
+   }
+}
+
+////////////////////////////////////////////////////////////////////////////////
+/// Get the batches of the real-valued columns in the range [first, first+len).
+/// The columns of the component datasets are lazily concatenated into
+/// internal buffers in composite row order.
+RooAbsData::RealSpans RooCompositeDataStore::getBatches(std::size_t first, std::size_t len) const
+{
+   fillBatchBuffers();
+
+   // Clamp against the actual number of entries, because the default
+   // arguments of RooAbsData::getBatches() ask for the maximum length.
+   first = std::min(first, static_cast<std::size_t>(numEntries()));
+   len = std::min(len, static_cast<std::size_t>(numEntries()) - first);
+
+   RooAbsData::RealSpans out;
+   for (auto const &item : _realBatchBuffers) {
+      out.emplace(item.first, std::span<const double>{item.second.data() + first, len});
+   }
+   return out;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+/// Get the batches of the category columns in the range [first, first+len),
+/// including the index category. See getBatches().
+RooAbsData::CategorySpans RooCompositeDataStore::getCategoryBatches(std::size_t first, std::size_t len) const
+{
+   fillBatchBuffers();
+
+   // See getBatches() for the clamping rationale.
+   first = std::min(first, static_cast<std::size_t>(numEntries()));
+   len = std::min(len, static_cast<std::size_t>(numEntries()) - first);
+
+   RooAbsData::CategorySpans out;
+   for (auto const &item : _catBatchBuffers) {
+      out.emplace(item.first, std::span<const RooAbsCategory::value_type>{item.second.data() + first, len});
+   }
+   return out;
+}