MB-69881: Re-architect vector search

c_api/IndexIVF_c_ex.cpp

@@ -23,30 +23,30 @@ int faiss_IndexIVF_set_direct_map(FaissIndexIVF* index, int direct_map_type) {
     CATCH_AND_HANDLE
 }
 
-int faiss_SearchParametersIVF_new_with_sel(
-        FaissSearchParametersIVF** p_sp,
-        FaissIDSelector* sel) {
+int faiss_count_ivf_list_vectors(
+        const FaissIndexIVF* index,
+        idx_t* list_counts,
+        size_t list_counts_size,
+        const FaissSearchParametersIVF* params) {
     try {
-        SearchParametersIVF* sp = new SearchParametersIVF;
-        sp->sel = reinterpret_cast<faiss::IDSelector*>(sel);
-        *p_sp = reinterpret_cast<FaissSearchParametersIVF*>(sp);
+        reinterpret_cast<const IndexIVF*>(index)->count_ivf_list_vectors(
+                list_counts,
+                list_counts_size,
+                reinterpret_cast<const faiss::SearchParameters*>(params));
     }
     CATCH_AND_HANDLE
 }
 
 int faiss_Search_closest_eligible_centroids(
-        FaissIndex* index,
+        const FaissIndexIVF* index,
         idx_t n,
         const float* query,
         idx_t k,
         float* centroid_distances,
         idx_t* centroid_ids,
         const FaissSearchParameters* params) {
     try {
-        faiss::IndexIVF* index_ivf = reinterpret_cast<IndexIVF*>(index);
-        assert(index_ivf);
-
-        index_ivf->quantizer->search(
+        reinterpret_cast<const IndexIVF*>(index)->quantizer->search(
                 n,
                 query,
                 k,
@@ -57,30 +57,6 @@ int faiss_Search_closest_eligible_centroids(
     CATCH_AND_HANDLE
 }
 
-int faiss_get_lists_for_keys(
-        FaissIndexIVF* index,
-        idx_t* keys,
-        size_t n_keys,
-        idx_t* lists) {
-    try {
-        reinterpret_cast<IndexIVF*>(index)->get_lists_for_keys(
-                keys, n_keys, lists);
-    }
-    CATCH_AND_HANDLE
-}
-
-int faiss_IndexIVF_get_centroids_and_cardinality(
-        const FaissIndexIVF* index,
-        float* centroid_vectors,
-        size_t* cardinalities,
-        idx_t* centroid_ids) {
-    try {
-        reinterpret_cast<const IndexIVF*>(index)->get_centroids_and_cardinality(
-            centroid_vectors, cardinalities, centroid_ids);
-    }
-    CATCH_AND_HANDLE
-}
-
 int faiss_IndexIVF_search_preassigned_with_params(
         const FaissIndexIVF* index,
         idx_t n,
@@ -123,6 +99,18 @@ int faiss_IndexIVF_compute_distance_to_codes_for_list(
     CATCH_AND_HANDLE
 }
 
+int faiss_IndexIVF_get_centroids_and_cardinality(
+        const FaissIndexIVF* index,
+        float* centroid_vectors,
+        size_t* cardinalities,
+        idx_t* centroid_ids) {
+    try {
+        reinterpret_cast<const IndexIVF*>(index)->get_centroids_and_cardinality(
+            centroid_vectors, cardinalities, centroid_ids);
+    }
+    CATCH_AND_HANDLE
+}
+
 int faiss_IndexIVF_compute_distance_table(
         FaissIndexIVF* index,
         const float* x,

c_api/IndexIVF_c_ex.h

@@ -19,11 +19,38 @@
 #ifdef __cplusplus
 extern "C" {
 #endif
+
+/*
+    Set the direct map type for the IVF index.
+
+    @param index            - Pointer to the Faiss IVF index
+    @param direct_map_type  - Type of direct map to set (faiss::DirectMap::Type)
+*/
 int faiss_IndexIVF_set_direct_map(FaissIndexIVF* index, int direct_map_type);
 
-int faiss_SearchParametersIVF_new_with_sel(
-        FaissSearchParametersIVF** p_sp,
-        FaissIDSelector* sel);
+/*
+    Count vectors per IVF list (cluster) for a given selector.
+
+    This function iterates over the vectors selected by the provided
+    search parameters and increments a counter for each IVF inverted
+    list (cluster) they belong to. The result is a per-list vector
+    count for the IVF index.
+
+    @param index             - Pointer to the Faiss IVF index
+    @param list_counts       - Output array of size index->nlist. Must be
+                               initialized to zero by the caller. On return,
+                               list_counts[i] contains the number of selected
+                               vectors assigned to IVF list i.
+    @param list_counts_size  - Size of list_counts array (must equal index->nlist)
+    @param params            - IVF search parameters containing the selector
+                               that defines which vectors are included
+*/
+int faiss_count_ivf_list_vectors(
+        const FaissIndexIVF* index,
+        idx_t* list_counts,
+        size_t list_counts_size,
+        const FaissSearchParametersIVF* params);
+
 /*
     Return 'k' centroids in the index closest to the query vector.
 
@@ -34,7 +61,7 @@ int faiss_SearchParametersIVF_new_with_sel(
     @param centroid_ids: output centroid IDs, size n * k.
 */
 int faiss_Search_closest_eligible_centroids(
-        FaissIndex* index,
+        const FaissIndexIVF* index,
         idx_t n,
         const float* query,
         idx_t k,
@@ -90,23 +117,6 @@ int faiss_IndexIVF_compute_distance_to_codes_for_list(
         float* dists,
         float* dist_table);
 
-/*
-    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
-*/
-
-int faiss_get_lists_for_keys(
-        FaissIndexIVF* index,
-        idx_t* keys,
-        size_t n_keys,
-        idx_t* lists);
-
 /*
     Get centroid information and cardinality for all centroids in an IVF index.
 

faiss/IndexIVF.cpp

@@ -936,9 +936,44 @@ void IndexIVF::reconstruct(idx_t key, float* recons) const {
     reconstruct_from_offset(lo_listno(lo), lo_offset(lo), recons);
 }
 
-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::count_ivf_list_vectors(
+        idx_t* list_counts,
+        size_t list_counts_size,
+        const faiss::SearchParameters* params) const {
+    FAISS_ASSERT(list_counts != nullptr);
+    FAISS_ASSERT(list_counts_size > 0);
+    FAISS_ASSERT(list_counts_size == nlist);
+    FAISS_ASSERT(params != nullptr && params->sel != nullptr);
+    FAISS_ASSERT(direct_map.type != DirectMap::NoMap);
+    const IDSelector* sel = params->sel;
+    // Optimized for bitmap selectors only
+    const IDSelectorBitmap* bitmap_sel = dynamic_cast<const IDSelectorBitmap*>(sel);
+    if (!bitmap_sel) {
+        FAISS_THROW_MSG("count_ivf_list_vectors supports only IDSelectorBitmap");
+    }
+    const uint8_t* bitmap = bitmap_sel->bitmap;
+    const size_t nbytes = bitmap_sel->n;
+    // Iterate over bitmap bytes
+    for (size_t byte_idx = 0; byte_idx < nbytes; ++byte_idx) {
+        uint8_t byte = bitmap[byte_idx];
+        if (byte == 0) {
+            continue; // fast skip
+        }
+        // Iterate over bits in the byte
+        for (uint8_t bit = 0; bit < 8; ++bit) {
+            if ((byte & (1 << bit)) == 0) {
+                continue;
+            }
+            idx_t id = (byte_idx << 3) + bit;
+            if (id >= ntotal) {
+                continue; // Safety check: skip invalid ids
+            }
+            uint64_t list_no = lo_listno(direct_map.get(id));
+            if (list_no >= nlist) {
+                continue; // Safety check: skip invalid list numbers
+            }
+            list_counts[list_no]++;
+        }
     }
 }
 

faiss/IndexIVF.h

@@ -361,7 +361,31 @@ struct IndexIVF : Index, IndexIVFInterface {
      */
     void reconstruct_n(idx_t i0, idx_t ni, float* recons) const override;
 
-    void get_lists_for_keys(idx_t* keys, size_t n_keys, idx_t* lists);
+    /** Count vectors per IVF list (cluster) for a given selector.
+     *
+     * This function iterates over the vectors selected by the provided
+     * search parameters and increments a counter for each IVF inverted
+     * list (cluster) they belong to. The result is a per-list vector
+     * count for the IVF index.
+     *
+     * @param list_counts      - Output array of size list_counts_size (must be == nlist).
+     *                          On return, list_counts[i] contains the number
+     *                          of selected vectors assigned to IVF list i.
+     *                          Precondition: every element of list_counts must
+     *                          be initialized to 0 by the caller before
+     *                          calling this function; otherwise the resulting
+     *                          per-list counts will be incorrect.
+     * @param list_counts_size - Size of list_counts array (must equal nlist)
+     * @param params           - Search parameters containing the selector
+     *                          that defines which vectors are included.
+     *                          Currently only IDSelectorBitmap is supported.
+     *
+     * @note Requires direct_map to be set (not NoMap)
+     */
+    void count_ivf_list_vectors(
+            idx_t* list_counts,
+            size_t list_counts_size,
+            const SearchParameters* params) const;
 
     /** Similar to search, but also reconstructs the stored vectors (or an
      * approximation in the case of lossy coding) for the search results.