added knn log ctx

Author: meiravgri

src/VecSim/vec_sim_index.h

@@ -98,4 +98,14 @@ struct VecSimIndexAbstract : public VecSimIndexInterface {
             delete[] buf;
         }
     }
+
+#ifdef BUILD_TESTS
+    // Set new log context to be sent to the log callback.
+    // Returns the previous logctx.
+    inline void *setLogCtx(void *new_logCtx) {
+        void *prev_logCtx = this->logCallbackCtx;
+        this->logCallbackCtx = new_logCtx;
+        return prev_logCtx;
+    }
+#endif
 };

src/VecSim/vec_sim_interface.cpp

@@ -14,3 +14,11 @@ void Vecsim_Log(void *ctx, const char *message) { std::cout << message << std::e
 timeoutCallbackFunction VecSimIndexInterface::timeoutCallback = [](void *ctx) { return 0; };
 logCallbackFunction VecSimIndexInterface::logCallback = Vecsim_Log;
 VecSimWriteMode VecSimIndexInterface::asyncWriteMode = VecSim_WriteAsync;
+
+#ifdef BUILD_TESTS
+static inline void Vecsim_Log_DO_NOTHING(void *ctx, const char *message) {}
+
+void VecSimIndexInterface::resetLogCallbackFunction() {
+    VecSimIndexInterface::logCallback = Vecsim_Log_DO_NOTHING;
+}
+#endif

src/VecSim/vec_sim_interface.h

@@ -186,6 +186,10 @@ struct VecSimIndexInterface : public VecsimBaseObject {
         VecSimIndexInterface::logCallback = callback;
     }
 
+#ifdef BUILD_TESTS
+    static void resetLogCallbackFunction();
+#endif
+
     /**
      * @brief Allow 3rd party to set the write mode for tiered index - async insert/delete using
      * background jobs, or insert/delete inplace.

src/VecSim/vec_sim_tiered_index.h

@@ -62,14 +62,20 @@ class VecSimTieredIndex : public VecSimIndexInterface {
                                      VecSimQueryParams *queryParams) override;
 
     static VecSimWriteMode getWriteMode() { return asyncWriteMode; }
+
+#ifdef BUILD_TESTS
+    inline VecSimIndexAbstract<DistType> *getFlatbufferIndex() { return this->frontendIndex; }
+#endif
 };
 
 template <typename DataType, typename DistType>
 VecSimQueryResult_List
 VecSimTieredIndex<DataType, DistType>::topKQuery(const void *queryBlob, size_t k,
                                                  VecSimQueryParams *queryParams) {
     this->flatIndexGuard.lock_shared();
-
+#ifdef BUILD_TESTS
+    this->getFlatbufferIndex()->log("");
+#endif
     // If the flat buffer is empty, we can simply query the main index.
     if (this->frontendIndex->indexSize() == 0) {
         // Release the flat lock and acquire the main lock.

src/python_bindings/bindings.cpp

@@ -381,7 +381,19 @@ class PyHNSWLibIndex : public PyVecSimIndex {
     }
 };
 
+template <typename DistType>
+struct KNNLogCtx {
+    VecSimIndexAbstract<DistType> *flat_index;
+    size_t curr_flat_size;
+    KNNLogCtx() : flat_index(nullptr), curr_flat_size(0) {}
+};
+
 class PyTIEREDIndex : public PyVecSimIndex {
+private:
+    VecSimIndexAbstract<float> *getFlatBuffer() {
+        return reinterpret_cast<VecSimTieredIndex<float, float> *>(this->index.get())
+            ->getFlatbufferIndex();
+    }
 
 protected:
     JobQueue jobQueue;          // External queue that holds the jobs.
@@ -395,6 +407,8 @@ class PyTIEREDIndex : public PyVecSimIndex {
     bool run_thread;
     std::bitset<MAX_POOL_SIZE> executions_status;
 
+    KNNLogCtx<float> knnLogCtx;
+
     TieredIndexParams TieredIndexParams_Init() {
         TieredIndexParams ret = {
             .jobQueue = &this->jobQueue,
@@ -409,14 +423,16 @@ class PyTIEREDIndex : public PyVecSimIndex {
     }
 
 public:
-    explicit PyTIEREDIndex(size_t BufferLimit = 20000000)
-        : submitCb(submit_callback), memoryCtx(0), UpdateMemCb(update_mem_callback), flatBufferLimit(BufferLimit),
-          run_thread(true) {
+    explicit PyTIEREDIndex(size_t BufferLimit = 1000)
+        : submitCb(submit_callback), memoryCtx(0), UpdateMemCb(update_mem_callback),
+          flatBufferLimit(BufferLimit), run_thread(true) {
 
         for (size_t i = 0; i < THREAD_POOL_SIZE; i++) {
             ThreadParams params(run_thread, executions_status, i, jobQueue);
             thread_pool.emplace_back(thread_main_loop, params);
         }
+
+        ResetLogCB();
     }
 
     virtual ~PyTIEREDIndex() = 0;
@@ -438,13 +454,34 @@ class PyTIEREDIndex : public PyVecSimIndex {
         }
     }
 
+    static void log_flat_buffer_size(void *ctx, const char *msg) {
+        auto *knnLogCtx = reinterpret_cast<KNNLogCtx<float> *>(ctx);
+        knnLogCtx->curr_flat_size = knnLogCtx->flat_index->indexLabelCount();
+    }
+    void SetKNNLogCtx() {
+        knnLogCtx.flat_index = getFlatBuffer();
+        knnLogCtx.curr_flat_size = 0;
+        knnLogCtx.flat_index->setLogCtx(&knnLogCtx);
+        this->index->setLogCallbackFunction(log_flat_buffer_size);
+    }
+    size_t getFlatIndexSize(const char *mode = "None") {
+        if (!strcmp(mode, "insert_and_knn")) {
+            return knnLogCtx.curr_flat_size;
+        }
 
+        return getFlatBuffer()->indexLabelCount();
+    }
+
+    void ResetLogCB() { this->index->resetLogCallbackFunction(); }
     static size_t GetThreadsNum() { return THREAD_POOL_SIZE; }
 
-    size_t getBufferLimit() {return flatBufferLimit; }
+    size_t getBufferLimit() { return flatBufferLimit; }
 };
 
-PyTIEREDIndex::~PyTIEREDIndex() { thread_pool_terminate(jobQueue, run_thread); }
+PyTIEREDIndex::~PyTIEREDIndex() {
+    thread_pool_terminate(jobQueue, run_thread);
+    ResetLogCB();
+}
 class PyTIERED_HNSWIndex : public PyTIEREDIndex {
 public:
     explicit PyTIERED_HNSWIndex(const HNSWParams &hnsw_params,
@@ -578,8 +615,11 @@ PYBIND11_MODULE(VecSim, m) {
 
     py::class_<PyTIEREDIndex, PyVecSimIndex>(m, "TIEREDIndex")
         .def("wait_for_index", &PyTIERED_HNSWIndex::WaitForIndex, py::arg("waiting_duration") = 10)
+        .def("get_curr_bf_size", &PyTIERED_HNSWIndex::getFlatIndexSize, py::arg("mode") = "None")
         .def("get_buffer_limit", &PyTIERED_HNSWIndex::getBufferLimit)
-        .def_static("get_threads_num", &PyTIEREDIndex::GetThreadsNum);
+        .def_static("get_threads_num", &PyTIEREDIndex::GetThreadsNum)
+        .def("reset_log", &PyTIERED_HNSWIndex::ResetLogCB)
+        .def("start_knn_log", &PyTIERED_HNSWIndex::SetKNNLogCtx);
 
     py::class_<PyTIERED_HNSWIndex, PyTIEREDIndex>(m, "TIERED_HNSWIndex")
         .def(

src/python_bindings/tiered_index_mock.h

@@ -1,4 +1,4 @@
- /*
+/*
  *Copyright Redis Ltd. 2021 - present
  *Licensed under your choice of the Redis Source Available License 2.0 (RSALv2) or
  *the Server Side Public License v1 (SSPLv1).
@@ -21,21 +21,6 @@ typedef struct RefManagedJob {
     std::weak_ptr<VecSimIndex> index_weak_ref;
 } RefManagedJob;
 
-struct SearchJobMock : public AsyncJob {
-    void *query; // The query vector. ownership is passed to the job in the constructor.
-    size_t k;    // The number of results to return.
-    size_t n;    // The number of vectors in the index (might be useful for the mock)
-    size_t dim;  // The dimension of the vectors in the index (might be useful for the mock)
-    std::atomic_int &successful_searches; // A reference to a shared counter that counts the number
-                                          // of successful searches.
-    SearchJobMock(std::shared_ptr<VecSimAllocator> allocator, JobCallback searchCB,
-                  VecSimIndex *index_, void *query_, size_t k_, size_t n_, size_t dim_,
-                  std::atomic_int &successful_searches_)
-        : AsyncJob(allocator, HNSW_SEARCH_JOB, searchCB, index_), query(query_), k(k_), n(n_),
-          dim(dim_), successful_searches(successful_searches_) {}
-    ~SearchJobMock() { this->allocator->free_allocation(query); }
-};
-
 using JobQueue = std::queue<RefManagedJob>;
 int submit_callback(void *job_queue, AsyncJob **jobs, size_t len, void *index_ctx);
 int update_mem_callback(void *mem_ctx, size_t mem);

tests/flow/common.py

@@ -41,4 +41,13 @@ def create_hnsw_index(dim, num_elements, metric, data_type, ef_construction=200,
 
     return HNSWIndex(hnsw_params)
 
+def bytes_to_mega(bytes, ndigits = 3):
+    return round(bytes/pow(10,6), ndigits)
+
+def round_(f_value, ndigits = 2):
+    return round(f_value, ndigits)
+
+
+def round_ms(f_value, ndigits = 2):
+    return round(f_value * 1000, ndigits)  
 

tests/flow/test_bm_hnsw_tiered_dataset.py

@@ -75,18 +75,19 @@ def __init__(self, data_size = 0, initialCap = 0, M = 32, ef_c = 512, ef_r = 10,
 
         data = load_data("dbpedia-768")
         self.num_elements = data_size if data_size != 0 else data.shape[0]
-        self.initialCap = initialCap if initialCap != 0 else 2 * self.num_elements
+        #self.initialCap = initialCap if initialCap != 0 else 2 * self.num_elements
+        self.initialCap = initialCap if initialCap != 0 else self.num_elements
 
         self.data = data[:self.num_elements]
         self.dim = len(self.data[0])
         self.metric = metric
-        self.type = data_type
+        self.data_type = data_type
         self.is_multi = is_multi
 
         self.hnsw_params = create_hnsw_params(dim=self.dim, 
                                               num_elements=self.initialCap, 
                                               metric=self.metric,
-                                              data_type=self.type,
+                                              data_type=self.data_type,
                                               ef_construction=ef_c,
                                               m=M,
                                               ef_runtime=ef_r,
@@ -102,22 +103,18 @@ def create_tiered(self):
 
     def create_hnsw(self):
         return HNSWIndex(self.hnsw_params)
-    
-    def set_num_vectors_per_label(self, num_per_label = 1):
-        self.num_per_label = num_per_label
 
     def init_and_populate_flat_index(self):
         bfparams = BFParams()
         bfparams.initialCapacity = self.num_elements
         bfparams.dim =self.dim
-        bfparams.type =self.type
+        bfparams.type =self.data_type
         bfparams.metric =self.metric
         bfparams.multi = self.is_multi
         self.flat_index = BFIndex(bfparams)
 
         for i, vector in enumerate(self.data):
-            for _ in range(self.num_per_label):
-                self.flat_index.add_vector(vector, i)
+            self.flat_index.add_vector(vector, i)
 
         return self.flat_index
 
@@ -129,6 +126,16 @@ def init_and_populate_hnsw_index(self):
         self.hnsw_index = hnsw_index
         return hnsw_index
 
+    def populate_index(self, index):
+        start = time.time()
+        duration = 0
+        for label, vector in enumerate(self.data):
+            start_add = time.time()
+            index.add_vector(vector, label)
+            duration += time.time() - start_add
+        end = time.time()
+        return (start, duration, end)
+    
     def generate_random_vectors(self, num_vectors):
         vectors = 0
         np_file_path = os.path.join(f'np_{num_vectors}vec_dim{self.dim}.npy')
@@ -154,7 +161,12 @@ def insert_in_batch(self, index, data, data_first_idx, batch_size, first_label):
             duration += time.time() - start_add
         end = time.time()
         return (duration, end)
+   
+    def generate_queries(self, num_queries):
+        self.rng = np.random.default_rng(seed=47)
 
+        queries = self.rng.random((num_queries, self.dim)) 
+        return np.float32(queries) if self.data_type == VecSimType_FLOAT32 else queries
 
 def create_dbpedia():
     indices_ctx = DBPediaIndexCtx()
@@ -192,7 +204,7 @@ def create_tiered():
     create_tiered()
 
 def create_dbpedia_graph():
-    indices_ctx = DBPediaIndexCtx(data_size = 100000)
+    indices_ctx = DBPediaIndexCtx()
 
     threads_num = TIEREDIndex.get_threads_num()
     print(f"thread num = {threads_num}")
@@ -283,9 +295,68 @@ def create_hnsw():
     print(f"Start hnsw creation")
 
     create_hnsw()
+
+def search_insert(is_multi: bool, num_per_label = 1):
+    indices_ctx = DBPediaIndexCtx(data_size=1000, mode=CreationMode.CREATE_TIERED_INDEX, is_multi=is_multi)
+    index = indices_ctx.tiered_index
+
+    num_elements = indices_ctx.num_elements
+    
+    query_data = indices_ctx.generate_queries(num_queries=1)
+    
+    # Add vectors to the flat index.
+    bf_index = indices_ctx.init_and_populate_flat_index()
+    
+    # Start background insertion to the tiered index.
+    index_start, _, _ = indices_ctx.populate_index(index)
+    
+    correct = 0
+    k = 10
+    searches_number = 0
+    
+    # config knn log
+    index.start_knn_log()
+    
+    # run knn query every 1 s. 
+    total_tiered_search_time = 0
+    prev_bf_size = num_elements
+    while index.hnsw_label_count() < num_elements:
+        # For each run get the current hnsw size and the query time.
+        bf_curr_size = index.get_curr_bf_size(mode = 'insert_and_knn')
+        query_start = time.time()
+        tiered_labels, _ = index.knn_query(query_data, k)
+        query_dur = time.time() - query_start
+        total_tiered_search_time += query_dur
+        
+        print(f"query time = {round_ms(query_dur)} ms")
+        
+        # BF size should decrease.
+        print(f"bf size = {bf_curr_size}")
+        assert bf_curr_size < prev_bf_size
+        
+        # Run the query also in the bf index to get the ground truth results.
+        bf_labels, _ = bf_index.knn_query(query_data, k)
+        correct += len(np.intersect1d(tiered_labels[0], bf_labels[0]))    
+        time.sleep(1)
+        searches_number += 1
+        prev_bf_size = bf_curr_size
+     
+    index.reset_log()
+    
+    # HNSW labels count updates before the job is done, so we need to wait for the queue to be empty.
+    index.wait_for_index(1)
+    index_dur = time.time() - index_start
+    print(f"indexing during search in tiered took {round_(index_dur)} s")
+    
+    # Measure recall.
+    recall = float(correct)/(k*searches_number)
+    print("Average recall is:", round_(recall, 3))
+    print("tiered query per seconds: ", round_(searches_number/total_tiered_search_time)) 
 
 def test_main():
     print("Test creation")
- #  create_dbpedia()
-    create_dbpedia_graph()
+    create_dbpedia()
+  #  create_dbpedia_graph()
+    print(f"\nStart insert & search test")
+   # search_insert(is_multi=False)