wire in new prepare/complete API

2 files changed, 32 insertions, 2 deletions

diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
index 719144efa41..2b908dd2bd4 100644
--- a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
@@ -11,6 +11,9 @@
 #include <vespa/vespalib/data/slime/inserter.h>
 #include <vespa/vespalib/datastore/array_store.hpp>
 #include <vespa/vespalib/util/rcuvector.hpp>
+#include <vespa/log/log.h>
+
+LOG_SETUP(".searchlib.tensor.hnsw_index");
 
 namespace search::tensor {
 
@@ -274,7 +277,6 @@ HnswIndex::add_document(uint32_t docid)
 {
     AddDocOperation op = internal_prepare_add(docid, get_vector(docid));
     internal_complete_add(docid, op);
-    auto up = std::make_unique<AddDocOperation>(std::move(op));
 }
 
 HnswIndex::AddDocOperation
@@ -337,6 +339,29 @@ HnswIndex::internal_complete_add(uint32_t docid, AddDocOperation &op)
     }
 }
 
+std::unique_ptr<PrepareResult>
+HnswIndex::prepare_add_document(uint32_t docid, 
+            TypedCells vector,
+            vespalib::GenerationHandler::Guard read_guard) const
+{
+    AddDocOperation op = internal_prepare_add(docid, vector);
+    (void) read_guard; // must keep guard until this point
+    return std::make_unique<AddDocOperation>(std::move(op));
+}
+
+void
+HnswIndex::complete_add_document(uint32_t docid, std::unique_ptr<PrepareResult> prepare_result)
+{
+    auto prepared = dynamic_cast<AddDocOperation *>(prepare_result.get());
+    if (prepared && (prepared->docid == docid)) {
+        internal_complete_add(docid, *prepared);
+    } else {
+        LOG(warning, "complete_add_document called with invalid prepare_result");
+        // fallback to normal add
+        add_document(docid);
+    }
+}
+
 void
 HnswIndex::mutual_reconnect(const LinkArrayRef &cluster, uint32_t level)
 {
diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index.h b/searchlib/src/vespa/searchlib/tensor/hnsw_index.h
index 1479e41b999..47712f56c18 100644
--- a/searchlib/src/vespa/searchlib/tensor/hnsw_index.h
+++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index.h
@@ -120,13 +120,14 @@ protected:
     std::vector<Neighbor> top_k_by_docid(uint32_t k, TypedCells vector,
                                          const BitVector *filter, uint32_t explore_k) const;
 
-    struct AddDocOperation {
+    struct AddDocOperation : public PrepareResult {
         uint32_t docid;
         int32_t max_level;
         std::vector<LinkArray> connections;
         AddDocOperation(uint32_t docid_in, int32_t max_level_in)
           : docid(docid_in), max_level(max_level_in), connections(max_level+1)
         {}
+        ~AddDocOperation() = default;
         AddDocOperation(AddDocOperation&& other) = default;
     };
     AddDocOperation internal_prepare_add(uint32_t docid, TypedCells input_vector) const;
@@ -141,6 +142,10 @@ public:
 
     // Implements NearestNeighborIndex
     void add_document(uint32_t docid) override;
+    std::unique_ptr<PrepareResult> prepare_add_document(uint32_t docid,
+            TypedCells vector,
+            vespalib::GenerationHandler::Guard read_guard) const override;
+    void complete_add_document(uint32_t docid, std::unique_ptr<PrepareResult> prepare_result) override;
     void remove_document(uint32_t docid) override;
     void transfer_hold_lists(generation_t current_gen) override;
     void trim_hold_lists(generation_t first_used_gen) override;