MB-65473: Batch converter for vector to cluster IDs (#49)

CascadingRadium · web-flow · commit 352484e0fc9d · 2025-04-01T13:06:10.000+05:30
- An API was introduced in #48 that returns the cluster containing a given vector in an IVF index. - This PR modifies the API to support batch processing, allowing multiple vector IDs to be provided and retrieving their corresponding cluster (list) IDs from an IVF index. This function efficiently assigns vector IDs to their respective inverted lists/clusters in a batch operation. - Fix typecasts in `faiss_Search_closest_eligible_centroids` - Ensure errors are correctly handled in the introduced API
diff --git a/c_api/IndexIVF_c_ex.cpp b/c_api/IndexIVF_c_ex.cpp
@@ -36,9 +36,9 @@ int faiss_SearchParametersIVF_new_with_sel(
 
 int faiss_Search_closest_eligible_centroids(
         FaissIndex* index,
-        int n,
-        float* query,
-        int k, 
+        idx_t n,
+        const float* query,
+        idx_t k, 
         float* centroid_distances,
         idx_t* centroid_ids,
         const FaissSearchParameters* params) {
@@ -52,11 +52,13 @@ int faiss_Search_closest_eligible_centroids(
     CATCH_AND_HANDLE
 }
 
-idx_t faiss_get_list_for_key(
-        FaissIndexIVF* index, 
-        idx_t key) {
+int faiss_get_lists_for_keys(
+        FaissIndexIVF* index,
+        idx_t* keys,
+        size_t n_keys,
+        idx_t* lists) {
     try {
-        return reinterpret_cast<IndexIVF*>(index)->get_list_for_key(key);
+        reinterpret_cast<IndexIVF*>(index)->get_lists_for_keys(keys, n_keys, lists);
     }
     CATCH_AND_HANDLE
 }
diff --git a/c_api/IndexIVF_c_ex.h b/c_api/IndexIVF_c_ex.h
@@ -37,9 +37,9 @@ int faiss_SearchParametersIVF_new_with_sel(
 */
 int faiss_Search_closest_eligible_centroids(
         FaissIndex* index,
-        int n,
-        float* query,
-        int k,
+        idx_t n,
+        const float* query,
+        idx_t k,
         float* centroid_distances,
         idx_t* centroid_ids,
         const FaissSearchParameters* params
@@ -91,17 +91,22 @@ int faiss_IndexIVF_compute_distance_to_codes_for_list(
         const uint8_t* codes,
         float* dists);
 
-/* 
-    Given a query vector ID `key`, return the list number  
-    where the vector is stored. In the context of an IVF index,  
-    this corresponds to the cluster that contains the vector.  
-
-    @param key - vector ID  
+/*
+    Given multiple vector IDs, retrieve the corresponding list (cluster) IDs  
+    from an IVF index. This function efficiently assigns vector IDs to their  
+    respective inverted lists/clusters in a batch operation.  
+
+    @param index  - Pointer to the Faiss IVF index  
+    @param keys   - Input array of vector IDs (keys)  
+    @param n_keys - Number of vector keys in the input array  
+    @param lists  - Output array where corresponding cluster (list) IDs are stored  
 */
 
-idx_t faiss_get_list_for_key(
+int faiss_get_lists_for_keys(
         FaissIndexIVF* index, 
-        idx_t key);
+        idx_t* keys,
+        size_t n_keys,
+        idx_t* lists);
 
 #ifdef __cplusplus
 }
diff --git a/faiss/IndexIVF.cpp b/faiss/IndexIVF.cpp
@@ -922,9 +922,10 @@ void IndexIVF::reconstruct(idx_t key, float* recons) const {
     reconstruct_from_offset(lo_listno(lo), lo_offset(lo), recons);
 }
 
-idx_t IndexIVF::get_list_for_key(idx_t key) {
-    idx_t lo = direct_map.get(key);
-    return lo_listno(lo);
+void IndexIVF::get_lists_for_keys(idx_t* keys, size_t n_keys, idx_t* lists) {
+    for (int i = 0; i < n_keys; i++) {
+        lists[i] = lo_listno(direct_map.get(keys[i]));
+    }
 }
 
 void IndexIVF::reconstruct_n(idx_t i0, idx_t ni, float* recons) const {
diff --git a/faiss/IndexIVF.h b/faiss/IndexIVF.h
@@ -343,7 +343,7 @@ struct IndexIVF : Index, IndexIVFInterface {
      */
     void reconstruct_n(idx_t i0, idx_t ni, float* recons) const override;
 
-    idx_t get_list_for_key(idx_t key);
+    void get_lists_for_keys(idx_t* keys, size_t n_keys, idx_t* lists);
 
     /** Similar to search, but also reconstructs the stored vectors (or an
      * approximation in the case of lossy coding) for the search results.

Original file line number	Diff line number	Diff line change
`@@ -922,9 +922,10 @@ void IndexIVF::reconstruct(idx_t key, float* recons) const {`
`922`	`922`	`reconstruct_from_offset(lo_listno(lo), lo_offset(lo), recons);`
`923`	`923`	`}`
`924`	`924`
`925`		`-idx_t IndexIVF::get_list_for_key(idx_t key) {`
`926`		`- idx_t lo = direct_map.get(key);`
`927`		`- return lo_listno(lo);`
	`925`	`+void IndexIVF::get_lists_for_keys(idx_t* keys, size_t n_keys, idx_t* lists) {`
	`926`	`+ for (int i = 0; i < n_keys; i++) {`
	`927`	`+ lists[i] = lo_listno(direct_map.get(keys[i]));`
	`928`	`+ }`
`928`	`929`	`}`
`929`	`930`
`930`	`931`	`void IndexIVF::reconstruct_n(idx_t i0, idx_t ni, float* recons) const {`