faiss 0.2.0 → 0.2.4

data/vendor/faiss/faiss/IndexBinaryIVF.h

@@ -10,15 +10,13 @@
 #ifndef FAISS_INDEX_BINARY_IVF_H
 #define FAISS_INDEX_BINARY_IVF_H
 
-
 #include <vector>
 
+#include <faiss/Clustering.h>
 #include <faiss/IndexBinary.h>
 #include <faiss/IndexIVF.h>
-#include <faiss/Clustering.h>
 #include <faiss/utils/Heap.h>
 
-
 namespace faiss {
 
 struct BinaryInvertedListScanner;
@@ -34,11 +32,11 @@ struct BinaryInvertedListScanner;
  */
 struct IndexBinaryIVF : IndexBinary {
     /// Access to the actual data
-    InvertedLists *invlists;
+    InvertedLists* invlists;
     bool own_invlists;
 
-    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
 
     /** Select between using a heap or counting to select the k smallest values
      * when scanning inverted lists.
@@ -48,20 +46,20 @@ struct IndexBinaryIVF : IndexBinary {
     /// map for direct access to the elements. Enables reconstruct().
     DirectMap direct_map;
 
-    IndexBinary *quantizer;   ///< quantizer that maps vectors to inverted lists
-    size_t nlist;             ///< number of possible key values
+    IndexBinary* quantizer; ///< quantizer that maps vectors to inverted lists
+    size_t nlist;           ///< number of possible key values
 
-    bool own_fields;          ///< whether object owns the quantizer
+    bool own_fields; ///< whether object owns the quantizer
 
     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
 
     /** The Inverted file takes a quantizer (an IndexBinary) 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 IndexBinaryIVF is in use.
      */
-    IndexBinaryIVF(IndexBinary *quantizer, size_t d, size_t nlist);
+    IndexBinaryIVF(IndexBinary* quantizer, size_t d, size_t nlist);
 
     IndexBinaryIVF();
 
@@ -70,15 +68,23 @@ struct IndexBinaryIVF : IndexBinary {
     void reset() override;
 
     /// Trains the quantizer
-    void train(idx_t n, const uint8_t *x) override;
+    void train(idx_t n, const uint8_t* x) override;
 
-    void add(idx_t n, const uint8_t *x) override;
+    void add(idx_t n, const uint8_t* x) override;
 
-    void add_with_ids(idx_t n, const uint8_t *x, const idx_t *xids) override;
+    void add_with_ids(idx_t n, const uint8_t* x, const idx_t* xids) override;
 
-    /// same as add_with_ids, with precomputed coarse quantizer
-    void add_core (idx_t n, const uint8_t * x, const idx_t *xids,
-                   const idx_t *precomputed_idx);
+    /** Implementation of vector addition where the vector assignments are
+     * predefined.
+     *
+     * @param precomputed_idx    quantization indices for the input vectors
+     * (size n)
+     */
+    void add_core(
+            idx_t n,
+            const uint8_t* x,
+            const idx_t* xids,
+            const idx_t* precomputed_idx);
 
     /** Search a set of vectors, that are pre-quantized by the IVF
      *  quantizer. Fill in the corresponding heaps with the query
@@ -97,25 +103,43 @@ struct IndexBinaryIVF : IndexBinary {
      *                     instead of ids (used for reranking).
      * @param params used to override the object's search parameters
      */
-    void search_preassigned(idx_t n, const uint8_t *x, idx_t k,
-                            const idx_t *assign,
-                            const int32_t *centroid_dis,
-                            int32_t *distances, idx_t *labels,
-                            bool store_pairs,
-                            const IVFSearchParameters *params=nullptr
-                            ) const;
-
-    virtual BinaryInvertedListScanner *get_InvertedListScanner (
-                                         bool store_pairs=false) const;
+    void search_preassigned(
+            idx_t n,
+            const uint8_t* x,
+            idx_t k,
+            const idx_t* assign,
+            const int32_t* centroid_dis,
+            int32_t* distances,
+            idx_t* labels,
+            bool store_pairs,
+            const IVFSearchParameters* params = nullptr) const;
+
+    virtual BinaryInvertedListScanner* get_InvertedListScanner(
+            bool store_pairs = false) const;
 
     /** assign the vectors, then call search_preassign */
-    void search(idx_t n, const uint8_t *x, idx_t k,
-                int32_t *distances, idx_t *labels) const override;
-
-    void range_search(idx_t n, const uint8_t *x, int radius,
-                      RangeSearchResult *result) const override;
-
-    void reconstruct(idx_t key, uint8_t *recons) const override;
+    void search(
+            idx_t n,
+            const uint8_t* x,
+            idx_t k,
+            int32_t* distances,
+            idx_t* labels) const override;
+
+    void range_search(
+            idx_t n,
+            const uint8_t* x,
+            int radius,
+            RangeSearchResult* result) const override;
+
+    void range_search_preassigned(
+            idx_t n,
+            const uint8_t* x,
+            int radius,
+            const idx_t* assign,
+            const int32_t* centroid_dis,
+            RangeSearchResult* result) const;
+
+    void reconstruct(idx_t key, uint8_t* recons) const override;
 
     /** Reconstruct a subset of the indexed vectors.
      *
@@ -126,7 +150,7 @@ struct IndexBinaryIVF : IndexBinary {
      * @param ni     nb of vectors to reconstruct
      * @param recons output array of reconstructed vectors, size ni * d / 8
      */
-    void reconstruct_n(idx_t i0, idx_t ni, uint8_t *recons) const override;
+    void reconstruct_n(idx_t i0, idx_t ni, uint8_t* recons) const override;
 
     /** Similar to search, but also reconstructs the stored vectors (or an
      * approximation in the case of lossy coding) for the search results.
@@ -137,9 +161,13 @@ struct IndexBinaryIVF : IndexBinary {
      *
      * @param recons      reconstructed vectors size (n, k, d / 8)
      */
-    void search_and_reconstruct(idx_t n, const uint8_t *x, idx_t k,
-                                int32_t *distances, idx_t *labels,
-                                uint8_t *recons) const override;
+    void search_and_reconstruct(
+            idx_t n,
+            const uint8_t* x,
+            idx_t k,
+            int32_t* distances,
+            idx_t* labels,
+            uint8_t* recons) const override;
 
     /** Reconstruct a vector given the location in terms of (inv list index +
      * inv list offset) instead of the id.
@@ -148,9 +176,10 @@ struct IndexBinaryIVF : IndexBinary {
      * the inv list offset is computed by search_preassigned() with
      * `store_pairs` set.
      */
-    virtual void reconstruct_from_offset(idx_t list_no, idx_t offset,
-                                         uint8_t* recons) const;
-
+    virtual void reconstruct_from_offset(
+            idx_t list_no,
+            idx_t offset,
+            uint8_t* recons) const;
 
     /// Dataset manipulation functions
     size_t remove_ids(const IDSelector& sel) override;
@@ -160,34 +189,33 @@ struct IndexBinaryIVF : IndexBinary {
      * sequential ids, this would be this->ntotal */
     virtual void merge_from(IndexBinaryIVF& other, idx_t add_id);
 
-    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 make_direct_map(bool new_maintain_direct_map = true);
 
-    void set_direct_map_type (DirectMap::Type type);
+    void set_direct_map_type(DirectMap::Type type);
 
-    void replace_invlists(InvertedLists *il, bool own=false);
+    void replace_invlists(InvertedLists* il, bool own = false);
 };
 
-
 struct BinaryInvertedListScanner {
-
     using idx_t = Index::idx_t;
 
     /// from now on we handle this query.
-    virtual void set_query (const uint8_t *query_vector) = 0;
+    virtual void set_query(const uint8_t* query_vector) = 0;
 
     /// following codes come from this inverted list
-    virtual void set_list (idx_t list_no, uint8_t coarse_dis) = 0;
+    virtual void set_list(idx_t list_no, uint8_t coarse_dis) = 0;
 
     /// compute a single query-to-code distance
-    virtual uint32_t distance_to_code (const uint8_t *code) const = 0;
+    virtual uint32_t distance_to_code(const uint8_t* code) const = 0;
 
     /** compute the distances to codes. (distances, labels) should be
      * organized as a min- or max-heap
@@ -199,23 +227,24 @@ struct BinaryInvertedListScanner {
      * @param labels     heap labels (size k)
      * @param k          heap size
      */
-    virtual size_t scan_codes (size_t n,
-                               const uint8_t *codes,
-                               const idx_t *ids,
-                               int32_t *distances, idx_t *labels,
-                               size_t k) const = 0;
-
-    virtual void scan_codes_range (size_t n,
-                                   const uint8_t *codes,
-                                   const idx_t *ids,
-                                   int radius,
-                                   RangeQueryResult &result) const = 0;
-
-    virtual ~BinaryInvertedListScanner () {}
-
+    virtual size_t scan_codes(
+            size_t n,
+            const uint8_t* codes,
+            const idx_t* ids,
+            int32_t* distances,
+            idx_t* labels,
+            size_t k) const = 0;
+
+    virtual void scan_codes_range(
+            size_t n,
+            const uint8_t* codes,
+            const idx_t* ids,
+            int radius,
+            RangeQueryResult& result) const = 0;
+
+    virtual ~BinaryInvertedListScanner() {}
 };
 
+} // namespace faiss
 
-}  // namespace faiss
-
-#endif  // FAISS_INDEX_BINARY_IVF_H
+#endif // FAISS_INDEX_BINARY_IVF_H