diff options
author | Geir Storli <geirst@verizonmedia.com> | 2020-06-15 12:00:36 +0000 |
---|---|---|
committer | Geir Storli <geirst@verizonmedia.com> | 2020-06-15 12:00:36 +0000 |
commit | f0ce676f7de72f90cd5919c33d2fa3fb6ec76c4f (patch) | |
tree | f32c6cab89e8ec0bbfbd37e0a18a86afc1d3862c /searchcore | |
parent | 8c0fa90a06966d7f6411e915b7c0d6906c53b130 (diff) |
Test that attribute writer can handle put in two phases.
Diffstat (limited to 'searchcore')
-rw-r--r-- | searchcore/src/tests/proton/attribute/attribute_test.cpp | 95 |
1 files changed, 86 insertions, 9 deletions
diff --git a/searchcore/src/tests/proton/attribute/attribute_test.cpp b/searchcore/src/tests/proton/attribute/attribute_test.cpp index c101c3e2bd5..feebe63f01a 100644 --- a/searchcore/src/tests/proton/attribute/attribute_test.cpp +++ b/searchcore/src/tests/proton/attribute/attribute_test.cpp @@ -34,6 +34,7 @@ #include <vespa/searchlib/index/dummyfileheadercontext.h> #include <vespa/searchlib/predicate/predicate_hash.h> #include <vespa/searchlib/predicate/predicate_index.h> +#include <vespa/searchlib/tensor/dense_tensor_attribute.h> #include <vespa/searchlib/tensor/tensor_attribute.h> #include <vespa/searchlib/test/directory_handler.h> #include <vespa/vespalib/btree/btreeroot.hpp> @@ -71,6 +72,8 @@ using search::index::DummyFileHeaderContext; using search::index::schema::CollectionType; using search::predicate::PredicateHash; using search::predicate::PredicateIndex; +using search::tensor::DenseTensorAttribute; +using search::tensor::PrepareResult; using search::tensor::TensorAttribute; using search::test::DirectoryHandler; using std::string; @@ -153,10 +156,14 @@ public: } AttributeVector::SP addAttribute(const AttributeSpec &spec) { auto ret = _mgr->addAttribute(spec.getName(), - AttributeFactory::createAttribute(spec.getName(), spec.getConfig())); + AttributeFactory::createAttribute(spec.getName(), spec.getConfig())); allocAttributeWriter(); return ret; } + void add_attribute(AttributeVector::SP attr) { + _mgr->addAttribute(attr->getName(), std::move(attr)); + allocAttributeWriter(); + } void put(SerialNum serialNum, const Document &doc, DocumentIdT lid, bool immediateCommit = true) { _aw->put(serialNum, doc, lid, immediateCommit, emptyCallback); @@ -625,18 +632,20 @@ Tensor::UP make_tensor(const TensorSpec &spec) { return Tensor::UP(dynamic_cast<Tensor*>(tensor.release())); } +const vespalib::string sparse_tensor = "tensor(x{},y{})"; + AttributeVector::SP createTensorAttribute(AttributeWriterTest &t) { AVConfig cfg(AVBasicType::TENSOR); - cfg.setTensorType(ValueType::from_spec("tensor(x{},y{})")); + cfg.setTensorType(ValueType::from_spec(sparse_tensor)); auto ret = t.addAttribute({"a1", cfg}); return ret; } Schema -createTensorSchema() { +createTensorSchema(const vespalib::string& tensor_spec = sparse_tensor) { Schema schema; - schema.addAttributeField(Schema::AttributeField("a1", schema::DataType::TENSOR, CollectionType::SINGLE)); + schema.addAttributeField(Schema::AttributeField("a1", schema::DataType::TENSOR, CollectionType::SINGLE, tensor_spec)); return schema; } @@ -654,7 +663,7 @@ TEST_F(AttributeWriterTest, can_write_to_tensor_attribute) auto a1 = createTensorAttribute(*this); Schema s = createTensorSchema(); DocBuilder builder(s); - auto tensor = make_tensor(TensorSpec("tensor(x{},y{})") + auto tensor = make_tensor(TensorSpec(sparse_tensor) .add({{"x", "4"}, {"y", "5"}}, 7)); Document::UP doc = createTensorPutDoc(builder, *tensor); put(1, *doc, 1); @@ -671,7 +680,7 @@ TEST_F(AttributeWriterTest, handles_tensor_assign_update) auto a1 = createTensorAttribute(*this); Schema s = createTensorSchema(); DocBuilder builder(s); - auto tensor = make_tensor(TensorSpec("tensor(x{},y{})") + auto tensor = make_tensor(TensorSpec(sparse_tensor) .add({{"x", "6"}, {"y", "7"}}, 9)); auto doc = createTensorPutDoc(builder, *tensor); put(1, *doc, 1); @@ -684,9 +693,9 @@ TEST_F(AttributeWriterTest, handles_tensor_assign_update) const document::DocumentType &dt(builder.getDocumentType()); DocumentUpdate upd(*builder.getDocumentTypeRepo(), dt, DocumentId("id:ns:searchdocument::1")); - auto new_tensor = make_tensor(TensorSpec("tensor(x{},y{})") + auto new_tensor = make_tensor(TensorSpec(sparse_tensor) .add({{"x", "8"}, {"y", "9"}}, 11)); - TensorDataType xySparseTensorDataType(vespalib::eval::ValueType::from_spec("tensor(x{},y{})")); + TensorDataType xySparseTensorDataType(vespalib::eval::ValueType::from_spec(sparse_tensor)); TensorFieldValue new_value(xySparseTensorDataType); new_value = new_tensor->clone(); upd.addUpdate(FieldUpdate(upd.getType().getField("a1")) @@ -762,15 +771,55 @@ TEST_F(AttributeWriterTest, spreads_write_over_3_write_contexts) putAttributes(*this, {0, 1, 2}); } +struct MockPrepareResult : public PrepareResult { + uint32_t docid; + const Tensor& tensor; + MockPrepareResult(uint32_t docid_in, const Tensor& tensor_in) : docid(docid_in), tensor(tensor_in) {} +}; + +class MockDenseTensorAttribute : public DenseTensorAttribute { +public: + mutable size_t prepare_set_tensor_cnt; + mutable size_t complete_set_tensor_cnt; + + MockDenseTensorAttribute(vespalib::stringref name, const AVConfig& cfg) + : DenseTensorAttribute(name, cfg), + prepare_set_tensor_cnt(0), + complete_set_tensor_cnt(0) + {} + std::unique_ptr<PrepareResult> prepare_set_tensor(uint32_t docid, const Tensor& tensor) const override { + ++prepare_set_tensor_cnt; + return std::make_unique<MockPrepareResult>(docid, tensor); + } + + virtual void complete_set_tensor(DocId docid, const Tensor& tensor, std::unique_ptr<PrepareResult> prepare_result) override { + ++complete_set_tensor_cnt; + assert(prepare_result); + auto* mock_result = dynamic_cast<MockPrepareResult*>(prepare_result.get()); + assert(mock_result); + EXPECT_EQ(docid, mock_result->docid); + EXPECT_EQ(tensor, mock_result->tensor); + } +}; + +const vespalib::string dense_tensor = "tensor(x[2])"; + AVConfig get_tensor_config(bool allow_multi_threaded_indexing) { AVConfig cfg(AVBasicType::TENSOR); - cfg.setTensorType(ValueType::from_spec("tensor(x[2])")); + cfg.setTensorType(ValueType::from_spec(dense_tensor)); cfg.set_hnsw_index_params(HnswIndexParams(4, 4, DistanceMetric::Euclidean, allow_multi_threaded_indexing)); return cfg; } +std::shared_ptr<MockDenseTensorAttribute> +make_mock_tensor_attribute(const vespalib::string& name, bool allow_multi_threaded_indexing) +{ + auto cfg = get_tensor_config(allow_multi_threaded_indexing); + return std::make_shared<MockDenseTensorAttribute>(name, cfg); +} + TEST_F(AttributeWriterTest, tensor_attributes_using_two_phase_put_are_in_separate_write_contexts) { addAttribute("a1"); @@ -793,6 +842,34 @@ TEST_F(AttributeWriterTest, tensor_attributes_using_two_phase_put_are_in_separat EXPECT_EQ("t2", ctx[2].getFields()[0].getAttribute().getName()); } +TEST_F(AttributeWriterTest, handles_put_in_two_phases_when_specified_for_tensor_attribute) +{ + setup(2); + auto a1 = make_mock_tensor_attribute("a1", true); + add_attribute(a1); + Schema schema = createTensorSchema(dense_tensor); + DocBuilder builder(schema); + auto tensor = make_tensor(TensorSpec(dense_tensor) + .add({{"x", 0}}, 3).add({{"x", 1}}, 5)); + auto doc = createTensorPutDoc(builder, *tensor); + + put(1, *doc, 1); + EXPECT_EQ(1, a1->prepare_set_tensor_cnt); + EXPECT_EQ(1, a1->complete_set_tensor_cnt); + assertExecuteHistory({1, 0}); + + put(2, *doc, 2); + EXPECT_EQ(2, a1->prepare_set_tensor_cnt); + EXPECT_EQ(2, a1->complete_set_tensor_cnt); + assertExecuteHistory({1, 0, 0, 0}); + + put(3, *doc, 3); + EXPECT_EQ(3, a1->prepare_set_tensor_cnt); + EXPECT_EQ(3, a1->complete_set_tensor_cnt); + // Note that the prepare step is executed round-robin between the 2 threads. + assertExecuteHistory({1, 0, 0, 0, 1, 0}); +} + ImportedAttributeVector::SP createImportedAttribute(const vespalib::string &name) |