Bläddra i källkod

feat(storage): hard-flip SimilaritySearch handler to LMDB-first

Last of the four read handlers migrated. Vectors are iterated from
the _vectors_<collection> sub-db via a new DocumentStore::scan_vectors
cursor callback; cosine similarity uses the SIMD primitives extracted
from MemoryStore into service/src/storage/cosine_simd.hpp (AVX2 /
SSE4.1 / scalar fallback paths preserved verbatim). A top-K min-heap
gates iteration so we never materialise more than topK candidates.
Documents are fetched via doc_store_->get() for each surviving result.
Gated on mirror_healthy_ + zero drift + non-system collection +
doc_store_ available. On LMDB throw, fall back to
store_.similaritySearch(). Dimension mismatch surfaces as
INVALID_ARGUMENT (matches MemoryStore semantics).

Test: tests/test_dual_write_mirror.cpp gains scan_vectors round-trip
coverage (5 vectors, dim=4, byte-identical readback + missing sub-db
no-op path). 12/12 ctest suite green. load_test_mixed 20s/4w/8r
PASS (2183 writes/s, 5242 reads/s, p99 ~100ms, 0 hard failures).
fszontagh 2 månader sedan
förälder
incheckning
0642caedaa

+ 126 - 2
service/src/database_grpc_impl.cpp

@@ -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

+ 116 - 0
service/src/storage/cosine_simd.hpp

@@ -0,0 +1,116 @@
+// v2.0 storage engine — SIMD-accelerated cosine similarity primitives.
+//
+// Extracted verbatim from service/src/memory_store.cpp::computeNorm and
+// ::cosineSimilaritySIMD (Stage 5 of the storage engine rewrite) so the
+// LMDB-first SimilaritySearch handler can reuse the AVX2 / SSE4.1 / scalar
+// fallback paths without copy-pasting. MemoryStore retains its own
+// implementations to keep the fallback path untouched.
+//
+// Compilation:
+//   __AVX2__ / __SSE4_1__ are predefined by the compiler when the
+//   corresponding instruction set is enabled at build time. No CMake-level
+//   intervention required; the same macros gate the original code in
+//   memory_store.cpp.
+//
+// Numeric semantics:
+//   - cosine_sim returns 0.0 when either norm is zero (matches
+//     MemoryStore::cosineSimilaritySIMD).
+//   - compute_norm returns sqrt(sum_of_squares) over the input range.
+//   - All inputs are read-only; no allocations.
+
+#pragma once
+
+#include <cmath>
+#include <cstddef>
+
+#if defined(__AVX2__)
+#include <immintrin.h>
+#endif
+#if defined(__SSE4_1__)
+#include <smmintrin.h>
+#endif
+
+namespace smartbotic::db::storage::cosine_simd {
+
+inline float compute_norm(const float* data, size_t dim) {
+    float sum = 0.0f;
+    size_t i = 0;
+
+#if defined(__AVX2__)
+    __m256 vsum = _mm256_setzero_ps();
+    for (; i + 8 <= dim; i += 8) {
+        __m256 v = _mm256_loadu_ps(data + i);
+        vsum = _mm256_fmadd_ps(v, v, vsum);
+    }
+    // Horizontal sum of 8 floats
+    __m128 hi = _mm256_extractf128_ps(vsum, 1);
+    __m128 lo = _mm256_castps256_ps128(vsum);
+    __m128 sum128 = _mm_add_ps(lo, hi);
+    sum128 = _mm_hadd_ps(sum128, sum128);
+    sum128 = _mm_hadd_ps(sum128, sum128);
+    sum = _mm_cvtss_f32(sum128);
+#elif defined(__SSE4_1__)
+    __m128 vsum = _mm_setzero_ps();
+    for (; i + 4 <= dim; i += 4) {
+        __m128 v = _mm_loadu_ps(data + i);
+        vsum = _mm_add_ps(vsum, _mm_mul_ps(v, v));
+    }
+    // Horizontal sum of 4 floats
+    vsum = _mm_hadd_ps(vsum, vsum);
+    vsum = _mm_hadd_ps(vsum, vsum);
+    sum = _mm_cvtss_f32(vsum);
+#endif
+
+    // Scalar remainder
+    for (; i < dim; ++i) {
+        sum += data[i] * data[i];
+    }
+
+    return std::sqrt(sum);
+}
+
+inline float cosine_sim(const float* a, const float* b, size_t dim,
+                         float normA, float normB) {
+    if (normA == 0.0f || normB == 0.0f) {
+        return 0.0f;
+    }
+
+    float dot = 0.0f;
+    size_t i = 0;
+
+#if defined(__AVX2__)
+    __m256 vdot = _mm256_setzero_ps();
+    for (; i + 8 <= dim; i += 8) {
+        __m256 va = _mm256_loadu_ps(a + i);
+        __m256 vb = _mm256_loadu_ps(b + i);
+        vdot = _mm256_fmadd_ps(va, vb, vdot);
+    }
+    // Horizontal sum
+    __m128 hi = _mm256_extractf128_ps(vdot, 1);
+    __m128 lo = _mm256_castps256_ps128(vdot);
+    __m128 sum128 = _mm_add_ps(lo, hi);
+    sum128 = _mm_hadd_ps(sum128, sum128);
+    sum128 = _mm_hadd_ps(sum128, sum128);
+    dot = _mm_cvtss_f32(sum128);
+#elif defined(__SSE4_1__)
+    __m128 vdot = _mm_setzero_ps();
+    for (; i + 4 <= dim; i += 4) {
+        __m128 va = _mm_loadu_ps(a + i);
+        __m128 vb = _mm_loadu_ps(b + i);
+        vdot = _mm_add_ps(vdot, _mm_mul_ps(va, vb));
+    }
+    // Horizontal sum
+    vdot = _mm_hadd_ps(vdot, vdot);
+    vdot = _mm_hadd_ps(vdot, vdot);
+    dot = _mm_cvtss_f32(vdot);
+#endif
+
+    // Scalar remainder
+    for (; i < dim; ++i) {
+        dot += a[i] * b[i];
+    }
+
+    return dot / (normA * normB);
+}
+
+}  // namespace smartbotic::db::storage::cosine_simd

