instead of renumbering implicitly, let caller provide remapper

Author: Jonathan Ellis

jvector-base/src/main/java/io/github/jbellis/jvector/disk/OnDiskGraphIndex.java

@@ -26,6 +26,8 @@
 import java.io.DataOutput;
 import java.io.IOException;
 import java.io.UncheckedIOException;
+import java.util.HashSet;
+import java.util.function.Function;
 import java.util.stream.IntStream;
 
 public class OnDiskGraphIndex<T> implements GraphIndex<T>, AutoCloseable, Accountable
@@ -147,19 +149,28 @@ public void close() throws IOException {
         readerSupplier.close();
     }
 
-    // takes Graph and Vectors separately since I'm reluctant to introduce a Vectors reference
-    // to OnHeapGraphIndex just for this method.  Maybe that will end up the best solution,
-    // but I'm not sure yet.
-    public static <T> void write(GraphIndex<T> graph, RandomAccessVectorValues<T> vectors, DataOutput out) throws IOException {
+    /**
+     * @param graph the graph to write
+     * @param vectors the vectors associated with each node
+     * @param ordinalMapper A function that maps from the graph's ordinals to the ordinals in the output.
+     *                      For simple use cases this can be the identity function.  To deal with deleted nodes,
+     *                      or to renumber the nodes to match rows or documents elsewhere, you may provide
+     *                      something custom.  The mapper must map from the graph's ordinals
+     *                      [0..getMaxNodeId()], to the output's ordinals [0..size()), with no gaps and
+     *                      no duplicates.
+     * @param out the output to write to
+     */
+    public static <T> void write(GraphIndex<T> graph,
+                                 RandomAccessVectorValues<T> vectors,
+                                 Function<Integer, Integer> ordinalMapper,
+                                 DataOutput out)
+            throws IOException
+    {
         if (graph instanceof OnHeapGraphIndex) {
             var ohgi = (OnHeapGraphIndex<T>) graph;
             if (ohgi.getDeletedNodes().cardinality() > 0) {
                 throw new IllegalArgumentException("Run builder.cleanup() before writing the graph");
             }
-            if (ohgi.hasPurgedNodes()) {
-                vectors = new RenumberingVectorValues<>(ohgi, vectors);
-                graph = new RenumberingGraphIndex<>(ohgi);
-            }
         }
         var view = graph.getView();
 
@@ -170,15 +181,22 @@ public static <T> void write(GraphIndex<T> graph, RandomAccessVectorValues<T> ve
         out.writeInt(graph.maxDegree());
 
         // for each graph node, write the associated vector and its neighbors
-        for (int node = 0; node < graph.size(); node++) {
-            out.writeInt(node); // unnecessary, but a reasonable sanity check
-            Io.writeFloats(out, (float[]) vectors.vectorValue(node));
+        var newOrdinals = new HashSet<Integer>();
+        for (int originalOrdinal = 0; originalOrdinal <= graph.getMaxNodeId(); originalOrdinal++) {
+            if (!graph.containsNode(originalOrdinal)) {
+                continue;
+            }
+
+            int newOrdinal = ordinalMapper.apply(originalOrdinal);
+            newOrdinals.add(newOrdinal);
+            out.writeInt(newOrdinal); // unnecessary, but a reasonable sanity check
+            Io.writeFloats(out, (float[]) vectors.vectorValue(originalOrdinal));
 
-            var neighbors = view.getNeighborsIterator(node);
+            var neighbors = view.getNeighborsIterator(originalOrdinal);
             out.writeInt(neighbors.size());
             int n = 0;
             for ( ; n < neighbors.size(); n++) {
-                out.writeInt(neighbors.nextInt());
+                out.writeInt(ordinalMapper.apply(neighbors.nextInt()));
             }
             assert !neighbors.hasNext();
 
@@ -187,5 +205,16 @@ public static <T> void write(GraphIndex<T> graph, RandomAccessVectorValues<T> ve
                 out.writeInt(-1);
             }
         }
+
+        // verify that the provided mapper was well-behaved
+        if (newOrdinals.size() > graph.size()) {
+            throw new IllegalArgumentException("graph modified during write");
+        }
+        if (newOrdinals.size() < graph.size()) {
+            throw new IllegalArgumentException("ordinalMapper resulted in duplicate entries");
+        }
+        if (graph.size() > 0 && newOrdinals.stream().mapToInt(i -> i).max().getAsInt() != graph.size() - 1)  {
+            throw new IllegalArgumentException("ordinalMapper produced out-of-range entries");
+        }
     }
 }

jvector-base/src/main/java/io/github/jbellis/jvector/disk/RenumberingGraphIndex.java

@@ -68,6 +68,13 @@ default int getMaxNodeId() {
     return size();
   }
 
+  /**
+   * @return true iff the graph contains the node with the given ordinal id
+   */
+  default boolean containsNode(int nodeId) {
+    return nodeId >= 0 && nodeId < size();
+  }
+
   @Override
   void close() throws IOException;
 

jvector-base/src/main/java/io/github/jbellis/jvector/disk/RenumberingVectorValues.java

@@ -261,12 +261,8 @@ int[] rawNodes() {
     return nodes.keySet().stream().mapToInt(i -> (int) (long) i).toArray();
   }
 
-  public boolean hasPurgedNodes() {
-    return hasPurgedNodes;
-  }
-
-  public boolean containsNode(int i) {
-    return nodes.containsKey(i);
+  public boolean containsNode(int nodeId) {
+    return nodes.containsKey(nodeId);
   }
 
   public double getAverageShortEdges() {

jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndex.java

@@ -38,6 +38,7 @@
 import java.util.List;
 import java.util.Set;
 import java.util.concurrent.atomic.LongAdder;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
@@ -59,7 +60,7 @@ private static void testRecall(int M, int efConstruction, List<Boolean> diskOpti
         var graphPath = testDirectory.resolve("graph" + M + efConstruction + ds.name);
         try {
             try (var outputStream = new DataOutputStream(new BufferedOutputStream(Files.newOutputStream(graphPath)))) {
-                OnDiskGraphIndex.write(onHeapGraph, floatVectors, outputStream);
+                OnDiskGraphIndex.write(onHeapGraph, floatVectors, Function.identity(), outputStream);
             }
             try (var onDiskGraph = new CachingGraphIndex(new OnDiskGraphIndex<>(ReaderSupplierFactory.open(graphPath), 0))) {
                 int queryRuns = 2;

jvector-base/src/main/java/io/github/jbellis/jvector/graph/OnHeapGraphIndex.java

@@ -35,6 +35,7 @@
 import java.util.HashSet;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
 import java.util.stream.IntStream;
 
 public class SiftSmall {
@@ -62,7 +63,7 @@ public static void testRecall(ArrayList<float[]> baseVectors, ArrayList<float[]>
         var graphPath = testDirectory.resolve("graph_test");
         try {
             DataOutputStream outputFile = new DataOutputStream(new FileOutputStream(graphPath.toFile()));
-            OnDiskGraphIndex.write(onHeapGraph, ravv, outputFile);
+            OnDiskGraphIndex.write(onHeapGraph, ravv, Function.identity(), outputFile);
 
             var onDiskGraph = new CachingGraphIndex(new OnDiskGraphIndex<>(ReaderSupplierFactory.open(graphPath), 0));
 

jvector-examples/src/main/java/io/github/jbellis/jvector/example/Bench.java

@@ -35,6 +35,7 @@
 import java.nio.file.attribute.BasicFileAttributes;
 import java.util.*;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -117,7 +118,7 @@ public static byte[] randomVector8(Random random, int dim) {
     public static <T> void writeGraph(GraphIndex<T> graph, RandomAccessVectorValues<T> vectors, Path outputPath) throws IOException {
         try (var indexOutputWriter = openFileForWriting(outputPath))
         {
-            OnDiskGraphIndex.write(graph, vectors, indexOutputWriter);
+            OnDiskGraphIndex.write(graph, vectors, Function.identity(), indexOutputWriter);
             indexOutputWriter.flush();
         }
     }

jvector-examples/src/main/java/io/github/jbellis/jvector/example/SiftSmall.java

@@ -32,7 +32,10 @@
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
 
 import static io.github.jbellis.jvector.TestUtil.getNeighborNodes;
 import static org.junit.Assert.assertArrayEquals;
@@ -98,24 +101,34 @@ public void testRenumberingOnDelete() throws IOException {
         assertEquals(1, getNeighborNodes(originalView, 2).size());
         assertTrue(getNeighborNodes(originalView, 2).contains(1));
 
-        // check renumbered
-        var renumbered = new RenumberingGraphIndex<>(original);
-        var renumberedView = renumbered.getView();
-        assertEquals(2, renumbered.size());
-        // 0 -> 1
-        assertEquals(1, getNeighborNodes(renumberedView, 0).size());
-        assertTrue(getNeighborNodes(renumberedView, 0).contains(1));
-        // 1 -> 0
-        assertEquals(1, getNeighborNodes(renumberedView, 1).size());
-        assertTrue(getNeighborNodes(renumberedView, 1).contains(0));
-
-        // writing to disk should be the same as the renumbered
-        var outputPath = testDirectory.resolve("large_graph");
-        TestUtil.writeGraph(original, ravv, outputPath);
+        // create renumbering map
+        Map<Integer, Integer> oldToNewMap = new HashMap<>();
+        int nextOrdinal = 0;
+        for (int i = 0; i <= originalView.getMaxNodeId(); i++) {
+            if (original.containsNode(i)) {
+                oldToNewMap.put(i, nextOrdinal++);
+            }
+        }
+        assertEquals(2, oldToNewMap.size());
+        assertEquals(0, (int) oldToNewMap.get(1));
+        assertEquals(1, (int) oldToNewMap.get(2));
+
+        // write the graph
+        var outputPath = testDirectory.resolve("renumbered_graph");
+        try (var indexOutputWriter = TestUtil.openFileForWriting(outputPath))
+        {
+            OnDiskGraphIndex.write(original, ravv, oldToNewMap::get, indexOutputWriter);
+            indexOutputWriter.flush();
+        }
+        // check that written graph ordinals match the new ones
         try (var marr = new SimpleMappedReader(outputPath.toAbsolutePath().toString());
-             var onDiskGraph = new OnDiskGraphIndex<float[]>(marr::duplicate, 0))
+             var onDiskGraph = new OnDiskGraphIndex<float[]>(marr::duplicate, 0);
+             var onDiskView = onDiskGraph.getView())
         {
-            TestUtil.assertGraphEquals(renumbered, onDiskGraph);
+            // 0 -> 1
+            assertTrue(getNeighborNodes(onDiskView, 0).contains(1));
+            // 1 -> 0
+            assertTrue(getNeighborNodes(onDiskView, 1).contains(0));
         }
     }