Merge pull request #24622 from vespa-engine/geirst/attribute-builder-helper

Add a helper class to build and fill AttributeVector instances in uni…

5 files changed, 180 insertions, 58 deletions

diff --git a/searchlib/src/tests/attribute/searchable/CMakeLists.txt b/searchlib/src/tests/attribute/searchable/CMakeLists.txt
index 29a4c122bf7..c3af20cc673 100644
--- a/searchlib/src/tests/attribute/searchable/CMakeLists.txt
+++ b/searchlib/src/tests/attribute/searchable/CMakeLists.txt
@@ -19,6 +19,7 @@ vespa_add_executable(searchlib_attribute_blueprint_test_app TEST
     attributeblueprint_test.cpp
     DEPENDS
     searchlib
+    searchlib_test
     GTest::GTest
 )
 vespa_add_test(NAME searchlib_attribute_blueprint_test_app COMMAND searchlib_attribute_blueprint_test_app)
diff --git a/searchlib/src/tests/attribute/searchable/attributeblueprint_test.cpp b/searchlib/src/tests/attribute/searchable/attributeblueprint_test.cpp
index 29dbf33d29c..504cc6506b5 100644
--- a/searchlib/src/tests/attribute/searchable/attributeblueprint_test.cpp
+++ b/searchlib/src/tests/attribute/searchable/attributeblueprint_test.cpp
@@ -1,10 +1,15 @@
 // Copyright Yahoo. 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/eval/value.h>
+#include <vespa/eval/eval/value_codec.h>
+#include <vespa/searchcommon/attribute/config.h>
+#include <vespa/searchcommon/attribute/iattributecontext.h>
+#include <vespa/searchlib/attribute/attribute.h>
 #include <vespa/searchlib/attribute/attribute_blueprint_factory.h>
 #include <vespa/searchlib/attribute/attribute_read_guard.h>
 #include <vespa/searchlib/attribute/attributecontext.h>
 #include <vespa/searchlib/attribute/attributefactory.h>
-#include <vespa/searchlib/attribute/attribute.h>
 #include <vespa/searchlib/fef/matchdata.h>
 #include <vespa/searchlib/query/tree/location.h>
 #include <vespa/searchlib/query/tree/point.h>
@@ -14,11 +19,7 @@
 #include <vespa/searchlib/queryeval/leaf_blueprints.h>
 #include <vespa/searchlib/queryeval/nearest_neighbor_blueprint.h>
 #include <vespa/searchlib/tensor/dense_tensor_attribute.h>
-#include <vespa/searchcommon/attribute/iattributecontext.h>
-#include <vespa/searchcommon/attribute/config.h>
-#include <vespa/eval/eval/tensor_spec.h>
-#include <vespa/eval/eval/value.h>
-#include <vespa/eval/eval/value_codec.h>
+#include <vespa/searchlib/test/attribute_builder.h>
 #include <vespa/vespalib/gtest/gtest.h>
 
 #include <vespa/log/log.h>
@@ -28,6 +29,7 @@ using search::AttributeGuard;
 using search::AttributeVector;
 using search::IAttributeManager;
 using search::attribute::IAttributeContext;
+using search::attribute::test::AttributeBuilder;
 using search::fef::MatchData;
 using search::fef::TermFieldMatchData;
 using search::query::Location;
@@ -144,53 +146,11 @@ downcast(ParentType& parent)
     return *result;
 }
 
-struct StringAttributeFiller {
-    using ValueType = vespalib::string;
-    static void add(AttributeVector& attr, const vespalib::string& value) {
-        auto& real = downcast<StringAttribute>(attr);
-        real.update(attr.getNumDocs() - 1, value);
-        real.commit();
-    }
-};
-
-struct WsetStringAttributeFiller {
-    using ValueType = vespalib::string;
-    static void add(AttributeVector& attr, const vespalib::string& value) {
-        auto& real = downcast<StringAttribute>(attr);
-        uint32_t docid = attr.getNumDocs() - 1;
-        real.append(docid, value, 1);
-        real.commit();
-    }
-};
-
-struct IntegerAttributeFiller {
-    using ValueType = int64_t;
-    static void add(AttributeVector& attr, int64_t value) {
-        auto& real = downcast<IntegerAttribute>(attr);
-        real.update(attr.getNumDocs() - 1, value);
-        real.commit();
-    }
-};
-
-template <typename FillerType>
-void
-fill(AttributeVector& attr, typename FillerType::ValueType value)
-{
-    AttributeVector::DocId docid;
-    attr.addDoc(docid);
-    attr.addDoc(docid);
-    attr.addDoc(docid);
-    assert(DOCID_LIMIT-1 == docid);
-    FillerType::add(attr, value);
-}
-
 AttributeVector::SP
 make_string_attribute(const std::string& value)
 {
     Config cfg(BasicType::STRING, CollectionType::SINGLE);
-    auto attr = AttributeFactory::createAttribute(field, cfg);
-    fill<StringAttributeFiller>(*attr, value);
-    return attr;
+    return AttributeBuilder(field, cfg).fill({"", value}).get();
 }
 
 AttributeVector::SP