+ 19 - 0
service/src/storage/document_store.hpp

@@ -19,6 +19,7 @@
 #pragma once
 
 #include <cstdint>
+#include <functional>
 #include <optional>
 #include <string>
 #include <string_view>
@@ -107,6 +108,24 @@ public:
     get_vector(std::string_view /*collection*/, std::string_view /*id*/) {
         return std::nullopt;
     }
+
+    // Iterate every (id, vector) pair in the collection's
+    // `_vectors_<collection>` sub-db. The callback receives the doc id
+    // and a read-only view of the float bytes — pointers MUST NOT outlive
+    // the callback return, because they reference mmap'd LMDB pages
+    // inside the cursor's read txn.
+    //
+    // Used by the v2.0 SimilaritySearch handler to compute cosine scores
+    // without round-tripping each vector through std::vector<float>.
+    //
+    // Default impl: no-op (backends without a vector sub-db produce no
+    // hits, which is consistent with get_vector() returning nullopt
+    // on the same collection).
+    virtual void scan_vectors(
+        std::string_view /*collection*/,
+        std::function<void(std::string_view id,
+                           const float* data,
+                           size_t count)> /*callback*/) {}
 };
 
 }  // namespace smartbotic::db::storage

+ 39 - 0
service/src/storage/document_store_lmdb.cpp

@@ -469,6 +469,45 @@ LmdbDocumentStore::get_vector(std::string_view collection, std::string_view id)
     return out;
 }
 
