From b582543566d4323359fc4694f8f1e56693178a87 Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Sat, 5 Sep 2026 06:52:33 +0000 Subject: [PATCH] [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. --- roofit/roofitcore/inc/RooCompositeDataStore.h | 33 ++++---- .../roofitcore/src/RooCompositeDataStore.cxx | 84 +++++++++++++++++++ 2 files changed, 101 insertions(+), 16 deletions(-) diff --git a/roofit/roofitcore/inc/RooCompositeDataStore.h b/roofit/roofitcore/inc/RooCompositeDataStore.h index b720c9d2ff62c..0f4bd75f963d2 100644 --- a/roofit/roofitcore/inc/RooCompositeDataStore.h +++ b/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::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 getWeightBatch(std::size_t first, std::size_t len) const override; protected: - - std::map _dataMap ; - RooCategory* _indexCat = nullptr; - mutable RooAbsDataStore* _curStore = nullptr; ///> _weightBuffer; /// _dataMap; + RooCategory *_indexCat = nullptr; + mutable RooAbsDataStore *_curStore = nullptr; ///> + _weightBuffer; ///> + _realBatchBuffers; ///> + _catBatchBuffers; /// #include #include @@ -408,3 +409,86 @@ std::span 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(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 *>> realCols; + std::vector *>> catCols; + + for (RooAbsArg const *arg : _vars) { + if (auto cat = dynamic_cast(arg)) { + auto &buf = _catBatchBuffers[arg]; + buf.reserve(n); + catCols.emplace_back(cat, &buf); + } else if (auto real = dynamic_cast(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(numEntries())); + len = std::min(len, static_cast(numEntries()) - first); + + RooAbsData::RealSpans out; + for (auto const &item : _realBatchBuffers) { + out.emplace(item.first, std::span{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(numEntries())); + len = std::min(len, static_cast(numEntries()) - first); + + RooAbsData::CategorySpans out; + for (auto const &item : _catBatchBuffers) { + out.emplace(item.first, std::span{item.second.data() + first, len}); + } + return out; +}