Merge pull request #11369 from vespa-engine/geirst/expose-query-tensors-in-request-context

Expose tensors passed with the query in IRequestContext.

6 files changed, 138 insertions, 7 deletions

diff --git a/searchcore/CMakeLists.txt b/searchcore/CMakeLists.txt
index 53435780d9a..7dbbb012b11 100644
--- a/searchcore/CMakeLists.txt
+++ b/searchcore/CMakeLists.txt
@@ -112,6 +112,7 @@ vespa_define_module(
     src/tests/proton/matching/match_loop_communicator
     src/tests/proton/matching/match_phase_limiter
     src/tests/proton/matching/partial_result
+    src/tests/proton/matching/request_context
     src/tests/proton/matching/same_element_builder
     src/tests/proton/matching/unpacking_iterators_optimizer
     src/tests/proton/metrics/documentdb_job_trackers
diff --git a/searchcore/src/tests/proton/matching/request_context/CMakeLists.txt b/searchcore/src/tests/proton/matching/request_context/CMakeLists.txt
new file mode 100644
index 00000000000..54696112302
--- /dev/null
+++ b/searchcore/src/tests/proton/matching/request_context/CMakeLists.txt
@@ -0,0 +1,9 @@
+# Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+vespa_add_executable(searchcore_matching_request_context_test_app TEST
+    SOURCES
+    request_context_test.cpp
+    DEPENDS
+    searchcore_matching
+    gtest
+)
+vespa_add_test(NAME searchcore_matching_request_context_test_app COMMAND searchcore_matching_request_context_test_app)
diff --git a/searchcore/src/tests/proton/matching/request_context/request_context_test.cpp b/searchcore/src/tests/proton/matching/request_context/request_context_test.cpp
new file mode 100644
index 00000000000..be70bacb4b1
--- /dev/null
+++ b/searchcore/src/tests/proton/matching/request_context/request_context_test.cpp
@@ -0,0 +1,81 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include <vespa/eval/eval/tensor_spec.h>
+#include <vespa/eval/tensor/default_tensor_engine.h>
+#include <vespa/searchcore/proton/matching/requestcontext.h>
+#include <vespa/searchlib/fef/properties.h>
+#include <vespa/vespalib/gtest/gtest.h>
+#include <vespa/vespalib/objects/nbostream.h>
+
+using search::attribute::IAttributeContext;
+using search::attribute::IAttributeFunctor;
+using search::attribute::IAttributeVector;
+using search::fef::Properties;
+using vespalib::eval::TensorSpec;
+using vespalib::eval::Value;
+using vespalib::tensor::DefaultTensorEngine;
+using namespace proton;
+
+class MyAttributeContext : public search::attribute::IAttributeContext {
+public:
+    const IAttributeVector* getAttribute(const vespalib::string&) const override { abort(); }
+    const IAttributeVector* getAttributeStableEnum(const vespalib::string&) const override { abort(); }
+    void getAttributeList(std::vector<const IAttributeVector*>&) const override { abort(); }
+    void asyncForAttribute(const vespalib::string&, std::unique_ptr<IAttributeFunctor>) const override { abort(); }
+};
+
+class RequestContextTest : public ::testing::Test {
+private:
+    vespalib::Clock _clock;
+    vespalib::Doom _doom;
+    MyAttributeContext _attr_ctx;
+    Properties _props;
+    RequestContext _request_ctx;
+    Value::UP _query_tensor;
+
+    void insert_tensor_in_properties(const vespalib::string& tensor_name, const Value& tensor_value) {
+        vespalib::nbostream stream;
+        DefaultTensorEngine::ref().encode(tensor_value, stream);
+        _props.add(tensor_name, vespalib::stringref(stream.c_str(), stream.size()));
+    }
+
+public:
+    RequestContextTest()
+        : _clock(),
+          _doom(_clock, fastos::SteadyTimeStamp()),
+          _attr_ctx(),
+          _props(),
+          _request_ctx(_doom, _attr_ctx, _props),
+          _query_tensor(DefaultTensorEngine::ref().from_spec(TensorSpec("tensor(x[2])")
+                                                                     .add({{"x", 0}}, 3).add({{"x", 1}}, 5)))
+    {
+        insert_tensor_in_properties("my_tensor", *_query_tensor);
+        _props.add("my_string", "foo bar");
+    }
+    TensorSpec expected_query_tensor() const { return DefaultTensorEngine::ref().to_spec(*_query_tensor); }
+    Value::UP get_query_tensor(const vespalib::string& tensor_name) const {
+        return _request_ctx.get_query_tensor(tensor_name);
+    }
+};
+
+TEST_F(RequestContextTest, query_tensor_can_be_retrieved)
+{
+    auto tensor = get_query_tensor("my_tensor");
+    ASSERT_TRUE(tensor);
+    EXPECT_TRUE(tensor->is_tensor());
+    EXPECT_EQ(expected_query_tensor(), DefaultTensorEngine::ref().to_spec(*tensor));
+}
+
+TEST_F(RequestContextTest, non_existing_query_tensor_returns_nullptr)
+{
+    auto tensor = get_query_tensor("non_existing");
+    EXPECT_FALSE(tensor);
+}
+
+TEST_F(RequestContextTest, rank_property_of_non_tensor_type_returns_nullptr)
+{
+    auto tensor = get_query_tensor("my_string");
+    EXPECT_FALSE(tensor);
+}
+
+GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/searchcore/src/vespa/searchcore/proton/matching/match_tools.cpp b/searchcore/src/vespa/searchcore/proton/matching/match_tools.cpp
index 22912671def..302765e152b 100644
--- a/searchcore/src/vespa/searchcore/proton/matching/match_tools.cpp
+++ b/searchcore/src/vespa/searchcore/proton/matching/match_tools.cpp
@@ -155,7 +155,7 @@ MatchToolsFactory(QueryLimiter               & queryLimiter,
                   const Properties           & rankProperties,
                   const Properties           & featureOverrides)
     : _queryLimiter(queryLimiter),
