Access graph structure for NSG (#2984)

Summary:
Pull Request resolved: #2984

It is not entirely trivial to access the NSG graph structure from Python (although it is a fixed size N-by-K matrix of vector ids).
This diff adds an inspect_tools function to do that.

Differential Revision: D48026775

fbshipit-source-id: 198202dfad4e77909c08761e6772cbaf7292283d

contrib/inspect_tools.py

@@ -96,3 +96,16 @@ def get_flat_data(index):
     """ copy and return the data matrix in an IndexFlat """
     xb = faiss.vector_to_array(index.codes).view("float32")
     return xb.reshape(index.ntotal, index.d)
+
+
+def get_NSG_neighbors(nsg):
+    """ get the neighbor list for the vectors stored in the NSG structure, as
+    a N-by-K matrix of indices """
+    graph = nsg.get_final_graph()
+    neighbors = np.zeros((graph.N, graph.K), dtype='int32')
+    faiss.memcpy(
+        faiss.swig_ptr(neighbors),
+        graph.data,
+        neighbors.nbytes
+    )
+    return neighbors

faiss/impl/NNDescent.cpp

@@ -374,6 +374,10 @@ void NNDescent::init_graph(DistanceComputer& qdis) {
 
 void NNDescent::build(DistanceComputer& qdis, const int n, bool verbose) {
     FAISS_THROW_IF_NOT_MSG(L >= K, "L should be >= K in NNDescent.build");
+    FAISS_THROW_IF_NOT_FMT(
+            n > NUM_EVAL_POINTS,
+            "NNDescent.build cannot build a graph smaller than %d",
+            int(NUM_EVAL_POINTS));
 
     if (verbose) {
         printf("Parameters: K=%d, S=%d, R=%d, L=%d, iter=%d\n",
@@ -403,7 +407,7 @@ void NNDescent::build(DistanceComputer& qdis, const int n, bool verbose) {
     has_built = true;
 
     if (verbose) {
-        printf("Addes %d points into the index\n", ntotal);
+        printf("Added %d points into the index\n", ntotal);
     }
 }
 

faiss/impl/NSG.h

@@ -54,7 +54,7 @@ namespace nsg {
 
 template <class node_t>
 struct Graph {
-    node_t* data;    ///< the flattened adjacency matrix
+    node_t* data;    ///< the flattened adjacency matrix, size N-by-K
     int K;           ///< nb of neighbors per node
     int N;           ///< total nb of nodes
     bool own_fields; ///< the underlying data owned by itself or not

faiss/python/swigfaiss.swig

@@ -454,7 +454,20 @@ void gpu_sync_all_devices()
 %include  <faiss/IndexNNDescent.h>
 
 %include  <faiss/IndexIVFFlat.h>
+
+%warnfilter(509) faiss::nsg::Graph< int >::at(int,int);
+
 %include  <faiss/impl/NSG.h>
+
+%template(NSG_Graph_int) faiss::nsg::Graph<int>;
+
+// not using %shared_ptr to avoid mem leaks
+%extend faiss::NSG {
+    faiss::nsg::Graph<int>* get_final_graph() {
+        return $self->final_graph.get();
+    }
+}
+
 %include  <faiss/IndexNSG.h>
 
 #ifndef SWIGWIN

tests/test_build_blocks.py

@@ -542,6 +542,23 @@ def subtest(self, d, K, metric):
         print('Metric: {}, knng accuracy: {}'.format(metric_names[metric], recall))
         assert recall > 0.99
 
+    def test_small_nndescent(self):
+        """ building a too small graph used to crash, make sure it raises
+        an exception instead.
+        TODO: build the exact knn graph for small cases
+        """
+        d = 32
+        K = 10
+        index = faiss.IndexNNDescentFlat(d, K, faiss.METRIC_L2)
+        index.nndescent.S = 10
+        index.nndescent.R = 32
+        index.nndescent.L = K + 20
+        index.nndescent.iter = 5
+        index.verbose = True
+
+        xb = np.zeros((78, d), dtype='float32')
+        self.assertRaises(RuntimeError, index.add, xb)
+
 
 class TestResultHeap(unittest.TestCase):
 

tests/test_contrib.py

@@ -219,6 +219,16 @@ def test_make_LT(self):
         Ynew = lt.apply(X)
         np.testing.assert_equal(Yref, Ynew)
 
+    def test_NSG_neighbors(self):
+        # FIXME number of elements to add should be >> 100
+        ds = datasets.SyntheticDataset(32, 0, 200, 10)
+        index = faiss.index_factory(ds.d, "NSG")
+        index.add(ds.get_database())
+        neighbors = inspect_tools.get_NSG_neighbors(index.nsg)
+        # neighbors should be either valid indexes or -1
+        np.testing.assert_array_less(-2, neighbors)
+        np.testing.assert_array_less(neighbors, ds.nb)
+
 
 class TestRangeEval(unittest.TestCase):