faiss 0.1.5 → 0.2.2

data/vendor/faiss/faiss/IndexResidual.cpp

@@ -0,0 +1,291 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <faiss/IndexResidual.h>
+
+#include <algorithm>
+#include <cmath>
+#include <cstring>
+
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/ResultHandler.h>
+#include <faiss/utils/distances.h>
+#include <faiss/utils/extra_distances.h>
+#include <faiss/utils/utils.h>
+
+namespace faiss {
+
+/**************************************************************************************
+ * IndexResidual
+ **************************************************************************************/
+
+IndexResidual::IndexResidual(
+        int d,        ///< dimensionality of the input vectors
+        size_t M,     ///< number of subquantizers
+        size_t nbits, ///< number of bit per subvector index
+        MetricType metric,
+        Search_type_t search_type_in)
+        : Index(d, metric), rq(d, M, nbits), search_type(ST_decompress) {
+    is_trained = false;
+    norm_max = norm_min = NAN;
+    set_search_type(search_type_in);
+}
+
+IndexResidual::IndexResidual(
+        int d,
+        const std::vector<size_t>& nbits,
+        MetricType metric,
+        Search_type_t search_type_in)
+        : Index(d, metric), rq(d, nbits), search_type(ST_decompress) {
+    is_trained = false;
+    norm_max = norm_min = NAN;
+    set_search_type(search_type_in);
+}
+
+IndexResidual::IndexResidual() : IndexResidual(0, 0, 0) {}
+
+void IndexResidual::set_search_type(Search_type_t new_search_type) {
+    int norm_bits = new_search_type == ST_norm_float ? 32
+            : new_search_type == ST_norm_qint8       ? 8
+                                                     : 0;
+
+    FAISS_THROW_IF_NOT(ntotal == 0);
+
+    search_type = new_search_type;
+    code_size = (rq.tot_bits + norm_bits + 7) / 8;
+}
+
+void IndexResidual::train(idx_t n, const float* x) {
+    rq.train(n, x);
+
+    std::vector<float> norms(n);
+    fvec_norms_L2sqr(norms.data(), x, d, n);
+
+    norm_min = HUGE_VALF;
+    norm_max = -HUGE_VALF;
+    for (idx_t i = 0; i < n; i++) {
+        if (norms[i] < norm_min) {
+            norm_min = norms[i];
+        }
+        if (norms[i] > norm_min) {
+            norm_max = norms[i];
+        }
+    }
+
+    is_trained = true;
+}
+
+void IndexResidual::add(idx_t n, const float* x) {
+    FAISS_THROW_IF_NOT(is_trained);
+    codes.resize((n + ntotal) * rq.code_size);
+    if (search_type == ST_decompress || search_type == ST_LUT_nonorm) {
+        rq.compute_codes(x, &codes[ntotal * rq.code_size], n);
+    } else {
+        // should compute codes + compute and quantize norms
+        FAISS_THROW_MSG("not implemented");
+    }
+    ntotal += n;
+}
+
+namespace {
+
+template <class VectorDistance, class ResultHandler>
+void search_with_decompress(
+        const IndexResidual& ir,
+        const float* xq,
+        VectorDistance& vd,
+        ResultHandler& res) {
+    const uint8_t* codes = ir.codes.data();
+    size_t ntotal = ir.ntotal;
+    size_t code_size = ir.code_size;
+
+    using SingleResultHandler = typename ResultHandler::SingleResultHandler;
+
+#pragma omp parallel for
+    for (int64_t q = 0; q < res.nq; q++) {
+        SingleResultHandler resi(res);
+        resi.begin(q);
+        std::vector<float> tmp(ir.d);
+        const float* x = xq + ir.d * q;
+        for (size_t i = 0; i < ntotal; i++) {
+            ir.rq.decode(codes + i * code_size, tmp.data(), 1);
+            float dis = vd(x, tmp.data());
+            resi.add_result(dis, i);
+        }
+        resi.end();
+    }
+}
+
+} // anonymous namespace
+
+void IndexResidual::search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels) const {
+    if (search_type == ST_decompress) {
+        if (metric_type == METRIC_L2) {
+            using VD = VectorDistance<METRIC_L2>;
+            VD vd = {size_t(d), metric_arg};
+            HeapResultHandler<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};
+            HeapResultHandler<VD::C> rh(n, distances, labels, k);
+            search_with_decompress(*this, x, vd, rh);
+        }
+    } else {
+        FAISS_THROW_MSG("not implemented");
+    }
+}
+
+void IndexResidual::reset() {
+    codes.clear();
+    ntotal = 0;
+}
+
+size_t IndexResidual::sa_code_size() const {
+    return code_size;
+}
+
+void IndexResidual::sa_encode(idx_t n, const float* x, uint8_t* bytes) const {
+    return rq.compute_codes(x, bytes, n);
+}
+
+void IndexResidual::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
+    return rq.decode(bytes, x, n);
+}
+
+/**************************************************************************************
+ * ResidualCoarseQuantizer
+ **************************************************************************************/
+
+ResidualCoarseQuantizer::ResidualCoarseQuantizer(
+        int d,        ///< dimensionality of the input vectors
+        size_t M,     ///< number of subquantizers
+        size_t nbits, ///< number of bit per subvector index
+        MetricType metric)
+        : Index(d, metric), rq(d, M, nbits), beam_factor(4.0) {
+    FAISS_THROW_IF_NOT(rq.tot_bits <= 63);
+    is_trained = false;
+}
+
+ResidualCoarseQuantizer::ResidualCoarseQuantizer(
+        int d,
+        const std::vector<size_t>& nbits,
+        MetricType metric)
+        : Index(d, metric), rq(d, nbits), beam_factor(4.0) {
+    FAISS_THROW_IF_NOT(rq.tot_bits <= 63);
+    is_trained = false;
+}
+
+ResidualCoarseQuantizer::ResidualCoarseQuantizer() {}
+
+void ResidualCoarseQuantizer::train(idx_t n, const float* x) {
+    rq.train(n, x);
+    is_trained = true;
+    ntotal = (idx_t)1 << rq.tot_bits;
+}
+
+void ResidualCoarseQuantizer::add(idx_t, const float*) {
+    FAISS_THROW_MSG("not applicable");
+}
+
+void ResidualCoarseQuantizer::set_beam_factor(float new_beam_factor) {
+    centroid_norms.resize(0);
+    beam_factor = new_beam_factor;
+    if (new_beam_factor > 0) {
+        FAISS_THROW_IF_NOT(new_beam_factor >= 1.0);
+        return;
+    }
+
+    if (metric_type == METRIC_L2) {
+        centroid_norms.resize((size_t)1 << rq.tot_bits);
+        rq.compute_centroid_norms(centroid_norms.data());
+    }
+}
+
+void ResidualCoarseQuantizer::search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels) const {
+    if (beam_factor < 0) {
+        if (metric_type == METRIC_INNER_PRODUCT) {
+            rq.knn_exact_inner_product(n, x, k, distances, labels);
+        } else if (metric_type == METRIC_L2) {
+            FAISS_THROW_IF_NOT(centroid_norms.size() == ntotal);
+            rq.knn_exact_L2(n, x, k, distances, labels, centroid_norms.data());
+        }
+        return;
+    }
+
+    int beam_size = int(k * beam_factor);
+
+    size_t memory_per_point = rq.memory_per_point(beam_size);
+
+    /*
+
+    printf("mem per point %ld n=%d max_mem_distance=%ld mem_kb=%zd\n",
+        memory_per_point, int(n), rq.max_mem_distances, get_mem_usage_kb());
+    */
+    if (n > 1 && memory_per_point * n > rq.max_mem_distances) {
+        // then split queries to reduce temp memory
+        idx_t bs = rq.max_mem_distances / memory_per_point;
+        if (bs == 0) {
+            bs = 1; // otherwise we can't do much
+        }
+        if (verbose) {
+            printf("ResidualCoarseQuantizer::search: run %d searches in batches of size %d\n",
+                   int(n),
+                   int(bs));
+        }
+        for (idx_t i0 = 0; i0 < n; i0 += bs) {
+            idx_t i1 = std::min(n, i0 + bs);
+            search(i1 - i0, x + i0 * d, k, distances + i0 * k, labels + i0 * k);
+            InterruptCallback::check();
+        }
+        return;
+    }
+
+    std::vector<int32_t> codes(beam_size * rq.M * n);
+    std::vector<float> beam_distances(n * beam_size);
+
+    rq.refine_beam(
+            n, 1, x, beam_size, codes.data(), nullptr, beam_distances.data());
+
+#pragma omp parallel for if (n > 4000)
+    for (idx_t i = 0; i < n; i++) {
+        memcpy(distances + i * k,
+               beam_distances.data() + beam_size * i,
+               k * sizeof(distances[0]));
+
+        const int32_t* codes_i = codes.data() + beam_size * i * rq.M;
+        for (idx_t j = 0; j < k; j++) {
+            idx_t l = 0;
+            int shift = 0;
+            for (int m = 0; m < rq.M; m++) {
+                l |= (*codes_i++) << shift;
+                shift += rq.nbits[m];
+            }
+            labels[i * k + j] = l;
+        }
+    }
+}
+
+void ResidualCoarseQuantizer::reconstruct(idx_t key, float* recons) const {
+    rq.decode_64bit(key, recons);
+}
+
+void ResidualCoarseQuantizer::reset() {
+    FAISS_THROW_MSG("not applicable");
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/IndexResidual.h

@@ -0,0 +1,152 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#ifndef FAISS_INDEX_RESIDUAL_H
+#define FAISS_INDEX_RESIDUAL_H
+
+#include <stdint.h>
+
+#include <vector>
+
+#include <faiss/Index.h>
+#include <faiss/impl/ResidualQuantizer.h>
+#include <faiss/impl/platform_macros.h>
+
+namespace faiss {
+
+/** Index based on a residual quantizer. Stored vectors are
+ * approximated by residual quantization codes.
+ * Can also be used as a codec
+ */
+struct IndexResidual : Index {
+    /// The residual quantizer used to encode the vectors
+    ResidualQuantizer rq;
+
+    enum Search_type_t {
+        ST_decompress, ///< decompress database vector
+        ST_LUT_nonorm, ///< use a LUT, don't include norms (OK for IP or
+                       ///< normalized vectors)
+        ST_norm_float, ///< use a LUT, and store float32 norm with the vectors
+        ST_norm_qint8, ///< use a LUT, and store 8bit-quantized norm
+    };
+    Search_type_t search_type;
+
+    /// min/max for quantization of norms
+    float norm_min, norm_max;
+
+    /// size of residual quantizer codes + norms
+    size_t code_size;
+
+    /// Codes. Size ntotal * rq.code_size
+    std::vector<uint8_t> codes;
+
+    /** Constructor.
+     *
+     * @param d      dimensionality of the input vectors
+     * @param M      number of subquantizers
+     * @param nbits  number of bit per subvector index
+     */
+    IndexResidual(
+            int d,        ///< dimensionality of the input vectors
+            size_t M,     ///< number of subquantizers
+            size_t nbits, ///< number of bit per subvector index
+            MetricType metric = METRIC_L2,
+            Search_type_t search_type = ST_decompress);
+
+    IndexResidual(
+            int d,
+            const std::vector<size_t>& nbits,
+            MetricType metric = METRIC_L2,
+            Search_type_t search_type = ST_decompress);
+
+    IndexResidual();
+
+    /// set search type and update parameters
+    void set_search_type(Search_type_t search_type);
+
+    void train(idx_t n, const float* x) override;
+
+    void add(idx_t n, const float* x) override;
+
+    /// not implemented
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels) const override;
+
+    void reset() override;
+
+    /* The standalone codec interface */
+    size_t sa_code_size() const override;
+
+    void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override;
+
+    void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
+
+    //    DistanceComputer* get_distance_computer() const override;
+};
+
+/** A "virtual" index where the elements are the residual quantizer centroids.
+ *
+ * Intended for use as a coarse quantizer in an IndexIVF.
+ */
+struct ResidualCoarseQuantizer : Index {
+    /// The residual quantizer used to encode the vectors
+    ResidualQuantizer rq;
+
+    /// factor between the beam size and the search k
+    /// if negative, use exact search-to-centroid
+    float beam_factor;
+
+    /// norms of centroids, useful for knn-search
+    std::vector<float> centroid_norms;
+
+    /// computes centroid norms if required
+    void set_beam_factor(float new_beam_factor);
+
+    /** Constructor.
+     *
+     * @param d      dimensionality of the input vectors
+     * @param M      number of subquantizers
+     * @param nbits  number of bit per subvector index
+     */
+    ResidualCoarseQuantizer(
+            int d,        ///< dimensionality of the input vectors
+            size_t M,     ///< number of subquantizers
+            size_t nbits, ///< number of bit per subvector index
+            MetricType metric = METRIC_L2);
+
+    ResidualCoarseQuantizer(
+            int d,
+            const std::vector<size_t>& nbits,
+            MetricType metric = METRIC_L2);
+
+    ResidualCoarseQuantizer();
+
+    void train(idx_t n, const float* x) override;
+
+    /// N/A
+    void add(idx_t n, const float* x) override;
+
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels) const override;
+
+    void reconstruct(idx_t key, float* recons) const override;
+
+    /// N/A
+    void reset() override;
+};
+
+} // namespace faiss
+
+#endif

