faiss 0.3.1 → 0.3.3

data/vendor/faiss/faiss/impl/simd_result_handlers.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -15,6 +15,7 @@
 #include <faiss/utils/simdlib.h>
 
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/IDSelector.h>
 #include <faiss/impl/ResultHandler.h>
 #include <faiss/impl/platform_macros.h>
 #include <faiss/utils/AlignedTable.h>
@@ -137,6 +138,7 @@ struct FixedStorageHandler : SIMDResultHandler {
             }
         }
     }
+
     virtual ~FixedStorageHandler() {}
 };
 
@@ -150,8 +152,10 @@ struct ResultHandlerCompare : SIMDResultHandlerToFloat {
     int64_t i0 = 0; // query origin
     int64_t j0 = 0; // db origin
 
-    ResultHandlerCompare(size_t nq, size_t ntotal)
-            : SIMDResultHandlerToFloat(nq, ntotal) {
+    const IDSelector* sel;
+
+    ResultHandlerCompare(size_t nq, size_t ntotal, const IDSelector* sel_in)
+            : SIMDResultHandlerToFloat(nq, ntotal), sel{sel_in} {
         this->is_CMax = C::is_max;
         this->sizeof_ids = sizeof(typename C::TI);
         this->with_fields = with_id_map;
@@ -232,9 +236,14 @@ struct SingleResultHandler : ResultHandlerCompare<C, with_id_map> {
     float* dis;
     int64_t* ids;
 
-    SingleResultHandler(size_t nq, size_t ntotal, float* dis, int64_t* ids)
-            : RHC(nq, ntotal), idis(nq), dis(dis), ids(ids) {
-        for (int i = 0; i < nq; i++) {
+    SingleResultHandler(
+            size_t nq,
+            size_t ntotal,
+            float* dis,
+            int64_t* ids,
+            const IDSelector* sel_in)
+            : RHC(nq, ntotal, sel_in), idis(nq), dis(dis), ids(ids) {
+        for (size_t i = 0; i < nq; i++) {
             ids[i] = -1;
             idis[i] = C::neutral();
         }
@@ -256,20 +265,36 @@ struct SingleResultHandler : ResultHandlerCompare<C, with_id_map> {
         d0.store(d32tab);
         d1.store(d32tab + 16);
 
-        while (lt_mask) {
-            // find first non-zero
-            int j = __builtin_ctz(lt_mask);
-            lt_mask -= 1 << j;
-            T d = d32tab[j];
-            if (C::cmp(idis[q], d)) {
-                idis[q] = d;
-                ids[q] = this->adjust_id(b, j);
+        if (this->sel != nullptr) {
+            while (lt_mask) {
+                // find first non-zero
+                int j = __builtin_ctz(lt_mask);
+                auto real_idx = this->adjust_id(b, j);
+                lt_mask -= 1 << j;
+                if (this->sel->is_member(real_idx)) {
+                    T d = d32tab[j];
+                    if (C::cmp(idis[q], d)) {
+                        idis[q] = d;
+                        ids[q] = real_idx;
+                    }
+                }
+            }
+        } else {
+            while (lt_mask) {
+                // find first non-zero
+                int j = __builtin_ctz(lt_mask);
+                lt_mask -= 1 << j;
+                T d = d32tab[j];
+                if (C::cmp(idis[q], d)) {
+                    idis[q] = d;
+                    ids[q] = this->adjust_id(b, j);
+                }
             }
         }
     }
 
     void end() {
-        for (int q = 0; q < this->nq; q++) {
+        for (size_t q = 0; q < this->nq; q++) {
             if (!normalizers) {
                 dis[q] = idis[q];
             } else {
@@ -296,8 +321,14 @@ struct HeapHandler : ResultHandlerCompare<C, with_id_map> {
 
     int64_t k; // number of results to keep
 
-    HeapHandler(size_t nq, size_t ntotal, int64_t k, float* dis, int64_t* ids)
-            : RHC(nq, ntotal),
+    HeapHandler(
+            size_t nq,
+            size_t ntotal,
+            int64_t k,
+            float* dis,
+            int64_t* ids,
+            const IDSelector* sel_in)
+            : RHC(nq, ntotal, sel_in),
               idis(nq * k),
               iids(nq * k),
               dis(dis),
@@ -330,21 +361,36 @@ struct HeapHandler : ResultHandlerCompare<C, with_id_map> {
         d0.store(d32tab);
         d1.store(d32tab + 16);
 
-        while (lt_mask) {
-            // find first non-zero
-            int j = __builtin_ctz(lt_mask);
-            lt_mask -= 1 << j;
-            T dis = d32tab[j];
-            if (C::cmp(heap_dis[0], dis)) {
-                int64_t idx = this->adjust_id(b, j);
-                heap_pop<C>(k, heap_dis, heap_ids);
-                heap_push<C>(k, heap_dis, heap_ids, dis, idx);
+        if (this->sel != nullptr) {
+            while (lt_mask) {
+                // find first non-zero
+                int j = __builtin_ctz(lt_mask);
+                auto real_idx = this->adjust_id(b, j);
+                lt_mask -= 1 << j;
+                if (this->sel->is_member(real_idx)) {
+                    T dis_2 = d32tab[j];
+                    if (C::cmp(heap_dis[0], dis_2)) {
+                        heap_replace_top<C>(
+                                k, heap_dis, heap_ids, dis_2, real_idx);
+                    }
+                }
+            }
+        } else {
+            while (lt_mask) {
+                // find first non-zero
+                int j = __builtin_ctz(lt_mask);
+                lt_mask -= 1 << j;
+                T dis_2 = d32tab[j];
+                if (C::cmp(heap_dis[0], dis_2)) {
+                    int64_t idx = this->adjust_id(b, j);
+                    heap_replace_top<C>(k, heap_dis, heap_ids, dis_2, idx);
+                }
             }
         }
     }
 
     void end() override {
-        for (int q = 0; q < this->nq; q++) {
+        for (size_t q = 0; q < this->nq; q++) {
             T* heap_dis_in = idis.data() + q * k;
             TI* heap_ids_in = iids.data() + q * k;
             heap_reorder<C>(k, heap_dis_in, heap_ids_in);
@@ -393,8 +439,12 @@ struct ReservoirHandler : ResultHandlerCompare<C, with_id_map> {
             size_t k,
             size_t cap,
             float* dis,
-            int64_t* ids)
-            : RHC(nq, ntotal), capacity((cap + 15) & ~15), dis(dis), ids(ids) {
+            int64_t* ids,
+            const IDSelector* sel_in)
+            : RHC(nq, ntotal, sel_in),
+              capacity((cap + 15) & ~15),
+              dis(dis),
+              ids(ids) {
         assert(capacity % 16 == 0);
         all_ids.resize(nq * capacity);
         all_vals.resize(nq * capacity);
@@ -423,12 +473,25 @@ struct ReservoirHandler : ResultHandlerCompare<C, with_id_map> {
         d0.store(d32tab);
         d1.store(d32tab + 16);
 
-        while (lt_mask) {
-            // find first non-zero
-            int j = __builtin_ctz(lt_mask);
-            lt_mask -= 1 << j;
-            T dis = d32tab[j];
-            res.add(dis, this->adjust_id(b, j));
+        if (this->sel != nullptr) {
+            while (lt_mask) {
+                // find first non-zero
+                int j = __builtin_ctz(lt_mask);
+                auto real_idx = this->adjust_id(b, j);
+                lt_mask -= 1 << j;
+                if (this->sel->is_member(real_idx)) {
+                    T dis_2 = d32tab[j];
+                    res.add(dis_2, real_idx);
+                }
+            }
+        } else {
+            while (lt_mask) {
+                // find first non-zero
+                int j = __builtin_ctz(lt_mask);
+                lt_mask -= 1 << j;
+                T dis_2 = d32tab[j];
+                res.add(dis_2, this->adjust_id(b, j));
+            }
         }
     }
 
@@ -439,7 +502,7 @@ struct ReservoirHandler : ResultHandlerCompare<C, with_id_map> {
                 CMin<float, int64_t>>::type;
 
         std::vector<int> perm(reservoirs[0].n);
-        for (int q = 0; q < reservoirs.size(); q++) {
+        for (size_t q = 0; q < reservoirs.size(); q++) {
             ReservoirTopN<C>& res = reservoirs[q];
             size_t n = res.n;
 
@@ -454,14 +517,14 @@ struct ReservoirHandler : ResultHandlerCompare<C, with_id_map> {
                 one_a = 1 / normalizers[2 * q];
                 b = normalizers[2 * q + 1];
             }
-            for (int i = 0; i < res.i; i++) {
+            for (size_t i = 0; i < res.i; i++) {
                 perm[i] = i;
             }
             // indirect sort of result arrays
             std::sort(perm.begin(), perm.begin() + res.i, [&res](int i, int j) {
                 return C::cmp(res.vals[j], res.vals[i]);
             });
-            for (int i = 0; i < res.i; i++) {
+            for (size_t i = 0; i < res.i; i++) {
                 heap_dis[i] = res.vals[perm[i]] * one_a + b;
                 heap_ids[i] = res.ids[perm[i]];
             }
@@ -472,7 +535,7 @@ struct ReservoirHandler : ResultHandlerCompare<C, with_id_map> {
     }
 };
 
-/** Result hanlder for range search. The difficulty is that the range distances
+/** Result handler for range search. The difficulty is that the range distances
  * have to be scaled using the scaler.
  */
 
@@ -499,13 +562,17 @@ struct RangeHandler : ResultHandlerCompare<C, with_id_map> {
     };
     std::vector<Triplet> triplets;
 
-    RangeHandler(RangeSearchResult& rres, float radius, size_t ntotal)
-            : RHC(rres.nq, ntotal), rres(rres), radius(radius) {
+    RangeHandler(
+            RangeSearchResult& rres,
+            float radius,
+            size_t ntotal,
+            const IDSelector* sel_in)
+            : RHC(rres.nq, ntotal, sel_in), rres(rres), radius(radius) {
         thresholds.resize(nq);
         n_per_query.resize(nq + 1);
     }
 
-    virtual void begin(const float* norms) {
+    virtual void begin(const float* norms) override {
         normalizers = norms;
         for (int q = 0; q < nq; ++q) {
             thresholds[q] =
@@ -528,13 +595,28 @@ struct RangeHandler : ResultHandlerCompare<C, with_id_map> {
         d0.store(d32tab);
         d1.store(d32tab + 16);
 
-        while (lt_mask) {
-            // find first non-zero
-            int j = __builtin_ctz(lt_mask);
-            lt_mask -= 1 << j;
-            T dis = d32tab[j];
-            n_per_query[q]++;
-            triplets.push_back({idx_t(q + q0), this->adjust_id(b, j), dis});
+        if (this->sel != nullptr) {
+            while (lt_mask) {
+                // find first non-zero
+                int j = __builtin_ctz(lt_mask);
+                lt_mask -= 1 << j;
+
+                auto real_idx = this->adjust_id(b, j);
+                if (this->sel->is_member(real_idx)) {
+                    T dis = d32tab[j];
+                    n_per_query[q]++;
+                    triplets.push_back({idx_t(q + q0), real_idx, dis});
+                }
+            }
+        } else {
+            while (lt_mask) {
+                // find first non-zero
+                int j = __builtin_ctz(lt_mask);
+                lt_mask -= 1 << j;
+                T dis = d32tab[j];
+                n_per_query[q]++;
+                triplets.push_back({idx_t(q + q0), this->adjust_id(b, j), dis});
+            }
         }
     }
 
@@ -578,8 +660,9 @@ struct PartialRangeHandler : RangeHandler<C, with_id_map> {
             float radius,
             size_t ntotal,
             size_t q0,
-            size_t q1)
-            : RangeHandler<C, with_id_map>(*pres.res, radius, ntotal),
+            size_t q1,
+            const IDSelector* sel_in)
+            : RangeHandler<C, with_id_map>(*pres.res, radius, ntotal, sel_in),
               pres(pres) {
         nq = q1 - q0;
         this->q0 = q0;
@@ -630,16 +713,16 @@ struct PartialRangeHandler : RangeHandler<C, with_id_map> {
  */
 
 template <class C, bool W, class Consumer, class... Types>
-void dispatch_SIMDResultHanlder_fixedCW(
+void dispatch_SIMDResultHandler_fixedCW(
         SIMDResultHandler& res,
         Consumer& consumer,
         Types... args) {
     if (auto resh = dynamic_cast<SingleResultHandler<C, W>*>(&res)) {
         consumer.template f<SingleResultHandler<C, W>>(*resh, args...);
-    } else if (auto resh = dynamic_cast<HeapHandler<C, W>*>(&res)) {
-        consumer.template f<HeapHandler<C, W>>(*resh, args...);
-    } else if (auto resh = dynamic_cast<ReservoirHandler<C, W>*>(&res)) {
-        consumer.template f<ReservoirHandler<C, W>>(*resh, args...);
+    } else if (auto resh_2 = dynamic_cast<HeapHandler<C, W>*>(&res)) {
+        consumer.template f<HeapHandler<C, W>>(*resh_2, args...);
+    } else if (auto resh_2 = dynamic_cast<ReservoirHandler<C, W>*>(&res)) {
+        consumer.template f<ReservoirHandler<C, W>>(*resh_2, args...);
     } else { // generic handler -- will not be inlined
         FAISS_THROW_IF_NOT_FMT(
                 simd_result_handlers_accept_virtual,
@@ -650,27 +733,27 @@ void dispatch_SIMDResultHanlder_fixedCW(
 }
 
 template <class C, class Consumer, class... Types>
-void dispatch_SIMDResultHanlder_fixedC(
+void dispatch_SIMDResultHandler_fixedC(
         SIMDResultHandler& res,
         Consumer& consumer,
         Types... args) {
     if (res.with_fields) {
-        dispatch_SIMDResultHanlder_fixedCW<C, true>(res, consumer, args...);
+        dispatch_SIMDResultHandler_fixedCW<C, true>(res, consumer, args...);
     } else {
-        dispatch_SIMDResultHanlder_fixedCW<C, false>(res, consumer, args...);
+        dispatch_SIMDResultHandler_fixedCW<C, false>(res, consumer, args...);
     }
 }
 
 template <class Consumer, class... Types>
-void dispatch_SIMDResultHanlder(
+void dispatch_SIMDResultHandler(
         SIMDResultHandler& res,
         Consumer& consumer,
         Types... args) {
     if (res.sizeof_ids == 0) {
         if (auto resh = dynamic_cast<StoreResultHandler*>(&res)) {
             consumer.template f<StoreResultHandler>(*resh, args...);
-        } else if (auto resh = dynamic_cast<DummyResultHandler*>(&res)) {
-            consumer.template f<DummyResultHandler>(*resh, args...);
+        } else if (auto resh_2 = dynamic_cast<DummyResultHandler*>(&res)) {
+            consumer.template f<DummyResultHandler>(*resh_2, args...);
         } else { // generic path
             FAISS_THROW_IF_NOT_FMT(
                     simd_result_handlers_accept_virtual,
@@ -680,24 +763,25 @@ void dispatch_SIMDResultHanlder(
         }
     } else if (res.sizeof_ids == sizeof(int)) {
         if (res.is_CMax) {
-            dispatch_SIMDResultHanlder_fixedC<CMax<uint16_t, int>>(
+            dispatch_SIMDResultHandler_fixedC<CMax<uint16_t, int>>(
                     res, consumer, args...);
         } else {
-            dispatch_SIMDResultHanlder_fixedC<CMin<uint16_t, int>>(
+            dispatch_SIMDResultHandler_fixedC<CMin<uint16_t, int>>(
                     res, consumer, args...);
         }
     } else if (res.sizeof_ids == sizeof(int64_t)) {
         if (res.is_CMax) {
-            dispatch_SIMDResultHanlder_fixedC<CMax<uint16_t, int64_t>>(
+            dispatch_SIMDResultHandler_fixedC<CMax<uint16_t, int64_t>>(
                     res, consumer, args...);
         } else {
-            dispatch_SIMDResultHanlder_fixedC<CMin<uint16_t, int64_t>>(
+            dispatch_SIMDResultHandler_fixedC<CMin<uint16_t, int64_t>>(
                     res, consumer, args...);
         }
     } else {
         FAISS_THROW_FMT("Unknown id size %d", res.sizeof_ids);
     }
 }
+
 } // namespace simd_result_handlers
 
 } // namespace faiss

data/vendor/faiss/faiss/index_factory.cpp

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -20,7 +20,6 @@
 
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/random.h>
-#include <faiss/utils/utils.h>
 
 #include <faiss/Index2Layer.h>
 #include <faiss/IndexAdditiveQuantizer.h>
@@ -140,8 +139,12 @@ std::map<std::string, ScalarQuantizer::QuantizerType> sq_types = {
         {"SQ4", ScalarQuantizer::QT_4bit},
         {"SQ6", ScalarQuantizer::QT_6bit},
         {"SQfp16", ScalarQuantizer::QT_fp16},
+        {"SQbf16", ScalarQuantizer::QT_bf16},
+        {"SQ8_direct_signed", ScalarQuantizer::QT_8bit_direct_signed},
+        {"SQ8_direct", ScalarQuantizer::QT_8bit_direct},
 };
-const std::string sq_pattern = "(SQ4|SQ8|SQ6|SQfp16)";
+const std::string sq_pattern =
+        "(SQ4|SQ8|SQ6|SQfp16|SQbf16|SQ8_direct_signed|SQ8_direct)";
 
 std::map<std::string, AdditiveQuantizer::Search_type_t> aq_search_type = {
         {"_Nfloat", AdditiveQuantizer::ST_norm_float},
@@ -222,6 +225,19 @@ VectorTransform* parse_VectorTransform(const std::string& description, int d) {
  * Parse IndexIVF
  */
 
+size_t parse_nlist(std::string s) {
+    size_t multiplier = 1;
+    if (s.back() == 'k') {
+        s.pop_back();
+        multiplier = 1024;
+    }
+    if (s.back() == 'M') {
+        s.pop_back();
+        multiplier = 1024 * 1024;
+    }
+    return std::stoi(s) * multiplier;
+}
+
 // parsing guard + function
 Index* parse_coarse_quantizer(
         const std::string& description,
@@ -236,8 +252,8 @@ Index* parse_coarse_quantizer(
     };
     use_2layer = false;
 
-    if (match("IVF([0-9]+)")) {
-        nlist = std::stoi(sm[1].str());
+    if (match("IVF([0-9]+[kM]?)")) {
+        nlist = parse_nlist(sm[1].str());
         return new IndexFlat(d, mt);
     }
     if (match("IMI2x([0-9]+)")) {
@@ -248,18 +264,18 @@ Index* parse_coarse_quantizer(
         nlist = (size_t)1 << (2 * nbit);
         return new MultiIndexQuantizer(d, 2, nbit);
     }
-    if (match("IVF([0-9]+)_HNSW([0-9]*)")) {
-        nlist = std::stoi(sm[1].str());
+    if (match("IVF([0-9]+[kM]?)_HNSW([0-9]*)")) {
+        nlist = parse_nlist(sm[1].str());
         int hnsw_M = sm[2].length() > 0 ? std::stoi(sm[2]) : 32;
         return new IndexHNSWFlat(d, hnsw_M, mt);
     }
-    if (match("IVF([0-9]+)_NSG([0-9]+)")) {
-        nlist = std::stoi(sm[1].str());
+    if (match("IVF([0-9]+[kM]?)_NSG([0-9]+)")) {
+        nlist = parse_nlist(sm[1].str());
         int R = std::stoi(sm[2]);
         return new IndexNSGFlat(d, R, mt);
     }
-    if (match("IVF([0-9]+)\\(Index([0-9])\\)")) {
-        nlist = std::stoi(sm[1].str());
+    if (match("IVF([0-9]+[kM]?)\\(Index([0-9])\\)")) {
+        nlist = parse_nlist(sm[1].str());
         int no = std::stoi(sm[2].str());
         FAISS_ASSERT(no >= 0 && no < parenthesis_indexes.size());
         return parenthesis_indexes[no].release();
@@ -526,11 +542,12 @@ Index* parse_other_indexes(
     }
 
     // IndexLSH
-    if (match("LSH(r?)(t?)")) {
-        bool rotate_data = sm[1].length() > 0;
-        bool train_thresholds = sm[2].length() > 0;
+    if (match("LSH([0-9]*)(r?)(t?)")) {
+        int nbits = sm[1].length() > 0 ? std::stoi(sm[1].str()) : d;
+        bool rotate_data = sm[2].length() > 0;
+        bool train_thresholds = sm[3].length() > 0;
         FAISS_THROW_IF_NOT(metric == METRIC_L2);
-        return new IndexLSH(d, d, rotate_data, train_thresholds);
+        return new IndexLSH(d, nbits, rotate_data, train_thresholds);
     }
 
     // IndexLattice
@@ -661,6 +678,24 @@ std::unique_ptr<Index> index_factory_sub(
     // for the current match
     std::smatch sm;
 
+    // IndexIDMap -- it turns out is was used both as a prefix and a suffix, so
+    // support both
+    if (re_match(description, "(.+),IDMap2", sm) ||
+        re_match(description, "IDMap2,(.+)", sm)) {
+        IndexIDMap2* idmap2 = new IndexIDMap2(
+                index_factory_sub(d, sm[1].str(), metric).release());
+        idmap2->own_fields = true;
+        return std::unique_ptr<Index>(idmap2);
+    }
+
+    if (re_match(description, "(.+),IDMap", sm) ||
+        re_match(description, "IDMap,(.+)", sm)) {
+        IndexIDMap* idmap = new IndexIDMap(
+                index_factory_sub(d, sm[1].str(), metric).release());
+        idmap->own_fields = true;
+        return std::unique_ptr<Index>(idmap);
+    }
+
     // handle refines
     if (re_match(description, "(.+),RFlat", sm) ||
         re_match(description, "(.+),Refine\\((.+)\\)", sm)) {
@@ -737,24 +772,6 @@ std::unique_ptr<Index> index_factory_sub(
                d);
     }
 
-    // IndexIDMap -- it turns out is was used both as a prefix and a suffix, so
-    // support both
-    if (re_match(description, "(.+),IDMap2", sm) ||
-        re_match(description, "IDMap2,(.+)", sm)) {
-        IndexIDMap2* idmap2 = new IndexIDMap2(
-                index_factory_sub(d, sm[1].str(), metric).release());
-        idmap2->own_fields = true;
-        return std::unique_ptr<Index>(idmap2);
-    }
-
-    if (re_match(description, "(.+),IDMap", sm) ||
-        re_match(description, "IDMap,(.+)", sm)) {
-        IndexIDMap* idmap = new IndexIDMap(
-                index_factory_sub(d, sm[1].str(), metric).release());
-        idmap->own_fields = true;
-        return std::unique_ptr<Index>(idmap);
-    }
-
     { // handle basic index types
         Index* index = parse_other_indexes(description, d, metric);
         if (index) {

data/vendor/faiss/faiss/index_factory.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.

data/vendor/faiss/faiss/index_io.h

@@ -1,12 +1,10 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
  */
 
-// -*- c++ -*-
-
 // I/O code for indexes
 
 #ifndef FAISS_INDEX_IO_H
@@ -35,9 +33,12 @@ struct IOReader;
 struct IOWriter;
 struct InvertedLists;
 
-void write_index(const Index* idx, const char* fname);
-void write_index(const Index* idx, FILE* f);
-void write_index(const Index* idx, IOWriter* writer);
+/// skip the storage for graph-based indexes
+const int IO_FLAG_SKIP_STORAGE = 1;
+
+void write_index(const Index* idx, const char* fname, int io_flags = 0);
+void write_index(const Index* idx, FILE* f, int io_flags = 0);
+void write_index(const Index* idx, IOWriter* writer, int io_flags = 0);
 
 void write_index_binary(const IndexBinary* idx, const char* fname);
 void write_index_binary(const IndexBinary* idx, FILE* f);
@@ -52,6 +53,12 @@ const int IO_FLAG_ONDISK_SAME_DIR = 4;
 const int IO_FLAG_SKIP_IVF_DATA = 8;
 // don't initialize precomputed table after loading
 const int IO_FLAG_SKIP_PRECOMPUTE_TABLE = 16;
+// don't compute the sdc table for PQ-based indices
+// this will prevent distances from being computed
+// between elements in the index. For indices like HNSWPQ,
+// this will prevent graph building because sdc
+// computations are required to construct the graph
+const int IO_FLAG_PQ_SKIP_SDC_TABLE = 32;
 // try to memmap data (useful to load an ArrayInvertedLists as an
 // OnDiskInvertedLists)
 const int IO_FLAG_MMAP = IO_FLAG_SKIP_IVF_DATA | 0x646f0000;

data/vendor/faiss/faiss/invlists/BlockInvertedLists.cpp

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
@@ -9,6 +9,7 @@
 
 #include <faiss/impl/CodePacker.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/IDSelector.h>
 
 #include <faiss/impl/io.h>
 #include <faiss/impl/io_macros.h>
@@ -54,7 +55,9 @@ size_t BlockInvertedLists::add_entries(
     codes[list_no].resize(n_block * block_size);
     if (o % block_size == 0) {
         // copy whole blocks
-        memcpy(&codes[list_no][o * code_size], code, n_block * block_size);
+        memcpy(&codes[list_no][o * packer->code_size],
+               code,
+               n_block * block_size);
     } else {
         FAISS_THROW_IF_NOT_MSG(packer, "missing code packer");
         std::vector<uint8_t> buffer(packer->code_size);
@@ -76,6 +79,29 @@ const uint8_t* BlockInvertedLists::get_codes(size_t list_no) const {
     return codes[list_no].get();
 }
 
+size_t BlockInvertedLists::remove_ids(const IDSelector& sel) {
+    idx_t nremove = 0;
+#pragma omp parallel for
+    for (idx_t i = 0; i < nlist; i++) {
+        std::vector<uint8_t> buffer(packer->code_size);
+        idx_t l = ids[i].size(), j = 0;
+        while (j < l) {
+            if (sel.is_member(ids[i][j])) {
+                l--;
+                ids[i][j] = ids[i][l];
+                packer->unpack_1(codes[i].data(), l, buffer.data());
+                packer->pack_1(buffer.data(), j, codes[i].data());
+            } else {
+                j++;
+            }
+        }
+        resize(i, l);
+        nremove += ids[i].size() - l;
+    }
+
+    return nremove;
+}
+
 const idx_t* BlockInvertedLists::get_ids(size_t list_no) const {
     assert(list_no < nlist);
     return ids[list_no].data();
@@ -101,13 +127,7 @@ void BlockInvertedLists::update_entries(
         size_t,
         const idx_t*,
         const uint8_t*) {
-    FAISS_THROW_MSG("not impemented");
-    /*
-    assert (list_no < nlist);
-    assert (n_entry + offset <= ids[list_no].size());
-    memcpy (&ids[list_no][offset], ids_in, sizeof(ids_in[0]) * n_entry);
-    memcpy (&codes[list_no][offset * code_size], codes_in, code_size * n_entry);
-    */
+    FAISS_THROW_MSG("not implemented");
 }
 
 BlockInvertedLists::~BlockInvertedLists() {