native-vector-store 0.3.8 → 0.4.0

package/src/vector_store.cpp

@@ -1,4 +1,7 @@
 #include "vector_store.h"
+#include "simple_tokenizer.h"
+#include <cctype>
+#include <algorithm>
 
 // ArenaAllocator implementation
 
@@ -109,7 +112,7 @@ void TopK::merge(const TopK& other) {
 
 // VectorStore implementation
 
-VectorStore::VectorStore(size_t dim) : dim_(dim) {
+VectorStore::VectorStore(size_t dim) : dim_(dim), postings_(), doc_freq_() {
     entries_.resize(1'000'000);  // Pre-size with default-constructed entries
     
     // Prepare per-thread arena allocators for zero-contention parallel loading
@@ -314,6 +317,45 @@ VectorStoreError VectorStore::add_document(simdjson::ondemand::object& json_doc)
     entry.doc = doc;
     entry.embedding = emb_ptr;
     
+    // Process text for BM25 - tokenize and build term frequencies
+    SimpleTokenizer tokenizer;
+    std::vector<std::string> tokens = tokenizer.split(std::string(text));
+    
+    // Build term frequency map
+    entry.tf.clear();
+    for (const std::string& token : tokens) {
+        // Convert to lowercase for case-insensitive matching
+        std::string lower_token = token;
+        std::transform(lower_token.begin(), lower_token.end(), lower_token.begin(), ::tolower);
+        entry.tf[lower_token]++;
+    }
+    
+    entry.length = tokens.size();
+    
+    // Update BM25 index structures using lock-free parallel hashmap operations
+    total_length_.fetch_add(entry.length, std::memory_order_relaxed);
+    
+    // Update postings and document frequencies
+    for (const auto& tf_pair : entry.tf) {
+        const std::string& term = tf_pair.first;
+        
+        // Update postings list using parallel hashmap's thread-safe lazy_emplace_l
+        postings_.lazy_emplace_l(term,
+            // If key exists, append to the vector
+            [&idx](auto& p) { p.second.push_back(idx); },
+            // If key doesn't exist, create new vector with this idx
+            [&term, &idx](const auto& ctor) { ctor(term, std::vector<size_t>{idx}); }
+        );
+        
+        // Update document frequency - parallel hashmap provides thread safety
+        doc_freq_.lazy_emplace_l(term,
+            // If key exists, increment the count
+            [](auto& p) { p.second++; },
+            // If key doesn't exist, create with value 1
+            [&term](const auto& ctor) { ctor(term, 1); }
+        );
+    }
+    
     entries_[idx] = entry;
     
     return VectorStoreError::SUCCESS;
@@ -360,7 +402,7 @@ void VectorStore::normalize_all() {
 }
 
 std::vector<std::pair<float, size_t>> 
-VectorStore::search(const float* query, size_t k) const {
+VectorStore::search(const float* __restrict__ query, size_t k) const {
     // Exclusive lock: prevent overlapping OpenMP teams
     // Since each search uses all threads via OpenMP, concurrent searches provide no benefit
     std::unique_lock<std::shared_mutex> lock(search_mutex_);
@@ -394,7 +436,7 @@ VectorStore::search(const float* query, size_t k) const {
         #pragma omp for  // default barrier kept - ensures all threads finish before merge
         for (int i = 0; i < static_cast<int>(n); ++i) {
             float score = 0.0f;
-            const float* emb = entries_[i].embedding;
+            const float* __restrict__ emb = entries_[i].embedding;
             
             #pragma omp simd reduction(+:score)
             for (size_t j = 0; j < dim_; ++j) {
@@ -430,3 +472,197 @@ size_t VectorStore::size() const {
 bool VectorStore::is_finalized() const {
     return is_finalized_.load(std::memory_order_acquire);
 }
+
+double VectorStore::avg_doc_length() const {
+    size_t n = size();
+    return n > 0 ? static_cast<double>(total_length_.load(std::memory_order_relaxed)) / n : 0.0;
+}
+
+void VectorStore::set_bm25_parameters(double k1, double b, double delta) {
+    k1_ = k1;
+    b_ = b; 
+    delta_ = delta;
+}
+
+std::vector<std::pair<size_t, double>> 
+VectorStore::search_bm25(const std::vector<std::string>& query_terms) const {
+    if (!is_finalized()) {
+        return {}; // Store must be finalized
+    }
+    
+    std::unordered_map<size_t, double> scores;
+    size_t N = size();
+    double avg_len = avg_doc_length();
+    
+    // Precompute IDF for each unique query term
+    std::unordered_map<std::string, double> idf_cache;
+    for (const auto& term : query_terms) {
+        if (idf_cache.find(term) == idf_cache.end()) {
+            auto df_it = doc_freq_.find(term);
+            int df = (df_it != doc_freq_.end()) ? df_it->second : 0;
+            idf_cache[term] = std::log((N - df + 0.5) / (df + 0.5) + 1.0);
+        }
+    }
+    
+    // For each unique term in the query:
+    for (const auto& term : query_terms) {
+        auto postings_it = postings_.find(term);
+        if (postings_it == postings_.end()) {
+            continue; // Term not found in corpus
+        }
+        
+        double idf_t = idf_cache[term];
+        for (size_t doc_id : postings_it->second) {
+            const Entry& entry = entries_[doc_id];
+            auto tf_it = entry.tf.find(term);
+            if (tf_it == entry.tf.end()) {
+                continue; // Should not happen if postings are consistent
+            }
+            
+            int tf = tf_it->second;
+            double norm = 1.0 - b_ + b_ * (entry.length / avg_len);
+            double tf_weight = (k1_ + 1) * tf / (tf + k1_ * norm);
+            scores[doc_id] += (tf_weight + delta_) * idf_t;
+        }
+    }
+    
+    // Collect and sort results
+    std::vector<std::pair<size_t, double>> results(scores.begin(), scores.end());
+    std::sort(results.begin(), results.end(),
+              [](const auto& a, const auto& b) { return a.second > b.second; });
+    
+    return results;
+}
+
+std::vector<std::pair<size_t, double>>
+VectorStore::search_hybrid(const float* __restrict__ query_vector, const std::vector<std::string>& query_terms, 
+                          double vector_weight, double bm25_weight, size_t k) const {
+    // Exclusive lock: prevent overlapping OpenMP teams
+    std::unique_lock<std::shared_mutex> lock(search_mutex_);
+    
+    if (!is_finalized()) {
+        return {}; // Store must be finalized
+    }
+    
+    size_t n = count_.load(std::memory_order_acquire);
+    if (n == 0 || k == 0) return {};
+    k = std::min(k, n);
+    
+    // Precompute BM25 IDF scores for query terms
+    std::unordered_map<std::string, double> idf_cache;
+    double avg_len = avg_doc_length();
+    
+    for (const auto& term : query_terms) {
+        auto df_it = doc_freq_.find(term);
+        int df = (df_it != doc_freq_.end()) ? df_it->second : 0;
+        idf_cache[term] = std::log((n - df + 0.5) / (df + 0.5) + 1.0);
+    }
+    
+    const int num_threads = omp_get_max_threads();
+    
+    // Each thread maintains TWO heaps - one for vector, one for BM25
+    struct DualTopK {
+        TopK vector_heap;
+        TopK bm25_heap;
+        DualTopK(size_t k) : vector_heap(k), bm25_heap(k) {}
+        
+        // Make DualTopK move-only like TopK
+        DualTopK(const DualTopK&) = delete;
+        DualTopK& operator=(const DualTopK&) = delete;
+        DualTopK(DualTopK&&) = default;
+        DualTopK& operator=(DualTopK&&) = default;
+    };
+    
+    std::vector<DualTopK> thread_heaps;
+    thread_heaps.reserve(num_threads);
+    for (int i = 0; i < num_threads; ++i) {
+        thread_heaps.emplace_back(k);
+    }
+    
+    #pragma omp parallel
+    {
+        const int tid = omp_get_thread_num();
+        DualTopK& local = thread_heaps[tid];
+        
+        #pragma omp for
+        for (int i = 0; i < static_cast<int>(n); ++i) {
+            // 1. Compute vector similarity score
+            float vector_score = 0.0f;
+            const float* __restrict__ emb = entries_[i].embedding;
+            
+            #pragma omp simd reduction(+:vector_score)
+            for (size_t j = 0; j < dim_; ++j) {
+                vector_score += emb[j] * query_vector[j];
+            }
+            
+            // 2. Compute BM25 score for this document
+            double bm25_score = 0.0;
+            const Entry& entry = entries_[i];
+            
+            for (const auto& term : query_terms) {
+                auto tf_it = entry.tf.find(term);
+                if (tf_it != entry.tf.end()) {
+                    int tf = tf_it->second;
+                    double norm = 1.0 - b_ + b_ * (entry.length / avg_len);
+                    double tf_weight = (k1_ + 1) * tf / (tf + k1_ * norm);
+                    bm25_score += (tf_weight + delta_) * idf_cache.at(term);
+                }
+            }
+            
+            // 3. Push to both heaps
+            local.vector_heap.push(vector_score, i);
+            local.bm25_heap.push(static_cast<float>(bm25_score), i);
+        }
+        
+        #pragma omp barrier
+    }
+    
+    // Merge thread-local heaps to get global top-k for each score type
+    TopK global_vector_heap(k);
+    TopK global_bm25_heap(k);
+    
+    for (auto& th : thread_heaps) {
+        global_vector_heap.merge(th.vector_heap);
+        global_bm25_heap.merge(th.bm25_heap);
+    }
+    
+    // Sort heaps to get ranking order
+    std::sort(global_vector_heap.heap.begin(), global_vector_heap.heap.end(),
+              [](const auto& a, const auto& b) { return a.first > b.first; });
+    std::sort(global_bm25_heap.heap.begin(), global_bm25_heap.heap.end(),
+              [](const auto& a, const auto& b) { return a.first > b.first; });
+    
+    // Apply Reciprocal Rank Fusion (RRF) with constant k=60 (typical value)
+    const double rrf_k = 60.0;
+    std::unordered_map<size_t, double> rrf_scores;
+    
+    // Add vector search rankings
+    for (size_t rank = 0; rank < global_vector_heap.heap.size(); ++rank) {
+        size_t doc_id = global_vector_heap.heap[rank].second;
+        // Weight the RRF contribution
+        rrf_scores[doc_id] += vector_weight * (1.0 / (rrf_k + rank + 1));
+    }
+    
+    // Add BM25 rankings
+    for (size_t rank = 0; rank < global_bm25_heap.heap.size(); ++rank) {
+        size_t doc_id = global_bm25_heap.heap[rank].second;
+        // Weight the RRF contribution
+        rrf_scores[doc_id] += bm25_weight * (1.0 / (rrf_k + rank + 1));
+    }
+    
+    // Sort by RRF score and return top-k
+    std::vector<std::pair<size_t, double>> results;
+    results.reserve(rrf_scores.size());
+    for (const auto& pair : rrf_scores) {
+        results.emplace_back(pair.first, pair.second);
+    }
+    
+    std::sort(results.begin(), results.end(),
+              [](const auto& a, const auto& b) { return a.second > b.second; });
+    
+    if (results.size() > k) {
+        results.resize(k);
+    }
+    
+    return results;
+}

package/src/vector_store.h

@@ -12,6 +12,9 @@
 #include <cassert>
 #include <algorithm>
 #include <functional>
+#include <unordered_map>
+#include <string>
+#include <parallel_hashmap/phmap.h>
 
 // Custom error codes for VectorStore
 enum class VectorStoreError {
@@ -196,6 +199,10 @@ public:
     struct Entry {
         Document doc;
         float* embedding;  // Extracted pointer for fast access
+        
+        // BM25 fields
+        size_t length;  // Total number of tokens in doc.text
+        phmap::flat_hash_map<std::string, int> tf;  // Term frequencies - better cache locality
     };
 
 private:
@@ -214,6 +221,35 @@ private:
     enum class TextFieldType { UNKNOWN, TEXT, CONTENT };
     std::atomic<TextFieldType> text_field_type_{TextFieldType::UNKNOWN};
     
+    // BM25 index structures - using parallel hashmap for lock-free concurrent updates
+    phmap::parallel_flat_hash_map<
+        std::string, 
+        std::vector<size_t>,
+        phmap::priv::hash_default_hash<std::string>,
+        phmap::priv::hash_default_eq<std::string>,
+        std::allocator<std::pair<const std::string, std::vector<size_t>>>,
+        4,  // 2^4 = 16 submaps for parallelism
+        std::mutex  // Use std::mutex for each submap
+    > postings_;  // term -> list of doc indices
+    
+    phmap::parallel_flat_hash_map<
+        std::string,
+        int,  // Regular int - parallel hashmap provides synchronization
+        phmap::priv::hash_default_hash<std::string>,
+        phmap::priv::hash_default_eq<std::string>,
+        std::allocator<std::pair<const std::string, int>>,
+        4,  // 16 submaps
+        std::mutex
+    > doc_freq_;  // document frequencies
+    
+    std::atomic<size_t> total_length_{0};  // sum of all document lengths - now atomic
+    // Note: bm25_index_mutex_ removed - no longer needed with parallel hashmap!
+    
+    // BM25 parameters
+    double k1_ = 1.2;
+    double b_ = 0.75; 
+    double delta_ = 1.0;
+    
 public:
     explicit VectorStore(size_t dim);
     
@@ -251,11 +287,26 @@ public:
     void normalize_all();
     
     std::vector<std::pair<float, size_t>> 
-    search(const float* query, size_t k) const;
+    search(const float* __restrict__ query, size_t k) const;
+    
+    // BM25 search
+    std::vector<std::pair<size_t, double>> 
+    search_bm25(const std::vector<std::string>& query_terms) const;
+    
+    // Hybrid search combining vector similarity and BM25
+    std::vector<std::pair<size_t, double>>
+    search_hybrid(const float* __restrict__ query_vector, const std::vector<std::string>& query_terms, 
+                  double vector_weight = 0.7, double bm25_weight = 0.3, size_t k = 10) const;
+    
+    // BM25 parameter setters
+    void set_bm25_parameters(double k1, double b, double delta);
     
     const Entry& get_entry(size_t idx) const;
     
     size_t size() const;
     
     bool is_finalized() const;
+    
+    // Get average document length for BM25
+    double avg_doc_length() const;
 };

package/src/vector_store_loader.cpp

@@ -93,7 +93,7 @@ void VectorStoreLoader::loadDirectory(VectorStore* store, const std::string& pat
     for (size_t w = 0; w < num_workers; ++w) {
         consumers.emplace_back([&]() {
             // Each thread needs its own parser with initial capacity
-            simdjson::ondemand::parser doc_parser(16 * 1024 * 1024); // 16MB initial capacity
+            simdjson::ondemand::parser doc_parser(1 * 1024 * 1024 * 1024); // 16MB initial capacity
             // Set a larger maximum capacity for very large files (up to 512MB)
             doc_parser.allocate(512 * 1024 * 1024);
             FileData* data = nullptr;

package/src/vector_store_loader_adaptive.cpp

@@ -133,7 +133,7 @@ void VectorStoreLoader::loadDirectoryAdaptive(VectorStore* store, const std::str
     for (size_t w = 0; w < num_workers; ++w) {
         consumers.emplace_back([&]() {
             // Each thread needs its own parser with initial capacity
-            simdjson::ondemand::parser doc_parser(16 * 1024 * 1024); // 16MB initial capacity
+            simdjson::ondemand::parser doc_parser(1 * 1024 * 1024 * 1024); // 16MB initial capacity
             // Set a larger maximum capacity for very large files (up to 512MB)
             doc_parser.allocate(512 * 1024 * 1024);
             MixedFileData* data = nullptr;

package/src/vector_store_loader_mmap.cpp

@@ -68,7 +68,7 @@ void VectorStoreLoader::loadDirectoryMMap(VectorStore* store, const std::string&
     for (size_t w = 0; w < num_workers; ++w) {
         consumers.emplace_back([&]() {
             // Each thread needs its own parser with initial capacity
-            simdjson::ondemand::parser doc_parser(16 * 1024 * 1024); // 16MB initial capacity
+            simdjson::ondemand::parser doc_parser(1 * 1024 * 1024 * 1024); // 16MB initial capacity
             // Set a larger maximum capacity for very large files (up to 512MB)
             doc_parser.set_max_capacity(512 * 1024 * 1024);
             MMapFileData* data = nullptr;
@@ -154,4 +154,4 @@ void VectorStoreLoader::loadDirectoryMMap(VectorStore* store, const std::string&
     
     // Finalize after batch load - normalize and switch to serving phase
     store->finalize();
-}
+}