data/vendor/faiss/faiss/IndexScalarQuantizer.cpp

@@ -9,231 +9,207 @@
 
 #include <faiss/IndexScalarQuantizer.h>
 
-#include <cstdio>
 #include <algorithm>
+#include <cstdio>
 
 #include <omp.h>
 
-#include <faiss/utils/utils.h>
-#include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/ScalarQuantizer.h>
+#include <faiss/utils/utils.h>
 
 namespace faiss {
 
-
-
 /*******************************************************************
  * IndexScalarQuantizer implementation
  ********************************************************************/
 
-IndexScalarQuantizer::IndexScalarQuantizer
-                      (int d, ScalarQuantizer::QuantizerType qtype,
-                       MetricType metric):
-          Index(d, metric),
-          sq (d, qtype)
-{
-    is_trained =
-        qtype == ScalarQuantizer::QT_fp16 ||
-        qtype == ScalarQuantizer::QT_8bit_direct;
+IndexScalarQuantizer::IndexScalarQuantizer(
+        int d,
+        ScalarQuantizer::QuantizerType qtype,
+        MetricType metric)
+        : Index(d, metric), sq(d, qtype) {
+    is_trained = qtype == ScalarQuantizer::QT_fp16 ||
+            qtype == ScalarQuantizer::QT_8bit_direct;
     code_size = sq.code_size;
 }
 
+IndexScalarQuantizer::IndexScalarQuantizer()
+        : IndexScalarQuantizer(0, ScalarQuantizer::QT_8bit) {}
 
-IndexScalarQuantizer::IndexScalarQuantizer ():
-    IndexScalarQuantizer(0, ScalarQuantizer::QT_8bit)
-{}
-
-void IndexScalarQuantizer::train(idx_t n, const float* x)
-{
+void IndexScalarQuantizer::train(idx_t n, const float* x) {
     sq.train(n, x);
     is_trained = true;
 }
 
-void IndexScalarQuantizer::add(idx_t n, const float* x)
-{
-    FAISS_THROW_IF_NOT (is_trained);
-    codes.resize ((n + ntotal) * code_size);
-    sq.compute_codes (x, &codes[ntotal * code_size], n);
+void IndexScalarQuantizer::add(idx_t n, const float* x) {
+    FAISS_THROW_IF_NOT(is_trained);
+    codes.resize((n + ntotal) * code_size);
+    sq.compute_codes(x, &codes[ntotal * code_size], n);
     ntotal += n;
 }
 
-
 void IndexScalarQuantizer::search(
         idx_t n,
         const float* x,
         idx_t k,
         float* distances,
-        idx_t* labels) const
-{
-    FAISS_THROW_IF_NOT (is_trained);
-    FAISS_THROW_IF_NOT (metric_type == METRIC_L2 ||
-                        metric_type == METRIC_INNER_PRODUCT);
+        idx_t* labels) const {
+    FAISS_THROW_IF_NOT(k > 0);
+
+    FAISS_THROW_IF_NOT(is_trained);
+    FAISS_THROW_IF_NOT(
+            metric_type == METRIC_L2 || metric_type == METRIC_INNER_PRODUCT);
 
 #pragma omp parallel
     {
-        InvertedListScanner* scanner = sq.select_InvertedListScanner
-            (metric_type, nullptr, true);
+        InvertedListScanner* scanner =
+                sq.select_InvertedListScanner(metric_type, nullptr, true);
         ScopeDeleter1<InvertedListScanner> del(scanner);
 
 #pragma omp for
         for (idx_t i = 0; i < n; i++) {
-            float * D = distances + k * i;
-            idx_t * I = labels + k * i;
+            float* D = distances + k * i;
+            idx_t* I = labels + k * i;
             // re-order heap
             if (metric_type == METRIC_L2) {
-                maxheap_heapify (k, D, I);
+                maxheap_heapify(k, D, I);
             } else {
-                minheap_heapify (k, D, I);
+                minheap_heapify(k, D, I);
             }
-            scanner->set_query (x + i * d);
-            scanner->scan_codes (ntotal, codes.data(),
-                                 nullptr, D, I, k);
+            scanner->set_query(x + i * d);
+            scanner->scan_codes(ntotal, codes.data(), nullptr, D, I, k);
 
             // re-order heap
             if (metric_type == METRIC_L2) {
-                maxheap_reorder (k, D, I);
+                maxheap_reorder(k, D, I);
             } else {
-                minheap_reorder (k, D, I);
+                minheap_reorder(k, D, I);
             }
         }
     }
-
 }
 
-
-DistanceComputer *IndexScalarQuantizer::get_distance_computer () const
-{
-    ScalarQuantizer::SQDistanceComputer *dc =
-        sq.get_distance_computer (metric_type);
+DistanceComputer* IndexScalarQuantizer::get_distance_computer() const {
+    ScalarQuantizer::SQDistanceComputer* dc =
+            sq.get_distance_computer(metric_type);
     dc->code_size = sq.code_size;
     dc->codes = codes.data();
     return dc;
 }
 
-
-void IndexScalarQuantizer::reset()
-{
+void IndexScalarQuantizer::reset() {
     codes.clear();
     ntotal = 0;
 }
 
-void IndexScalarQuantizer::reconstruct_n(
-             idx_t i0, idx_t ni, float* recons) const
-{
-    std::unique_ptr<ScalarQuantizer::Quantizer> squant(sq.select_quantizer ());
+void IndexScalarQuantizer::reconstruct_n(idx_t i0, idx_t ni, float* recons)
+        const {
+    std::unique_ptr<ScalarQuantizer::Quantizer> squant(sq.select_quantizer());
     for (size_t i = 0; i < ni; i++) {
         squant->decode_vector(&codes[(i + i0) * code_size], recons + i * d);
     }
 }
 
-void IndexScalarQuantizer::reconstruct(idx_t key, float* recons) const
-{
+void IndexScalarQuantizer::reconstruct(idx_t key, float* recons) const {
     reconstruct_n(key, 1, recons);
 }
 
 /* Codec interface */
-size_t IndexScalarQuantizer::sa_code_size () const
-{
+size_t IndexScalarQuantizer::sa_code_size() const {
     return sq.code_size;
 }
 
-void IndexScalarQuantizer::sa_encode (idx_t n, const float *x,
-                      uint8_t *bytes) const
-{
-    FAISS_THROW_IF_NOT (is_trained);
-    sq.compute_codes (x, bytes, n);
+void IndexScalarQuantizer::sa_encode(idx_t n, const float* x, uint8_t* bytes)
+        const {
+    FAISS_THROW_IF_NOT(is_trained);
+    sq.compute_codes(x, bytes, n);
 }
 
-void IndexScalarQuantizer::sa_decode (idx_t n, const uint8_t *bytes,
-                                              float *x) const
-{
-    FAISS_THROW_IF_NOT (is_trained);
+void IndexScalarQuantizer::sa_decode(idx_t n, const uint8_t* bytes, float* x)
+        const {
+    FAISS_THROW_IF_NOT(is_trained);
     sq.decode(bytes, x, n);
 }
 
-
-
 /*******************************************************************
  * IndexIVFScalarQuantizer implementation
  ********************************************************************/
 
-IndexIVFScalarQuantizer::IndexIVFScalarQuantizer (
-            Index *quantizer, size_t d, size_t nlist,
-            ScalarQuantizer::QuantizerType qtype,
-            MetricType metric, bool encode_residual)
-    : IndexIVF(quantizer, d, nlist, 0, metric),
-      sq(d, qtype),
-      by_residual(encode_residual)
-{
+IndexIVFScalarQuantizer::IndexIVFScalarQuantizer(
+        Index* quantizer,
+        size_t d,
+        size_t nlist,
+        ScalarQuantizer::QuantizerType qtype,
+        MetricType metric,
+        bool encode_residual)
+        : IndexIVF(quantizer, d, nlist, 0, metric),
+          sq(d, qtype),
+          by_residual(encode_residual) {
     code_size = sq.code_size;
     // was not known at construction time
     invlists->code_size = code_size;
     is_trained = false;
 }
 
-IndexIVFScalarQuantizer::IndexIVFScalarQuantizer ():
-    IndexIVF(),
-    by_residual(true)
-{
-}
+IndexIVFScalarQuantizer::IndexIVFScalarQuantizer()
+        : IndexIVF(), by_residual(true) {}
 
-void IndexIVFScalarQuantizer::train_residual (idx_t n, const float *x)
-{
+void IndexIVFScalarQuantizer::train_residual(idx_t n, const float* x) {
     sq.train_residual(n, x, quantizer, by_residual, verbose);
 }
 
-void IndexIVFScalarQuantizer::encode_vectors(idx_t n, const float* x,
-                                             const idx_t *list_nos,
-                                             uint8_t * codes,
-                                             bool include_listnos) const
-{
-    std::unique_ptr<ScalarQuantizer::Quantizer> squant (sq.select_quantizer ());
-    size_t coarse_size = include_listnos ? coarse_code_size () : 0;
+void IndexIVFScalarQuantizer::encode_vectors(
+        idx_t n,
+        const float* x,
+        const idx_t* list_nos,
+        uint8_t* codes,
+        bool include_listnos) const {
+    std::unique_ptr<ScalarQuantizer::Quantizer> squant(sq.select_quantizer());
+    size_t coarse_size = include_listnos ? coarse_code_size() : 0;
     memset(codes, 0, (code_size + coarse_size) * n);
 
-#pragma omp parallel if(n > 1000)
+#pragma omp parallel if (n > 1000)
     {
-        std::vector<float> residual (d);
+        std::vector<float> residual(d);
 
 #pragma omp for
         for (idx_t i = 0; i < n; i++) {
-            int64_t list_no = list_nos [i];
+            int64_t list_no = list_nos[i];
             if (list_no >= 0) {
-                const float *xi = x + i * d;
-                uint8_t *code = codes + i * (code_size + coarse_size);
+                const float* xi = x + i * d;
+                uint8_t* code = codes + i * (code_size + coarse_size);
                 if (by_residual) {
-                    quantizer->compute_residual (
-                          xi, residual.data(), list_no);
-                    xi = residual.data ();
+                    quantizer->compute_residual(xi, residual.data(), list_no);
+                    xi = residual.data();
                 }
                 if (coarse_size) {
-                    encode_listno (list_no, code);
+                    encode_listno(list_no, code);
                 }
-                squant->encode_vector (xi, code + coarse_size);
+                squant->encode_vector(xi, code + coarse_size);
             }
         }
     }
 }
 
-void IndexIVFScalarQuantizer::sa_decode (idx_t n, const uint8_t *codes,
-                                                 float *x) const
-{
-    std::unique_ptr<ScalarQuantizer::Quantizer> squant (sq.select_quantizer ());
-    size_t coarse_size = coarse_code_size ();
+void IndexIVFScalarQuantizer::sa_decode(idx_t n, const uint8_t* codes, float* x)
+        const {
+    std::unique_ptr<ScalarQuantizer::Quantizer> squant(sq.select_quantizer());
+    size_t coarse_size = coarse_code_size();
 
-#pragma omp parallel if(n > 1000)
+#pragma omp parallel if (n > 1000)
     {
-        std::vector<float> residual (d);
+        std::vector<float> residual(d);
 
 #pragma omp for
         for (idx_t i = 0; i < n; i++) {
-            const uint8_t *code = codes + i * (code_size + coarse_size);
-            int64_t list_no = decode_listno (code);
-            float *xi = x + i * d;
-            squant->decode_vector (code + coarse_size, xi);
+            const uint8_t* code = codes + i * (code_size + coarse_size);
+            int64_t list_no = decode_listno(code);
+            float* xi = x + i * d;
+            squant->decode_vector(code + coarse_size, xi);
             if (by_residual) {
-                quantizer->reconstruct (list_no, residual.data());
+                quantizer->reconstruct(list_no, residual.data());
                 for (size_t j = 0; j < d; j++) {
                     xi[j] += residual[j];
                 }
@@ -242,83 +218,72 @@ void IndexIVFScalarQuantizer::sa_decode (idx_t n, const uint8_t *codes,
     }
 }
 
+void IndexIVFScalarQuantizer::add_core(
+        idx_t n,
+        const float* x,
+        const idx_t* xids,
+        const idx_t* coarse_idx) {
+    FAISS_THROW_IF_NOT(is_trained);
 
-
-void IndexIVFScalarQuantizer::add_with_ids
-       (idx_t n, const float * x, const idx_t *xids)
-{
-    FAISS_THROW_IF_NOT (is_trained);
-    std::unique_ptr<int64_t []> idx (new int64_t [n]);
-    quantizer->assign (n, x, idx.get());
     size_t nadd = 0;
-    std::unique_ptr<ScalarQuantizer::Quantizer> squant(sq.select_quantizer ());
+    std::unique_ptr<ScalarQuantizer::Quantizer> squant(sq.select_quantizer());
 
-    DirectMapAdd dm_add (direct_map, n, xids);
+    DirectMapAdd dm_add(direct_map, n, xids);
 
-#pragma omp parallel reduction(+: nadd)
+#pragma omp parallel reduction(+ : nadd)
     {
-        std::vector<float> residual (d);
-        std::vector<uint8_t> one_code (code_size);
+        std::vector<float> residual(d);
+        std::vector<uint8_t> one_code(code_size);
         int nt = omp_get_num_threads();
         int rank = omp_get_thread_num();
 
         // each thread takes care of a subset of lists
         for (size_t i = 0; i < n; i++) {
-            int64_t list_no = idx [i];
+            int64_t list_no = coarse_idx[i];
             if (list_no >= 0 && list_no % nt == rank) {
                 int64_t id = xids ? xids[i] : ntotal + i;
 
-                const float * xi = x + i * d;
+                const float* xi = x + i * d;
                 if (by_residual) {
-                    quantizer->compute_residual (xi, residual.data(), list_no);
+                    quantizer->compute_residual(xi, residual.data(), list_no);
                     xi = residual.data();
                 }
 
-                memset (one_code.data(), 0, code_size);
-                squant->encode_vector (xi, one_code.data());
+                memset(one_code.data(), 0, code_size);
+                squant->encode_vector(xi, one_code.data());
 
-                size_t ofs = invlists->add_entry (list_no, id, one_code.data());
+                size_t ofs = invlists->add_entry(list_no, id, one_code.data());
 
-                dm_add.add (i, list_no, ofs);
+                dm_add.add(i, list_no, ofs);
                 nadd++;
 
             } else if (rank == 0 && list_no == -1) {
-                dm_add.add (i, -1, 0);
+                dm_add.add(i, -1, 0);
             }
         }
     }
 
-
     ntotal += n;
 }
 
-
-
-
-
-InvertedListScanner* IndexIVFScalarQuantizer::get_InvertedListScanner
-    (bool store_pairs) const
-{
-    return sq.select_InvertedListScanner (metric_type, quantizer, store_pairs,
-                                          by_residual);
+InvertedListScanner* IndexIVFScalarQuantizer::get_InvertedListScanner(
+        bool store_pairs) const {
+    return sq.select_InvertedListScanner(
+            metric_type, quantizer, store_pairs, by_residual);
 }
 
-
-void IndexIVFScalarQuantizer::reconstruct_from_offset (int64_t list_no,
-                                                       int64_t offset,
-                                                       float* recons) const
-{
+void IndexIVFScalarQuantizer::reconstruct_from_offset(
+        int64_t list_no,
+        int64_t offset,
+        float* recons) const {
     std::vector<float> centroid(d);
-    quantizer->reconstruct (list_no, centroid.data());
+    quantizer->reconstruct(list_no, centroid.data());
 
-    const uint8_t* code = invlists->get_single_code (list_no, offset);
-    sq.decode (code, recons, 1);
+    const uint8_t* code = invlists->get_single_code(list_no, offset);
+    sq.decode(code, recons, 1);
     for (int i = 0; i < d; ++i) {
         recons[i] += centroid[i];
     }
 }
 
-
-
-
 } // namespace faiss