From c62c2d98f5db547536ef003b89adf32798da9dab Mon Sep 17 00:00:00 2001
From: Aditi Ahuja <aditi.ahuja@couchbase.com>
Date: Thu, 5 Sep 2024 11:32:03 +0530
Subject: [PATCH 1/3] search only specific clusters

---
 c_api/IndexIVF_c_ex.cpp | 34 ++++++++++++++++++++++++++++++++++
 c_api/IndexIVF_c_ex.h   | 21 +++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/c_api/IndexIVF_c_ex.cpp b/c_api/IndexIVF_c_ex.cpp
index 5b72e15c94..959c08ff8f 100644
--- a/c_api/IndexIVF_c_ex.cpp
+++ b/c_api/IndexIVF_c_ex.cpp
@@ -31,4 +31,38 @@ int faiss_SearchParametersIVF_new_with_sel(
         *p_sp = reinterpret_cast<FaissSearchParametersIVF*>(sp);
     }
     CATCH_AND_HANDLE
+}
+
+int faiss_Search_closest_eligible_centroids(
+        FaissIndex* index,
+        float* query,
+        int n, // top n centroids.
+        float* centroid_distances,
+        idx_t* centroid_ids) {
+    try {
+        faiss::IndexIVF* index_ivf = reinterpret_cast<IndexIVF*>(index); 
+        assert(index_ivf);
+
+        index_ivf->quantizer->search(1, query, n, centroid_distances, centroid_ids);
+    }
+    CATCH_AND_HANDLE
+}
+
+int faiss_IndexIVF_search_preassigned_with_params(
+        const FaissIndexIVF* index,
+        idx_t n,
+        const float* x,
+        idx_t k,
+        const idx_t* assign,
+        const float* centroid_dis,
+        float* distances,
+        idx_t* labels,
+        int store_pairs,
+        const FaissSearchParametersIVF* params) {
+    try {
+        reinterpret_cast<const IndexIVF*>(index)->search_preassigned(
+                n, x, k, assign, centroid_dis, distances, labels, store_pairs,
+                reinterpret_cast<const faiss::SearchParametersIVF*>(params));
+    }
+    CATCH_AND_HANDLE
 }
\ No newline at end of file
diff --git a/c_api/IndexIVF_c_ex.h b/c_api/IndexIVF_c_ex.h
index e82c5c106f..8a0af8d3f9 100644
--- a/c_api/IndexIVF_c_ex.h
+++ b/c_api/IndexIVF_c_ex.h
@@ -27,6 +27,27 @@ int faiss_SearchParametersIVF_new_with_sel(
         FaissSearchParametersIVF** p_sp,
         FaissIDSelector* sel);
 
+int faiss_Search_closest_eligible_centroids(
+        FaissIndex* index,
+        float* query,
+        int n,
+        float* centroid_distances,
+        idx_t* centroid_ids
+);
+
+int faiss_IndexIVF_search_preassigned_with_params(
+        const FaissIndexIVF* index,
+        idx_t n,
+        const float* x,
+        idx_t k,
+        const idx_t* assign,
+        const float* centroid_dis,
+        float* distances,
+        idx_t* labels,
+        int store_pairs,
+        const FaissSearchParametersIVF* params);
+
+
 #ifdef __cplusplus
 }
 #endif

From 9d28768bfc9f3ac6377886463350cad7cb8e5ea8 Mon Sep 17 00:00:00 2001
From: Aditi Ahuja <aditi.ahuja@couchbase.com>
Date: Fri, 13 Sep 2024 14:33:18 +0530
Subject: [PATCH 2/3] updated comments

---
 c_api/IndexIVF_c_ex.cpp |  5 +++--
 c_api/IndexIVF_c_ex.h   | 22 +++++++++++++++++++++-
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/c_api/IndexIVF_c_ex.cpp b/c_api/IndexIVF_c_ex.cpp
index 959c08ff8f..39041a5265 100644
--- a/c_api/IndexIVF_c_ex.cpp
+++ b/c_api/IndexIVF_c_ex.cpp
@@ -35,15 +35,16 @@ int faiss_SearchParametersIVF_new_with_sel(
 
 int faiss_Search_closest_eligible_centroids(
         FaissIndex* index,
+        int n,
         float* query,
-        int n, // top n centroids.
+        int k, 
         float* centroid_distances,
         idx_t* centroid_ids) {
     try {
         faiss::IndexIVF* index_ivf = reinterpret_cast<IndexIVF*>(index); 
         assert(index_ivf);
 
-        index_ivf->quantizer->search(1, query, n, centroid_distances, centroid_ids);
+        index_ivf->quantizer->search(n, query, k, centroid_distances, centroid_ids);
     }
     CATCH_AND_HANDLE
 }
diff --git a/c_api/IndexIVF_c_ex.h b/c_api/IndexIVF_c_ex.h
index 8a0af8d3f9..33959e451c 100644
--- a/c_api/IndexIVF_c_ex.h
+++ b/c_api/IndexIVF_c_ex.h
@@ -26,15 +26,35 @@ int faiss_IndexIVF_set_direct_map(
 int faiss_SearchParametersIVF_new_with_sel(
         FaissSearchParametersIVF** p_sp,
         FaissIDSelector* sel);
+/*
+    Return 'k' centroids in the index closest to the query vector.
 
+    @param n: number of queries.
+    @param query: query vector.
+    @param k: count of closest number of vectors.
+    @param centroid_distances: output distances, size n * k.
+    @param centroid_ids: output centroid IDs, size n * k.    
+*/
 int faiss_Search_closest_eligible_centroids(
         FaissIndex* index,
-        float* query,
         int n,
+        float* query,
+        int k,
         float* centroid_distances,
         idx_t* centroid_ids
 );
 
+/*
+    Search the clusters whose IDs are in 'assign' and
+    return return the 'k' nearest neighbours from among them.     
+
+    @param n: number of queries.
+    @param x: query vector, size n * d.
+    @param k: count of nearest neighbours to be returned for each query.
+    @param centroid_ids: output centroid IDs, size n * k. 
+    @param distance: output distances, size n * k
+    @param labels: output labels, size n * k
+*/
 int faiss_IndexIVF_search_preassigned_with_params(
         const FaissIndexIVF* index,
         idx_t n,

From 89f8261a52694f5cfb37acf064a1d5916a35cea9 Mon Sep 17 00:00:00 2001
From: Aditi Ahuja <aditi.ahuja@couchbase.com>
Date: Tue, 17 Sep 2024 17:57:51 +0530
Subject: [PATCH 3/3] fixed comment

---
 c_api/IndexIVF_c_ex.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/c_api/IndexIVF_c_ex.h b/c_api/IndexIVF_c_ex.h
index 33959e451c..1aa9378b66 100644
--- a/c_api/IndexIVF_c_ex.h
+++ b/c_api/IndexIVF_c_ex.h
@@ -46,7 +46,7 @@ int faiss_Search_closest_eligible_centroids(
 
 /*
     Search the clusters whose IDs are in 'assign' and
-    return return the 'k' nearest neighbours from among them.     
+    return the 'k' nearest neighbours from among them.     
 
     @param n: number of queries.
     @param x: query vector, size n * d.