@@ -199,18 +159,14 @@ make_wset_string_attribute(const std::string& value)
     Config cfg(BasicType::STRING, CollectionType::WSET);
     // fast-search is needed to trigger use of DirectAttributeBlueprint.
     cfg.setFastSearch(true);
-    auto attr = AttributeFactory::createAttribute(field, cfg);
-    fill<WsetStringAttributeFiller>(*attr, value);
-    return attr;
+    return AttributeBuilder(field, cfg).fill_array({{}, {value}}).get();
 }
 
 AttributeVector::SP
 make_int_attribute(int64_t value)
 {
     Config cfg(BasicType::INT32, CollectionType::SINGLE);
-    auto attr = AttributeFactory::createAttribute(field, cfg);
-    fill<IntegerAttributeFiller>(*attr, value);
-    return attr;
+    return AttributeBuilder(field, cfg).fill({-1, value}).get();
 }
 
 AttributeVector::SP
@@ -218,9 +174,7 @@ make_fast_search_long_attribute(int64_t value)
 {
     Config cfg(BasicType::fromType(int64_t()), CollectionType::SINGLE);
     cfg.setFastSearch(true);
-    auto attr = AttributeFactory::createAttribute(field, cfg);
-    fill<IntegerAttributeFiller>(*attr, value);
-    return attr;
+    return AttributeBuilder(field, cfg).fill({-1, value}).get();
 }
 
 MyAttributeManager
diff --git a/searchlib/src/vespa/searchlib/test/CMakeLists.txt b/searchlib/src/vespa/searchlib/test/CMakeLists.txt
index b53b3097850..7decdb992e6 100644
--- a/searchlib/src/vespa/searchlib/test/CMakeLists.txt
+++ b/searchlib/src/vespa/searchlib/test/CMakeLists.txt
@@ -1,6 +1,7 @@
 # Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 vespa_add_library(searchlib_test
     SOURCES
+    attribute_builder.cpp
     document_weight_attribute_helper.cpp
     doc_builder.cpp
     imported_attribute_fixture.cpp