+void LmdbDocumentStore::scan_vectors(
+    std::string_view collection,
+    std::function<void(std::string_view, const float*, size_t)> callback) {
+    if (!callback) return;
+    const std::string subdb = vector_subdb_name(collection);
+    ReadTxn rtxn(env_);
+    auto dbi_opt = try_open_for_read(rtxn, subdb);
+    if (!dbi_opt) return;  // no _vectors_<coll> sub-db -> no entries
+
+    MDB_cursor* cursor = nullptr;
+    mdb_check(mdb_cursor_open(rtxn.raw(), *dbi_opt, &cursor),
+              "cursor_open (vectors)");
+    struct CursorGuard {
+        MDB_cursor* c;
+        ~CursorGuard() { if (c) mdb_cursor_close(c); }
+    } guard{cursor};
+
+    MDB_val k{0, nullptr};
+    MDB_val v{0, nullptr};
+    int rc = mdb_cursor_get(cursor, &k, &v, MDB_FIRST);
+    while (rc == MDB_SUCCESS) {
+        if (v.mv_size % sizeof(float) != 0) {
+            throw std::runtime_error(
+                "scan_vectors: stored bytes are not a multiple of sizeof(float)");
+        }
+        const size_t n = v.mv_size / sizeof(float);
+        // Pointer + length view into the mmap'd page. The callback contract
+        // (see DocumentStore::scan_vectors docstring) forbids keeping these
+        // pointers past return — they go stale when the cursor advances.
+        std::string_view id(static_cast<const char*>(k.mv_data), k.mv_size);
+        const float* data = static_cast<const float*>(v.mv_data);
+        callback(id, data, n);
+        rc = mdb_cursor_get(cursor, &k, &v, MDB_NEXT);
+    }
+    if (rc != MDB_NOTFOUND && rc != MDB_SUCCESS) {
+        throw_mdb(rc, "cursor_get (vectors)");
+    }
+}
+
 bool LmdbDocumentStore::drop_collection(std::string_view collection) {
     // Probe with a read txn first; if the sub-db doesn't exist, return
     // false without opening a write txn at all.

+ 5 - 0
service/src/storage/document_store_lmdb.hpp

@@ -60,6 +60,11 @@ public:
     bool del_vector(std::string_view collection, std::string_view id) override;
     std::optional<std::vector<float>>
     get_vector(std::string_view collection, std::string_view id) override;
+    void scan_vectors(
+        std::string_view collection,
+        std::function<void(std::string_view id,
+                           const float* data,
+                           size_t count)> callback) override;
 
 private:
     // Sub-db handle cache. MDB_dbi is typedef'd to unsigned int — held as

+ 57 - 0
tests/test_dual_write_mirror.cpp

@@ -375,6 +375,62 @@ void test_backfill_iteration_pattern() {
     check(healthy.load() && drift.load() == 0, "backfill: clean");
 }
 
+// v2.0 Stage 5 — LMDB scan_vectors round-trip. Populates a
+// `_vectors_<coll>` sub-db with 5 vectors of dim=4 via put_vector, then
+// iterates them with scan_vectors and verifies (a) count, (b) every
+// (id, vector) pair seen exactly once, (c) byte-identical contents.
+// Backbone for the LMDB-first SimilaritySearch handler.
+void test_scan_vectors_round_trip() {
+    TmpEnv tmp("scan-vectors");
+    LmdbDocumentStore store(tmp.env);
+
+    // 5 vectors of dimension 4. Distinct values so we can spot any
+    // cross-talk between iterations.
+    const std::vector<std::pair<std::string, std::vector<float>>> seed = {
+        {"u1", {1.0f, 0.0f, 0.0f, 0.0f}},
+        {"u2", {0.0f, 1.0f, 0.0f, 0.0f}},
+        {"u3", {0.0f, 0.0f, 1.0f, 0.0f}},
+        {"u4", {0.0f, 0.0f, 0.0f, 1.0f}},
+        {"u5", {0.5f, 0.5f, 0.5f, 0.5f}},
+    };
+    for (const auto& [id, vec] : seed) {
+        store.put_vector("embeddings", id, vec);
+    }
+
+    // Iterate and record observed pairs.
+    std::vector<std::pair<std::string, std::vector<float>>> observed;
+    store.scan_vectors("embeddings",
+        [&](std::string_view id, const float* data, size_t count) {
+            std::vector<float> copy(data, data + count);
+            observed.emplace_back(std::string(id), std::move(copy));
+        });
+
+    check(observed.size() == seed.size(),
+          "scan_vectors: count matches seed");
+
+    // Validate every seeded pair appears exactly once with byte-identical
+    // contents. Iteration order is LMDB's lexicographic key order, but we
+    // don't depend on that — we cross-check via lookup.
+    for (const auto& [id, vec] : seed) {
+        auto it = std::find_if(observed.begin(), observed.end(),
+            [&](const auto& obs) { return obs.first == id; });
+        bool found = (it != observed.end());
+        check(found, "scan_vectors: each seeded id observed");
+        if (!found) continue;
+        check(it->second.size() == vec.size(),
+              "scan_vectors: dim preserved");
+        bool bytes_match = std::equal(it->second.begin(), it->second.end(),
+                                       vec.begin());
+        check(bytes_match, "scan_vectors: byte-identical payload");
+    }
+
+    // Empty collection -> callback never invoked.
+    int empty_calls = 0;
+    store.scan_vectors("nonexistent",
+        [&](std::string_view, const float*, size_t) { ++empty_calls; });
+    check(empty_calls == 0, "scan_vectors: missing sub-db -> no callbacks");
+}
+
 }  // namespace
 
 int main() {
@@ -389,6 +445,7 @@ int main() {
     test_vector_mirror_system_collection_skipped();
     test_vector_mirror_via_memory_store();
     test_backfill_iteration_pattern();
+    test_scan_vectors_round_trip();
 
     std::cout << "dual_write_mirror: " << g_pass << " passed, " << g_fail << " failed\n";
     return g_fail == 0 ? 0 : 1;