From 95bdc81ea132c3d1375e6c0c7dbc1833534b208a Mon Sep 17 00:00:00 2001 From: Geir Storli Date: Wed, 20 Nov 2019 12:22:56 +0000 Subject: Expose tensors passed with the query in IRequestContext. --- searchcore/CMakeLists.txt | 1 + .../proton/matching/request_context/CMakeLists.txt | 9 +++ .../request_context/request_context_test.cpp | 83 ++++++++++++++++++++++ .../searchcore/proton/matching/match_tools.cpp | 2 +- .../searchcore/proton/matching/requestcontext.cpp | 40 +++++++++-- .../searchcore/proton/matching/requestcontext.h | 12 +++- .../searchlib/queryeval/fake_requestcontext.h | 8 ++- .../vespa/searchlib/queryeval/irequestcontext.h | 7 ++ 8 files changed, 154 insertions(+), 8 deletions(-) create mode 100644 searchcore/src/tests/proton/matching/request_context/CMakeLists.txt create mode 100644 searchcore/src/tests/proton/matching/request_context/request_context_test.cpp 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..98173066f93 --- /dev/null +++ b/searchcore/src/tests/proton/matching/request_context/request_context_test.cpp @@ -0,0 +1,83 @@ +// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +#include +#include +#include +#include +#include +#include +#include +#include + +using search::attribute::IAttributeContext; +using search::attribute::IAttributeFunctor; +using search::attribute::IAttributeVector; +using search::fef::Properties; +using vespalib::eval::TensorSpec; +using vespalib::tensor::Tensor; +using vespalib::tensor::TypedBinaryFormat; +using vespalib::tensor::test::makeTensor; +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 override { abort(); } + void asyncForAttribute(const vespalib::string&, std::unique_ptr) const override { abort(); } +}; + +class RequestContextTest : public ::testing::Test { +private: + vespalib::Clock _clock; + vespalib::Doom _doom; + MyAttributeContext _attr_ctx; + Properties _props; + RequestContext _request_ctx; + Tensor::UP _query_tensor; + + void insert_tensor_in_properties(const vespalib::string& tensor_name, const Tensor& tensor) { + vespalib::nbostream stream; + TypedBinaryFormat::serialize(stream, tensor); + _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(makeTensor(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"); + } + const Tensor& expected_query_tensor() const { return *_query_tensor; } + Tensor::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_EQUAL(expected_query_tensor(), *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..1669b90f97e 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 #include +#include +#include +#include + +#include +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 func) const { +void +RequestContext::asyncForAttribute(const vespalib::string &name, std::unique_ptr func) const +{ _attributeContext.asyncForAttribute(name, std::move(func)); } +vespalib::tensor::Tensor::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::TypedBinaryFormat::deserialize(stream); + } catch (vespalib::IllegalArgumentException& ex) { + LOG(warning, "Query tensor '%s' could not be deserialized", tensor_name.c_str()); + return vespalib::tensor::Tensor::UP(); + } + } + return vespalib::tensor::Tensor::UP(); +} + } diff --git a/searchcore/src/vespa/searchcore/proton/matching/requestcontext.h b/searchcore/src/vespa/searchcore/proton/matching/requestcontext.h index 19e6331d14d..02dd96a9771 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 #include #include +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 func) const override; const search::attribute::IAttributeVector *getAttributeStableEnum(const vespalib::string &name) const override; + + vespalib::tensor::Tensor::UP get_query_tensor(const vespalib::string& tensor_name) const override; + + private: const Doom _softDoom; IAttributeContext & _attributeContext; + const search::fef::Properties& _rank_properties; }; } diff --git a/searchlib/src/vespa/searchlib/queryeval/fake_requestcontext.h b/searchlib/src/vespa/searchlib/queryeval/fake_requestcontext.h index fefa7f5b6ff..6940438ce43 100644 --- a/searchlib/src/vespa/searchlib/queryeval/fake_requestcontext.h +++ b/searchlib/src/vespa/searchlib/queryeval/fake_requestcontext.h @@ -2,9 +2,10 @@ #pragma once -#include +#include #include #include +#include #include namespace search::queryeval { @@ -24,6 +25,11 @@ public: ? _attributeContext->getAttribute(name) : nullptr; } + vespalib::tensor::Tensor::UP get_query_tensor(const vespalib::string& tensor_name) const override { + (void) tensor_name; + return vespalib::tensor::Tensor::UP(); + } + private: vespalib::Clock _clock; const vespalib::Doom _doom; diff --git a/searchlib/src/vespa/searchlib/queryeval/irequestcontext.h b/searchlib/src/vespa/searchlib/queryeval/irequestcontext.h index b779ac400bf..88399406292 100644 --- a/searchlib/src/vespa/searchlib/queryeval/irequestcontext.h +++ b/searchlib/src/vespa/searchlib/queryeval/irequestcontext.h @@ -6,6 +6,7 @@ #include namespace search::attribute { class IAttributeVector; } +namespace vespalib::tensor { class Tensor; } namespace search::queryeval { @@ -29,6 +30,12 @@ public: */ virtual const attribute::IAttributeVector *getAttribute(const vespalib::string &name) const = 0; virtual const attribute::IAttributeVector *getAttributeStableEnum(const vespalib::string &name) const = 0; + + /** + * Returns the tensor of the given name that was passed with the query. + * Returns nullptr if the tensor is not found or if it is not a tensor. + */ + virtual std::unique_ptr get_query_tensor(const vespalib::string& tensor_name) const = 0; }; } -- cgit v1.2.3