-      _requestContext(softDoom, attributeContext),
+      _requestContext(softDoom, attributeContext, rankProperties),
       _hardDoom(hardDoom),
       _query(),
       _match_limiter(),
diff --git a/searchcore/src/vespa/searchcore/proton/matching/requestcontext.cpp b/searchcore/src/vespa/searchcore/proton/matching/requestcontext.cpp
index 07c672a3fe1..918e4a14649 100644
--- a/searchcore/src/vespa/searchcore/proton/matching/requestcontext.cpp
+++ b/searchcore/src/vespa/searchcore/proton/matching/requestcontext.cpp
@@ -1,15 +1,26 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
 #include "requestcontext.h"
+#include <vespa/eval/tensor/default_tensor_engine.h>
 #include <vespa/searchlib/attribute/attributevector.h>
+#include <vespa/searchlib/fef/properties.h>
+#include <vespa/vespalib/objects/nbostream.h>
+#include <vespa/vespalib/util/exceptions.h>
+
+#include <vespa/log/log.h>
+LOG_SETUP(".proton.matching.requestcontext");
 
 namespace proton {
 
 using search::attribute::IAttributeVector;
 
-RequestContext::RequestContext(const Doom & softDoom, IAttributeContext & attributeContext) :
-    _softDoom(softDoom),
-    _attributeContext(attributeContext)
-{ }
+RequestContext::RequestContext(const Doom & softDoom, IAttributeContext & attributeContext,
+                               const search::fef::Properties& rank_properties)
+    : _softDoom(softDoom),
+      _attributeContext(attributeContext),
+      _rank_properties(rank_properties)
+{
+}
 
 const search::attribute::IAttributeVector *
 RequestContext::getAttribute(const vespalib::string &name) const
@@ -23,8 +34,27 @@ RequestContext::getAttributeStableEnum(const vespalib::string &name) const
     return _attributeContext.getAttributeStableEnum(name);
 }
 
-void RequestContext::asyncForAttribute(const vespalib::string &name, std::unique_ptr<IAttributeFunctor> func) const {
+void
+RequestContext::asyncForAttribute(const vespalib::string &name, std::unique_ptr<IAttributeFunctor> func) const
+{
     _attributeContext.asyncForAttribute(name, std::move(func));
 }
 
+vespalib::eval::Value::UP
+RequestContext::get_query_tensor(const vespalib::string& tensor_name) const
+{
+    auto property = _rank_properties.lookup(tensor_name);
+    if (property.found() && !property.get().empty()) {
+        const vespalib::string& value = property.get();
+        vespalib::nbostream stream(value.data(), value.size());
+        try {
+            return vespalib::tensor::DefaultTensorEngine::ref().decode(stream);
+        } catch (vespalib::IllegalArgumentException& ex) {
+            LOG(warning, "Query tensor '%s' could not be deserialized", tensor_name.c_str());
+            return vespalib::eval::Value::UP();
+        }
+    }
+    return vespalib::eval::Value::UP();
+}
+
 }
diff --git a/searchcore/src/vespa/searchcore/proton/matching/requestcontext.h b/searchcore/src/vespa/searchcore/proton/matching/requestcontext.h
index 19e6331d14d..0352e28eea2 100644
--- a/searchcore/src/vespa/searchcore/proton/matching/requestcontext.h
+++ b/searchcore/src/vespa/searchcore/proton/matching/requestcontext.h
@@ -2,9 +2,12 @@
 
 #pragma once
 
+#include <vespa/eval/tensor/tensor.h>
 #include <vespa/searchlib/queryeval/irequestcontext.h>
 #include <vespa/searchcommon/attribute/iattributecontext.h>
 
+namespace search::fef { class Properties; }
+
 namespace proton {
 
 class RequestContext : public search::queryeval::IRequestContext,
@@ -14,16 +17,23 @@ public:
     using IAttributeContext = search::attribute::IAttributeContext;
     using IAttributeFunctor = search::attribute::IAttributeFunctor;
     using Doom = vespalib::Doom;
-    RequestContext(const Doom & softDoom, IAttributeContext & attributeContext);
+    RequestContext(const Doom & softDoom, IAttributeContext & attributeContext,
+                   const search::fef::Properties& rank_properties);
+
     const Doom & getSoftDoom() const override { return _softDoom; }
     const search::attribute::IAttributeVector *getAttribute(const vespalib::string &name) const override;
 
     void asyncForAttribute(const vespalib::string &name, std::unique_ptr<IAttributeFunctor> func) const override;
 
     const search::attribute::IAttributeVector *getAttributeStableEnum(const vespalib::string &name) const override;
+
+    vespalib::eval::Value::UP get_query_tensor(const vespalib::string& tensor_name) const override;
+
+
 private:
     const Doom          _softDoom;
     IAttributeContext & _attributeContext;
+    const search::fef::Properties& _rank_properties;
 };
 
 }