Setup ranking assets repo for streaming search.

6 files changed, 124 insertions, 40 deletions

diff --git a/streamingvisitors/src/vespa/searchvisitor/indexenvironment.cpp b/streamingvisitors/src/vespa/searchvisitor/indexenvironment.cpp
index 81df2b5492f..1242195c9df 100644
--- a/streamingvisitors/src/vespa/searchvisitor/indexenvironment.cpp
+++ b/streamingvisitors/src/vespa/searchvisitor/indexenvironment.cpp
@@ -1,6 +1,7 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
 #include "indexenvironment.h"
+#include <vespa/searchlib/fef/i_ranking_assets_repo.h>
 
 using namespace search::fef;
 
@@ -13,7 +14,8 @@ IndexEnvironment::IndexEnvironment(const ITableManager & tableManager) :
     _fieldNames(),
     _motivation(RANK),
     _rankAttributes(),
-    _dumpAttributes()
+    _dumpAttributes(),
+    _ranking_assets_repo()
 {
 }
 
@@ -50,5 +52,28 @@ IndexEnvironment::hintAttributeAccess(const string & name) const {
     }
 }
 
-} // namespace streaming
+void
+IndexEnvironment::set_ranking_assets_repo(std::shared_ptr<const IRankingAssetsRepo> ranking_assets_repo)
+{
+    _ranking_assets_repo = std::move(ranking_assets_repo);
+}
 
+vespalib::eval::ConstantValue::UP
+IndexEnvironment::getConstantValue(const vespalib::string& name) const
+{
+    return _ranking_assets_repo->getConstant(name);
+}
+
+vespalib::string
+IndexEnvironment::getRankingExpression(const vespalib::string& name) const
+{
+    return _ranking_assets_repo->getExpression(name);
+}
+
+const search::fef::OnnxModel*
+IndexEnvironment::getOnnxModel(const vespalib::string& name) const
+{
+    return _ranking_assets_repo->getOnnxModel(name);
+}
+
+}
diff --git a/streamingvisitors/src/vespa/searchvisitor/indexenvironment.h b/streamingvisitors/src/vespa/searchvisitor/indexenvironment.h
index ef679cacdf0..88ea5a5ada8 100644
--- a/streamingvisitors/src/vespa/searchvisitor/indexenvironment.h
+++ b/streamingvisitors/src/vespa/searchvisitor/indexenvironment.h
@@ -12,6 +12,8 @@
 #include <vespa/vespalib/stllike/hash_map.h>
 #include <set>
 
