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