lucene/core/src/test/org/apache/lucene/index/TestKnnGraph.java - lucene-solr - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one or more
  * contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.
  * The ASF licenses this file to You under the Apache License, Version 2.0
  * (the "License"); you may not use this file except in compliance with
  * the License.  You may obtain a copy of the License at
  *
  *     http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
 package org.apache.lucene.index;


 import java.io.IOException;
 import java.util.Arrays;
 import java.util.HashSet;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Set;

 import org.apache.lucene.document.Document;
 import org.apache.lucene.document.Field;
 import org.apache.lucene.document.KnnGraphField;
 import org.apache.lucene.document.StringField;
 import org.apache.lucene.search.GraphSearch;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.util.LuceneTestCase;

 /** Tests indexing of a knn-graph by KnnGraphWriter */
 public class TestKnnGraph extends LuceneTestCase {

   private static final String KNN_GRAPH_FIELD = "vector";
   private static final String KNN_GRAPH_NBR_FIELD = "vector$nbr";

   /**
    * Basic test of creating documents in a graph
    */
   public void testBasic() throws Exception {
     try (Directory dir = newDirectory();
          IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null))) {
       int numDoc = atLeast(10);
       int dimension = atLeast(3);
       float[][] values = new float[numDoc][];
       for (int i = 0; i < numDoc; i++) {
         if (random().nextBoolean()) {
           values[i] = new float[dimension];
           for (int j = 0; j < dimension; j++) {
             values[i][j] = random().nextFloat();
           }
         }
         add(iw, i, values[i]);
       }
       assertConsistentGraph(iw, dimension, values);
     }
   }

   /**
    * Verify that the graph properties are preserved when merging
    */
   public void testMerge() throws Exception {
     try (Directory dir = newDirectory();
          IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null))) {
       int numDoc = atLeast(100);
       int dimension = atLeast(10);
       float[][] values = new float[numDoc][];
       for (int i = 0; i < numDoc; i++) {
         if (random().nextBoolean()) {
           values[i] = new float[dimension];
           for (int j = 0; j < dimension; j++) {
             // FIXME why do all the distances look identical?
             values[i][j] = random().nextFloat();
           }
         }
         add(iw, i, values[i]);
         if (random().nextInt(10) == 3) {
           //System.out.println("commit");
           iw.commit();
         }
       }
       if (random().nextBoolean()) {
         iw.forceMerge(1);
       }
       assertConsistentGraph(iw, dimension, values);
     }
   }

   // TODO: testSorted
   // TODO: testDeletions

   /**
    * Verify that searching does something reasonable
    */
   public void testSearch() throws Exception {
     try (Directory dir = newDirectory();
          IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null))) {
       // Add a document for every cartesian point  in an NxN square so we can
       // easily know which are the nearest neighbors to every point. Insert by iterating
       // using a prime number that is not a divisor of N*N so that we will hit each point once,
       // and chosen so that points will be inserted in a deterministic
       // but somewhat distributed pattern
       int n = 5, stepSize = 17;
       float[][] values = new float[n * n][];
       int index = 0;
       for (int i = 0; i < values.length; i++) {
         // System.out.printf("%d: (%d, %d)\n", i, index % n, index / n);
         values[i] = new float[]{index % n, index / n};
         index = (index + stepSize) % (n * n);
         add(iw, i, values[i]);
         if (i == 13) {
           // create 2 segments
           iw.commit();
         }
       }
       //System.out.println("");
       // TODO: enable this randomness
       if (random().nextBoolean()) {
         iw.forceMerge(1);
       }
       assertConsistentGraph(iw, 2, values);
       try (DirectoryReader dr = DirectoryReader.open(iw)) {
         IndexSearcher searcher = new IndexSearcher(dr);
         // results are ordered by distance (descending) and docid (ascending);
         // This is the docid ordering:
         // column major, origin at upper left
         //  0 15  5 20 10
         //  3 18  8 23 13
         //  6 21 11  1 16
         //  9 24 14  4 19
         // 12  2 17  7 22

         // For this small graph it seems we can always get exact results with 2 probes
         assertResults(new int[]{11, 1, 8, 14, 21},
             GraphSearch.search(searcher, KNN_GRAPH_FIELD, 5, 2, new float[]{2, 2}));
         assertResults(new int[]{0, 3, 15, 18, 5},
             GraphSearch.search(searcher, KNN_GRAPH_FIELD, 5, 2, new float[]{0, 0}));
         assertResults(new int[]{15, 18, 0, 3, 5},
             GraphSearch.search(searcher, KNN_GRAPH_FIELD, 5, 2, new float[]{0.3f, 0.8f}));
       }
     }
   }

   private void assertResults(int[] expected, TopDocs topDocs) {
     assertEquals(expected.length, topDocs.scoreDocs.length);
     for (int i = 0; i < expected.length; i++) {
       assertEquals(expected[i], topDocs.scoreDocs[i].doc);
     }
   }

   private void assertConsistentGraph(IndexWriter iw, int dimension, float[][] values) throws IOException {
     float[] scratch = new float[dimension];
     try (DirectoryReader dr = DirectoryReader.open(iw)) {
       for (LeafReaderContext ctx: dr.leaves()) {
         LeafReader reader = ctx.reader();
         VectorDocValues vectorDocValues = VectorDocValues.get(reader, KNN_GRAPH_FIELD);
         SortedNumericDocValues neighbors = DocValues.getSortedNumeric(reader, KNN_GRAPH_NBR_FIELD);
         int[][] graph = new int[reader.maxDoc()][];
         boolean singleNodeGraph = false;
         int graphSize = 0;
         for (int i = 0; i < reader.maxDoc(); i++) {
           int id = Integer.parseInt(reader.document(i).get("id"));
           if (values[id] == null) {
             // documents without KnnGraphValues have no vectors or neighbors
             assertFalse("document " + id + " was not expected to have values", vectorDocValues.advanceExact(i));
             assertFalse(neighbors.advanceExact(i));
           } else {
             ++graphSize;
             // documents with KnnGraphValues have the expected vectors
             assertTrue("doc " + i + " has no vector value", vectorDocValues.advanceExact(i));
             vectorDocValues.vector(scratch);
             assertArrayEquals(values[id], scratch, 0f);
             // We collect neighbors for analysis below
             if (neighbors.advanceExact(i)) {
               graph[i] = new int[neighbors.docValueCount()];
               for (int j = 0; j < graph[i].length; j++) {
                 graph[i][j] = (int) neighbors.nextValue();
                 //System.out.println("" + i + " -> " + graph[i][j]);
               }
             } else {
               // graph must have a single node
               singleNodeGraph = true;
             }
           }
         }
         assertTrue(singleNodeGraph || graphSize != 1);
         if (graphSize > 0) {
           assertEquals(dimension, vectorDocValues.dimension());
         }
         // assert that the graph in each leaf is connected and undirected (ie links are reciprocated)
         assertReciprocal(graph);
         assertConnected(graph);
       }
     }
   }

   private void assertReciprocal(int[][] graph) {
     // The graph is undirected: if a -> b then b -> a.
     for (int i = 0; i < graph.length; i++) {
       if (graph[i] != null) {
         for (int j = 0; j < graph[i].length; j++) {
           int k = graph[i][j];
           assertTrue("" + i + "->" + k + " is not reciprocated", Arrays.binarySearch(graph[k], i) >= 0);
         }
       }
     }
   }

   private void assertConnected(int[][] graph) {
     // every node in the graph is reachable from every other node
     Set<Integer> visited = new HashSet<>();
     List<Integer> queue = new LinkedList<>();
     int count = 0;
     for (int[] entry : graph) {
       if (entry != null) {
         if (queue.isEmpty()) {
           queue.add(entry[0]); // start from any node
           //System.out.println("start at " + entry[0]);
         }
         ++count;
       }
     }
     while(queue.isEmpty() == false) {
       int i = queue.remove(0);
       assertNotNull("expected neighbors of " + i, graph[i]);
       visited.add(i);
       for (int j : graph[i]) {
         if (visited.contains(j) == false) {
           //System.out.println("  ... " + j);
           queue.add(j);
         }
       }
     }
     // we visited each node exactly once
     assertEquals(count, visited.size());
   }


   private void add(IndexWriter iw, int id, float[] vector) throws IOException {
     Document doc = new Document();
     if (vector != null) {
       doc.add(new KnnGraphField(KNN_GRAPH_FIELD, vector));
     }
     doc.add(new StringField("id", Integer.toString(id), Field.Store.YES));
     //System.out.println("add " + id + " " + vector);
     iw.addDocument(doc);
   }

 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one or more
	* contributor license agreements. See the NOTICE file distributed with
	* this work for additional information regarding copyright ownership.
	* The ASF licenses this file to You under the Apache License, Version 2.0
	* (the "License"); you may not use this file except in compliance with
	* the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software
	* distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	* See the License for the specific language governing permissions and
	* limitations under the License.
	*/
	package org.apache.lucene.index;


	import java.io.IOException;
	import java.util.Arrays;
	import java.util.HashSet;
	import java.util.LinkedList;
	import java.util.List;
	import java.util.Set;

	import org.apache.lucene.document.Document;
	import org.apache.lucene.document.Field;
	import org.apache.lucene.document.KnnGraphField;
	import org.apache.lucene.document.StringField;
	import org.apache.lucene.search.GraphSearch;
	import org.apache.lucene.search.IndexSearcher;
	import org.apache.lucene.search.TopDocs;
	import org.apache.lucene.store.Directory;
	import org.apache.lucene.util.LuceneTestCase;

	/** Tests indexing of a knn-graph by KnnGraphWriter */
	public class TestKnnGraph extends LuceneTestCase {

	private static final String KNN_GRAPH_FIELD = "vector";
	private static final String KNN_GRAPH_NBR_FIELD = "vector$nbr";

	/**
	* Basic test of creating documents in a graph
	*/
	public void testBasic() throws Exception {
	try (Directory dir = newDirectory();
	IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null))) {
	int numDoc = atLeast(10);
	int dimension = atLeast(3);
	float[][] values = new float[numDoc][];
	for (int i = 0; i < numDoc; i++) {
	if (random().nextBoolean()) {
	values[i] = new float[dimension];
	for (int j = 0; j < dimension; j++) {
	values[i][j] = random().nextFloat();
	}
	}
	add(iw, i, values[i]);
	}
	assertConsistentGraph(iw, dimension, values);
	}
	}

	/**
	* Verify that the graph properties are preserved when merging
	*/
	public void testMerge() throws Exception {
	try (Directory dir = newDirectory();
	IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null))) {
	int numDoc = atLeast(100);
	int dimension = atLeast(10);
	float[][] values = new float[numDoc][];
	for (int i = 0; i < numDoc; i++) {
	if (random().nextBoolean()) {
	values[i] = new float[dimension];
	for (int j = 0; j < dimension; j++) {
	// FIXME why do all the distances look identical?
	values[i][j] = random().nextFloat();
	}
	}
	add(iw, i, values[i]);
	if (random().nextInt(10) == 3) {
	//System.out.println("commit");
	iw.commit();
	}
	}
	if (random().nextBoolean()) {
	iw.forceMerge(1);
	}
	assertConsistentGraph(iw, dimension, values);
	}
	}

	// TODO: testSorted
	// TODO: testDeletions

	/**
	* Verify that searching does something reasonable
	*/
	public void testSearch() throws Exception {
	try (Directory dir = newDirectory();
	IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null))) {
	// Add a document for every cartesian point in an NxN square so we can
	// easily know which are the nearest neighbors to every point. Insert by iterating
	// using a prime number that is not a divisor of N*N so that we will hit each point once,
	// and chosen so that points will be inserted in a deterministic
	// but somewhat distributed pattern
	int n = 5, stepSize = 17;
	float[][] values = new float[n * n][];
	int index = 0;
	for (int i = 0; i < values.length; i++) {
	// System.out.printf("%d: (%d, %d)\n", i, index % n, index / n);
	values[i] = new float[]{index % n, index / n};
	index = (index + stepSize) % (n * n);
	add(iw, i, values[i]);
	if (i == 13) {
	// create 2 segments
	iw.commit();
	}
	}
	//System.out.println("");
	// TODO: enable this randomness
	if (random().nextBoolean()) {
	iw.forceMerge(1);
	}
	assertConsistentGraph(iw, 2, values);
	try (DirectoryReader dr = DirectoryReader.open(iw)) {
	IndexSearcher searcher = new IndexSearcher(dr);
	// results are ordered by distance (descending) and docid (ascending);
	// This is the docid ordering:
	// column major, origin at upper left
	// 0 15 5 20 10
	// 3 18 8 23 13
	// 6 21 11 1 16
	// 9 24 14 4 19
	// 12 2 17 7 22

	// For this small graph it seems we can always get exact results with 2 probes
	assertResults(new int[]{11, 1, 8, 14, 21},
	GraphSearch.search(searcher, KNN_GRAPH_FIELD, 5, 2, new float[]{2, 2}));
	assertResults(new int[]{0, 3, 15, 18, 5},
	GraphSearch.search(searcher, KNN_GRAPH_FIELD, 5, 2, new float[]{0, 0}));
	assertResults(new int[]{15, 18, 0, 3, 5},
	GraphSearch.search(searcher, KNN_GRAPH_FIELD, 5, 2, new float[]{0.3f, 0.8f}));
	}
	}
	}

	private void assertResults(int[] expected, TopDocs topDocs) {
	assertEquals(expected.length, topDocs.scoreDocs.length);
	for (int i = 0; i < expected.length; i++) {
	assertEquals(expected[i], topDocs.scoreDocs[i].doc);
	}
	}

	private void assertConsistentGraph(IndexWriter iw, int dimension, float[][] values) throws IOException {
	float[] scratch = new float[dimension];
	try (DirectoryReader dr = DirectoryReader.open(iw)) {
	for (LeafReaderContext ctx: dr.leaves()) {
	LeafReader reader = ctx.reader();
	VectorDocValues vectorDocValues = VectorDocValues.get(reader, KNN_GRAPH_FIELD);
	SortedNumericDocValues neighbors = DocValues.getSortedNumeric(reader, KNN_GRAPH_NBR_FIELD);
	int[][] graph = new int[reader.maxDoc()][];
	boolean singleNodeGraph = false;
	int graphSize = 0;
	for (int i = 0; i < reader.maxDoc(); i++) {
	int id = Integer.parseInt(reader.document(i).get("id"));
	if (values[id] == null) {
	// documents without KnnGraphValues have no vectors or neighbors
	assertFalse("document " + id + " was not expected to have values", vectorDocValues.advanceExact(i));
	assertFalse(neighbors.advanceExact(i));
	} else {
	++graphSize;
	// documents with KnnGraphValues have the expected vectors
	assertTrue("doc " + i + " has no vector value", vectorDocValues.advanceExact(i));
	vectorDocValues.vector(scratch);
	assertArrayEquals(values[id], scratch, 0f);
	// We collect neighbors for analysis below
	if (neighbors.advanceExact(i)) {
	graph[i] = new int[neighbors.docValueCount()];
	for (int j = 0; j < graph[i].length; j++) {
	graph[i][j] = (int) neighbors.nextValue();
	//System.out.println("" + i + " -> " + graph[i][j]);
	}
	} else {
	// graph must have a single node
	singleNodeGraph = true;
	}
	}
	}
	assertTrue(singleNodeGraph \|\| graphSize != 1);
	if (graphSize > 0) {
	assertEquals(dimension, vectorDocValues.dimension());
	}
	// assert that the graph in each leaf is connected and undirected (ie links are reciprocated)
	assertReciprocal(graph);
	assertConnected(graph);
	}
	}
	}

	private void assertReciprocal(int[][] graph) {
	// The graph is undirected: if a -> b then b -> a.
	for (int i = 0; i < graph.length; i++) {
	if (graph[i] != null) {
	for (int j = 0; j < graph[i].length; j++) {
	int k = graph[i][j];
	assertTrue("" + i + "->" + k + " is not reciprocated", Arrays.binarySearch(graph[k], i) >= 0);
	}
	}
	}
	}

	private void assertConnected(int[][] graph) {
	// every node in the graph is reachable from every other node
	Set<Integer> visited = new HashSet<>();
	List<Integer> queue = new LinkedList<>();
	int count = 0;
	for (int[] entry : graph) {
	if (entry != null) {
	if (queue.isEmpty()) {
	queue.add(entry[0]); // start from any node
	//System.out.println("start at " + entry[0]);
	}
	++count;
	}
	}
	while(queue.isEmpty() == false) {
	int i = queue.remove(0);
	assertNotNull("expected neighbors of " + i, graph[i]);
	visited.add(i);
	for (int j : graph[i]) {
	if (visited.contains(j) == false) {
	//System.out.println(" ... " + j);
	queue.add(j);
	}
	}
	}
	// we visited each node exactly once
	assertEquals(count, visited.size());
	}


	private void add(IndexWriter iw, int id, float[] vector) throws IOException {
	Document doc = new Document();
	if (vector != null) {
	doc.add(new KnnGraphField(KNN_GRAPH_FIELD, vector));
	}
	doc.add(new StringField("id", Integer.toString(id), Field.Store.YES));
	//System.out.println("add " + id + " " + vector);
	iw.addDocument(doc);
	}

	}