diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/NodeRecordTask.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/NodeRecordTask.java index f58a9deb5..a9c665512 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/NodeRecordTask.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/NodeRecordTask.java @@ -189,19 +189,11 @@ private void callLegacy() throws Exception { // never called for them, so nothing here is a gap: the whole record extends // the current run with no flush needed. writer.writeInt(newOrdinal); - for (var feature : inlineFeatures) { - for (int i = 0; i < feature.featureSize(); i++) writer.writeByte(0); - } - writer.writeInt(0); // neighbor count - for (int n = 0; n < graph.getDegree(0); n++) writer.writeInt(-1); + writeOmittedFeaturesAndEmptyNeighbors(writer); continue; } - if (!graph.containsNode(originalOrdinal)) { - throw new IllegalStateException(String.format( - "Ordinal mapper mapped new ordinal %d to non-existing node %d", - newOrdinal, originalOrdinal)); - } + checkNodeExists(newOrdinal, originalOrdinal); // Ordinal: always owned. writer.writeInt(newOrdinal); @@ -225,24 +217,7 @@ private void callLegacy() throws Exception { } // Neighbor section: always owned — extends the current run. - var neighbors = view.getNeighborsIterator(0, originalOrdinal); - if (neighbors.size() > graph.getDegree(0)) { - throw new IllegalStateException(String.format( - "Node %d has more neighbors %d than max degree %d -- run Builder.cleanup()!", - originalOrdinal, neighbors.size(), graph.getDegree(0))); - } - writer.writeInt(neighbors.size()); - int n = 0; - for (; n < neighbors.size(); n++) { - int newNeighbor = ordinalMapper.oldToNew(neighbors.nextInt()); - if (newNeighbor < 0 || newNeighbor > ordinalMapper.maxOrdinal()) { - throw new IllegalStateException(String.format( - "Neighbor ordinal out of bounds: %d/%d", - newNeighbor, ordinalMapper.maxOrdinal())); - } - writer.writeInt(newNeighbor); - } - for (; n < graph.getDegree(0); n++) writer.writeInt(-1); + writeNeighborSection(writer, originalOrdinal); } // Final trailing run. @@ -281,39 +256,64 @@ private void buildFullRecord(ByteBufferIndexWriter writer, int newOrdinal) throw writer.writeInt(newOrdinal); if (originalOrdinal == OrdinalMapper.OMITTED) { - for (var feature : inlineFeatures) { - for (int i = 0; i < feature.featureSize(); i++) writer.writeByte(0); - } - writer.writeInt(0); - for (int n = 0; n < graph.getDegree(0); n++) writer.writeInt(-1); + writeOmittedFeaturesAndEmptyNeighbors(writer); } else { - if (!graph.containsNode(originalOrdinal)) { - throw new IllegalStateException(String.format( - "Ordinal mapper mapped new ordinal %d to non-existing node %d", - newOrdinal, originalOrdinal)); - } + checkNodeExists(newOrdinal, originalOrdinal); for (var feature : inlineFeatures) { feature.writeInline(writer, featureStateSuppliers.get(feature.id()).apply(originalOrdinal)); } - var neighbors = view.getNeighborsIterator(0, originalOrdinal); - if (neighbors.size() > graph.getDegree(0)) { + writeNeighborSection(writer, originalOrdinal); + } + } + + /** + * Writes zero bytes for every inline feature followed by an empty (all-padding) neighbor + * section, for an {@link OrdinalMapper#OMITTED} "hole" ordinal. The caller has already + * written the new ordinal itself. + */ + private void writeOmittedFeaturesAndEmptyNeighbors(ByteBufferIndexWriter writer) { + for (var feature : inlineFeatures) { + for (int i = 0; i < feature.featureSize(); i++) writer.writeByte(0); + } + writer.writeInt(0); // neighbor count + for (int n = 0; n < graph.getDegree(0); n++) writer.writeInt(-1); + } + + /** + * Validates that {@code originalOrdinal} (the value the ordinal mapper produced for + * {@code newOrdinal}) actually exists in the graph. + */ + private void checkNodeExists(int newOrdinal, int originalOrdinal) { + if (!graph.containsNode(originalOrdinal)) { + throw new IllegalStateException(String.format( + "Ordinal mapper mapped new ordinal %d to non-existing node %d", + newOrdinal, originalOrdinal)); + } + } + + /** + * Writes the neighbor count followed by each neighbor's remapped ordinal, padded with -1 + * up to the graph's max degree at level 0. + */ + private void writeNeighborSection(ByteBufferIndexWriter writer, int originalOrdinal) { + var neighbors = view.getNeighborsIterator(0, originalOrdinal); + if (neighbors.size() > graph.getDegree(0)) { + throw new IllegalStateException(String.format( + "Node %d has more neighbors %d than max degree %d -- run Builder.cleanup()!", + originalOrdinal, neighbors.size(), graph.getDegree(0))); + } + writer.writeInt(neighbors.size()); + int n = 0; + for (; n < neighbors.size(); n++) { + int newNeighbor = ordinalMapper.oldToNew(neighbors.nextInt()); + if (newNeighbor < 0 || newNeighbor > ordinalMapper.maxOrdinal()) { throw new IllegalStateException(String.format( - "Node %d has more neighbors %d than max degree %d -- run Builder.cleanup()!", - originalOrdinal, neighbors.size(), graph.getDegree(0))); - } - writer.writeInt(neighbors.size()); - int n = 0; - for (; n < neighbors.size(); n++) { - int newNeighbor = ordinalMapper.oldToNew(neighbors.nextInt()); - if (newNeighbor < 0 || newNeighbor > ordinalMapper.maxOrdinal()) { - throw new IllegalStateException(String.format( - "Neighbor ordinal out of bounds: %d/%d", - newNeighbor, ordinalMapper.maxOrdinal())); - } - writer.writeInt(newNeighbor); + "Neighbor ordinal out of bounds: %d/%d", + newNeighbor, ordinalMapper.maxOrdinal())); } - for (; n < graph.getDegree(0); n++) writer.writeInt(-1); + writer.writeInt(newNeighbor); } + for (; n < graph.getDegree(0); n++) writer.writeInt(-1); } /** A not-yet-fully-written buffer destined for a fixed file offset. */ diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java index b9287fffe..b19413ebc 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java @@ -17,6 +17,7 @@ package io.github.jbellis.jvector.example; import io.github.jbellis.jvector.disk.BufferedRandomAccessWriter; +import io.github.jbellis.jvector.disk.RandomAccessWriter; import io.github.jbellis.jvector.disk.ReaderSupplierFactory; import io.github.jbellis.jvector.example.benchmarks.AccuracyBenchmark; import io.github.jbellis.jvector.example.benchmarks.BenchmarkTablePrinter; @@ -494,18 +495,13 @@ private static BuilderWithSuppliers builderWithSuppliers(Set features var identityMapper = new OrdinalMapper.IdentityMapper(floatVectors.size() - 1); Map> suppliers = new EnumMap<>(FeatureId.class); - RandomAccessOnDiskGraphIndexWriter writer; - if (useParallelConstruction) { - var builder = new OnDiskParallelGraphIndexWriter.Builder(onHeapGraph, outPath); - builder.withMapper(identityMapper); - addFeaturesToBuilder(builder, features, onHeapGraph, floatVectors, pq, constructionMetrics, suppliers); - writer = builder.build(); - } else { - var builder = new OnDiskGraphIndexWriter.Builder(onHeapGraph, outPath); - builder.withMapper(identityMapper); - addFeaturesToBuilder(builder, features, onHeapGraph, floatVectors, pq, constructionMetrics, suppliers); - writer = builder.build(); - } + AbstractGraphIndexWriter.Builder builder = + useParallelConstruction + ? new OnDiskParallelGraphIndexWriter.Builder(onHeapGraph, outPath) + : new OnDiskGraphIndexWriter.Builder(onHeapGraph, outPath); + builder.withMapper(identityMapper); + addFeaturesToBuilder(builder, features, onHeapGraph, floatVectors, pq, constructionMetrics, suppliers); + RandomAccessOnDiskGraphIndexWriter writer = builder.build(); return new BuilderWithSuppliers(writer, suppliers); } diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestRandomAccessOnDiskGraphIndexWriter.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestRandomAccessOnDiskGraphIndexWriter.java index 23a8b7659..302d5e77a 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestRandomAccessOnDiskGraphIndexWriter.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestRandomAccessOnDiskGraphIndexWriter.java @@ -27,6 +27,8 @@ import io.github.jbellis.jvector.graph.disk.feature.FeatureId; import io.github.jbellis.jvector.graph.disk.feature.InlineVectors; import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; import org.junit.After; import org.junit.Before; import org.junit.Test; @@ -35,10 +37,18 @@ import java.nio.file.Files; import java.nio.file.Path; import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; import java.util.Map; +import java.util.Set; + +import static org.junit.Assert.assertEquals; @ThreadLeakScope(ThreadLeakScope.Scope.NONE) public class TestRandomAccessOnDiskGraphIndexWriter extends LuceneTestCase { + private static final VectorTypeSupport vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport(); + private Path testDirectory; @@ -132,4 +142,105 @@ private void testGraphWrite(boolean addHierarchy, boolean twoPhase) throws IOExc } } + @Test + public void testOrdinalHolesOnePhase() throws IOException { + testGraphWriteWithHoles(false); + } + + @Test + public void testOrdinalHolesTwoPhase() throws IOException { + testGraphWriteWithHoles(true); + } + + /** + * Exercises the OMITTED-ordinal ("hole") branch of NodeRecordTask for both the batched + * (one-phase) and legacy (two-phase / writeFeaturesInline) parallel write paths, using an + * OrdinalMapper that leaves gaps in the new ordinal space. + */ + private void testGraphWriteWithHoles(boolean twoPhase) throws IOException { + int dimension = 8; + int size = 6; + int maxConnections = 4; + int beamWidth = 20; + float alpha = 1.2f; + float neighborOverflow = 1.2f; + + var ravv = new ListRandomAccessVectorValues( + new ArrayList<>(TestUtil.createRandomVectors(size, dimension)), + dimension + ); + + var builder = new GraphIndexBuilder( + ravv, + VectorSimilarityFunction.COSINE, + maxConnections, + beamWidth, + neighborOverflow, + alpha, + false + ); + ImmutableGraphIndex graph = TestUtil.buildSequentially(builder, ravv); + + // Map old ordinals 0..size-1 onto a sparser new ordinal space, leaving holes at new + // ordinals 3, 6 and 7 (maxOrdinal is 8, so the disk-resident ordinal space is 0..8). + int[] newOrdinals = {0, 1, 2, 4, 5, 8}; + Map oldToNew = new HashMap<>(); + for (int i = 0; i < size; i++) { + oldToNew.put(i, newOrdinals[i]); + } + var mapper = new OrdinalMapper.MapMapper(oldToNew); + int maxOrdinal = mapper.maxOrdinal(); + Set holeOrdinals = new HashSet<>(List.of(3, 6, 7)); + + Path indexPath = testDirectory.resolve("graph_index_holes"); + + var suppliers = Feature.singleStateFactory( + FeatureId.INLINE_VECTORS, + nodeId -> new InlineVectors.State(ravv.getVector(nodeId)) + ); + + try (var writer = new OnDiskParallelGraphIndexWriter.Builder(graph, indexPath) + .withParallelDirectBuffers(true) + .with(new InlineVectors(ravv.dimension())) + .withMapper(mapper) + .build()) { + if (twoPhase) { + for (int oldOrdinal = 0; oldOrdinal < graph.size(0); oldOrdinal++) { + Map stateMap = Map.of( + FeatureId.INLINE_VECTORS, + new InlineVectors.State(ravv.getVector(oldOrdinal)) + ); + writer.writeFeaturesInline(mapper.oldToNew(oldOrdinal), stateMap); + } + writer.write(Map.of()); + } else { + writer.write(suppliers); + } + } + + try (var readerSupplier = new SimpleMappedReader.Supplier(indexPath)) { + var onDiskGraph = OnDiskGraphIndex.load(readerSupplier); + // The on-disk L0 record range is [0, maxOrdinal], regardless of the header's + // declared layer-0 size (which reflects the source graph's node count). + var view = onDiskGraph.getView(); + + for (int oldOrdinal = 0; oldOrdinal < size; oldOrdinal++) { + int newOrdinal = mapper.oldToNew(oldOrdinal); + assertEquals("vector mismatch at new ordinal " + newOrdinal, + ravv.getVector(oldOrdinal), view.getVector(newOrdinal)); + } + + var zero = vectorTypeSupport.createFloatVector(dimension); + for (int newOrdinal = 0; newOrdinal <= maxOrdinal; newOrdinal++) { + if (holeOrdinals.contains(newOrdinal)) { + assertEquals("hole at new ordinal " + newOrdinal + " should have no neighbors", + 0, view.getNeighborsIterator(0, newOrdinal).size()); + assertEquals("hole at new ordinal " + newOrdinal + " should have a zero vector", + zero, view.getVector(newOrdinal)); + } + } + view.close(); + } + } + }