[RF] Implement batch data access for RooCompositeDataStore

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;
+}