Added getNumElements + getVector

java/src/main/java/com/spotify/voyager/jni/StringIndex.java

@@ -253,9 +253,15 @@ public void saveIndex(OutputStream indexOutputStream, OutputStream namesListOutp
   }
 
   public void addItem(String name, float[] vector) {
-    int nextIndex = names.size();
-    index.addItem(vector, nextIndex);
-    names.add(name);
+    int nextIndex = names.indexOf(name);
+    if (nextIndex == -1) {
+      // we add it
+      index.addItem(vector, names.size());
+      names.add(name);
+    } else {
+      // up just update vector
+      index.addItem(vector, nextIndex);
+    }
   }
 
   public void addItem(String name, List<Float> vector) {
@@ -272,15 +278,28 @@ public void addItems(Map<String, List<Float>> vectors) {
     Iterator<Entry<String, List<Float>>> iterator = vectors.entrySet().iterator();
     for (int i = 0; i < numVectors; i++) {
       Entry<String, List<Float>> nextVector = iterator.next();
-      newNames.add(nextVector.getKey());
+      // Check if we already have this vector - and update it if needed
+      int nextIndex = names.indexOf(nextVector.getKey());
+      if (nextIndex == -1) {
+        newNames.add(nextVector.getKey());
+        nextIndex = names.size() + i;
+      }
       assignPrimitive(nextVector.getValue(), primitiveVectors[i]);
-      labels[i] = names.size() + i;
+      labels[i] = nextIndex;
     }
 
     names.addAll(newNames);
     index.addItems(primitiveVectors, labels, -1);
   }
 
+  public long getNumElements() {
+    return index.getNumElements();
+  }
+
+  public float[] getVector(String name) {
+    return index.getVector(names.indexOf(name));
+  }
+
   private float[] toPrimitive(List<Float> vector) {
     float[] vectorValues = new float[vector.size()];
     assignPrimitive(vector, vectorValues);

java/src/test/java/com/spotify/voyager/jni/StringIndexTest.java

@@ -261,6 +261,55 @@ public void itResizesIndex() throws Exception {
     }
   }
 
+  @Test
+  public void itUpsertItem() throws Exception {
+    List<Vector> testVectors = TestUtils.getTestVectors();
+    try (final StringIndex index =
+        new StringIndex(
+            SpaceType.Cosine,
+            testVectors.get(0).vector.length,
+            20,
+            2,
+            0,
+            2,
+            StorageDataType.E4M3)) {
+      Vector v1 = testVectors.get(0);
+      Vector v2 = testVectors.get(1);
+      // Add a couple of vectors//
+      index.addItem(v1.name, v1.vector);
+      index.addItem(v2.name, v2.vector);
+      // Again add the first one which should be an update not an insert with the same name
+      index.addItem(v1.name, v1.vector);
+      // We should only have 2 entries in it.
+      assertEquals(2, index.getNumElements());
+    }
+  }
+
+  @Test
+  public void itUpsertItems() throws Exception {
+    int testSize = 15;
+    List<Vector> testVectors = TestUtils.getTestVectors();
+    try (final StringIndex index =
+        new StringIndex(
+            SpaceType.Cosine,
+            testVectors.get(0).vector.length,
+            20,
+            testSize,
+            0,
+            testSize,
+            StorageDataType.E4M3)) {
+      Map<String, List<Float>> vectors =
+          testVectors.stream()
+              .limit(testSize)
+              .collect(Collectors.toMap(vec -> vec.name, vec -> convert(vec.vector)));
+      index.addItems(vectors);
+      assertEquals(vectors.size(), index.getNumElements());
+      // If we add then again number of elements  should stay the same//
+      index.addItems(vectors);
+      assertEquals(vectors.size(), index.getNumElements());
+    }
+  }
+
   public static class CustomResult {
     private final String name;
     private final float distance;