+namespace search::fef { struct IRankingAssetsRepo; }
+
 namespace streaming {
 
 /**
@@ -29,6 +31,7 @@ private:
     mutable FeatureMotivation            _motivation;
     mutable std::set<vespalib::string>   _rankAttributes;
     mutable std::set<vespalib::string>   _dumpAttributes;
+    std::shared_ptr<const search::fef::IRankingAssetsRepo> _ranking_assets_repo;
 
 public:
     IndexEnvironment(const search::fef::ITableManager & tableManager);
@@ -71,17 +74,11 @@ public:
 
     void hintAttributeAccess(const string & name) const override;
 
-    vespalib::eval::ConstantValue::UP getConstantValue(const vespalib::string &) const override {
-        return vespalib::eval::ConstantValue::UP();
-    }
+    vespalib::eval::ConstantValue::UP getConstantValue(const vespalib::string& name) const override;
 
-    vespalib::string getRankingExpression(const vespalib::string &) const override {
-        return {};
-    }
+    vespalib::string getRankingExpression(const vespalib::string& name) const override;
 
-    const search::fef::OnnxModel *getOnnxModel(const vespalib::string &) const override {
-        return nullptr;
-    }
+    const search::fef::OnnxModel *getOnnxModel(const vespalib::string& name) const override;
 
     bool addField(const vespalib::string& name,
                   bool isAttribute,
@@ -89,6 +86,8 @@ public:
 
     search::fef::Properties & getProperties() { return _properties; }
 
+    void set_ranking_assets_repo(std::shared_ptr<const search::fef::IRankingAssetsRepo> ranking_assets_repo);
+
     const std::set<vespalib::string> & getHintedRankAttributes() const { return _rankAttributes; }
 
     const std::set<vespalib::string> & getHintedDumpAttributes() const { return _dumpAttributes; }
diff --git a/streamingvisitors/src/vespa/searchvisitor/rankmanager.cpp b/streamingvisitors/src/vespa/searchvisitor/rankmanager.cpp
index 81a2a48fb4d..1340ecd1e70 100644
--- a/streamingvisitors/src/vespa/searchvisitor/rankmanager.cpp
+++ b/streamingvisitors/src/vespa/searchvisitor/rankmanager.cpp
@@ -17,6 +17,7 @@ using vespa::config::search::vsm::VsmfieldsConfig;
 using search::fef::Blueprint;
 using search::fef::BlueprintFactory;
 using search::fef::FieldInfo;
+using search::fef::IRankingAssetsRepo;
 using search::fef::Properties;
 using search::fef::RankSetup;
 using vsm::VsmfieldsHandle;
@@ -162,13 +163,6 @@ RankManager::Snapshot::Snapshot() :
 }
 
 RankManager::Snapshot::~Snapshot() = default;
-    
-bool
-RankManager::Snapshot::setup(const RankManager & rm, const std::vector<NamedPropertySet> & properties)
-{
-    _properties = properties;
-    return setup(rm);
-}
 
 bool
 RankManager::Snapshot::setup(const RankManager & rm)
@@ -183,26 +177,27 @@ RankManager::Snapshot::setup(const RankManager & rm)
 }
 
 bool
-RankManager::Snapshot::setup(const RankManager & rm, const RankProfilesConfig & cfg)
+RankManager::Snapshot::setup(const RankManager & rm, const RankProfilesConfig & cfg, std::shared_ptr<const IRankingAssetsRepo> ranking_assets_repo)
 {
+    _protoEnv.set_ranking_assets_repo(std::move(ranking_assets_repo));
     addProperties(cfg);
     return setup(rm);
 }
 
 void
-RankManager::notify(const vsm::VSMConfigSnapshot & snap)
+RankManager::notify(const vsm::VSMConfigSnapshot & snap, std::shared_ptr<const IRankingAssetsRepo> ranking_assets_repo)
 {
-    configureRankProfiles(*snap.getConfig<RankProfilesConfig>());
+    configureRankProfiles(*snap.getConfig<RankProfilesConfig>(), std::move(ranking_assets_repo));
 }
 
 
 void
-RankManager::configureRankProfiles(const RankProfilesConfig & cfg)
+RankManager::configureRankProfiles(const RankProfilesConfig & cfg, std::shared_ptr<const IRankingAssetsRepo> ranking_assets_repo)
 {
     LOG(debug, "configureRankProfiles(): Size of cfg rankprofiles: %zd", cfg.rankprofile.size());
 
     auto snapshot = std::make_unique<Snapshot>();
-    if (snapshot->setup(*this, cfg)) {
+    if (snapshot->setup(*this, cfg, std::move(ranking_assets_repo))) {
         _snapshot.set(snapshot.release());
         _snapshot.latch(); // switch to the new config object
     } else {
@@ -224,9 +219,9 @@ RankManager::RankManager(VSMAdapter * const vsmAdapter) :
 RankManager::~RankManager() = default;
 
 void
-RankManager::configure(const vsm::VSMConfigSnapshot & snap)
+RankManager::configure(const vsm::VSMConfigSnapshot & snap, std::shared_ptr<const IRankingAssetsRepo> ranking_assets_repo)
 {
-    notify(snap);
+    notify(snap, std::move(ranking_assets_repo));
 }
     
 }
diff --git a/streamingvisitors/src/vespa/searchvisitor/rankmanager.h b/streamingvisitors/src/vespa/searchvisitor/rankmanager.h
index ad437bf3bb3..71910c65242 100644
--- a/streamingvisitors/src/vespa/searchvisitor/rankmanager.h
+++ b/streamingvisitors/src/vespa/searchvisitor/rankmanager.h
@@ -19,6 +19,7 @@ class RankManager
 public:
     /** collection of field ids for an index **/
     using View = std::vector<uint32_t>;
+    using IRankingAssetsRepo = search::fef::IRankingAssetsRepo;
 
     /**
      * This class represents a snapshot of the rank-profiles config with associated setup per rank profile.
@@ -51,8 +52,7 @@ public:
         Snapshot();
         ~Snapshot();
         const std::vector<NamedPropertySet> & getProperties() const { return _properties; }
-        bool setup(const RankManager & manager, const vespa::config::search::RankProfilesConfig & cfg);
-        bool setup(const RankManager & manager, const std::vector<NamedPropertySet> & properties);
+        bool setup(const RankManager & manager, const vespa::config::search::RankProfilesConfig & cfg, std::shared_ptr<const IRankingAssetsRepo> ranking_assets_repo);
         const search::fef::RankSetup & getRankSetup(const vespalib::string &rankProfile) const {
             return *(_rankSetup[getIndex(rankProfile)]);
         }
@@ -73,14 +73,14 @@ private:
     vespalib::PtrHolder<const Snapshot> _snapshot;
     const vsm::VSMAdapter       * _vsmAdapter;
 
-    void configureRankProfiles(const vespa::config::search::RankProfilesConfig & cfg);
-    virtual void notify(const vsm::VSMConfigSnapshot & snapshot);
+    void configureRankProfiles(const vespa::config::search::RankProfilesConfig & cfg, std::shared_ptr<const IRankingAssetsRepo> ranking_assets_repo);
+    virtual void notify(const vsm::VSMConfigSnapshot & snapshot, std::shared_ptr<const IRankingAssetsRepo> ranking_assets_repo);
 
 public:
     RankManager(vsm::VSMAdapter * const vsmAdapter);
     virtual ~RankManager();
 
-    void configure(const vsm::VSMConfigSnapshot & snap);
+    void configure(const vsm::VSMConfigSnapshot & snap, std::shared_ptr<const IRankingAssetsRepo> ranking_assets_repo);
 
     /**
      * Retrieves the current snapshot of the rank-profiles config.
diff --git a/streamingvisitors/src/vespa/searchvisitor/searchenvironment.cpp b/streamingvisitors/src/vespa/searchvisitor/searchenvironment.cpp
index 7dab92370ae..75e07615bd9 100644
--- a/streamingvisitors/src/vespa/searchvisitor/searchenvironment.cpp
+++ b/streamingvisitors/src/vespa/searchvisitor/searchenvironment.cpp
@@ -2,13 +2,26 @@
 
 #include "searchenvironment.h"
 #include "search_environment_snapshot.h"
+#include <vespa/config-onnx-models.h>
+#include <vespa/config-ranking-constants.h>
+#include <vespa/config-ranking-expressions.h>
+#include <vespa/config/retriever/configsnapshot.hpp>
+#include <vespa/eval/eval/fast_value.h>
+#include <vespa/searchlib/fef/ranking_assets_builder.h>
+#include <vespa/searchlib/fef/ranking_assets_repo.h>
 #include <vespa/vespalib/stllike/hash_map.hpp>
 #include <vespa/searchsummary/config/config-juniperrc.h>
+#include <cassert>
 
 #include <vespa/log/log.h>
 LOG_SETUP(".visitor.instance.searchenvironment");
 
 using search::docsummary::JuniperProperties;
+using search::fef::OnnxModels;
+using search::fef::RankingAssetsBuilder;
+using search::fef::RankingAssetsRepo;
+using search::fef::RankingConstants;
+using search::fef::RankingExpressions;
 using vsm::VSMAdapter;
 
 namespace streaming {
@@ -22,6 +35,13 @@ SearchEnvironment::Env::Env(const config::ConfigUri& configUri, const Fast_Norma
       _rankManager(std::make_unique<RankManager>(_vsmAdapter.get())),
       _snapshot(),
       _lock(),
+      _tensor_loader(vespalib::eval::FastValueBuilderFactory::get()),
+      _constant_value_cache(_tensor_loader),
+      _generation(-1),
+      _onnx_models(),
+      _ranking_constants(),
+      _ranking_expressions(),
+      _ranking_assets_repo(),
       _transport(transport),
       _file_distributor_connection_spec(file_distributor_connection_spec)
 {
@@ -36,7 +56,10 @@ SearchEnvironment::Env::createKeySet(const vespalib::string & configId)
             vespa::config::search::SummaryConfig,
             vespa::config::search::vsm::VsmsummaryConfig,
             vespa::config::search::summary::JuniperrcConfig,
-            vespa::config::search::RankProfilesConfig>(configId);
+            vespa::config::search::RankProfilesConfig,
+            vespa::config::search::core::OnnxModelsConfig,
+            vespa::config::search::core::RankingConstantsConfig,
+            vespa::config::search::core::RankingExpressionsConfig>(configId);
     return set;
 }
 
@@ -44,8 +67,14 @@ void
 SearchEnvironment::Env::configure(const config::ConfigSnapshot & snapshot)
 {
     vsm::VSMConfigSnapshot snap(_configId, snapshot);
+    RankingAssetsBuilder builder(_transport, _file_distributor_connection_spec);
+    configure_ranking_asset<vespa::config::search::core::OnnxModelsConfig, OnnxModels>(_onnx_models, snapshot, builder);
+    configure_ranking_asset<vespa::config::search::core::RankingConstantsConfig, RankingConstants>(_ranking_constants, snapshot, builder);
+    configure_ranking_asset<vespa::config::search::core::RankingExpressionsConfig, RankingExpressions>(_ranking_expressions, snapshot, builder);
+    _ranking_assets_repo = std::make_shared<const RankingAssetsRepo>(_constant_value_cache, _ranking_constants, _ranking_expressions, _onnx_models);
+    _generation = snapshot.getGeneration();
     _vsmAdapter->configure(snap);
-    _rankManager->configure(snap);
+    _rankManager->configure(snap, _ranking_assets_repo);
     auto se_snapshot = std::make_shared<const SearchEnvironmentSnapshot>(*_rankManager, *_vsmAdapter);
     std::lock_guard guard(_lock);
     std::swap(se_snapshot, _snapshot);
@@ -58,6 +87,19 @@ SearchEnvironment::Env::get_snapshot()
     return _snapshot;
 }
 
+template <typename ConfigType, typename RankingAssetType>
+void
+SearchEnvironment::Env::configure_ranking_asset(std::shared_ptr<const RankingAssetType> &ranking_asset,
+                                                const config::ConfigSnapshot& snapshot,
+                                                RankingAssetsBuilder& builder)
+{
+    if (snapshot.isChanged<ConfigType>(_configId, _generation)) {
+        ranking_asset = builder.build(*snapshot.getConfig<ConfigType>(_configId));
+    } else {
+        assert(ranking_asset);
+    }
+}
+
 SearchEnvironment::Env::~Env()
 {
     _configurer.close();
diff --git a/streamingvisitors/src/vespa/searchvisitor/searchenvironment.h b/streamingvisitors/src/vespa/searchvisitor/searchenvironment.h
index caae62cac58..f3bbfddd76c 100644
--- a/streamingvisitors/src/vespa/searchvisitor/searchenvironment.h
+++ b/streamingvisitors/src/vespa/searchvisitor/searchenvironment.h
@@ -3,6 +3,8 @@
 #pragma once
 
 #include "rankmanager.h"
+#include <vespa/eval/eval/value_cache/constant_tensor_loader.h>
+#include <vespa/eval/eval/value_cache/constant_value_cache.h>
 #include <vespa/searchsummary/docsummary/juniperproperties.h>
 #include <vespa/storage/visiting/visitor.h>
 #include <vespa/config/retriever/simpleconfigurer.h>
@@ -13,6 +15,16 @@
 
 class FNET_Transport;
 
+namespace search::fef {
+
+struct IRankingAssetsRepo;
+class OnnxModels;
+class RankingAssetsBuilder;
+class RankingConstants;
+class RankingExpressions;
+
+}
+
 namespace streaming {
 
 class SearchEnvironmentSnapshot;
@@ -30,14 +42,25 @@ private:
         static config::ConfigKeySet createKeySet(const vespalib::string & configId);
         std::shared_ptr<const SearchEnvironmentSnapshot> get_snapshot();
     private:
-        const vespalib::string                           _configId;
-        config::SimpleConfigurer                         _configurer;
-        std::unique_ptr<vsm::VSMAdapter>                 _vsmAdapter;
-        std::unique_ptr<RankManager>                     _rankManager;
-        std::shared_ptr<const SearchEnvironmentSnapshot> _snapshot;
-        std::mutex                                       _lock;
-        FNET_Transport* const                            _transport;
-        const vespalib::string                           _file_distributor_connection_spec;
+        template <typename ConfigType, typename RankingAssetType>
+        void configure_ranking_asset(std::shared_ptr<const RankingAssetType> &ranking_asset,
+                                     const config::ConfigSnapshot& snapshot,
+                                     search::fef::RankingAssetsBuilder& builder);
+        const vespalib::string                                 _configId;
+        config::SimpleConfigurer                               _configurer;
+        std::unique_ptr<vsm::VSMAdapter>                       _vsmAdapter;
+        std::unique_ptr<RankManager>                           _rankManager;
+        std::shared_ptr<const SearchEnvironmentSnapshot>       _snapshot;
+        std::mutex                                             _lock;
+        vespalib::eval::ConstantTensorLoader                   _tensor_loader;
+        vespalib::eval::ConstantValueCache                     _constant_value_cache;
+        uint64_t                                               _generation;
+        std::shared_ptr<const search::fef::OnnxModels>         _onnx_models;
+        std::shared_ptr<const search::fef::RankingConstants>   _ranking_constants;
+        std::shared_ptr<const search::fef::RankingExpressions> _ranking_expressions;
+        std::shared_ptr<const search::fef::IRankingAssetsRepo> _ranking_assets_repo;
+        FNET_Transport* const                                  _transport;
+        const vespalib::string                                 _file_distributor_connection_spec;
     };
     using EnvMap = vespalib::hash_map<vespalib::string, Env::SP>;
     using EnvMapUP = std::unique_ptr<EnvMap>;