faiss 0.5.3 → 0.6.0

data/lib/faiss/version.rb

@@ -1,3 +1,3 @@
 module Faiss
-  VERSION = "0.5.3"
+  VERSION = "0.6.0"
 end

data/lib/faiss.rb

@@ -1,5 +1,5 @@
 # dependencies
-require "numo/narray"
+require "numo/narray/alt"
 
 # ext
 require "faiss/ext"

data/vendor/faiss/faiss/AutoTune.cpp

@@ -19,6 +19,8 @@
 #include <faiss/utils/random.h>
 #include <faiss/utils/utils.h>
 
+#include <faiss/IndexBinaryHNSW.h>
+#include <faiss/IndexBinaryIVF.h>
 #include <faiss/IndexHNSW.h>
 #include <faiss/IndexIDMap.h>
 #include <faiss/IndexIVF.h>
@@ -408,21 +410,24 @@ void ParameterSpace::initialize(const Index* index) {
 
 #undef DC
 
-/// set a combination of parameters on an index
-void ParameterSpace::set_index_parameters(Index* index, size_t cno) const {
-    for (int i = 0; i < parameter_ranges.size(); i++) {
-        const ParameterRange& pr = parameter_ranges[i];
+template <typename SetParamFunc>
+static void set_index_parameters_common(
+        const ParameterSpace* ps,
+        size_t cno,
+        SetParamFunc set_param) {
+    for (int i = 0; i < ps->parameter_ranges.size(); i++) {
+        const ParameterRange& pr = ps->parameter_ranges[i];
         size_t j = cno % pr.values.size();
         cno /= pr.values.size();
         double val = pr.values[j];
-        set_index_parameter(index, pr.name, val);
+        set_param(pr.name, val);
     }
 }
 
-/// set a combination of parameters on an index
-void ParameterSpace::set_index_parameters(
-        Index* index,
-        const char* description_in) const {
+template <typename SetParamFunc>
+static void set_index_parameters_string_common(
+        const char* description_in,
+        SetParamFunc set_param) {
     std::string description(description_in);
     char* ptr;
 
@@ -433,10 +438,47 @@ void ParameterSpace::set_index_parameters(
         int ret = sscanf(tok, "%99[^=]=%lf", name, &val);
         FAISS_THROW_IF_NOT_FMT(
                 ret == 2, "could not interpret parameters %s", tok);
-        set_index_parameter(index, name, val);
+        set_param(name, val);
     }
 }
 
+/// set a combination of parameters on an index
+void ParameterSpace::set_index_parameters(Index* index, size_t cno) const {
+    set_index_parameters_common(
+            this, cno, [this, index](const std::string& name, double val) {
+                this->set_index_parameter(index, name, val);
+            });
+}
+
+/// set a combination of parameters on an index
+void ParameterSpace::set_index_parameters(
+        Index* index,
+        const char* description_in) const {
+    set_index_parameters_string_common(
+            description_in, [this, index](const std::string& name, double val) {
+                this->set_index_parameter(index, name, val);
+            });
+}
+
+/// set a combination of parameters on a binary index
+void ParameterSpace::set_index_parameters(IndexBinary* index, size_t cno)
+        const {
+    set_index_parameters_common(
+            this, cno, [this, index](const std::string& name, double val) {
+                this->set_index_parameter(index, name, val);
+            });
+}
+
+/// set a combination of parameters on a binary index
+void ParameterSpace::set_index_parameters(
+        IndexBinary* index,
+        const char* description_in) const {
+    set_index_parameters_string_common(
+            description_in, [this, index](const std::string& name, double val) {
+                this->set_index_parameter(index, name, val);
+            });
+}
+
 // non-const version
 // Do not use this macro if ix will be unused
 #define DC(classname) classname* ix = dynamic_cast<classname*>(index)
@@ -490,7 +532,6 @@ void ParameterSpace::set_index_parameter(
     }
 
     if (name == "verbose") {
-        index->verbose = int(val);
         return; // last verbose that we could find
     }
 
@@ -573,6 +614,84 @@ void ParameterSpace::set_index_parameter(
             name.c_str());
 }
 
+void ParameterSpace::set_index_parameter(
+        IndexBinary* index,
+        const std::string& name,
+        double val) const {
+    if (verbose > 1) {
+        printf("    set_index_parameter (binary) %s=%g\n", name.c_str(), val);
+    }
+
+    if (name == "verbose") {
+        index->verbose = int(val);
+        // and fall through to also enable it on sub-indexes
+    }
+
+    if (DC(IndexBinaryIDMap)) {
+        set_index_parameter(ix->index, name, val);
+        return;
+    }
+
+    if (name == "verbose") {
+        return; // last verbose that we could find
+    }
+
+    if (name == "nprobe") {
+        if (DC(IndexBinaryIVF)) {
+            ix->nprobe = int(val);
+            return;
+        }
+    }
+
+    if (name == "max_codes") {
+        if (DC(IndexBinaryIVF)) {
+            ix->max_codes = std::isfinite(val) ? size_t(val) : 0;
+            return;
+        }
+    }
+
+    if (name == "efConstruction") {
+        if (DC(IndexBinaryHNSW)) {
+            ix->hnsw.efConstruction = int(val);
+            return;
+        }
+        if (DC(IndexBinaryIVF)) {
+            if (IndexBinaryHNSW* cq =
+                        dynamic_cast<IndexBinaryHNSW*>(ix->quantizer)) {
+                cq->hnsw.efConstruction = int(val);
+                return;
+            }
+        }
+    }
+
+    if (name == "efSearch") {
+        if (DC(IndexBinaryHNSW)) {
+            ix->hnsw.efSearch = int(val);
+            return;
+        }
+        if (DC(IndexBinaryIVF)) {
+            if (IndexBinaryHNSW* cq =
+                        dynamic_cast<IndexBinaryHNSW*>(ix->quantizer)) {
+                cq->hnsw.efSearch = int(val);
+                return;
+            }
+        }
+    }
+
+    if (name.find("quantizer_") == 0) {
+        if (DC(IndexBinaryIVF)) {
+            std::string sub_name = name.substr(strlen("quantizer_"));
+            set_index_parameter(ix->quantizer, sub_name, val);
+            return;
+        }
+    }
+
+    FAISS_THROW_FMT(
+            "ParameterSpace::set_index_parameter:"
+            "could not set parameter %s on binary index",
+            name.c_str());
+}
+
 #undef DC
 
 void ParameterSpace::display() const {

data/vendor/faiss/faiss/AutoTune.h

@@ -177,12 +177,25 @@ struct ParameterSpace {
     /// set a combination of parameters described by a string
     void set_index_parameters(Index* index, const char* param_string) const;
 
-    /// set one of the parameters, returns whether setting was successful
+    /// set one of the parameters
     virtual void set_index_parameter(
             Index* index,
             const std::string& name,
             double val) const;
 
+    /// set a combination of parameters on a binary index
+    void set_index_parameters(IndexBinary* index, size_t cno) const;
+
+    /// set a combination of parameters described by a string on a binary index
+    void set_index_parameters(IndexBinary* index, const char* param_string)
+            const;
+
+    /// set one of the parameters on a binary index
+    virtual void set_index_parameter(
+            IndexBinary* index,
+            const std::string& name,
+            double val) const;
+
     /** find an upper bound on the performance and a lower bound on t
      * for configuration cno given another operating point op */
     void update_bounds(

data/vendor/faiss/faiss/Clustering.cpp

@@ -407,19 +407,52 @@ void Clustering::train_encoded(
             printf("Outer iteration %d / %d\n", redo, nredo);
         }
 
-        // initialize (remaining) centroids with random points from the dataset
+        // initialize centroids using the selected method
         centroids.resize(d * k);
-        std::vector<int> perm(nx);
 
-        rand_perm(perm.data(), nx, actual_seed + 1 + redo * 15486557L);
+        size_t k_to_init = k - n_input_centroids;
+        if (k_to_init > 0) {
+            // Fast path for RANDOM initialization - preserves exact original
+            // behavior
+            if (init_method == ClusteringInitMethod::RANDOM) {
+                std::vector<int> perm(nx);
+                rand_perm(perm.data(), nx, actual_seed + 1 + redo * 15486557L);
+                for (size_t i = 0; i < k_to_init; i++) {
+                    if (!codec) {
+                        memcpy(centroids.data() + (n_input_centroids + i) * d,
+                               x + perm[n_input_centroids + i] * line_size,
+                               line_size);
+                    } else {
+                        codec->sa_decode(
+                                1,
+                                x + perm[n_input_centroids + i] * line_size,
+                                centroids.data() + (n_input_centroids + i) * d);
+                    }
+                }
+            } else {
+                // For k-means++ and AFK-MC², we need all vectors decoded
+                const float* x_float = nullptr;
+                std::vector<float> x_decoded;
 
-        if (!codec) {
-            for (int i = n_input_centroids; i < k; i++) {
-                memcpy(&centroids[i * d], x + perm[i] * line_size, line_size);
-            }
-        } else {
-            for (int i = n_input_centroids; i < k; i++) {
-                codec->sa_decode(1, x + perm[i] * line_size, &centroids[i * d]);
+                if (!codec) {
+                    x_float = reinterpret_cast<const float*>(x);
+                } else {
+                    // Decode all vectors for initialization
+                    x_decoded.resize(nx * d);
+                    codec->sa_decode(nx, x, x_decoded.data());
+                    x_float = x_decoded.data();
+                }
+
+                ClusteringInitialization initializer(d, k_to_init);
+                initializer.method = init_method;
+                initializer.seed = actual_seed + 1 + redo * 15486557L;
+                initializer.afkmc2_chain_length = afkmc2_chain_length;
+                initializer.init_centroids(
+                        nx,
+                        x_float,
+                        centroids.data() + n_input_centroids * d,
+                        n_input_centroids,
+                        n_input_centroids > 0 ? centroids.data() : nullptr);
             }
         }
 
@@ -529,6 +562,22 @@ void Clustering::train_encoded(
 
             index.add(k, centroids.data());
             InterruptCallback::check();
+
+            // Early stopping: if objective didn't change, we've converged.
+            // Safe to access iteration_stats[size - 2] because we push_back
+            // above, so size >= i + 1, and when i > 0 we have size >= 2.
+            if (i > 0) {
+                float prev_obj =
+                        iteration_stats[iteration_stats.size() - 2].obj;
+                if (obj == prev_obj) {
+                    if (verbose) {
+                        printf("\n  Converged at iteration %d: "
+                               "objective did not change\n",
+                               i);
+                    }
+                    break;
+                }
+            }
         }
 
         if (verbose) {

data/vendor/faiss/faiss/Clustering.h

@@ -10,6 +10,7 @@
 #ifndef FAISS_CLUSTERING_H
 #define FAISS_CLUSTERING_H
 #include <faiss/Index.h>
+#include <faiss/impl/ClusteringInitialization.h>
 
 #include <vector>
 
@@ -57,6 +58,17 @@ struct ClusteringParameters {
     /// Whether to use splitmix64-based random number generator for subsampling,
     /// which is faster, but may pick duplicate points.
     bool use_faster_subsampling = false;
+
+    /// Initialization method for centroids.
+    /// RANDOM: uniform random sampling (default, current behavior)
+    /// KMEANS_PLUS_PLUS: k-means++ (O(nkd), better quality)
+    /// AFK_MC2: Assumption-Free K-MC² (O(nd) + O(mk²d), fast approximation)
+    ClusteringInitMethod init_method = ClusteringInitMethod::RANDOM;
+
+    /// Chain length for AFK-MC² initialization.
+    /// Only used when init_method = AFK_MC2.
+    /// Longer chains give better approximation but are slower.
+    uint16_t afkmc2_chain_length = 50;
 };
 
 struct ClusteringIterationStats {

data/vendor/faiss/faiss/IVFlib.cpp

@@ -18,6 +18,7 @@
 #include <faiss/MetaIndexes.h>
 #include <faiss/clone_index.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/index_io.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/hamming.h>
@@ -512,39 +513,41 @@ void ivf_residual_add_from_flat_codes(
     const ResidualQuantizer& rq = index->rq;
 
     // populate inverted lists
-#pragma omp parallel if (nb > 10000)
-    {
-        std::vector<uint8_t> tmp_code(index->code_size);
-        std::vector<float> tmp(rq.d);
-        int nt = omp_get_num_threads();
-        int rank = omp_get_thread_num();
+    with_simd_level([&]<SIMDLevel SL>() {
+#pragma omp parallel
+        {
+            std::vector<uint8_t> tmp_code(index->rq.code_size);
+            std::vector<float> tmp(rq.d);
+            int nt = omp_get_num_threads();
+            int rank = omp_get_thread_num();
 
 #pragma omp for
-        for (idx_t i = 0; i < nb; i++) {
-            const uint8_t* code = &raw_codes[i * code_size];
-            BitstringReader rd(code, code_size);
-            idx_t list_no = rd.read(rcq->rq.tot_bits);
-
-            if (list_no % nt ==
-                rank) { // each thread takes care of 1/nt of the invlists
-                // copy AQ indexes one by one
-                BitstringWriter wr(tmp_code.data(), tmp_code.size());
-                for (int j = 0; j < rq.M; j++) {
-                    int nbit = rq.nbits[j];
-                    wr.write(rd.read(nbit), nbit);
+            for (idx_t i = 0; i < nb; i++) {
+                const uint8_t* code = &raw_codes[i * code_size];
+                BitstringReader rd(code, code_size);
+                idx_t list_no = rd.read(rcq->rq.tot_bits);
+
+                if (list_no % nt ==
+                    rank) { // each thread takes care of 1/nt of the invlists
+                    // copy AQ indexes one by one
+                    BitstringWriter wr(tmp_code.data(), tmp_code.size());
+                    for (int j = 0; j < rq.M; j++) {
+                        int nbit = rq.nbits[j];
+                        wr.write(rd.read(nbit), nbit);
+                    }
+                    // we need to recompute the norm
+                    // decode first, does not use the norm component, so that's
+                    // ok
+                    index->rq.decode(tmp_code.data(), tmp.data(), 1);
+                    float norm = fvec_norm_L2sqr<SL>(tmp.data(), rq.d);
+                    wr.write(rq.encode_norm(norm), rq.norm_bits);
+
+                    // add code to the inverted list
+                    invlists.add_entry(list_no, i, tmp_code.data());
                 }
-                // we need to recompute the norm
-                // decode first, does not use the norm component, so that's
-                // ok
-                index->rq.decode(tmp_code.data(), tmp.data(), 1);
-                float norm = fvec_norm_L2sqr(tmp.data(), rq.d);
-                wr.write(rq.encode_norm(norm), rq.norm_bits);
-
-                // add code to the inverted list
-                invlists.add_entry(list_no, i, tmp_code.data());
             }
         }
-    }
+    });
     index->ntotal += nb;
 }
 

data/vendor/faiss/faiss/Index.cpp

@@ -24,12 +24,20 @@ void Index::train(idx_t /*n*/, const float* /*x*/) {
     // does nothing by default
 }
 
+void Index::train(
+        idx_t /*n*/,
+        const float* /*x*/,
+        idx_t /*n_train_q*/,
+        const float* /*xq_train*/) {
+    // does nothing by default
+}
+
 void Index::range_search(
         idx_t,
         const float*,
         float,
         RangeSearchResult*,
-        const SearchParameters* params) const {
+        const SearchParameters* /*params*/) const {
     FAISS_THROW_MSG("range search not implemented");
 }
 
@@ -105,16 +113,20 @@ void Index::search_and_reconstruct(
 }
 
 void Index::search_subset(
-        idx_t n,
-        const float* x,
-        idx_t k_base,
-        const idx_t* base_labels,
-        idx_t k,
-        float* distances,
-        idx_t* labels) const {
+        idx_t /*n*/,
+        const float* /*x*/,
+        idx_t /*k_base*/,
+        const idx_t* /*base_labels*/,
+        idx_t /*k*/,
+        float* /*distances*/,
+        idx_t* /*labels*/) const {
     FAISS_THROW_MSG("search_subset not implemented for this type of index");
 }
 
+void Index::search1(const float*, ResultHandler&, SearchParameters*) const {
+    FAISS_THROW_MSG("search1 not implemented for this type of index");
+}
+
 void Index::compute_residual(const float* x, float* residual, idx_t key) const {
     reconstruct(key, residual);
     for (size_t i = 0; i < d; i++) {

data/vendor/faiss/faiss/Index.h

@@ -14,11 +14,10 @@
 #include <faiss/impl/FaissAssert.h>
 
 #include <cstdio>
-#include <sstream>
 
 #define FAISS_VERSION_MAJOR 1
-#define FAISS_VERSION_MINOR 13
-#define FAISS_VERSION_PATCH 2
+#define FAISS_VERSION_MINOR 14
+#define FAISS_VERSION_PATCH 1
 
 // Macro to combine the version components into a single string
 #ifndef FAISS_STRINGIFY
@@ -55,6 +54,9 @@ namespace faiss {
 struct IDSelector;
 struct RangeSearchResult;
 struct DistanceComputer;
+template <typename T, typename TI>
+struct ResultHandlerUnordered;
+using ResultHandler = ResultHandlerUnordered<float, idx_t>;
 
 enum NumericType {
     Float32,
@@ -129,6 +131,20 @@ struct Index {
      */
     virtual void train(idx_t n, const float* x);
 
+    /** Perfrom training on a representative set of vectors and a representative
+     * set of queries
+     *
+     * @param n         nb of training vectors
+     * @param x         training vectors, size n * d
+     * @param n_train_q nb of training queries
+     * @param xq_train  training queries, size n_train_q * d
+     */
+    virtual void train(
+            idx_t n,
+            const float* x,
+            idx_t n_train_q,
+            const float* xq_train);
+
     virtual void train_ex(idx_t n, const void* x, NumericType numeric_type) {
         if (numeric_type == NumericType::Float32) {
             train(n, static_cast<const float*>(x));
@@ -216,6 +232,12 @@ struct Index {
         }
     }
 
+    /** search one vector with a custom result handler */
+    virtual void search1(
+            const float* x,
+            ResultHandler& handler,
+            SearchParameters* params = nullptr) const;
+
     /** query n vectors of dimension d to the index.
      *
      * return all vectors with distance < radius. Note that many

data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp

@@ -8,13 +8,11 @@
 #include <faiss/IndexAdditiveQuantizer.h>
 
 #include <algorithm>
-#include <cmath>
 #include <cstring>
 
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/ResidualQuantizer.h>
 #include <faiss/impl/ResultHandler.h>
-#include <faiss/utils/distances.h>
 #include <faiss/utils/extra_distances.h>
 
 namespace faiss {
@@ -189,17 +187,14 @@ void search_with_LUT(
 FlatCodesDistanceComputer* IndexAdditiveQuantizer::
         get_FlatCodesDistanceComputer() const {
     if (aq->search_type == AdditiveQuantizer::ST_decompress) {
-        if (metric_type == METRIC_L2) {
-            using VD = VectorDistance<METRIC_L2>;
-            VD vd = {size_t(d), metric_arg};
-            return new AQDistanceComputerDecompress<VD>(*this, vd);
-        } else if (metric_type == METRIC_INNER_PRODUCT) {
-            using VD = VectorDistance<METRIC_INNER_PRODUCT>;
-            VD vd = {size_t(d), metric_arg};
-            return new AQDistanceComputerDecompress<VD>(*this, vd);
-        } else {
-            FAISS_THROW_MSG("unsupported metric");
-        }
+        return with_VectorDistance(
+                d,
+                metric_type,
+                metric_arg,
+                [&](auto vd) -> FlatCodesDistanceComputer* {
+                    return new AQDistanceComputerDecompress<decltype(vd)>(
+                            *this, vd);
+                });
     } else {
         if (metric_type == METRIC_INNER_PRODUCT) {
             return new AQDistanceComputerLUT<
@@ -242,17 +237,17 @@ void IndexAdditiveQuantizer::search(
             !params, "search params not supported for this index");
 
     if (aq->search_type == AdditiveQuantizer::ST_decompress) {
-        if (metric_type == METRIC_L2) {
-            using VD = VectorDistance<METRIC_L2>;
-            VD vd = {size_t(d), metric_arg};
-            HeapBlockResultHandler<VD::C> rh(n, distances, labels, k);
-            search_with_decompress(*this, x, vd, rh);
-        } else if (metric_type == METRIC_INNER_PRODUCT) {
-            using VD = VectorDistance<METRIC_INNER_PRODUCT>;
-            VD vd = {size_t(d), metric_arg};
-            HeapBlockResultHandler<VD::C> rh(n, distances, labels, k);
-            search_with_decompress(*this, x, vd, rh);
-        }
+        with_VectorDistance(d, metric_type, metric_arg, [&](auto vd) {
+            if constexpr (decltype(vd)::is_similarity) {
+                HeapBlockResultHandler<CMin<float, idx_t>> rh(
+                        n, distances, labels, k);
+                search_with_decompress(*this, x, vd, rh);
+            } else {
+                HeapBlockResultHandler<CMax<float, idx_t>> rh(
+                        n, distances, labels, k);
+                search_with_decompress(*this, x, vd, rh);
+            }
+        });
     } else {
         if (metric_type == METRIC_INNER_PRODUCT) {
             HeapBlockResultHandler<CMin<float, idx_t>> rh(

data/vendor/faiss/faiss/IndexBinary.cpp

@@ -12,6 +12,7 @@
 
 #include <cinttypes>
 #include <cstring>
+#include <typeinfo>
 
 namespace faiss {
 

data/vendor/faiss/faiss/IndexBinaryHNSW.cpp

@@ -22,6 +22,7 @@
 #include <faiss/impl/DistanceComputer.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/ResultHandler.h>
+#include <faiss/impl/VisitedTable.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/hamming.h>
 #include <faiss/utils/random.h>
@@ -205,10 +206,14 @@ void IndexBinaryHNSW::search(
         idx_t k,
         int32_t* distances,
         idx_t* labels,
-        const SearchParameters* params) const {
-    FAISS_THROW_IF_NOT_MSG(
-            !params, "search params not supported for this index");
+        const SearchParameters* params_in) const {
     FAISS_THROW_IF_NOT(k > 0);
+    const SearchParametersHNSW* params = nullptr;
+    if (params_in) {
+        params = dynamic_cast<const SearchParametersHNSW*>(params_in);
+        FAISS_THROW_IF_NOT_MSG(
+                params, "IndexBinaryHNSW params have incorrect type");
+    }
 
     // we use the buffer for distances as float but convert them back
     // to int in the end
@@ -231,7 +236,7 @@ void IndexBinaryHNSW::search(
             // as the index parameter. This state does not get used in the
             // search function, as it is merely there to to enable Panorama
             // execution for IndexHNSWFlatPanorama.
-            hnsw.search(*dis, nullptr, res, vt);
+            hnsw.search(*dis, nullptr, res, vt, params_in);
             res.end();
         }
     }