Add preassigned functions to contrib

Summary:
Adds the preassigned add and search python wrappers to contrib.
Adds the preassigned search for the binary case (was missing before).
Also adds a real test for that functionality.

Reviewed By: beauby

Differential Revision: D26560021

fbshipit-source-id: 330b715a9ed0073cfdadbfbcb1c23b10bed963a5

contrib/README.md

@@ -44,6 +44,9 @@ Tested in `tests/test_contrib_torch.py` (CPU) and `gpu/test/test_contrib_torch_g
 Functions to inspect C++ objects wrapped by SWIG. Most often this just means reading
 fields and converting them to the proper python array.
 
+### ivf_tools.py
+
+A few functions to override the coarse quantizer in IVF, providing additional flexibility for assignment.
 
 ### datasets.py
 

contrib/ivf_tools.py

@@ -0,0 +1,91 @@
+# 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.
+
+import numpy as np
+import faiss
+
+def add_preassigned(index_ivf, x, a, ids=None):
+    """
+    Add elements to an IVF index, where the assignment is already computed
+    """
+    n, d = x.shape
+    assert a.shape == (n, )
+    if isinstance(index_ivf, faiss.IndexBinaryIVF):
+        d *= 8
+    assert d == index_ivf.d
+    if ids is not None:
+        assert ids.shape == (n, )
+        ids = faiss.swig_ptr(ids)
+    index_ivf.add_core(
+        n, faiss.swig_ptr(x), ids, faiss.swig_ptr(a)
+    )
+
+
+def search_preassigned(index_ivf, xq, k, list_nos, coarse_dis=None):
+    """
+    Perform a search in the IVF index, with predefined lists to search into
+    """
+    n, d = xq.shape
+    if isinstance(index_ivf, faiss.IndexBinaryIVF):
+        d *= 8
+        dis_type = "int32"
+    else:
+        dis_type = "float32"
+
+    assert d == index_ivf.d
+    assert list_nos.shape == (n, index_ivf.nprobe)
+
+    # the coarse distances are used in IVFPQ with L2 distance and by_residual=True
+    # otherwise we provide dummy coarse_dis
+    if coarse_dis is None:
+        coarse_dis = np.zeros((n, index_ivf.nprobe), dtype=dis_type)
+    else:
+        assert coarse_dis.shape == (n, index_ivf.nprobe)
+
+    D = np.empty((n, k), dtype=dis_type)
+    I = np.empty((n, k), dtype='int64')
+
+    sp = faiss.swig_ptr
+    index_ivf.search_preassigned(
+        n, sp(xq), k,
+        sp(list_nos), sp(coarse_dis), sp(D), sp(I), False)
+    return D, I
+
+
+def range_search_preassigned(index_ivf, x, radius, list_nos, coarse_dis=None):
+    """
+    Perform a range search in the IVF index, with predefined lists to search into
+    """
+    n, d = x.shape
+    if isinstance(index_ivf, faiss.IndexBinaryIVF):
+        d *= 8
+        dis_type = "int32"
+    else:
+        dis_type = "float32"
+
+    # the coarse distances are used in IVFPQ with L2 distance and by_residual=True
+    # otherwise we provide dummy coarse_dis
+    if coarse_dis is None:
+        coarse_dis = np.empty((n, index_ivf.nprobe), dtype=dis_type)
+    else:
+        assert coarse_dis.shape == (n, index_ivf.nprobe)
+
+    assert d == index_ivf.d
+    assert list_nos.shape == (n, index_ivf.nprobe)
+
+    res = faiss.RangeSearchResult(n)
+    sp = faiss.swig_ptr
+
+    index_ivf.range_search_preassigned(
+        n, sp(x), radius,
+        sp(list_nos), sp(coarse_dis),
+        res
+    )
+    # get pointers and copy them
+    lims = faiss.rev_swig_ptr(res.lims, n + 1).copy()
+    num_results = int(lims[-1])
+    dist = faiss.rev_swig_ptr(res.distances, num_results).copy()
+    indices = faiss.rev_swig_ptr(res.labels, num_results).copy()
+    return lims, dist, indices

faiss/IndexBinaryIVF.cpp

@@ -656,6 +656,19 @@ void IndexBinaryIVF::range_search(
     t0 = getmillisecs();
     invlists->prefetch_lists(idx.get(), n * nprobe);
 
+    range_search_preassigned(n, x, radius, idx.get(), coarse_dis.get(), res);
+
+    indexIVF_stats.search_time += getmillisecs() - t0;
+}
+
+void IndexBinaryIVF::range_search_preassigned(
+        idx_t n,
+        const uint8_t* x,
+        int radius,
+        const idx_t* assign,
+        const int32_t* centroid_dis,
+        RangeSearchResult* res) const {
+    const size_t nprobe = std::min(nlist, this->nprobe);
     bool store_pairs = false;
     size_t nlistv = 0, ndis = 0;
 
@@ -671,7 +684,7 @@ void IndexBinaryIVF::range_search(
         all_pres[omp_get_thread_num()] = &pres;
 
         auto scan_list_func = [&](size_t i, size_t ik, RangeQueryResult& qres) {
-            idx_t key = idx[i * nprobe + ik]; /* select the list  */
+            idx_t key = assign[i * nprobe + ik]; /* select the list  */
             if (key < 0)
                 return;
             FAISS_THROW_IF_NOT_FMT(
@@ -688,7 +701,7 @@ void IndexBinaryIVF::range_search(
             InvertedLists::ScopedCodes scodes(invlists, key);
             InvertedLists::ScopedIds ids(invlists, key);
 
-            scanner->set_list(key, coarse_dis[i * nprobe + ik]);
+            scanner->set_list(key, assign[i * nprobe + ik]);
             nlistv++;
             ndis += list_size;
             scanner->scan_codes_range(
@@ -711,7 +724,6 @@ void IndexBinaryIVF::range_search(
     indexIVF_stats.nq += n;
     indexIVF_stats.nlist += nlistv;
     indexIVF_stats.ndis += ndis;
-    indexIVF_stats.search_time += getmillisecs() - t0;
 }
 
 IndexBinaryIVF::~IndexBinaryIVF() {

faiss/IndexBinaryIVF.h

@@ -131,6 +131,14 @@ struct IndexBinaryIVF : IndexBinary {
             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.

tests/test_contrib.py

@@ -11,6 +11,7 @@ import platform
 from faiss.contrib import datasets
 from faiss.contrib import inspect_tools
 from faiss.contrib import evaluation
+from faiss.contrib import ivf_tools
 
 from common import get_dataset_2
 try:
@@ -265,3 +266,118 @@ class TestRangeEval(unittest.TestCase):
 
             np.testing.assert_array_almost_equal(ref_precisions, precisions)
             np.testing.assert_array_almost_equal(ref_recalls, recalls)
+
+
+class TestPreassigned(unittest.TestCase):
+
+    def test_float(self):
+        ds = datasets.SyntheticDataset(128, 2000, 2000, 200)
+
+        d = ds.d
+        xt = ds.get_train()
+        xq = ds.get_queries()
+        xb = ds.get_database()
+
+        # define alternative quantizer on the 20 first dims of vectors
+        km = faiss.Kmeans(20, 50)
+        km.train(xt[:, :20].copy())
+        alt_quantizer = km.index
+
+        index = faiss.index_factory(d, "IVF50,PQ16np")
+        index.by_residual = False
+
+        # (optional) fake coarse quantizer
+        fake_centroids = np.zeros((index.nlist, index.d), dtype="float32")
+        index.quantizer.add(fake_centroids)
+
+        # train the PQ part
+        index.train(xt)
+
+        # add elements xb
+        a = alt_quantizer.search(xb[:, :20].copy(), 1)[1].ravel()
+        ivf_tools.add_preassigned(index, xb, a)
+
+        # search elements xq, increase nprobe, check 4 first results w/ groundtruth
+        prev_inter_perf = 0
+        for nprobe in 1, 10, 20:
+
+            index.nprobe = nprobe
+            a = alt_quantizer.search(xq[:, :20].copy(), index.nprobe)[1]
+            D, I = ivf_tools.search_preassigned(index, xq, 4, a)
+            inter_perf = (I == ds.get_groundtruth()[:, :4]).sum() / I.size
+            self.assertTrue(inter_perf >= prev_inter_perf)
+            prev_inter_perf = inter_perf
+
+        # test range search
+
+        index.nprobe = 20
+
+        a = alt_quantizer.search(xq[:, :20].copy(), index.nprobe)[1]
+
+        # just to find a reasonable radius
+        D, I = ivf_tools.search_preassigned(index, xq, 4, a)
+        radius = D.max() * 1.01
+
+        lims, DR, IR = ivf_tools.range_search_preassigned(index, xq, radius, a)
+
+        # with that radius the k-NN results are a subset of the range search results
+        for q in range(len(xq)):
+            l0, l1 = lims[q], lims[q + 1]
+            self.assertTrue(set(I[q]) <= set(IR[l0:l1]))
+
+    def test_binary(self):
+        ds = datasets.SyntheticDataset(128, 2000, 2000, 200)
+
+        d = ds.d
+        xt = ds.get_train()
+        xq = ds.get_queries()
+        xb = ds.get_database()
+
+        # define alternative quantizer on the 20 first dims of vectors (will be in float)
+        km = faiss.Kmeans(20, 50)
+        km.train(xt[:, :20].copy())
+        alt_quantizer = km.index
+
+        binarizer = faiss.index_factory(d, "ITQ,LSHt")
+        binarizer.train(xt)
+
+        xb_bin = binarizer.sa_encode(xb)
+        xq_bin = binarizer.sa_encode(xq)
+
+        index = faiss.index_binary_factory(d, "BIVF200")
+
+        fake_centroids = np.zeros((index.nlist, index.d // 8), dtype="uint8")
+        index.quantizer.add(fake_centroids)
+        index.is_trained = True
+
+        # add elements xb
+        a = alt_quantizer.search(xb[:, :20].copy(), 1)[1].ravel()
+        ivf_tools.add_preassigned(index, xb_bin, a)
+
+        # search elements xq, increase nprobe, check 4 first results w/ groundtruth
+        prev_inter_perf = 0
+        for nprobe in 1, 10, 20:
+
+            index.nprobe = nprobe
+            a = alt_quantizer.search(xq[:, :20].copy(), index.nprobe)[1]
+            D, I = ivf_tools.search_preassigned(index, xq_bin, 4, a)
+            inter_perf = (I == ds.get_groundtruth()[:, :4]).sum() / I.size
+            self.assertTrue(inter_perf >= prev_inter_perf)
+            prev_inter_perf = inter_perf
+
+        # test range search
+
+        index.nprobe = 20
+
+        a = alt_quantizer.search(xq[:, :20].copy(), index.nprobe)[1]
+
+        # just to find a reasonable radius
+        D, I = ivf_tools.search_preassigned(index, xq_bin, 4, a)
+        radius = int(D.max() + 1)
+
+        lims, DR, IR = ivf_tools.range_search_preassigned(index, xq_bin, radius, a)
+
+        # with that radius the k-NN results are a subset of the range search results
+        for q in range(len(xq)):
+            l0, l1 = lims[q], lims[q + 1]
+            self.assertTrue(set(I[q]) <= set(IR[l0:l1]))