Fix OnDiskGraphIndex when ordinal remapping with holes

Solves #399. Increasing OnDiskGraphIndex to version 4 to serialize the needed elements.

Author: marianotepper

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

@@ -30,13 +30,15 @@ class CommonHeader {
     public final int dimension;
     public final int entryNode;
     public final int maxDegree;
+    public final int idUpperBound;
 
-    CommonHeader(int version, int size, int dimension, int entryNode, int maxDegree) {
+    CommonHeader(int version, int size, int dimension, int entryNode, int maxDegree, int idUpperBound) {
         this.version = version;
         this.size = size;
         this.dimension = dimension;
         this.entryNode = entryNode;
         this.maxDegree = maxDegree;
+        this.idUpperBound = idUpperBound;
     }
 
     void write(DataOutput out) throws IOException {
@@ -48,6 +50,10 @@ void write(DataOutput out) throws IOException {
         out.writeInt(dimension);
         out.writeInt(entryNode);
         out.writeInt(maxDegree);
+
+        if (version >= 4) {
+            out.writeInt(idUpperBound);
+        }
     }
 
     static CommonHeader load(RandomAccessReader reader) throws IOException {
@@ -65,10 +71,22 @@ static CommonHeader load(RandomAccessReader reader) throws IOException {
         int entryNode = reader.readInt();
         int maxDegree = reader.readInt();
 
-        return new CommonHeader(version, size, dimension, entryNode, maxDegree);
+        int idUpperBound = size;
+        if (version >= 4) {
+            idUpperBound = reader.readInt();
+        }
+
+        return new CommonHeader(version, size, dimension, entryNode, maxDegree, idUpperBound);
     }
 
     int size() {
-        return ((version >= 3 ? 2 : 0) + 4) * Integer.BYTES;
+        int size = 4;
+        if (version >= 3) {
+            size += 2;
+        }
+        if (version >= 4) {
+            size += 1;
+        }
+        return size * Integer.BYTES;
     }
 }

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

@@ -38,7 +38,6 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
-import java.util.stream.IntStream;
 
 /**
  * A class representing a graph index stored on disk. The base graph contains only graph structure.
@@ -49,7 +48,7 @@
  */
 public class OnDiskGraphIndex implements GraphIndex, AutoCloseable, Accountable
 {
-    public static final int CURRENT_VERSION = 3;
+    public static final int CURRENT_VERSION = 4;
     static final int MAGIC = 0xFFFF0D61; // FFFF to distinguish from old graphs, which should never start with a negative size "ODGI"
     static final VectorTypeSupport vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport();
     final ReaderSupplier readerSupplier;
@@ -58,6 +57,7 @@ public class OnDiskGraphIndex implements GraphIndex, AutoCloseable, Accountable
     final int maxDegree;
     final int dimension;
     final int entryNode;
+    final int idUpperBound;
     final int inlineBlockSize; // total size of all inline elements contributed by features
     final EnumMap<FeatureId, ? extends Feature> features;
     final EnumMap<FeatureId, Integer> inlineOffsets;
@@ -71,6 +71,7 @@ public class OnDiskGraphIndex implements GraphIndex, AutoCloseable, Accountable
         this.dimension = header.common.dimension;
         this.entryNode = header.common.entryNode;
         this.maxDegree = header.common.maxDegree;
+        this.idUpperBound = header.common.idUpperBound;
         this.features = header.features;
         this.neighborsOffset = neighborsOffset;
         var inlineBlockSize = 0;
@@ -120,9 +121,30 @@ public int maxDegree() {
     }
 
     @Override
-    public NodesIterator getNodes()
-    {
-        return NodesIterator.fromPrimitiveIterator(IntStream.range(0, size).iterator(), size);
+    public int getIdUpperBound() {
+        return idUpperBound;
+    }
+
+    @Override
+    public NodesIterator getNodes() {
+        try (var reader = readerSupplier.get()) {
+            int[] valid_nodes = new int[size];
+            int pos = 0;
+            for (int node = 0; node < getIdUpperBound(); node++) {
+                long node_offset = neighborsOffset +
+                        (node * ((long) Integer.BYTES // ids
+                                + inlineBlockSize // inline elements
+                                + (Integer.BYTES * (long) (maxDegree + 1)) // neighbor count + neighbors)
+                        ));
+                reader.seek(node_offset);
+                if (reader.readInt() != -1) {
+                    valid_nodes[pos++] = node;
+                }
+            }
+            return new NodesIterator.ArrayNodesIterator(valid_nodes, size);
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
     }
 
     @Override
@@ -250,11 +272,17 @@ public int entryNode() {
             return entryNode;
         }
 
+        @Override
+        public int getIdUpperBound() {
+            return idUpperBound;
+        }
+
         @Override
         public Bits liveNodes() {
             return Bits.ALL;
         }
 
+
         @Override
         public void close() throws IOException {
             reader.close();

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

@@ -69,7 +69,7 @@ private OnDiskGraphIndexWriter(RandomAccessWriter out,
         this.startOffset = startOffset;
 
         // create a mock Header to determine the correct size
-        var ch = new CommonHeader(version, 0, dimension, view.entryNode(), graph.maxDegree());
+        var ch = new CommonHeader(version, 0, dimension, view.entryNode(), graph.maxDegree(), 0);
         var placeholderHeader = new Header(ch, featureMap);
         this.headerSize = placeholderHeader.size();
     }
@@ -241,7 +241,8 @@ public synchronized void writeHeader() throws IOException {
                                             graph.size(),
                                             dimension,
                                             ordinalMapper.oldToNew(view.entryNode()),
-                                            graph.maxDegree());
+                                            graph.maxDegree(),
+                                            ordinalMapper.maxOrdinal() + 1);
         var header = new Header(commonHeader, featureMap);
         header.write(out);
         out.flush();

jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndex.java

@@ -22,6 +22,7 @@
 import io.github.jbellis.jvector.disk.SimpleMappedReader;
 import io.github.jbellis.jvector.graph.GraphIndexBuilder;
 import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues;
+import io.github.jbellis.jvector.graph.NodesIterator;
 import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
 import io.github.jbellis.jvector.graph.TestVectorGraph;
 import io.github.jbellis.jvector.quantization.PQVectors;
@@ -35,11 +36,7 @@
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
-import java.util.EnumSet;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 
 import static io.github.jbellis.jvector.TestUtil.getNeighborNodes;
 import static org.junit.Assert.*;
@@ -180,6 +177,13 @@ public void testReorderingWithHoles() throws IOException {
              var onDiskView = onDiskGraph.getView())
         {
             assertEquals(11, onDiskGraph.getIdUpperBound());
+
+            Set<Integer> nodesInGraph = new HashSet<>();
+            for (NodesIterator it = onDiskGraph.getNodes(); it.hasNext(); ) {
+                nodesInGraph.add(it.next());
+            }
+            assertEquals(nodesInGraph, Set.of(0, 2, 10));
+
             assertEquals(onDiskView.getVector(0), ravv.getVector(2));
             assertEquals(onDiskView.getVector(10), ravv.getVector(1));
             assertEquals(onDiskView.getVector(2), ravv.getVector(0));