Merge pull request #17 from yatsukav/master

added dot metric index support

Author: Erik Bernhardsson

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

@@ -64,8 +64,8 @@ public ANNIndex(final int dimension,
  final int blockSize) throws IOException {
  DIMENSION = dimension;
  INDEX_TYPE = indexType;
- INDEX_TYPE_OFFSET = (INDEX_TYPE == IndexType.ANGULAR) ? 4 : 8;
- K_NODE_HEADER_STYLE = (INDEX_TYPE == IndexType.ANGULAR) ? 12 : 16;
+ INDEX_TYPE_OFFSET = INDEX_TYPE.getOffset();
+ K_NODE_HEADER_STYLE = INDEX_TYPE.getkNodeHeaderStyle();
  // we can store up to MIN_LEAF_SIZE children in leaf nodes (we put
  // them where the separating plane normally goes)
  this.MIN_LEAF_SIZE = DIMENSION + 2;
@@ -151,6 +151,10 @@ private float getNodeBias(final long nodeOffset) { // euclidean-only
  return getFloatInAnnBuf(nodeOffset + 4);
   }
 
+ private float getDotFactor(final long nodeOffset) { // dot-only
+ return getFloatInAnnBuf(nodeOffset + 12);
+  }
+
  public final float[] getItemVector(final int itemIndex) {
  return getNodeVector(itemIndex * NODE_SIZE);
   }
@@ -175,11 +179,19 @@ private static float euclideanDistance(final float[] u, final float[] v) {
  return norm(diff);
   }
 
- public static float cosineMargin(final float[] u, final float[] v) {
+ public static float dot(final float[] u, final float[] v) {
  double d = 0;
  for (int i = 0; i < u.length; i++)
  d += u[i] * v[i];
- return (float) (d / (norm(u) * norm(v)));
+ return (float) d;
+  }
+
+ public static float cosineMargin(final float[] u, final float[] v) {
+ return dot(u, v) / (norm(u) * norm(v));
+  }
+
+ public static float dotMargin(final float[] u, final float[] v, final float norm) {
+ return dot(u, v) + norm * norm;
   }
 
  public static float euclideanMargin(final float[] u,
@@ -271,9 +283,9 @@ public final List<Integer> getNearest(final float[] queryVector,
  nearestNeighbors.add(j);
         }
       } else {
- float margin = (INDEX_TYPE == IndexType.ANGULAR) ?
- cosineMargin(v, queryVector) :
- euclideanMargin(v, queryVector, getNodeBias(topNodeOffset));
+ float margin = (INDEX_TYPE == IndexType.ANGULAR) ? cosineMargin(v, queryVector)
+      : (INDEX_TYPE == IndexType.DOT) ? dotMargin(v, queryVector, getDotFactor(topNodeOffset))
+      : euclideanMargin(v, queryVector, getNodeBias(topNodeOffset));
  long childrenMemOffset = topNodeOffset + INDEX_TYPE_OFFSET;
  long lChild = NODE_SIZE * getIntInAnnBuf(childrenMemOffset);
  long rChild = NODE_SIZE * getIntInAnnBuf(childrenMemOffset + 4);
@@ -286,11 +298,10 @@ public final List<Integer> getNearest(final float[] queryVector,
  for (int nn : nearestNeighbors) {
  float[] v = getItemVector(nn);
  if (!isZeroVec(v)) {
- sortedNNs.add(
- new PQEntry((INDEX_TYPE == IndexType.ANGULAR) ?
- cosineMargin(v, queryVector) :
-                        -euclideanDistance(v, queryVector),
- nn));
+ float margin = (INDEX_TYPE == IndexType.ANGULAR) ? cosineMargin(v, queryVector)
+                     : (INDEX_TYPE == IndexType.DOT) ? dot(v, queryVector)
+                     : -euclideanDistance(v, queryVector);
+ sortedNNs.add(new PQEntry(margin, nn));
       }
     }
  Collections.sort(sortedNNs);
@@ -317,6 +328,8 @@ public static void main(final String[] args) throws IOException {
  IndexType indexType = null;                 // 2
  if (args[2].toLowerCase().equals("angular"))
  indexType = IndexType.ANGULAR;
+ else if (args[2].toLowerCase().equals("dot"))
+ indexType = IndexType.DOT;
  else if (args[2].toLowerCase().equals("euclidean"))
  indexType = IndexType.EUCLIDEAN;
  else throw new RuntimeException("wrong index type specified");

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

@@ -1,5 +1,23 @@
 package com.spotify.annoy;
 
 public enum IndexType {
- ANGULAR, EUCLIDEAN
+ ANGULAR(4, 12),
+ EUCLIDEAN(8, 16),
+ DOT(4, 16);
+
+ private final int offset;
+ private final int kNodeHeaderStyle;
+
+ IndexType(int offset, int kNodeHeaderStyle) {
+ this.offset = offset;
+ this.kNodeHeaderStyle = kNodeHeaderStyle;
+  }
+
+ public int getOffset() {
+ return offset;
+  }
+
+ public int getkNodeHeaderStyle() {
+ return kNodeHeaderStyle;
+  }
 }

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

@@ -70,6 +70,16 @@ public void testAngular() throws IOException {
  testIndex(IndexType.ANGULAR, 0, false);
   }
 
+ @Test
+ /**
+   Make sure that the NNs retrieved by the Java version match the
+   ones pre-computed by the C++ version of the Angular index
+   using the default block size (for files up to 2GB).
+   */
+ public void testDot() throws IOException {
+ testIndex(IndexType.DOT, 0, false);
+  }
+
 
  @Test
  /**
@@ -92,6 +102,17 @@ public void testAngularBlocks() throws IOException {
  testIndex(IndexType.ANGULAR, 1, false);
   }
 
+ @Test
+ /**
+   Make sure that the NNs retrieved by the Java version match the
+   ones pre-computed by the C++ version of the Angular index
+   simulating files larger than 2GB.
+   */
+ public void testDotBlocks() throws IOException {
+ testIndex(IndexType.DOT, 10, false);
+ testIndex(IndexType.DOT, 1, false);
+  }
+
 
  @Test
  /**

src/test/resources/makeindex.py

@@ -3,13 +3,17 @@
 from annoy import AnnoyIndex
 
 a = AnnoyIndex(8, 'a')
+d = AnnoyIndex(8, 'd')
 e = AnnoyIndex(8, 'e')
 for n, l in enumerate(open('points.csv')):
  x = [float(f) for f in l.strip().split(',')]
  a.add_item(n, x)
+ d.add_item(n, x)
  e.add_item(n, x)
 
 a.build(-1)
 a.save('points.angular.annoy')
+d.build(-1)
+d.save('points.dot.annoy')
 e.build(-1)
 e.save('points.euclidean.annoy')

src/test/resources/points.dot.ann.txt

@@ -0,0 +1,10 @@
+1443	1379,161,2218,3294,1343,1507,340,529,140,2100
+1240	3610,430,2995,10,2848,200,6,822,3673,2168
+818	430,2995,3610,10,200,429,822,3774,1250,3182
+1725	627,696,1081,440,1353,2779,2364,6,1088,429
+1290	1147,4066,3673,1343,1003,1566,340,3476,1353,2779
+2031	1450,1272,440,627,3920,886,3227,1396,1081,4117
+1117	430,2995,3610,200,10,822,429,3774,3776,2435
+1211	1195,1365,1493,3263,2246,1081,2826,2064,221,627
+1902	430,2995,3610,200,10,822,3774,429,1250,3776
+603	3294,2100,2218,1854,4066,1343,1379,140,161,3913

src/test/resources/points.dot.annoy

@@ -12,4 +12,5 @@ def do(indextype):
  print >> out, '%s\t%s' % (q_index, ','.join([str(n) for n in nns]))
 
 do('angular')
+do('dot')
 do('euclidean')