faiss 0.2.0 → 0.2.4

data/vendor/faiss/faiss/IndexIVF.h

@@ -10,30 +10,27 @@
 #ifndef FAISS_INDEX_IVF_H
 #define FAISS_INDEX_IVF_H
 
-
-#include <vector>
-#include <unordered_map>
 #include <stdint.h>
+#include <unordered_map>
+#include <vector>
 
-#include <faiss/Index.h>
-#include <faiss/invlists/InvertedLists.h>
-#include <faiss/invlists/DirectMap.h>
 #include <faiss/Clustering.h>
+#include <faiss/Index.h>
 #include <faiss/impl/platform_macros.h>
+#include <faiss/invlists/DirectMap.h>
+#include <faiss/invlists/InvertedLists.h>
 #include <faiss/utils/Heap.h>
 
-
 namespace faiss {
 
-
 /** Encapsulates a quantizer object for the IndexIVF
  *
  * The class isolates the fields that are independent of the storage
  * of the lists (especially training)
  */
 struct Level1Quantizer {
-    Index * quantizer;        ///< quantizer that maps vectors to inverted lists
-    size_t nlist;             ///< number of possible key values
+    Index* quantizer; ///< quantizer that maps vectors to inverted lists
+    size_t nlist;     ///< number of possible key values
 
     /**
      * = 0: use the quantizer as index in a kmeans training
@@ -41,40 +38,37 @@ struct Level1Quantizer {
      * = 2: kmeans training on a flat index + add the centroids to the quantizer
      */
     char quantizer_trains_alone;
-    bool own_fields;          ///< whether object owns the quantizer
+    bool own_fields; ///< whether object owns the quantizer (false by default)
 
     ClusteringParameters cp; ///< to override default clustering params
-    Index *clustering_index; ///< to override index used during clustering
+    Index* clustering_index; ///< to override index used during clustering
 
     /// Trains the quantizer and calls train_residual to train sub-quantizers
-    void train_q1 (size_t n, const float *x, bool verbose,
-                   MetricType metric_type);
-
+    void train_q1(
+            size_t n,
+            const float* x,
+            bool verbose,
+            MetricType metric_type);
 
     /// compute the number of bytes required to store list ids
-    size_t coarse_code_size () const;
-    void encode_listno (Index::idx_t list_no, uint8_t *code) const;
-    Index::idx_t decode_listno (const uint8_t *code) const;
-
-    Level1Quantizer (Index * quantizer, size_t nlist);
+    size_t coarse_code_size() const;
+    void encode_listno(Index::idx_t list_no, uint8_t* code) const;
+    Index::idx_t decode_listno(const uint8_t* code) const;
 
-    Level1Quantizer ();
+    Level1Quantizer(Index* quantizer, size_t nlist);
 
-    ~Level1Quantizer ();
+    Level1Quantizer();
 
+    ~Level1Quantizer();
 };
 
-
-
 struct IVFSearchParameters {
-    size_t nprobe;            ///< number of probes at query time
-    size_t max_codes;         ///< max nb of codes to visit to do a query
-    IVFSearchParameters(): nprobe(1), max_codes(0) {}
-    virtual ~IVFSearchParameters () {}
+    size_t nprobe;    ///< number of probes at query time
+    size_t max_codes; ///< max nb of codes to visit to do a query
+    IVFSearchParameters() : nprobe(1), max_codes(0) {}
+    virtual ~IVFSearchParameters() {}
 };
 
-
-
 struct InvertedListScanner;
 struct IndexIVFStats;
 
@@ -98,15 +92,15 @@ struct IndexIVFStats;
  * Sub-classes implement a post-filtering of the index that refines
  * the distance estimation from the query to databse vectors.
  */
-struct IndexIVF: Index, Level1Quantizer {
+struct IndexIVF : Index, Level1Quantizer {
     /// Access to the actual data
-    InvertedLists *invlists;
+    InvertedLists* invlists;
     bool own_invlists;
 
-    size_t code_size;              ///< code size per vector in bytes
+    size_t code_size; ///< code size per vector in bytes
 
-    size_t nprobe;            ///< number of probes at query time
-    size_t max_codes;         ///< max nb of codes to visit to do a query
+    size_t nprobe;    ///< number of probes at query time
+    size_t max_codes; ///< max nb of codes to visit to do a query
 
     /** Parallel mode determines how queries are parallelized with OpenMP
      *
@@ -127,12 +121,14 @@ struct IndexIVF: Index, Level1Quantizer {
 
     /** The Inverted file takes a quantizer (an Index) on input,
      * which implements the function mapping a vector to a list
-     * identifier. The pointer is borrowed: the quantizer should not
-     * be deleted while the IndexIVF is in use.
+     * identifier.
      */
-    IndexIVF (Index * quantizer, size_t d,
-              size_t nlist, size_t code_size,
-              MetricType metric = METRIC_L2);
+    IndexIVF(
+            Index* quantizer,
+            size_t d,
+            size_t nlist,
+            size_t code_size,
+            MetricType metric = METRIC_L2);
 
     void reset() override;
 
@@ -145,6 +141,19 @@ struct IndexIVF: Index, Level1Quantizer {
     /// default implementation that calls encode_vectors
     void add_with_ids(idx_t n, const float* x, const idx_t* xids) override;
 
+    /** Implementation of vector addition where the vector assignments are
+     * predefined. The default implementation hands over the code extraction to
+     * encode_vectors.
+     *
+     * @param precomputed_idx    quantization indices for the input vectors
+     * (size n)
+     */
+    virtual void add_core(
+            idx_t n,
+            const float* x,
+            const idx_t* xids,
+            const idx_t* precomputed_idx);
+
     /** Encodes a set of vectors as they would appear in the inverted lists
      *
      * @param list_nos   inverted list ids as returned by the
@@ -154,14 +163,23 @@ struct IndexIVF: Index, Level1Quantizer {
      *                   include the list ids in the code (in this case add
      *                   ceil(log8(nlist)) to the code size)
      */
-    virtual void encode_vectors(idx_t n, const float* x,
-                                const idx_t *list_nos,
-                                uint8_t * codes,
-                                bool include_listno = false) const = 0;
+    virtual void encode_vectors(
+            idx_t n,
+            const float* x,
+            const idx_t* list_nos,
+            uint8_t* codes,
+            bool include_listno = false) const = 0;
+
+    /** Add vectors that are computed with the standalone codec
+     *
+     * @param codes  codes to add size n * sa_code_size()
+     * @param xids   corresponding ids, size n
+     */
+    void add_sa_codes(idx_t n, const uint8_t* codes, const idx_t* xids);
 
     /// Sub-classes that encode the residuals can train their encoders here
     /// does nothing by default
-    virtual void train_residual (idx_t n, const float *x);
+    virtual void train_residual(idx_t n, const float* x);
 
     /** search a set of vectors, that are pre-quantized by the IVF
      *  quantizer. Fill in the corresponding heaps with the query
@@ -182,36 +200,53 @@ struct IndexIVF: Index, Level1Quantizer {
      * @param params used to override the object's search parameters
      * @param stats  search stats to be updated (can be null)
      */
-    virtual void search_preassigned (
-            idx_t n, const float *x, idx_t k,
-            const idx_t *assign, const float *centroid_dis,
-            float *distances, idx_t *labels,
+    virtual void search_preassigned(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            const idx_t* assign,
+            const float* centroid_dis,
+            float* distances,
+            idx_t* labels,
             bool store_pairs,
-            const IVFSearchParameters *params=nullptr,
-            IndexIVFStats *stats=nullptr
-            ) const;
+            const IVFSearchParameters* params = nullptr,
+            IndexIVFStats* stats = nullptr) const;
 
     /** assign the vectors, then call search_preassign */
-    void search (idx_t n, const float *x, idx_t k,
-                 float *distances, idx_t *labels) const override;
-
-    void range_search (idx_t n, const float* x, float radius,
-                       RangeSearchResult* result) const override;
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels) const override;
+
+    void range_search(
+            idx_t n,
+            const float* x,
+            float radius,
+            RangeSearchResult* result) const override;
 
     void range_search_preassigned(
-            idx_t nx, const float *x, float radius,
-            const idx_t *keys, const float *coarse_dis,
-            RangeSearchResult *result,
-            bool store_pairs=false,
-            const IVFSearchParameters *params=nullptr,
-            IndexIVFStats *stats=nullptr) const;
-
-    /// get a scanner for this index (store_pairs means ignore labels)
-    virtual InvertedListScanner *get_InvertedListScanner (
-        bool store_pairs=false) const;
+            idx_t nx,
+            const float* x,
+            float radius,
+            const idx_t* keys,
+            const float* coarse_dis,
+            RangeSearchResult* result,
+            bool store_pairs = false,
+            const IVFSearchParameters* params = nullptr,
+            IndexIVFStats* stats = nullptr) const;
+
+    /** Get a scanner for this index (store_pairs means ignore labels)
+     *
+     * The default search implementation uses this to compute the distances
+     */
+    virtual InvertedListScanner* get_InvertedListScanner(
+            bool store_pairs = false) const;
 
-    /** reconstruct a vector. Works only if maintain_direct_map is set to 1 or 2 */
-    void reconstruct (idx_t key, float* recons) const override;
+    /** reconstruct a vector. Works only if maintain_direct_map is set to 1 or 2
+     */
+    void reconstruct(idx_t key, float* recons) const override;
 
     /** Update a subset of vectors.
      *
@@ -221,7 +256,7 @@ struct IndexIVF: Index, Level1Quantizer {
      * @param idx    vector indices to update, size nv
      * @param v      vectors of new values, size nv*d
      */
-    virtual void update_vectors (int nv, const idx_t *idx, const float *v);
+    virtual void update_vectors(int nv, const idx_t* idx, const float* v);
 
     /** Reconstruct a subset of the indexed vectors.
      *
@@ -243,9 +278,13 @@ struct IndexIVF: Index, Level1Quantizer {
      *
      * @param recons      reconstructed vectors size (n, k, d)
      */
-    void search_and_reconstruct (idx_t n, const float *x, idx_t k,
-                                 float *distances, idx_t *labels,
-                                 float *recons) const override;
+    void search_and_reconstruct(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            float* recons) const override;
 
     /** Reconstruct a vector given the location in terms of (inv list index +
      * inv list offset) instead of the id.
@@ -254,9 +293,10 @@ struct IndexIVF: Index, Level1Quantizer {
      * the inv list offset is computed by search_preassigned() with
      * `store_pairs` set.
      */
-    virtual void reconstruct_from_offset (int64_t list_no, int64_t offset,
-                                          float* recons) const;
-
+    virtual void reconstruct_from_offset(
+            int64_t list_no,
+            int64_t offset,
+            float* recons) const;
 
     /// Dataset manipulation functions
 
@@ -265,12 +305,12 @@ struct IndexIVF: Index, Level1Quantizer {
     /** check that the two indexes are compatible (ie, they are
      * trained in the same way and have the same
      * parameters). Otherwise throw. */
-    void check_compatible_for_merge (const IndexIVF &other) const;
+    void check_compatible_for_merge(const IndexIVF& other) const;
 
     /** moves the entries from another dataset to self. On output,
      * other is empty. add_id is added to all moved ids (for
      * sequential ids, this would be this->ntotal */
-    virtual void merge_from (IndexIVF &other, idx_t add_id);
+    virtual void merge_from(IndexIVF& other, idx_t add_id);
 
     /** copy a subset of the entries index to the other index
      *
@@ -279,34 +319,36 @@ struct IndexIVF: Index, Level1Quantizer {
      * if subset_type == 2: copies inverted lists such that a1
      *                      elements are left before and a2 elements are after
      */
-    virtual void copy_subset_to (IndexIVF & other, int subset_type,
-                                 idx_t a1, idx_t a2) const;
+    virtual void copy_subset_to(
+            IndexIVF& other,
+            int subset_type,
+            idx_t a1,
+            idx_t a2) const;
 
     ~IndexIVF() override;
 
-    size_t get_list_size (size_t list_no) const
-    { return invlists->list_size(list_no); }
+    size_t get_list_size(size_t list_no) const {
+        return invlists->list_size(list_no);
+    }
 
     /** intialize a direct map
      *
      * @param new_maintain_direct_map    if true, create a direct map,
      *                                   else clear it
      */
-    void make_direct_map (bool new_maintain_direct_map=true);
-
-    void set_direct_map_type (DirectMap::Type type);
+    void make_direct_map(bool new_maintain_direct_map = true);
 
+    void set_direct_map_type(DirectMap::Type type);
 
     /// replace the inverted lists, old one is deallocated if own_invlists
-    void replace_invlists (InvertedLists *il, bool own=false);
+    void replace_invlists(InvertedLists* il, bool own = false);
 
     /* The standalone codec interface (except sa_decode that is specific) */
-    size_t sa_code_size () const override;
+    size_t sa_code_size() const override;
 
-    void sa_encode (idx_t n, const float *x,
-                          uint8_t *bytes) const override;
+    void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override;
 
-    IndexIVF ();
+    IndexIVF();
 };
 
 struct RangeQueryResult;
@@ -316,20 +358,28 @@ struct RangeQueryResult;
  * distance_to_code and scan_codes can be called in multiple
  * threads */
 struct InvertedListScanner {
-
     using idx_t = Index::idx_t;
 
+    idx_t list_no = -1;    ///< remember current list
+    bool keep_max = false; ///< keep maximum instead of minimum
+    /// store positions in invlists rather than labels
+    bool store_pairs = false;
+
+    /// used in default implementation of scan_codes
+    size_t code_size = 0;
+
     /// from now on we handle this query.
-    virtual void set_query (const float *query_vector) = 0;
+    virtual void set_query(const float* query_vector) = 0;
 
     /// following codes come from this inverted list
-    virtual void set_list (idx_t list_no, float coarse_dis) = 0;
+    virtual void set_list(idx_t list_no, float coarse_dis) = 0;
 
     /// compute a single query-to-code distance
-    virtual float distance_to_code (const uint8_t *code) const = 0;
+    virtual float distance_to_code(const uint8_t* code) const = 0;
 
     /** scan a set of codes, compute distances to current query and
-     * update heap of results if necessary.
+     * update heap of results if necessary. Default implemetation
+     * calls distance_to_code.
      *
      * @param n      number of codes to scan
      * @param codes  codes to scan (n * code_size)
@@ -339,45 +389,46 @@ struct InvertedListScanner {
      * @param k          heap size
      * @return number of heap updates performed
      */
-    virtual size_t scan_codes (size_t n,
-                               const uint8_t *codes,
-                               const idx_t *ids,
-                               float *distances, idx_t *labels,
-                               size_t k) const = 0;
+    virtual size_t scan_codes(
+            size_t n,
+            const uint8_t* codes,
+            const idx_t* ids,
+            float* distances,
+            idx_t* labels,
+            size_t k) const;
 
     /** scan a set of codes, compute distances to current query and
      * update results if distances are below radius
      *
      * (default implementation fails) */
-    virtual void scan_codes_range (size_t n,
-                                   const uint8_t *codes,
-                                   const idx_t *ids,
-                                   float radius,
-                                   RangeQueryResult &result) const;
-
-    virtual ~InvertedListScanner () {}
-
+    virtual void scan_codes_range(
+            size_t n,
+            const uint8_t* codes,
+            const idx_t* ids,
+            float radius,
+            RangeQueryResult& result) const;
+
+    virtual ~InvertedListScanner() {}
 };
 
-
 struct IndexIVFStats {
-    size_t nq;       // nb of queries run
-    size_t nlist;    // nb of inverted lists scanned
-    size_t ndis;     // nb of distancs computed
-    size_t nheap_updates; // nb of times the heap was updated
+    size_t nq;                // nb of queries run
+    size_t nlist;             // nb of inverted lists scanned
+    size_t ndis;              // nb of distances computed
+    size_t nheap_updates;     // nb of times the heap was updated
     double quantization_time; // time spent quantizing vectors (in ms)
     double search_time;       // time spent searching lists (in ms)
 
-    IndexIVFStats () {reset (); }
-    void reset ();
-    void add (const IndexIVFStats & other);
+    IndexIVFStats() {
+        reset();
+    }
+    void reset();
+    void add(const IndexIVFStats& other);
 };
 
 // global var that collects them all
 FAISS_API extern IndexIVFStats indexIVF_stats;
 
-
 } // namespace faiss
 
-
 #endif