faiss 0.6.2 → 0.6.3

data/vendor/faiss/faiss/IndexFastScan.cpp

@@ -8,6 +8,7 @@
 #include <faiss/IndexFastScan.h>
 
 #include <omp.h>
+#include <algorithm>
 #include <cstring>
 #include <memory>
 
@@ -266,8 +267,7 @@ void IndexFastScan::search(
         float* distances,
         idx_t* labels,
         const SearchParameters* params) const {
-    FAISS_THROW_IF_NOT_MSG(
-            !params, "search params not supported for this index");
+    FAISS_THROW_IF_MSG(params, "search params not supported for this index");
     FAISS_THROW_IF_NOT(k > 0);
 
     FastScanDistancePostProcessing empty_context{};
@@ -316,7 +316,7 @@ void IndexFastScan::search_dispatch_implem(
     if (implem == 1) {
         FAISS_THROW_MSG("not implemented");
     } else if (implem == 2 || implem == 3 || implem == 4) {
-        FAISS_THROW_IF_NOT(orig_codes != nullptr);
+        FAISS_THROW_IF_NOT(orig_codes);
         search_implem_234<Cfloat>(n, x, k, distances, labels, context);
     } else if (impl >= 12 && impl <= 15) {
         FAISS_THROW_IF_NOT(ntotal < INT_MAX);
@@ -617,8 +617,19 @@ template void IndexFastScan::search_dispatch_implem<false>(
         const FastScanDistancePostProcessing& context) const;
 
 void IndexFastScan::reconstruct(idx_t key, float* recons) const {
-    std::vector<uint8_t> code(code_size, 0);
+    FAISS_THROW_IF_NOT_FMT(
+            key >= 0 && key < ntotal,
+            "IndexFastScan::reconstruct: key %zd out of range (ntotal=%zd)",
+            (size_t)key,
+            (size_t)ntotal);
     std::unique_ptr<CodePacker> packer(get_CodePacker());
+    size_t block_no = (size_t)key / packer->nvec;
+    FAISS_THROW_IF_NOT_MSG(
+            mul_no_overflow(
+                    block_no + 1, packer->block_size, "IndexFastScan codes") <=
+                    codes.size(),
+            "IndexFastScan::reconstruct: packed codes buffer too small");
+    std::vector<uint8_t> code(std::max(code_size, packer->code_size), 0);
     packer->unpack_1(codes.data(), key, code.data());
     sa_decode(1, code.data(), recons);
 }

data/vendor/faiss/faiss/IndexFlat.cpp

@@ -77,7 +77,8 @@ void IndexFlat::range_search(
             range_search_L2sqr(x, get_xb(), d, n, ntotal, radius, result, sel);
             break;
         default:
-            FAISS_THROW_MSG("metric type not supported");
+            IndexFlatCodes::range_search(n, x, radius, result, params);
+            break;
     }
 }
 
@@ -105,13 +106,8 @@ namespace {
 template <SIMDLevel SL>
 struct FlatL2Dis : FlatCodesDistanceComputer {
     size_t d;
-    idx_t nb;
-    const float* b;
-    size_t ndis;
-    size_t npartial_dot_products;
 
     float distance_to_code(const uint8_t* code) final {
-        ndis++;
         return fvec_L2sqr<SL>(q, (float*)code, d);
     }
 
@@ -119,25 +115,21 @@ struct FlatL2Dis : FlatCodesDistanceComputer {
             const idx_t i,
             const uint32_t offset,
             const uint32_t num_components) final override {
-        npartial_dot_products++;
+        const float* b = (const float*)this->codes;
         return fvec_inner_product<SL>(
                 q + offset, b + i * d + offset, num_components);
     }
 
     float symmetric_dis(idx_t i, idx_t j) override {
+        const float* b = (const float*)this->codes;
         return fvec_L2sqr<SL>(b + j * d, b + i * d, d);
     }
 
-    explicit FlatL2Dis(const IndexFlat& storage, const float* q_ = nullptr)
+    explicit FlatL2Dis(const IndexFlat& storage)
             : FlatCodesDistanceComputer(
                       storage.codes.data(),
-                      storage.code_size,
-                      q_),
-              d(storage.d),
-              nb(storage.ntotal),
-              b(storage.get_xb()),
-              ndis(0),
-              npartial_dot_products(0) {}
+                      storage.code_size),
+              d(storage.d) {}
 
     void set_query(const float* x) override {
         q = x;
@@ -153,8 +145,6 @@ struct FlatL2Dis : FlatCodesDistanceComputer {
             float& dis1,
             float& dis2,
             float& dis3) final override {
-        ndis += 4;
-
         // compute first, assign next
         const float* __restrict y0 =
                 reinterpret_cast<const float*>(codes + idx0 * code_size);
@@ -187,8 +177,6 @@ struct FlatL2Dis : FlatCodesDistanceComputer {
             float& dp3,
             const uint32_t offset,
             const uint32_t num_components) final override {
-        npartial_dot_products += 4;
-
         // compute first, assign next
         const float* __restrict y0 =
                 reinterpret_cast<const float*>(codes + idx0 * code_size);
@@ -224,29 +212,21 @@ struct FlatL2Dis : FlatCodesDistanceComputer {
 template <SIMDLevel SL>
 struct FlatIPDis : FlatCodesDistanceComputer {
     size_t d;
-    idx_t nb;
-    const float* q;
-    const float* b;
-    size_t ndis;
 
     float symmetric_dis(idx_t i, idx_t j) final override {
+        const float* b = (const float*)this->codes;
         return fvec_inner_product<SL>(b + j * d, b + i * d, d);
     }
 
     float distance_to_code(const uint8_t* code) final override {
-        ndis++;
         return fvec_inner_product<SL>(q, (const float*)code, d);
     }
 
-    explicit FlatIPDis(const IndexFlat& storage, const float* q_in = nullptr)
+    explicit FlatIPDis(const IndexFlat& storage)
             : FlatCodesDistanceComputer(
                       storage.codes.data(),
                       storage.code_size),
-              d(storage.d),
-              nb(storage.ntotal),
-              q(q_in),
-              b(storage.get_xb()),
-              ndis(0) {}
+              d(storage.d) {}
 
     void set_query(const float* x) override {
         q = x;
@@ -262,8 +242,6 @@ struct FlatIPDis : FlatCodesDistanceComputer {
             float& dis1,
             float& dis2,
             float& dis3) final override {
-        ndis += 4;
-
         // compute first, assign next
         const float* __restrict y0 =
                 reinterpret_cast<const float*>(codes + idx0 * code_size);
@@ -296,8 +274,7 @@ FlatCodesDistanceComputer* IndexFlat::get_FlatCodesDistanceComputer() const {
     } else if (metric_type == METRIC_INNER_PRODUCT) {
         with_simd_level([&]<SIMDLevel SL>() { dc = new FlatIPDis<SL>(*this); });
     } else {
-        dc = get_extra_distance_computer(
-                d, metric_type, metric_arg, ntotal, get_xb());
+        dc = get_extra_distance_computer(d, metric_type, metric_arg, get_xb());
     }
     return dc;
 }
@@ -327,16 +304,11 @@ namespace {
 template <SIMDLevel SL>
 struct FlatL2WithNormsDis : FlatCodesDistanceComputer {
     size_t d;
-    idx_t nb;
-    const float* q;
-    const float* b;
-    size_t ndis;
 
     const float* l2norms;
     float query_l2norm;
 
     float distance_to_code(const uint8_t* code) final override {
-        ndis++;
         return fvec_L2sqr<SL>(q, (float*)code, d);
     }
 
@@ -361,17 +333,11 @@ struct FlatL2WithNormsDis : FlatCodesDistanceComputer {
         return l2norms[i] + l2norms[j] - 2 * dp0;
     }
 
-    explicit FlatL2WithNormsDis(
-            const IndexFlatL2& storage,
-            const float* q_in = nullptr)
+    explicit FlatL2WithNormsDis(const IndexFlatL2& storage)
             : FlatCodesDistanceComputer(
                       storage.codes.data(),
                       storage.code_size),
               d(storage.d),
-              nb(storage.ntotal),
-              q(q_in),
-              b(storage.get_xb()),
-              ndis(0),
               l2norms(storage.cached_l2norms.data()),
               query_l2norm(0) {}
 
@@ -390,8 +356,6 @@ struct FlatL2WithNormsDis : FlatCodesDistanceComputer {
             float& dis1,
             float& dis2,
             float& dis3) final override {
-        ndis += 4;
-
         // compute first, assign next
         const float* __restrict y0 =
                 reinterpret_cast<const float*>(codes + idx0 * code_size);
@@ -487,8 +451,7 @@ void IndexFlat1D::search(
         float* distances,
         idx_t* labels,
         const SearchParameters* params) const {
-    FAISS_THROW_IF_NOT_MSG(
-            !params, "search params not supported for this index");
+    FAISS_THROW_IF_MSG(params, "search params not supported for this index");
     FAISS_THROW_IF_NOT(k > 0);
     FAISS_THROW_IF_NOT_MSG(
             perm.size() == static_cast<size_t>(ntotal),

data/vendor/faiss/faiss/IndexHNSW.cpp

@@ -280,12 +280,12 @@ void hnsw_search(
 
 #pragma omp parallel if (i1 - i0 > 1)
         {
-            std::unique_ptr<VisitedTable> vt;
+            VisitedTable* vt = nullptr;
             std::unique_ptr<typename BlockResultHandler::SingleResultHandler>
                     res;
             std::unique_ptr<DistanceComputer> dis;
             try {
-                vt = VisitedTable::create(
+                vt = &VisitedTable::get_reusable(
                         index->ntotal, hnsw.use_visited_hashset);
                 res = std::make_unique<
                         typename BlockResultHandler::SingleResultHandler>(bres);
@@ -479,11 +479,11 @@ void IndexHNSW::search_level_0(
         {
             std::unique_ptr<DistanceComputer> qdis;
             HNSWStats search_stats;
-            std::unique_ptr<VisitedTable> vt;
+            VisitedTable* vt = nullptr;
             std::unique_ptr<typename RH::SingleResultHandler> res;
             try {
                 qdis.reset(storage_distance_computer(storage));
-                vt = VisitedTable::create(
+                vt = &VisitedTable::get_reusable(
                         hnsw_ntotal, hnsw.use_visited_hashset);
                 res = std::make_unique<typename RH::SingleResultHandler>(bres);
             } catch (...) {
@@ -903,8 +903,7 @@ void IndexHNSW2Level::search(
         idx_t* labels,
         const SearchParameters* params) const {
     FAISS_THROW_IF_NOT(k > 0);
-    FAISS_THROW_IF_NOT_MSG(
-            !params, "search params not supported for this index");
+    FAISS_THROW_IF_MSG(params, "search params not supported for this index");
 
     if (dynamic_cast<const Index2Layer*>(storage)) {
         IndexHNSW::search(n, x, k, distances, labels);
@@ -1084,8 +1083,8 @@ IndexHNSWCagra::IndexHNSWCagra(
 }
 
 void IndexHNSWCagra::add(idx_t n, const float* x) {
-    FAISS_THROW_IF_NOT_MSG(
-            !base_level_only,
+    FAISS_THROW_IF_MSG(
+            base_level_only,
             "Cannot add vectors when base_level_only is set to True");
 
     IndexHNSW::add(n, x);
@@ -1208,11 +1207,11 @@ void IndexHNSWCagra::range_search(
 
             RangeQueryResult& qres = pres.new_result(i);
             RangeResultHandler<C> res(&qres, radius);
-            std::unique_ptr<VisitedTable> vt =
-                    VisitedTable::create(ntotal, hnsw.use_visited_hashset);
+            VisitedTable& vt = VisitedTable::get_reusable(
+                    ntotal, hnsw.use_visited_hashset);
             HNSWStats stats;
             hnsw.search_level_0(
-                    *dis, res, 1, &nearest, &nearest_d, 1, stats, *vt, params);
+                    *dis, res, 1, &nearest, &nearest_d, 1, stats, vt, params);
             n1 += stats.n1;
             n2 += stats.n2;
             ndis += stats.ndis;

data/vendor/faiss/faiss/IndexIDMap.cpp

@@ -138,7 +138,8 @@ void IndexIDMapTemplate<IndexT>::add_sa_codes(
         idx_t n,
         const uint8_t* codes,
         const idx_t* xids) {
-    index->add_sa_codes(n, codes, xids);
+    // don't pass the ids to the sub-index, they are kept in id_map
+    index->add_sa_codes(n, codes, nullptr);
     for (idx_t i = 0; i < n; i++) {
         id_map.push_back(xids[i]);
     }
@@ -228,14 +229,14 @@ void IndexIDMapTemplate<IndexT>::range_search(
         typename IndexT::distance_t radius,
         RangeSearchResult* result,
         const SearchParameters* params) const {
-    if (params) {
+    if (params && params->sel) {
         SearchParameters internal_search_parameters;
         IDSelectorTranslated id_selector_translated(id_map, params->sel);
         internal_search_parameters.sel = &id_selector_translated;
 
         index->range_search(n, x, radius, result, &internal_search_parameters);
     } else {
-        index->range_search(n, x, radius, result);
+        index->range_search(n, x, radius, result, params);
     }
 
     const idx_t id_map_size = static_cast<idx_t>(id_map.size());
@@ -329,6 +330,18 @@ void IndexIDMap2Template<IndexT>::add_with_ids(
             xids);
 }
 
+template <typename IndexT>
+void IndexIDMap2Template<IndexT>::add_sa_codes(
+        idx_t n,
+        const uint8_t* codes,
+        const idx_t* xids) {
+    idx_t prev_ntotal = this->ntotal;
+    IndexIDMapTemplate<IndexT>::add_sa_codes(n, codes, xids);
+    for (idx_t i = prev_ntotal; i < this->ntotal; i++) {
+        rev_map[this->id_map[i]] = i;
+    }
+}
+
 template <typename IndexT>
 void IndexIDMap2Template<IndexT>::check_consistency() const {
     FAISS_THROW_IF_NOT(rev_map.size() == this->id_map.size());

data/vendor/faiss/faiss/IndexIDMap.h

@@ -110,6 +110,8 @@ struct IndexIDMap2Template : IndexIDMapTemplate<IndexT> {
             NumericType numeric_type,
             const idx_t* xids) override;
 
+    void add_sa_codes(idx_t n, const uint8_t* x, const idx_t* xids) override;
+
     size_t remove_ids(const IDSelector& sel) override;
 
     void reconstruct(idx_t key, component_t* recons) const override;

data/vendor/faiss/faiss/IndexIVF.cpp

@@ -17,7 +17,6 @@
 #include <cstdio>
 #include <limits>
 
-#include <faiss/utils/hamming.h>
 #include <faiss/utils/utils.h>
 
 #include <faiss/IndexFlat.h>
@@ -501,7 +500,7 @@ void IndexIVF::search_preassigned(
                 if (!do_heap_init) {
                     return;
                 }
-                if (metric_type == METRIC_INNER_PRODUCT) {
+                if (is_similarity_metric(metric_type)) {
                     heap_heapify<HeapForIP>(k, simi, idxi);
                 } else {
                     heap_heapify<HeapForL2>(k, simi, idxi);
@@ -512,7 +511,7 @@ void IndexIVF::search_preassigned(
                                          const idx_t* local_idx,
                                          float* simi,
                                          idx_t* idxi) {
-                if (metric_type == METRIC_INNER_PRODUCT) {
+                if (is_similarity_metric(metric_type)) {
                     heap_addn<HeapForIP>(
                             k, simi, idxi, local_dis, local_idx, k);
                 } else {
@@ -525,7 +524,7 @@ void IndexIVF::search_preassigned(
                 if (!do_heap_init) {
                     return;
                 }
-                if (metric_type == METRIC_INNER_PRODUCT) {
+                if (is_similarity_metric(metric_type)) {
                     heap_reorder<HeapForIP>(k, simi, idxi);
                 } else {
                     heap_reorder<HeapForL2>(k, simi, idxi);
@@ -599,7 +598,7 @@ void IndexIVF::search_preassigned(
 
                     size_t old_scan_cnt = 0;
                     size_t old_heap_updates = 0;
-                    if (metric_type == METRIC_INNER_PRODUCT) {
+                    if (is_similarity_metric(metric_type)) {
                         HeapResultHandler<HeapForIP, false> handler(
                                 k, simi, idxi);
                         old_scan_cnt = handler.stats.scan_cnt;
@@ -1056,7 +1055,19 @@ InvertedListScanner* IndexIVF::get_InvertedListScanner(
 
 void IndexIVF::reconstruct(idx_t key, float* recons) const {
     idx_t lo = direct_map.get(key);
-    reconstruct_from_offset(lo_listno(lo), lo_offset(lo), recons);
+    const size_t list_no = lo_listno(lo);
+    const size_t offset = lo_offset(lo);
+    FAISS_THROW_IF_NOT_FMT(
+            list_no < nlist,
+            "IndexIVF::reconstruct: list_no %zd out of range (nlist=%zd)",
+            list_no,
+            nlist);
+    FAISS_THROW_IF_NOT_FMT(
+            offset < invlists->list_size(list_no),
+            "IndexIVF::reconstruct: offset %zd out of range (list_size=%zd)",
+            offset,
+            invlists->list_size(list_no));
+    reconstruct_from_offset(list_no, offset, recons);
 }
 
 void IndexIVF::reconstruct_n(idx_t i0, idx_t ni, float* recons) const {

data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp

@@ -283,7 +283,7 @@ InvertedListScanner* IndexIVFAdditiveQuantizer::get_InvertedListScanner(
         bool store_pairs,
         const IDSelector* sel,
         const IVFSearchParameters*) const {
-    FAISS_THROW_IF_NOT(!sel);
+    FAISS_THROW_IF_MSG(sel, "id selector not supported for this index");
     if (metric_type == METRIC_INNER_PRODUCT) {
         if (aq->search_type == AdditiveQuantizer::ST_decompress) {
             return new AQInvertedListScannerDecompress<true>(

data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.cpp

@@ -48,8 +48,8 @@ void IndexIVFAdditiveQuantizerFastScan::init(
         MetricType metric,
         int bbs_,
         bool own_invlists_) {
-    FAISS_THROW_IF_NOT(aq_ != nullptr);
-    FAISS_THROW_IF_NOT(!aq_->nbits.empty());
+    FAISS_THROW_IF_NOT(aq_);
+    FAISS_THROW_IF_MSG(aq_->nbits.empty(), "quantizer nbits must not be empty");
     FAISS_THROW_IF_NOT(aq_->nbits[0] == 4);
     if (metric == METRIC_INNER_PRODUCT) {
         FAISS_THROW_IF_NOT_MSG(
@@ -310,8 +310,7 @@ void IndexIVFAdditiveQuantizerFastScan::search(
         float* distances,
         idx_t* labels,
         const SearchParameters* params) const {
-    FAISS_THROW_IF_NOT_MSG(
-            !params, "search params not supported for this index");
+    FAISS_THROW_IF_MSG(params, "search params not supported for this index");
 
     FAISS_THROW_IF_NOT(k > 0);
     bool rescale = (rescale_norm && norm_scale > 1 && metric_type == METRIC_L2);