use shared state for query tensors

prepare shared state in ftlib feature test setup

4 files changed, 60 insertions, 11 deletions

diff --git a/searchlib/src/vespa/searchlib/features/constant_tensor_executor.h b/searchlib/src/vespa/searchlib/features/constant_tensor_executor.h
index 05a2e1452f6..5babd0a5814 100644
--- a/searchlib/src/vespa/searchlib/features/constant_tensor_executor.h
+++ b/searchlib/src/vespa/searchlib/features/constant_tensor_executor.h
@@ -41,4 +41,20 @@ public:
     }
 };
 
+/**
+ * Feature executor that returns a constant tensor.
+ */
+class ConstantTensorRefExecutor : public fef::FeatureExecutor
+{
+private:
+    const vespalib::eval::Value &_tensor_ref;
+public:
+    ConstantTensorRefExecutor(const vespalib::eval::Value &tensor_ref)
+      : _tensor_ref(tensor_ref) {}
+    bool isPure() final override { return true; }
+    void execute(uint32_t) final override {
+        outputs().set_object(0, _tensor_ref);
+    }
+};
+
 }
diff --git a/searchlib/src/vespa/searchlib/features/queryfeature.cpp b/searchlib/src/vespa/searchlib/features/queryfeature.cpp
index 56375eecbad..c72e3e6b921 100644
--- a/searchlib/src/vespa/searchlib/features/queryfeature.cpp
+++ b/searchlib/src/vespa/searchlib/features/queryfeature.cpp
@@ -25,8 +25,13 @@ using namespace search::fef;
 using namespace search::fef::indexproperties;
 using document::TensorDataType;
 using vespalib::eval::ValueType;
+using vespalib::eval::Value;
 using vespalib::Issue;
 using search::fef::FeatureType;
+using search::fef::AnyWrapper;
+using search::fef::Anything;
+
+using ValueWrapper = AnyWrapper<Value::UP>;
 
 namespace search::features {
 
@@ -56,6 +61,15 @@ feature_t asFeature(const vespalib::string &str) {
     return val;
 }
 
+// query(foo):
+// query.value.foo -> decoded tensor value 'foo'
+vespalib::string make_value_key(const vespalib::string &base, const vespalib::string &sub_key) {
+    vespalib::string key(base);
+    key.append(".value.");
+    key.append(sub_key);
+    return key;
+}
+
 } // namespace search::features::<unnamed>
 
 QueryBlueprint::QueryBlueprint() :
@@ -86,6 +100,7 @@ QueryBlueprint::setup(const IIndexEnvironment &env, const ParameterList &params)
     _key = params[0].getValue();
     _key2 = "$";
     _key2.append(_key);
+    _stored_value_key = make_value_key(getBaseName(), _key);
 
     vespalib::string key3;
     key3.append("query(");
@@ -114,10 +129,9 @@ QueryBlueprint::setup(const IIndexEnvironment &env, const ParameterList &params)
 
 namespace {
 
-FeatureExecutor &
-createTensorExecutor(const IQueryEnvironment &env,
-                     const vespalib::string &queryKey,
-                     const ValueType &valueType, vespalib::Stash &stash)
+Value::UP make_tensor_value(const IQueryEnvironment &env,
+                            const vespalib::string &queryKey,
+                            const ValueType &valueType)
 {
     Property prop = env.getProperties().lookup(queryKey);
     if (prop.found() && !prop.get().empty()) {
@@ -125,27 +139,42 @@ createTensorExecutor(const IQueryEnvironment &env,
         vespalib::nbostream stream(value.data(), value.size());
         try {
             auto tensor = vespalib::eval::decode_value(stream, vespalib::eval::FastValueBuilderFactory::get());
-            if (!TensorDataType::isAssignableType(valueType, tensor->type())) {
+            if (TensorDataType::isAssignableType(valueType, tensor->type())) {
+                return tensor;
+            } else {
                 Issue::report("Query feature type is '%s' but other tensor type is '%s'",
                               valueType.to_spec().c_str(), tensor->type().to_spec().c_str());
-                return ConstantTensorExecutor::createEmpty(valueType, stash);
             }
-            return ConstantTensorExecutor::create(std::move(tensor), stash);
         } catch (const vespalib::eval::DecodeValueException &e) {
             Issue::report("Query feature has invalid binary format: %s", e.what());
-            return ConstantTensorExecutor::createEmpty(valueType, stash);
         }
     }
-    return ConstantTensorExecutor::createEmpty(valueType, stash);
+    return {};
+}
+
 }
 
+void
+QueryBlueprint::prepareSharedState(const fef::IQueryEnvironment &env, fef::IObjectStore &store) const
+{
+    if (!_stored_value_key.empty() && (store.get(_stored_value_key) == nullptr)) {
+        auto value = make_tensor_value(env, _key, _valueType);
+        if (value) {
+            store.add(_stored_value_key, std::make_unique<ValueWrapper>(std::move(value)));
+        }
+    }
 }
 
 FeatureExecutor &
 QueryBlueprint::createExecutor(const IQueryEnvironment &env, vespalib::Stash &stash) const
 {
     if (_valueType.has_dimensions()) {
-        return createTensorExecutor(env, _key, _valueType, stash);
+        if (const Anything *wrapped_value = env.getObjectStore().get(_stored_value_key)) {
+            if (const Value *value = ValueWrapper::getValue(*wrapped_value).get()) {
+                return stash.create<ConstantTensorRefExecutor>(*value);
+            }
+        }
+        return ConstantTensorExecutor::createEmpty(_valueType, stash);
     } else {
         std::vector<feature_t> values;
         Property p = env.getProperties().lookup(_key);
diff --git a/searchlib/src/vespa/searchlib/features/queryfeature.h b/searchlib/src/vespa/searchlib/features/queryfeature.h
index b097f5785ed..05e44bcd923 100644
--- a/searchlib/src/vespa/searchlib/features/queryfeature.h
+++ b/searchlib/src/vespa/searchlib/features/queryfeature.h
@@ -17,6 +17,7 @@ class QueryBlueprint : public fef::Blueprint {
 private:
     vespalib::string _key;  // 'foo'
     vespalib::string _key2; // '$foo'
+    vespalib::string _stored_value_key;
     feature_t _defaultValue;
     vespalib::eval::ValueType _valueType;
 
@@ -30,6 +31,7 @@ public:
         return fef::ParameterDescriptions().desc().string();
     }
     bool setup(const fef::IIndexEnvironment &env, const fef::ParameterList &params) override;
+    void prepareSharedState(const fef::IQueryEnvironment &env, fef::IObjectStore &store) const override;
     fef::FeatureExecutor &createExecutor(const fef::IQueryEnvironment &env, vespalib::Stash &stash) const override;
 };
 
diff --git a/searchlib/src/vespa/searchlib/fef/test/featuretest.cpp b/searchlib/src/vespa/searchlib/fef/test/featuretest.cpp
index 6a2feaf14fa..81c1666bed6 100644
--- a/searchlib/src/vespa/searchlib/fef/test/featuretest.cpp
+++ b/searchlib/src/vespa/searchlib/fef/test/featuretest.cpp
@@ -69,7 +69,9 @@ FeatureTest::setup()
         LOG(error, "Failed to compile blueprint resolver.");
         return false;
     }
-
+    for (const auto &spec: _resolver->getExecutorSpecs()) {
+        spec.blueprint->prepareSharedState(_queryEnv, _queryEnv.getObjectStore());
+    }
     _rankProgram->setup(*_match_data, _queryEnv, _overrides);
     _doneSetup = true;
     return true;