faiss 0.2.3 → 0.2.5

data/vendor/faiss/faiss/impl/IDSelector.h

@@ -0,0 +1,135 @@
+/**
+ * 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.
+ */
+
+#pragma once
+
+#include <unordered_set>
+#include <vector>
+
+#include <faiss/Index.h>
+
+/** IDSelector is intended to define a subset of vectors to handle (for removal
+ * or as subset to search) */
+
+namespace faiss {
+
+/** Encapsulates a set of ids to handle. */
+struct IDSelector {
+    using idx_t = Index::idx_t;
+    virtual bool is_member(idx_t id) const = 0;
+    virtual ~IDSelector() {}
+};
+
+/** ids between [imin, imax) */
+struct IDSelectorRange : IDSelector {
+    idx_t imin, imax;
+
+    /// Assume that the ids to handle are sorted. In some cases this can speed
+    /// up processing
+    bool assume_sorted;
+
+    IDSelectorRange(idx_t imin, idx_t imax, bool assume_sorted = false);
+
+    bool is_member(idx_t id) const final;
+
+    /// for sorted ids, find the range of list indices where the valid ids are
+    /// stored
+    void find_sorted_ids_bounds(
+            size_t list_size,
+            const idx_t* ids,
+            size_t* jmin,
+            size_t* jmax) const;
+
+    ~IDSelectorRange() override {}
+};
+
+/** Simple array of elements
+ *
+ * is_member calls are very inefficient, but some operations can use the ids
+ * directly.
+ */
+struct IDSelectorArray : IDSelector {
+    size_t n;
+    const idx_t* ids;
+
+    /** Construct with an array of ids to process
+     *
+     * @param n number of ids to store
+     * @param ids elements to store. The pointer should remain valid during
+     *            IDSelectorArray's lifetime
+     */
+    IDSelectorArray(size_t n, const idx_t* ids);
+    bool is_member(idx_t id) const final;
+    ~IDSelectorArray() override {}
+};
+
+/** Ids from a set.
+ *
+ * Repetitions of ids in the indices set passed to the constructor does not hurt
+ * performance.
+ *
+ * The hash function used for the bloom filter and GCC's implementation of
+ * unordered_set are just the least significant bits of the id. This works fine
+ * for random ids or ids in sequences but will produce many hash collisions if
+ * lsb's are always the same
+ */
+struct IDSelectorBatch : IDSelector {
+    std::unordered_set<idx_t> set;
+
+    // Bloom filter to avoid accessing the unordered set if it is unlikely
+    // to be true
+    std::vector<uint8_t> bloom;
+    int nbits;
+    idx_t mask;
+
+    /** Construct with an array of ids to process
+     *
+     * @param n number of ids to store
+     * @param ids elements to store. The pointer can be released after
+     *            construction
+     */
+    IDSelectorBatch(size_t n, const idx_t* indices);
+    bool is_member(idx_t id) const final;
+    ~IDSelectorBatch() override {}
+};
+
+/** One bit per element. Constructed with a bitmap, size ceil(n / 8).
+ */
+struct IDSelectorBitmap : IDSelector {
+    size_t n;
+    const uint8_t* bitmap;
+
+    /** Construct with a binary mask
+     *
+     * @param n size of the bitmap array
+     * @param bitmap id will be selected iff id / 8 < n and bit number
+     *               (i%8) of bitmap[floor(i / 8)] is 1.
+     */
+    IDSelectorBitmap(size_t n, const uint8_t* bitmap);
+    bool is_member(idx_t id) const final;
+    ~IDSelectorBitmap() override {}
+};
+
+/** reverts the membership test of another selector */
+struct IDSelectorNot : IDSelector {
+    const IDSelector* sel;
+    IDSelectorNot(const IDSelector* sel) : sel(sel) {}
+    bool is_member(idx_t id) const final {
+        return !sel->is_member(id);
+    }
+    virtual ~IDSelectorNot() {}
+};
+
+/// selects all entries (useful for benchmarking)
+struct IDSelectorAll : IDSelector {
+    bool is_member(idx_t id) const final {
+        return true;
+    }
+    virtual ~IDSelectorAll() {}
+};
+
+} // namespace faiss