Merge pull request #8 from wy4515/master

Throw exception if init with wrong dimension

src/main/java/com/spotify/annoy/ANNIndex.java

@@ -88,13 +88,18 @@ private void load(final String filename) throws IOException {
     int buffIndex =  (numNodes - 1) / MAX_NODES_IN_BUFFER;
     int rest = (int) (fileSize % BLOCK_SIZE);
     int blockSize = (rest > 0 ? rest : BLOCK_SIZE);
+    // Two valid relations between dimension and file size:
+    // 1) rest % NODE_SIZE == 0 makes sure either everything fits into buffer or rest is a multiple of NODE_SIZE;
+    // 2) (file_size - rest) % NODE_SIZE == 0 makes sure everything else is a multiple of NODE_SIZE.
+    if (rest % NODE_SIZE != 0 || (fileSize - rest) % NODE_SIZE != 0) {
+      throw new RuntimeException("ANNIndex initiated with wrong dimension size");
+    }
     long position = fileSize - blockSize;
-
     buffers = new MappedByteBuffer[buffIndex + 1];
     boolean process = true;
     int m = -1;
     long index = fileSize;
-    while(position >= 0) {
+    while (position >= 0) {
       MappedByteBuffer annBuf = memoryMappedFile.getChannel().map(
               FileChannel.MapMode.READ_ONLY, position, blockSize);
       annBuf.order(ByteOrder.LITTLE_ENDIAN);
@@ -232,6 +237,11 @@ private static boolean isZeroVec(float[] v) {
   public final List<Integer> getNearest(final float[] queryVector,
                                         final int nResults) {
 
+    if (queryVector.length != DIMENSION) {
+      throw new RuntimeException(String.format("queryVector must be size of %d, but was %d",
+              DIMENSION, queryVector.length));
+    }
+
     PriorityQueue<PQEntry> pq = new PriorityQueue<>(
             roots.size() * FLOAT_SIZE);
     final float kMaxPriority = 1e30f;
@@ -273,7 +283,6 @@ public final List<Integer> getNearest(final float[] queryVector,
     }
 
     ArrayList<PQEntry> sortedNNs = new ArrayList<PQEntry>();
-    int i = 0;
     for (int nn : nearestNeighbors) {
       float[] v = getItemVector(nn);
       if (!isZeroVec(v)) {
@@ -287,7 +296,7 @@ public final List<Integer> getNearest(final float[] queryVector,
     Collections.sort(sortedNNs);
 
     ArrayList<Integer> result = new ArrayList<>(nResults);
-    for (i = 0; i < nResults && i < sortedNNs.size(); i++) {
+    for (int i = 0; i < nResults && i < sortedNNs.size(); i++) {
       result.add((int) sortedNNs.get(i).nodeOffset);
     }
     return result;
@@ -315,6 +324,7 @@ else if (args[2].toLowerCase().equals("euclidean"))
 
     ANNIndex annIndex = new ANNIndex(dimension, indexPath, indexType);
 
+    // input vector
     float[] u = annIndex.getItemVector(queryItem);
     System.out.printf("vector[%d]: ", queryItem);
     for (float x : u) {

src/test/java/com/spotify/annoy/ANNIndexTest.java

@@ -104,4 +104,21 @@ public void testEuclideanMultipleBlocks() throws IOException {
     testIndex(IndexType.EUCLIDEAN, 1, false);
   }
 
+  @Test(expected = RuntimeException.class)
+  /**
+   Make sure wrong dimension size used to init ANNIndex will throw RuntimeException.
+   */
+  public void testLoadFileWithWrongDimension() throws IOException {
+    ANNIndex index = new ANNIndex(7, "src/test/resources/points.euclidean.annoy");
+  }
+
+  @Test(expected = RuntimeException.class)
+  /**
+   Make sure wrong dimension size throw exception in getNearest()
+   */
+  public void testGetNearesWithWrongDim() throws IOException {
+    ANNIndex index = new ANNIndex(8, "src/test/resources/points.angular.annoy", IndexType.ANGULAR);
+    float[] u = {0f, 1.0f, 0.2f, 0.1f, 0f, 1.0f, 0.2f, 0.1f, 1f};
+    index.getNearest(u, 10);
+  }
 }