|
|
@@ -1,5 +1,6 @@
|
|
|
#include "database_grpc_impl.hpp"
|
|
|
#include "database_service.hpp"
|
|
|
+#include "storage/cosine_simd.hpp"
|
|
|
#include "storage/document_store.hpp"
|
|
|
#include "json_parse.hpp"
|
|
|
#include "persistence/wal.hpp"
|
|
|
@@ -8,6 +9,12 @@
|
|
|
#include <spdlog/spdlog.h>
|
|
|
#include <nlohmann/json.hpp>
|
|
|
|
|
|
+#include <algorithm>
|
|
|
+#include <functional>
|
|
|
+#include <queue>
|
|
|
+#include <utility>
|
|
|
+#include <vector>
|
|
|
+
|
|
|
namespace smartbotic::database {
|
|
|
|
|
|
// Namespace alias for proto types
|
|
|
@@ -833,9 +840,126 @@ grpc::Status DatabaseGrpcImpl::SimilaritySearch(
|
|
|
|
|
|
uint32_t topK = request->top_k() > 0 ? request->top_k() : 10;
|
|
|
float minScore = request->min_score();
|
|
|
+ const std::string& collection = request->collection();
|
|
|
+
|
|
|
+ // v2.0 Stage 5 — LMDB-first SimilaritySearch. Same gate shape as
|
|
|
+ // Exists/Get/Find: non-system collection + mirror healthy + zero
|
|
|
+ // drift + doc_store_ available. We iterate the
|
|
|
+ // `_vectors_<collection>` sub-db via scan_vectors, score with the
|
|
|
+ // SIMD cosine primitives extracted into storage/cosine_simd.hpp,
|
|
|
+ // maintain a top-K min-heap, then fetch each surviving doc via
|
|
|
+ // doc_store_->get() for the response payload. On LMDB throw, fall
|
|
|
+ // back to MemoryStore::similaritySearch — semantic identity is
|
|
|
+ // guaranteed by the dual-write mirror.
|
|
|
+ bool used_lmdb = false;
|
|
|
+ std::vector<MemoryStore::SimilarityResult> results;
|
|
|
+ if (!collection.empty() && collection[0] != '_'
|
|
|
+ && service_.mirrorHealthy() && service_.mirrorDriftCount() == 0) {
|
|
|
+ if (auto* ds = service_.docStore()) {
|
|
|
+ try {
|
|
|
+ // Pre-compute query norm once outside the iteration.
|
|
|
+ const float queryNorm =
|
|
|
+ smartbotic::db::storage::cosine_simd::compute_norm(
|
|
|
+ queryVec.data(), queryVec.size());
|
|
|
+ if (queryNorm == 0.0f) {
|
|
|
+ // All-zero query vector -> all cosines are 0; matches
|
|
|
+ // MemoryStore behaviour which returns empty.
|
|
|
+ used_lmdb = true;
|
|
|
+ } else {
|
|
|
+ // Min-heap of (score, id) keyed on score-ascending so
|
|
|
+ // pop() evicts the weakest current top-K candidate.
|
|
|
+ using Scored = std::pair<float, std::string>;
|
|
|
+ auto cmp = [](const Scored& a, const Scored& b) {
|
|
|
+ return a.first > b.first; // min-heap on score
|
|
|
+ };
|
|
|
+ std::priority_queue<Scored, std::vector<Scored>,
|
|
|
+ decltype(cmp)> heap(cmp);
|
|
|
+ size_t expectedDim = queryVec.size();
|
|
|
+ bool dim_mismatch = false;
|
|
|
+ std::string mismatch_msg;
|
|
|
+
|
|
|
+ ds->scan_vectors(collection,
|
|
|
+ [&](std::string_view id, const float* data,
|
|
|
+ size_t count) {
|
|
|
+ if (count != expectedDim) {
|
|
|
+ // Capture the first mismatch and let the
|
|
|
+ // scan finish; we surface the error after
|
|
|
+ // the txn closes.
|
|
|
+ if (!dim_mismatch) {
|
|
|
+ dim_mismatch = true;
|
|
|
+ mismatch_msg =
|
|
|
+ "Query vector dimension mismatch: expected " +
|
|
|
+ std::to_string(count) + ", got " +
|
|
|
+ std::to_string(expectedDim);
|
|
|
+ }
|
|
|
+ return;
|
|
|
+ }
|
|
|
+ const float docNorm =
|
|
|
+ smartbotic::db::storage::cosine_simd::compute_norm(
|
|
|
+ data, count);
|
|
|
+ const float score =
|
|
|
+ smartbotic::db::storage::cosine_simd::cosine_sim(
|
|
|
+ queryVec.data(), data, count,
|
|
|
+ queryNorm, docNorm);
|
|
|
+ if (score < minScore) return;
|
|
|
+ if (heap.size() < topK) {
|
|
|
+ heap.emplace(score, std::string(id));
|
|
|
+ } else if (!heap.empty() && score > heap.top().first) {
|
|
|
+ heap.pop();
|
|
|
+ heap.emplace(score, std::string(id));
|
|
|
+ }
|
|
|
+ });
|
|
|
+
|
|
|
+ if (dim_mismatch) {
|
|
|
+ return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
|
|
|
+ mismatch_msg);
|
|
|
+ }
|
|
|
+
|
|
|
+ // Drain heap; results end up in ascending order, so
|
|
|
+ // reverse for descending-by-score response.
|
|
|
+ std::vector<Scored> ranked;
|
|
|
+ ranked.reserve(heap.size());
|
|
|
+ while (!heap.empty()) {
|
|
|
+ ranked.push_back(heap.top());
|
|
|
+ heap.pop();
|
|
|
+ }
|
|
|
+ std::reverse(ranked.begin(), ranked.end());
|
|
|
+
|
|
|
+ results.reserve(ranked.size());
|
|
|
+ for (auto& [score, id] : ranked) {
|
|
|
+ auto doc = ds->get(collection, id);
|
|
|
+ if (!doc) {
|
|
|
+ // Vector survived a doc delete that the
|
|
|
+ // mirror handled but somehow the doc isn't
|
|
|
+ // here — skip; matches MemoryStore which
|
|
|
+ // also drops missing docs from results.
|
|
|
+ continue;
|
|
|
+ }
|
|
|
+ MemoryStore::SimilarityResult r;
|
|
|
+ r.id = id;
|
|
|
+ r.score = score;
|
|
|
+ r.document = std::move(*doc);
|
|
|
+ results.push_back(std::move(r));
|
|
|
+ }
|
|
|
+ used_lmdb = true;
|
|
|
+ }
|
|
|
+ } catch (const std::invalid_argument&) {
|
|
|
+ // Argument-shape errors should propagate, not silently
|
|
|
+ // fall back; re-throw so the outer catch returns
|
|
|
+ // INVALID_ARGUMENT.
|
|
|
+ throw;
|
|
|
+ } catch (const std::exception& e) {
|
|
|
+ spdlog::warn("v2.0 SimilaritySearch LMDB scan failed "
|
|
|
+ "coll={}: {} -- falling back to MemoryStore",
|
|
|
+ collection, e.what());
|
|
|
+ results.clear(); // partial top-K is unsafe to keep
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
|
|
|
- auto results = store_.similaritySearch(
|
|
|
- request->collection(), queryVec, topK, minScore);
|
|
|
+ if (!used_lmdb) {
|
|
|
+ results = store_.similaritySearch(collection, queryVec, topK, minScore);
|
|
|
+ }
|
|
|
|
|
|
for (auto& result : results) {
|
|
|
// Decrypt sensitive fields before returning
|