feat: vector storage with SIMD cosine similarity search

service/src/memory_store.cpp

@@ -3,11 +3,19 @@
 #include <spdlog/spdlog.h>
 
 #include <algorithm>
+#include <cmath>
 #include <random>
 #include <regex>
 #include <sstream>
 #include <iomanip>
 
+#if defined(__AVX2__)
+#include <immintrin.h>
+#endif
+#if defined(__SSE4_1__)
+#include <smmintrin.h>
+#endif
+
 namespace smartbotic::database {
 
 MemoryStore::MemoryStore(Config config)
@@ -223,6 +231,9 @@ std::string MemoryStore::insert(const std::string& collection, Document doc) {
         throw std::runtime_error("Document with ID '" + doc.id + "' already exists");
     }
 
+    // Extract vector before storing document (strips _vector from JSON data)
+    auto vec = extractVector(*coll, doc.data);
+
     // Set metadata
     doc.collection = collection;
     doc.version = 1;
@@ -244,6 +255,9 @@ std::string MemoryStore::insert(const std::string& collection, Document doc) {
     coll->documents[docId] = doc;
     coll->updatedAt = doc.updatedAt;
 
+    // Store vector separately
+    storeVector(*coll, docId, vec);
+
     // Track memory increment
     uint64_t docSize = estimateDocumentSize(doc);
 
@@ -350,6 +364,10 @@ bool MemoryStore::update(const std::string& collection, const std::string& id, c
 
     // Update document
     Document updated = doc;
+
+    // Extract vector before storing document (strips _vector from JSON data)
+    auto vec = extractVector(*coll, updated.data);
+
     updated.id = id;
     updated.collection = collection;
     updated.version = it->second.version + 1;
@@ -367,6 +385,13 @@ bool MemoryStore::update(const std::string& collection, const std::string& id, c
     it->second = updated;
     coll->updatedAt = updated.updatedAt;
 
+    // Update vector (replace or remove if no new vector provided)
+    if (!vec.empty()) {
+        storeVector(*coll, id, vec);
+    } else {
+        removeVector(*coll, id);
+    }
+
     // Track memory change (new size)
     uint64_t newSize = estimateDocumentSize(updated);
 
@@ -412,6 +437,9 @@ std::string MemoryStore::upsert(const std::string& collection, Document doc) {
         }
     }
 
+    // Extract vector before storing document (strips _vector from JSON data)
+    auto vec = extractVector(*coll, doc.data);
+
     std::string docId = doc.id;
     bool isInsert = false;
     uint64_t oldSize = 0;
@@ -459,6 +487,14 @@ std::string MemoryStore::upsert(const std::string& collection, Document doc) {
         it->second = doc;
     }
 
+    // Store or update vector
+    if (!vec.empty()) {
+        storeVector(*coll, docId, vec);
+    } else if (!isInsert) {
+        // On update without new vector, remove old vector
+        removeVector(*coll, docId);
+    }
+
     coll->updatedAt = doc.updatedAt;
 
     // Track new document size
@@ -516,6 +552,7 @@ bool MemoryStore::remove(const std::string& collection, const std::string& id) {
 
     saveToHistory(*coll, it->second);
     coll->documents.erase(it);
+    removeVector(*coll, id);
     coll->updatedAt = currentTimeMs();
 
     lock.unlock();
@@ -1005,6 +1042,228 @@ uint64_t MemoryStore::bulkDelete(const std::string& collection, const std::vecto
     return deleted;
 }
 
+// ===== Vector Similarity Search =====
+
+float MemoryStore::computeNorm(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);
+}
+
+float MemoryStore::cosineSimilaritySIMD(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);
+}
+
+std::vector<float> MemoryStore::extractVector(CollectionData& coll, nlohmann::json& docData) {
+    if (coll.options.vectorDimension == 0) {
+        return {};
+    }
+
+    if (!docData.contains("_vector") || !docData["_vector"].is_array()) {
+        return {};
+    }
+
+    auto vec = docData["_vector"].get<std::vector<float>>();
+    docData.erase("_vector");
+
+    if (vec.size() != coll.options.vectorDimension) {
+        throw std::invalid_argument(
+            "Vector dimension mismatch: expected " +
+            std::to_string(coll.options.vectorDimension) +
+            ", got " + std::to_string(vec.size()));
+    }
+
+    return vec;
+}
+
+void MemoryStore::storeVector(CollectionData& coll, const std::string& docId,
+                               const std::vector<float>& vec) {
+    if (vec.empty()) return;
+    coll.vectors[docId] = vec;
+    coll.vectorNorms[docId] = computeNorm(vec.data(), vec.size());
+}
+
+void MemoryStore::removeVector(CollectionData& coll, const std::string& docId) {
+    coll.vectors.erase(docId);
+    coll.vectorNorms.erase(docId);
+}
+
+std::vector<MemoryStore::SimilarityResult> MemoryStore::similaritySearch(
+    const std::string& collection,
+    const std::vector<float>& queryVector,
+    uint32_t topK,
+    float minScore) {
+
+    std::shared_lock<std::shared_mutex> globalLock(globalMutex_);
+
+    auto it = collections_.find(collection);
+    if (it == collections_.end()) {
+        throw std::runtime_error("Collection '" + collection + "' not found");
+    }
+
+    auto& coll = *it->second;
+
+    if (coll.options.vectorDimension == 0) {
+        throw std::runtime_error("Collection '" + collection + "' does not support vector search");
+    }
+
+    if (queryVector.size() != coll.options.vectorDimension) {
+        throw std::invalid_argument(
+            "Query vector dimension mismatch: expected " +
+            std::to_string(coll.options.vectorDimension) +
+            ", got " + std::to_string(queryVector.size()));
+    }
+
+    std::shared_lock<std::shared_mutex> collLock(coll.mutex);
+
+    if (coll.vectors.empty()) {
+        return {};
+    }
+
+    float queryNorm = computeNorm(queryVector.data(), queryVector.size());
+    if (queryNorm == 0.0f) {
+        return {};
+    }
+
+    // Compute similarity for all vectors
+    struct ScoredEntry {
+        std::string id;
+        float score;
+    };
+    std::vector<ScoredEntry> scored;
+    scored.reserve(coll.vectors.size());
+
+    for (const auto& [docId, vec] : coll.vectors) {
+        auto normIt = coll.vectorNorms.find(docId);
+        float docNorm = (normIt != coll.vectorNorms.end()) ? normIt->second : 0.0f;
+
+        float score = cosineSimilaritySIMD(
+            queryVector.data(), vec.data(), vec.size(), queryNorm, docNorm);
+
+        if (score >= minScore) {
+            scored.push_back({docId, score});
+        }
+    }
+
+    // partial_sort for top-K
+    uint32_t k = std::min(topK, static_cast<uint32_t>(scored.size()));
+    if (k == 0) {
+        return {};
+    }
+
+    std::partial_sort(scored.begin(), scored.begin() + k, scored.end(),
+                      [](const ScoredEntry& a, const ScoredEntry& b) {
+                          return a.score > b.score;
+                      });
+
+    // Build results with documents
+    std::vector<SimilarityResult> results;
+    results.reserve(k);
+
+    for (uint32_t i = 0; i < k; ++i) {
+        auto docIt = coll.documents.find(scored[i].id);
+        if (docIt != coll.documents.end() && !docIt->second.isExpired()) {
+            results.push_back({scored[i].id, scored[i].score, docIt->second});
+        }
+    }
+
+    return results;
+}
+
+const std::unordered_map<std::string, std::vector<float>>*
+MemoryStore::getCollectionVectors(const std::string& collection) const {
+    std::shared_lock<std::shared_mutex> globalLock(globalMutex_);
+
+    auto it = collections_.find(collection);
+    if (it == collections_.end()) {
+        return nullptr;
+    }
+
+    return &it->second->vectors;
+}
+
+void MemoryStore::loadVector(const std::string& collection, const std::string& docId,
+                              std::vector<float> vec) {
+    CollectionData* coll = getOrCreateCollection(collection);
+
+    std::unique_lock<std::shared_mutex> lock(coll->mutex);
+
+    coll->vectors[docId] = std::move(vec);
+    coll->vectorNorms[docId] = computeNorm(coll->vectors[docId].data(),
+                                            coll->vectors[docId].size());
+}
+
 // ===== TTL Management =====
 
 uint64_t MemoryStore::expireDocuments() {

service/src/memory_store.hpp

@@ -229,6 +229,44 @@ public:
      */
     uint64_t bulkDelete(const std::string& collection, const std::vector<std::string>& ids);
 
+    // ===== Vector Similarity Search =====
+
+    /**
+     * Result of a similarity search.
+     */
+    struct SimilarityResult {
+        std::string id;
+        float score;
+        Document document;
+    };
+
+    /**
+     * Find documents by vector similarity (cosine similarity).
+     * @param collection Collection name (must have vectorDimension set)
+     * @param queryVector Query embedding vector
+     * @param topK Maximum number of results to return
+     * @param minScore Minimum cosine similarity score (0.0 to 1.0)
+     * @return Top-K documents sorted by descending similarity score
+     */
+    std::vector<SimilarityResult> similaritySearch(
+        const std::string& collection,
+        const std::vector<float>& queryVector,
+        uint32_t topK,
+        float minScore = 0.0f);
+
+    /**
+     * Get all vectors for a collection (for snapshotting/persistence).
+     * @return Pointer to vectors map, or nullptr if collection not found
+     */
+    [[nodiscard]] const std::unordered_map<std::string, std::vector<float>>*
+        getCollectionVectors(const std::string& collection) const;
+
+    /**
+     * Load a vector directly (for snapshot/WAL recovery, bypasses event callbacks).
+     */
+    void loadVector(const std::string& collection, const std::string& docId,
+                    std::vector<float> vec);
+
     // ===== TTL Management =====
 
     /**
@@ -419,6 +457,8 @@ private:
         std::unordered_map<std::string, Document> documents;
         std::map<uint64_t, std::unordered_set<std::string>> expirationIndex;  // expiresAt -> ids
         std::unordered_map<std::string, std::deque<DocumentVersion>> versionHistory;  // docId -> versions (oldest first)
+        std::unordered_map<std::string, std::vector<float>> vectors;       // docId -> embedding vector
+        std::unordered_map<std::string, float> vectorNorms;                // cached L2 norms
         uint64_t createdAt = 0;
         uint64_t updatedAt = 0;
         mutable std::shared_mutex mutex;
@@ -445,6 +485,13 @@ private:
     // Recursively search JSON for string matches (case-insensitive)
     [[nodiscard]] static bool searchInJson(const nlohmann::json& obj, const std::string& lowerSearchTerm);
 
+    // Vector storage helpers
+    void storeVector(CollectionData& coll, const std::string& docId, const std::vector<float>& vec);
+    void removeVector(CollectionData& coll, const std::string& docId);
+    std::vector<float> extractVector(CollectionData& coll, nlohmann::json& docData);
+    static float cosineSimilaritySIMD(const float* a, const float* b, size_t dim, float normA, float normB);
+    static float computeNorm(const float* data, size_t dim);
+
     // Emit an event
     void emitEvent(EventType type, const std::string& collection, const std::string& id,
                    const std::optional<nlohmann::json>& data = std::nullopt);