MB-62985 - Adding BQ vector-related utils.

c_api/CMakeLists.txt

@@ -24,6 +24,7 @@ set(FAISS_C_SRC
   Index_c.cpp
   Index_c_ex.cpp
   IndexBinary_c.cpp
+  IndexBinaryIVF_c.cpp
   IndexScalarQuantizer_c.cpp
   MetaIndexes_c.cpp
   clone_index_c.cpp

c_api/IndexBinaryIVF_c.cpp

@@ -0,0 +1,87 @@
+#include "IndexBinaryIVF_c.h"
+#include <faiss/IndexBinaryIVF.h>
+#include "macros_impl.h"
+
+using faiss::IndexBinaryIVF;
+
+DEFINE_GETTER_PERMISSIVE(IndexBinaryIVF, FaissIndexBinaryPtr, quantizer)
+
+extern "C" {
+
+int faiss_IndexBinaryIVF_set_direct_map(
+        FaissIndexBinaryIVF* index,
+        int direct_map_type) {
+    try {
+        reinterpret_cast<IndexBinaryIVF*>(index)->set_direct_map_type(
+                static_cast<faiss::DirectMap::Type>(direct_map_type));
+        return 0;
+    }
+    CATCH_AND_HANDLE
+}
+
+int faiss_get_lists_for_keys_binary(
+        FaissIndexBinaryIVF* index,
+        idx_t* keys,
+        size_t n_keys,
+        idx_t* lists) {
+    try {
+        reinterpret_cast<IndexBinaryIVF*>(index)->get_lists_for_keys(
+                keys, n_keys, lists);
+        return 0;
+    }
+    CATCH_AND_HANDLE
+}
+
+int faiss_Search_closest_eligible_centroids_binary(
+        FaissIndexBinaryIVF* index,
+        idx_t n,
+        const uint8_t* query,
+        idx_t k,
+        int32_t* centroid_distances,
+        idx_t* centroid_ids,
+        const FaissSearchParameters* params) {
+    try {
+        auto idx = reinterpret_cast<IndexBinaryIVF*>(index);
+        idx->quantizer->search(
+                n,
+                query,
+                k,
+                centroid_distances,
+                centroid_ids,
+                reinterpret_cast<const faiss::SearchParameters*>(params));
+
+        return 0;
+    }
+    CATCH_AND_HANDLE
+}
+
+int faiss_IndexBinaryIVF_set_is_trained(FaissIndexBinaryIVF* index, int is_trained) {
+    try {
+        reinterpret_cast<faiss::IndexBinaryIVF*>(index)->is_trained = static_cast<bool>(is_trained);
+    }
+    CATCH_AND_HANDLE
+}
+
+int faiss_IndexBinaryIVF_search_preassigned_with_params(
+        const FaissIndexBinaryIVF* index,
+        idx_t n,
+        const uint8_t* x,
+        idx_t k,
+        const idx_t* assign,
+        const int32_t* centroid_dis,
+        int32_t* distances,
+        idx_t* labels,
+        int store_pairs,
+        const FaissSearchParametersIVF* params) {
+    try {
+        auto idx = reinterpret_cast<const IndexBinaryIVF*>(index);
+        idx->search_preassigned(
+                n, x, k, assign, centroid_dis, distances, labels,
+                static_cast<bool>(store_pairs),
+                reinterpret_cast<const faiss::IVFSearchParameters*>(params));
+        return 0;
+    }
+    CATCH_AND_HANDLE
+}
+
+} // extern "C"

c_api/IndexBinaryIVF_c.h

@@ -0,0 +1,57 @@
+#ifndef FAISS_INDEX_BINARY_IVF_C_H
+#define FAISS_INDEX_BINARY_IVF_C_H
+
+#include "Index_c.h"
+#include "IndexBinary_c.h"
+#include "IndexIVF_c.h"
+#include "faiss_c.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+FAISS_DECLARE_GETTER_SETTER(IndexBinaryIVF, size_t, nlist)
+
+FAISS_DECLARE_GETTER_SETTER(IndexBinaryIVF, size_t, nprobe)
+
+int faiss_IndexBinaryIVF_set_direct_map(
+        FaissIndexBinaryIVF* index,
+        int direct_map_type);
+
+int faiss_get_lists_for_keys_binary(
+        FaissIndexBinaryIVF* index,
+        idx_t* keys,
+        size_t n_keys,
+        idx_t* lists);
+
+int faiss_Search_closest_eligible_centroids_binary(
+        FaissIndexBinaryIVF* index,
+        idx_t n,
+        const uint8_t* query,
+        idx_t k,
+        int32_t* centroid_distances,
+        idx_t* centroid_ids,
+        const FaissSearchParameters* params);
+
+int faiss_IndexBinaryIVF_set_is_trained(FaissIndexBinaryIVF* index, int is_trained);
+
+int faiss_IndexBinaryIVF_search_preassigned_with_params(
+        const FaissIndexBinaryIVF* index,
+        idx_t n,
+        const uint8_t* x,
+        idx_t k,
+        const idx_t* assign,
+        const int32_t* centroid_dis,
+        int32_t* distances,
+        idx_t* labels,
+        int store_pairs,
+        const FaissSearchParametersIVF* params);
+
+typedef FaissIndexBinary* FaissIndexBinaryPtr;
+FaissIndexBinaryPtr faiss_IndexBinaryIVF_quantizer(const FaissIndexBinaryIVF* index);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif

c_api/IndexBinary_c.cpp

@@ -9,6 +9,9 @@
 
 #include "IndexBinary_c.h"
 #include <faiss/IndexBinary.h>
+#include <faiss/IndexBinaryIVF.h>
+#include <faiss/IndexIVF.h>
+#include "IndexIVF_c_ex.h"
 #include "macros_impl.h"
 
 extern "C" {
@@ -18,6 +21,7 @@ DEFINE_DESTRUCTOR(IndexBinary)
 DEFINE_GETTER(IndexBinary, int, d)
 
 DEFINE_GETTER(IndexBinary, int, is_trained)
+DEFINE_SETTER(IndexBinary, int, is_trained)
 
 DEFINE_GETTER(IndexBinary, idx_t, ntotal)
 
@@ -26,6 +30,14 @@ DEFINE_GETTER(IndexBinary, FaissMetricType, metric_type)
 DEFINE_GETTER(IndexBinary, int, verbose);
 DEFINE_SETTER(IndexBinary, int, verbose);
 
+DEFINE_GETTER(IndexBinaryIVF, size_t, nprobe);
+DEFINE_SETTER(IndexBinaryIVF, size_t, nprobe);
+
+DEFINE_GETTER(IndexBinaryIVF, size_t, nlist);
+DEFINE_SETTER(IndexBinaryIVF, size_t, nlist);
+
+DEFINE_INDEXBINARY_DOWNCAST(IndexBinaryIVF)
+
 int faiss_IndexBinary_train(
         FaissIndexBinary* index,
         idx_t n,
@@ -68,6 +80,26 @@ int faiss_IndexBinary_search(
     CATCH_AND_HANDLE
 }
 
+int faiss_IndexBinary_search_with_params(
+        const FaissIndexBinary* index,
+        idx_t n,
+        const uint8_t* x,
+        idx_t k,
+        const FaissSearchParameters* params,
+        int32_t* distances,
+        idx_t* labels) {
+    try {
+        reinterpret_cast<const faiss::IndexBinary*>(index)->search(
+                n,
+                x,
+                k,
+                distances,
+                labels,
+                reinterpret_cast<const faiss::SearchParameters*>(params));
+    }
+    CATCH_AND_HANDLE
+}
+
 int faiss_IndexBinary_range_search(
         const FaissIndexBinary* index,
         idx_t n,

c_api/IndexBinary_c.h

@@ -28,11 +28,17 @@ typedef struct FaissIDSelector_H FaissIDSelector;
 FAISS_DECLARE_CLASS(IndexBinary)
 FAISS_DECLARE_DESTRUCTOR(IndexBinary)
 
+FAISS_DECLARE_CLASS_INHERITED(IndexBinaryIVF, IndexBinary)
+
+/// Cast function for IndexBinaryIVF
+FaissIndexBinaryIVF* faiss_IndexBinaryIVF_cast(FaissIndexBinary* index);
+
 /// Getter for d
 FAISS_DECLARE_GETTER(IndexBinary, int, d)
 
 /// Getter for is_trained
-FAISS_DECLARE_GETTER(IndexBinary, int, is_trained)
+
+FAISS_DECLARE_GETTER_SETTER(IndexBinary, int, is_trained)
 
 /// Getter for ntotal
 FAISS_DECLARE_GETTER(IndexBinary, idx_t, ntotal)
@@ -92,6 +98,27 @@ int faiss_IndexBinary_search(
         int32_t* distances,
         idx_t* labels);
 
+/** query n vectors of dimension d to the index.
+ *
+ * return at most k vectors. If there are not enough results for a
+ * query, the result array is padded with -1s.
+ *
+ * @param index       opaque pointer to index object
+ * @param x           input vectors to search, size n * d
+ * @param k           number of results to return
+ * @param params      search parameters
+ * @param distances   output distances, size n*k
+ * @param labels      output labels of the NNs, size n*k
+ */
+int faiss_IndexBinary_search_with_params(
+        const FaissIndexBinary* index,
+        idx_t n,
+        const uint8_t* x,
+        idx_t k,
+        const FaissSearchParameters* params,
+        int32_t* distances,
+        idx_t* labels);
+
 /** query n vectors of dimension d to the index.
  *
  * return all vectors with distance < radius. Note that many

c_api/IndexIVF_c.cpp

@@ -182,3 +182,27 @@ void faiss_IndexIVFStats_reset(FaissIndexIVFStats* stats) {
 FaissIndexIVFStats* faiss_get_indexIVF_stats() {
     return reinterpret_cast<FaissIndexIVFStats*>(&faiss::indexIVF_stats);
 }
+
+int faiss_IndexIVF_get_centroids(
+        const FaissIndexIVF* index,
+        float* centroids) {
+    try {
+        const IndexIVF* ivf = reinterpret_cast<const IndexIVF*>(index);
+        const faiss::Index* quantizer = ivf->quantizer;
+
+        if (!quantizer) {
+            FAISS_THROW_MSG("No quantizer found in IVF index");
+        }
+
+        // Get all centroids using reconstruct_n
+        quantizer->reconstruct_n(0, ivf->nlist, centroids);
+    }
+    CATCH_AND_HANDLE
+}
+
+int faiss_IndexIVF_free_centroids(float* centroids) {
+    try {
+        delete[] centroids;
+    }
+    CATCH_AND_HANDLE
+}

c_api/IndexIVF_c.h

@@ -177,6 +177,19 @@ inline void faiss_IndexIVFStats_init(FaissIndexIVFStats* stats) {
 /// global var that collects all statists
 FaissIndexIVFStats* faiss_get_indexIVF_stats();
 
+/** Get centroids from an IVF index
+ *
+ * @param index       IVF index
+ * @param centroids   output centroids, size nlist * d (allocated by caller)
+ * @return            0 if successful
+ */
+int faiss_IndexIVF_get_centroids(
+        const FaissIndexIVF* index,
+        float* centroids);
+
+// Free function for auto-allocated centroids
+int faiss_IndexIVF_free_centroids(float* centroids);
+
 #ifdef __cplusplus
 }
 #endif

c_api/IndexIVF_c_ex.cpp

@@ -38,16 +38,21 @@ int faiss_Search_closest_eligible_centroids(
         FaissIndex* index,
         idx_t n,
         const float* query,
-        idx_t k, 
+        idx_t k,
         float* centroid_distances,
         idx_t* centroid_ids,
         const FaissSearchParameters* params) {
     try {
-        faiss::IndexIVF* index_ivf = reinterpret_cast<IndexIVF*>(index); 
+        faiss::IndexIVF* index_ivf = reinterpret_cast<IndexIVF*>(index);
         assert(index_ivf);
 
-        index_ivf->quantizer->search(n, query, k, centroid_distances, centroid_ids,
-            reinterpret_cast<const faiss::SearchParameters*>(params));
+        index_ivf->quantizer->search(
+                n,
+                query,
+                k,
+                centroid_distances,
+                centroid_ids,
+                reinterpret_cast<const faiss::SearchParameters*>(params));
     }
     CATCH_AND_HANDLE
 }
@@ -58,7 +63,8 @@ int faiss_get_lists_for_keys(
         size_t n_keys,
         idx_t* lists) {
     try {
-        reinterpret_cast<IndexIVF*>(index)->get_lists_for_keys(keys, n_keys, lists);
+        reinterpret_cast<IndexIVF*>(index)->get_lists_for_keys(
+                keys, n_keys, lists);
     }
     CATCH_AND_HANDLE
 }
@@ -76,7 +82,14 @@ int faiss_IndexIVF_search_preassigned_with_params(
         const FaissSearchParametersIVF* params) {
     try {
         reinterpret_cast<const IndexIVF*>(index)->search_preassigned(
-                n, x, k, assign, centroid_dis, distances, labels, store_pairs,
+                n,
+                x,
+                k,
+                assign,
+                centroid_dis,
+                distances,
+                labels,
+                store_pairs,
                 reinterpret_cast<const faiss::SearchParametersIVF*>(params));
     }
     CATCH_AND_HANDLE
@@ -92,7 +105,7 @@ int faiss_IndexIVF_compute_distance_to_codes_for_list(
         float* dist_table) {
     try {
         reinterpret_cast<IndexIVF*>(index)->compute_distance_to_codes_for_list(
-               list_no, x, n, codes, dists, dist_table);
+                list_no, x, n, codes, dists, dist_table);
         return 0;
     }
     CATCH_AND_HANDLE
@@ -104,7 +117,7 @@ int faiss_IndexIVF_compute_distance_table(
         float* dist_table) {
     try {
         reinterpret_cast<IndexIVF*>(index)->compute_distance_table(
-               x, dist_table);
+                x, dist_table);
         return 0;
     }
     CATCH_AND_HANDLE