diff --git a/searchlib/src/vespa/searchlib/test/attribute_builder.cpp b/searchlib/src/vespa/searchlib/test/attribute_builder.cpp
new file mode 100644
index 00000000000..b804723d3be
--- /dev/null
+++ b/searchlib/src/vespa/searchlib/test/attribute_builder.cpp
@@ -0,0 +1,125 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "attribute_builder.h"
+#include <vespa/searchlib/attribute/attributefactory.h>
+#include <vespa/searchlib/attribute/attributevector.h>
+#include <vespa/searchlib/attribute/integerbase.h>
+#include <vespa/searchlib/attribute/stringbase.h>
+#include <cassert>
+
+namespace search::attribute::test {
+
+AttributeBuilder::AttributeBuilder(const vespalib::string& name, const Config& cfg)
+    : _attr_ptr(AttributeFactory::createAttribute(name, cfg)),
+      _attr(*_attr_ptr)
+{
+}
+
+AttributeBuilder::AttributeBuilder(AttributeVector& attr)
+    : _attr_ptr(),
+      _attr(attr)
+{
+}
+
+namespace {
+
+void
+add_docs(AttributeVector& attr, size_t num_docs)
+{
+    attr.addReservedDoc();
+    attr.addDocs(num_docs);
+}
+
+template <typename AttrType, typename ValueType>
+void
+fill_helper(AttributeVector& attr, const std::vector<ValueType>& values)
+{
+    assert(attr.getConfig().collectionType() == CollectionType::SINGLE);
+    add_docs(attr, values.size());
+    auto& real = dynamic_cast<AttrType&>(attr);
+    for (size_t i = 0; i < values.size(); ++i) {
+        uint32_t docid = (i + 1);
+        real.update(docid, values[i]);
+    }
+    attr.commit(true);
+}
+
+template <typename AttrType, typename ValueType>
+void
+fill_array_helper(AttributeVector& attr, const std::vector<std::vector<ValueType>>& values)
+{
+    assert((attr.getConfig().collectionType() == CollectionType::ARRAY) ||
+           (attr.getConfig().collectionType() == CollectionType::WSET));
+    add_docs(attr, values.size());
+    auto& real = dynamic_cast<AttrType&>(attr);
+    for (size_t i = 0; i < values.size(); ++i) {
+        uint32_t docid = (i + 1);
+        for (auto value : values[i]) {
+            real.append(docid, value, 1);
+        }
+    }
+    attr.commit(true);
+}
+
+template <typename AttrType, typename ValueType>
+void
+fill_wset_helper(AttributeVector& attr, const std::vector<std::vector<std::pair<ValueType, int32_t>>>& values)
+{
+    assert(attr.getConfig().collectionType() == CollectionType::WSET);
+    add_docs(attr, values.size());
+    auto& real = dynamic_cast<AttrType&>(attr);
+    for (size_t i = 0; i < values.size(); ++i) {
+        uint32_t docid = (i + 1);
+        for (auto value : values[i]) {
+            real.append(docid, value.first, value.second);
+        }
+    }
+    attr.commit(true);
+}
+
+}
+
+AttributeBuilder&
+AttributeBuilder::fill(const std::vector<int64_t>& values)
+{
+    fill_helper<IntegerAttribute, int64_t>(_attr, values);
+    return *this;
+}
+
+AttributeBuilder&
+AttributeBuilder::fill_array(const std::vector<std::vector<int64_t>>& values)
+{
+    fill_array_helper<IntegerAttribute, int64_t>(_attr, values);
+    return *this;
+}
+
+AttributeBuilder&
+AttributeBuilder::fill_wset(const std::vector<std::vector<std::pair<int64_t, int32_t>>>& values)
+{
+    fill_wset_helper<IntegerAttribute, int64_t>(_attr, values);
+    return *this;
+}
+
+AttributeBuilder&
+AttributeBuilder::fill(const std::vector<vespalib::string>& values)
+{
+    fill_helper<StringAttribute, vespalib::string>(_attr, values);
+    return *this;
+}
+
+AttributeBuilder&
+AttributeBuilder::fill_array(const std::vector<std::vector<vespalib::string>>& values)
+{
+    fill_array_helper<StringAttribute, vespalib::string>(_attr, values);
+    return *this;
+}
+
+AttributeBuilder&
+AttributeBuilder::fill_wset(const std::vector<std::vector<std::pair<vespalib::string, int32_t>>>& values)
+{
+    fill_wset_helper<StringAttribute, vespalib::string>(_attr, values);
+    return *this;
+}
+
+}
+
diff --git a/searchlib/src/vespa/searchlib/test/attribute_builder.h b/searchlib/src/vespa/searchlib/test/attribute_builder.h
new file mode 100644
index 00000000000..6d3099b5ea1
--- /dev/null
+++ b/searchlib/src/vespa/searchlib/test/attribute_builder.h
@@ -0,0 +1,41 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include <vespa/searchcommon/attribute/config.h>
+#include <vespa/vespalib/stllike/string.h>
+#include <memory>
+#include <utility>
+#include <vector>
+
+namespace search { class AttributeVector; }
+namespace search::attribute { class Config; }
+
+namespace search::attribute::test {
+
+/**
+ * Helper class used to build and fill AttributeVector instances in unit tests.
+ */
+class AttributeBuilder {
+private:
+    std::shared_ptr<AttributeVector> _attr_ptr;
+    AttributeVector& _attr;
+
+public:
+    AttributeBuilder(const vespalib::string& name, const Config& cfg);
+    AttributeBuilder(AttributeVector& attr);
+
+    // Fill functions for integer attributes
+    AttributeBuilder& fill(const std::vector<int64_t>& values);
+    AttributeBuilder& fill_array(const std::vector<std::vector<int64_t>>& values);
+    AttributeBuilder& fill_wset(const std::vector<std::vector<std::pair<int64_t, int32_t>>>& values);
+
+    // Fill functions for string attributes
+    AttributeBuilder& fill(const std::vector<vespalib::string>& values);
+    AttributeBuilder& fill_array(const std::vector<std::vector<vespalib::string>>& values);
+    AttributeBuilder& fill_wset(const std::vector<std::vector<std::pair<vespalib::string, int32_t>>>& values);
+
+    std::shared_ptr<AttributeVector> get() const { return _attr_ptr; }
+};
+
+}