Skip to content

Consolidate NodeRecordTask and test ordinal holes #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
jaipilot wants to merge 1 commit into
base: jaipilot-v8/jvector-695
Choose a base branch
from
Draft

Consolidate NodeRecordTask and test ordinal holes #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
108 changes: 54 additions & 54 deletions jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/NodeRecordTask.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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.
Expand Down Expand Up @@ -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. */
Expand Down
20 changes: 8 additions & 12 deletions jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -494,18 +495,13 @@ private static BuilderWithSuppliers builderWithSuppliers(Set<FeatureId> features
var identityMapper = new OrdinalMapper.IdentityMapper(floatVectors.size() - 1);
Map<FeatureId, IntFunction<Feature.State>> 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<? extends RandomAccessOnDiskGraphIndexWriter, RandomAccessWriter> 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);
}

Expand Down
111 changes: 111 additions & 0 deletions ...est/java/io/github/jbellis/jvector/graph/disk/TestRandomAccessOnDiskGraphIndexWriter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;


Expand Down Expand Up @@ -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<Integer, Integer> 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<Integer> 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<FeatureId, Feature.State> 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();
}
}

}