Răsfoiți Sursa

feat: SimilaritySearch gRPC handler + vector persistence in CRUD

- Implement SimilaritySearch RPC: converts query vector, calls
  store_.similaritySearch(), decrypts sensitive fields in results
- Wire vector_dimension in CreateCollection and GetCollectionInfo
- Add VEC_PUT persistence after Insert/Update/Upsert when document
  contains a _vector field
- Add VEC_DELETE persistence in Delete (checks for vector before
  remove() erases it) and Update/Upsert when vector is removed
fszontagh 4 luni în urmă
părinte
comite
35addbe0b4
2 a modificat fișierele cu 96 adăugiri și 0 ștergeri
  1. 90 0
      service/src/database_grpc_impl.cpp
  2. 6 0
      service/src/database_grpc_impl.hpp

+ 90 - 0
service/src/database_grpc_impl.cpp

@@ -55,6 +55,15 @@ grpc::Status DatabaseGrpcImpl::Insert(
 
         std::string id = store_.insert(request->collection(), doc);
 
+        // Persist vector if the document had a _vector field
+        auto* vectors = store_.getCollectionVectors(request->collection());
+        if (vectors) {
+            auto vit = vectors->find(id);
+            if (vit != vectors->end()) {
+                persistence_.logVecPut(request->collection(), id, vit->second);
+            }
+        }
+
         response->set_id(id);
         response->set_version(1);
 
@@ -130,6 +139,18 @@ grpc::Status DatabaseGrpcImpl::Update(
             }
         }
 
+        // Persist vector changes
+        auto* vectors = store_.getCollectionVectors(request->collection());
+        if (vectors) {
+            auto vit = vectors->find(request->id());
+            if (vit != vectors->end()) {
+                persistence_.logVecPut(request->collection(), request->id(), vit->second);
+            } else {
+                // Vector was removed (update without _vector field removes existing vector)
+                persistence_.logVecDelete(request->collection(), request->id());
+            }
+        }
+
         // Get updated version
         auto updated = store_.get(request->collection(), request->id());
         response->set_new_version(updated ? updated->version : 0);
@@ -185,6 +206,18 @@ grpc::Status DatabaseGrpcImpl::Upsert(
 
         std::string id = store_.upsert(request->collection(), doc);
 
+        // Persist vector changes
+        auto* vectors = store_.getCollectionVectors(request->collection());
+        if (vectors) {
+            auto vit = vectors->find(id);
+            if (vit != vectors->end()) {
+                persistence_.logVecPut(request->collection(), id, vit->second);
+            } else if (existed) {
+                // Existing doc updated without _vector — remove old vector
+                persistence_.logVecDelete(request->collection(), id);
+            }
+        }
+
         auto updated = store_.get(request->collection(), id);
         response->set_id(id);
         response->set_version(updated ? updated->version : 1);
@@ -205,7 +238,19 @@ grpc::Status DatabaseGrpcImpl::Delete(
     const pb::DeleteRequest* request,
     pb::DeleteResponse* response
 ) {
+    // Check if document has a vector before deleting (remove() erases it from memory)
+    bool hadVector = false;
+    auto* vectors = store_.getCollectionVectors(request->collection());
+    if (vectors) {
+        hadVector = vectors->count(request->id()) > 0;
+    }
+
     bool deleted = store_.remove(request->collection(), request->id());
+
+    if (deleted && hadVector) {
+        persistence_.logVecDelete(request->collection(), request->id());
+    }
+
     response->set_deleted(deleted);
     return grpc::Status::OK;
 }
@@ -415,6 +460,48 @@ grpc::Status DatabaseGrpcImpl::Count(
     return grpc::Status::OK;
 }
 
+grpc::Status DatabaseGrpcImpl::SimilaritySearch(
+    grpc::ServerContext* /*context*/,
+    const pb::SimilaritySearchRequest* request,
+    pb::SimilaritySearchResponse* response
+) {
+    try {
+        if (request->collection().empty()) {
+            return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "Collection name is required");
+        }
+
+        // Convert proto repeated float to std::vector<float>
+        std::vector<float> queryVec(request->query_vector().begin(),
+                                     request->query_vector().end());
+        if (queryVec.empty()) {
+            return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "Query vector is required");
+        }
+
+        uint32_t topK = request->top_k() > 0 ? request->top_k() : 10;
+        float minScore = request->min_score();
+
+        auto results = store_.similaritySearch(
+            request->collection(), queryVec, topK, minScore);
+
+        for (auto& result : results) {
+            // Decrypt sensitive fields before returning
+            encryption_.decryptSensitiveFields(result.document);
+
+            auto* protoResult = response->add_results();
+            protoResult->set_id(result.id);
+            protoResult->set_score(result.score);
+            protoResult->set_data(result.document.data.dump());
+        }
+
+        return grpc::Status::OK;
+
+    } catch (const std::invalid_argument& e) {
+        return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, e.what());
+    } catch (const std::exception& e) {
+        return grpc::Status(grpc::StatusCode::INTERNAL, e.what());
+    }
+}
+
 // ===== Set Operations =====
 
 grpc::Status DatabaseGrpcImpl::SetAdd(
@@ -474,6 +561,7 @@ grpc::Status DatabaseGrpcImpl::CreateCollection(
     opts.defaultTtlSeconds = request->options().default_ttl_seconds();
     opts.encrypted = request->options().encrypted();
     opts.maxVersions = request->options().max_versions();
+    opts.vectorDimension = request->options().vector_dimension();
     for (const auto& field : request->options().sensitive_fields()) {
         opts.sensitiveFields.push_back(field);
     }
@@ -832,11 +920,13 @@ pb::CollectionInfo DatabaseGrpcImpl::toProtoCollInfo(const struct CollectionInfo
     proto.mutable_options()->set_default_ttl_seconds(info.options.defaultTtlSeconds);
     proto.mutable_options()->set_encrypted(info.options.encrypted);
     proto.mutable_options()->set_max_versions(info.options.maxVersions);
+    proto.mutable_options()->set_vector_dimension(info.options.vectorDimension);
     for (const auto& field : info.options.sensitiveFields) {
         proto.mutable_options()->add_sensitive_fields(field);
     }
     proto.set_created_at(info.createdAt);
     proto.set_updated_at(info.updatedAt);
+    proto.set_vector_dimension(info.options.vectorDimension);
     return proto;
 }
 

+ 6 - 0
service/src/database_grpc_impl.hpp

@@ -132,6 +132,12 @@ public:
         pb::CountResponse* response
     ) override;
 
+    grpc::Status SimilaritySearch(
+        grpc::ServerContext* context,
+        const pb::SimilaritySearchRequest* request,
+        pb::SimilaritySearchResponse* response
+    ) override;
+
     // ===== Set Operations =====
 
     grpc::Status SetAdd(