From fbf1badb2b702777ed18e181ec6936b06f733525 Mon Sep 17 00:00:00 2001 From: Mark Wolters Date: Thu, 23 Jul 2026 07:46:34 -0400 Subject: [PATCH 1/5] initial commit --- .../disk/AbstractGraphIndexSerializer.java | 221 ++++++++++++++++++ .../graph/disk/AbstractGraphIndexWriter.java | 145 +----------- .../graph/disk/GraphIndexMetadata.java | 119 ++++++++++ .../graph/disk/GraphIndexSerializer.java | 132 +++++++++++ .../disk/GraphIndexSerializerFactory.java | 103 ++++++++ .../graph/disk/GraphIndexSerializerV2.java | 111 +++++++++ .../graph/disk/GraphIndexSerializerV3.java | 119 ++++++++++ .../graph/disk/GraphIndexSerializerV4.java | 162 +++++++++++++ .../graph/disk/GraphIndexSerializerV5.java | 47 ++++ .../graph/disk/GraphIndexSerializerV6.java | 181 ++++++++++++++ 10 files changed, 1206 insertions(+), 134 deletions(-) create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexSerializer.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexMetadata.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializer.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerFactory.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV2.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV3.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV4.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV5.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV6.java diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexSerializer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexSerializer.java new file mode 100644 index 000000000..423110911 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexSerializer.java @@ -0,0 +1,221 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph.disk; + +import io.github.jbellis.jvector.disk.IndexWriter; +import io.github.jbellis.jvector.graph.ImmutableGraphIndex; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.graph.disk.feature.SeparatedFeature; + +import java.io.IOException; +import java.util.EnumSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.function.IntFunction; + +/** + * Abstract base class for graph index serializers providing common functionality. + */ +abstract class AbstractGraphIndexSerializer implements GraphIndexSerializer { + private final int version; + private final Set supportedFeatures; + private final boolean supportsMultiLayer; + private final boolean usesFooter; + private final FeatureOrdering featureOrdering; + + /** A magic number to indicate the file footer */ + public static final int FOOTER_MAGIC = 0x4a564244; + /** The size of the offset in the footer. */ + public static final int FOOTER_OFFSET_SIZE = Long.BYTES; + /** The size of the magic number in the footer. */ + public static final int FOOTER_MAGIC_SIZE = Integer.BYTES; + /** The total size of the footer. */ + public static final int FOOTER_SIZE = FOOTER_MAGIC_SIZE + FOOTER_OFFSET_SIZE; + + protected AbstractGraphIndexSerializer(int version, + Set supportedFeatures, + boolean supportsMultiLayer, + boolean usesFooter, + FeatureOrdering featureOrdering) { + this.version = version; + this.supportedFeatures = supportedFeatures; + this.supportsMultiLayer = supportsMultiLayer; + this.usesFooter = usesFooter; + this.featureOrdering = featureOrdering; + } + + @Override + public int getVersion() { + return version; + } + + @Override + public boolean supportsFeature(FeatureId feature) { + return supportedFeatures.contains(feature); + } + + @Override + public Set getSupportedFeatures() { return supportedFeatures; } + + @Override + public boolean supportsMultiLayer() { + return supportsMultiLayer; + } + + @Override + public boolean usesFooter() { + return usesFooter; + } + + @Override + public FeatureOrdering getFeatureOrdering() { + return featureOrdering; + } + + @Override + public void writeSparseLevels(ImmutableGraphIndex graph, OrdinalMapper ordinalMapper, IndexWriter out, Map> featureStateSuppliers) throws IOException { + for (int level = 1; level <= graph.getMaxLevel(); level++) { + int layerSize = graph.size(level); + int layerDegree = graph.getDegree(level); + int nodesWritten = 0; + for (var it = graph.getNodes(level); it.hasNext(); ) { + int originalOrdinal = it.nextInt(); + // node id + final int newOrdinal = ordinalMapper.oldToNew(originalOrdinal); + out.writeInt(newOrdinal); + // neighbors + var view = graph.getView(); + var neighbors = view.getNeighborsIterator(level, originalOrdinal); + out.writeInt(neighbors.size()); + int n = 0; + for ( ; n < neighbors.size(); n++) { + out.writeInt(ordinalMapper.oldToNew(neighbors.nextInt())); + } + assert !neighbors.hasNext() : "Mismatch between neighbor's reported size and actual size"; + // pad out to degree + for (; n < layerDegree; n++) { + out.writeInt(-1); + } + nodesWritten++; + } + if (nodesWritten != layerSize) { + throw new IllegalStateException("Mismatch between layer size and nodes written"); + } + } + } + + @Override + public synchronized void writeHeader(ImmutableGraphIndex graph, OrdinalMapper ordinalMapper, Map featureMap, int dimension, long startOffset, IndexWriter out, long headerSize) throws IOException { + // graph-level properties + var layerInfo = CommonHeader.LayerInfo.fromGraph(graph, ordinalMapper); + var commonHeader = new CommonHeader(version, + dimension, + graph.getView().entryNode() == null ? ImmutableGraphIndex.ENTRY_NODE_ABSENT : ordinalMapper.oldToNew(graph.getView().entryNode().node), + layerInfo, + ordinalMapper.maxOrdinal() + 1); + var header = new Header(commonHeader, featureMap); + header.write(out); + assert out.position() == startOffset + headerSize : String.format("%d != %d", out.position(), startOffset + headerSize); + } + + @Override + public long featureOffsetForOrdinal(ImmutableGraphIndex graph, List inlineFeatures, long startOffset, int ordinal, long headerSize) { + int edgeSize = Integer.BYTES * (1 + graph.getDegree(0)); + long inlineBytes = ordinal * (long) (Integer.BYTES + inlineFeatures.stream().mapToInt(Feature::featureSize).sum() + edgeSize); + return startOffset + + headerSize + + inlineBytes // previous nodes + + Integer.BYTES; // the ordinal of the node whose features we're about to write + } + + @Override + public void writeFooter(ImmutableGraphIndex graph, OrdinalMapper ordinalMapper, long headerOffset, int dimension, IndexWriter out, Map featureMap, long headerSize) throws IOException { + var layerInfo = CommonHeader.LayerInfo.fromGraph(graph, ordinalMapper); + var commonHeader = new CommonHeader(version, + dimension, + graph.getView().entryNode() == null ? ImmutableGraphIndex.ENTRY_NODE_ABSENT : ordinalMapper.oldToNew(graph.getView().entryNode().node), + layerInfo, + ordinalMapper.maxOrdinal() + 1); + var header = new Header(commonHeader, featureMap); + header.write(out); // write the header + out.writeLong(headerOffset); // We write the offset of the header at the end of the file + out.writeInt(FOOTER_MAGIC); + final long expectedPosition = headerOffset + headerSize + FOOTER_SIZE; + assert out.position() == expectedPosition : String.format("%d != %d", out.position(), expectedPosition); + } + + public void writeSeparatedFeatures(Map> featureStateSuppliers, Map featureMap, IndexWriter out, OrdinalMapper ordinalMapper) throws IOException { + for (var featureEntry : featureMap.entrySet()) { + if (isSeparated(featureEntry.getValue())) { + var fid = featureEntry.getKey(); + var supplier = featureStateSuppliers.get(fid); + if (supplier == null) { + throw new IllegalStateException("Supplier for feature " + fid + " not found"); + } + + // Set the offset for this feature + var feature = (SeparatedFeature) featureEntry.getValue(); + feature.setOffset(out.position()); + + // Write separated data for each node + for (int newOrdinal = 0; newOrdinal <= ordinalMapper.maxOrdinal(); newOrdinal++) { + int originalOrdinal = ordinalMapper.newToOld(newOrdinal); + if (originalOrdinal != OrdinalMapper.OMITTED) { + feature.writeSeparately(out, supplier.apply(originalOrdinal)); + } else { + // write zeros for missing data as padding + for (int i = 0; i < feature.featureSize(); i++) { + out.writeByte(0); + } + } + } + } + } + } + + boolean isSeparated(Feature feature) { + return feature instanceof SeparatedFeature; + } + + /** + * Helper to create a set of all known features. + * Only correct for the current maximum version — do not use for older versions. + * When a new version is introduced with a new feature, the previous version's serializer + * must be updated to use an explicit set that excludes the new feature. + */ + protected static Set allFeatures() { + return EnumSet.allOf(FeatureId.class); + } + + /** + * Helper to create a set of all features except FUSED_PQ. + * Used by versions 3–5, which support all non-fused features but predate fused PQ hierarchy + * support introduced in version 6. + */ + protected static Set nonFusedFeatures() { + return EnumSet.complementOf(EnumSet.of(FeatureId.FUSED_PQ)); + } + + /** + * Helper to create a set with only inline vectors (for version 2). + */ + protected static Set inlineVectorsOnly() { + return EnumSet.of(FeatureId.INLINE_VECTORS); + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexWriter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexWriter.java index 09d0d0ec0..d8c8ed407 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexWriter.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexWriter.java @@ -42,7 +42,7 @@ * Abstract base class for writing graph indexes to disk. * @param the type of the output writer */ -public abstract class AbstractGraphIndexWriter implements GraphIndexWriter { +public abstract class AbstractGraphIndexWriter implements GraphIndexWriter { /** A magic number to indicate the file footer */ public static final int FOOTER_MAGIC = 0x4a564244; /** The size of the offset in the footer. */ @@ -57,9 +57,11 @@ public abstract class AbstractGraphIndexWriter implements final int dimension; final Map featureMap; final T out; /* output for graph nodes and inline features */ - final int headerSize; + private final long headerSize; + volatile int maxOrdinalWritten = -1; final List inlineFeatures; + final GraphIndexSerializer serializer; AbstractGraphIndexWriter(T out, int version, @@ -68,7 +70,8 @@ public abstract class AbstractGraphIndexWriter implements int dimension, EnumMap features) { - if (graph.getMaxLevel() > 0 && version < 4) { + serializer = GraphIndexSerializerFactory.forVersion(version); + if (!serializer.supportsMultiLayer()) { throw new IllegalArgumentException("Multilayer graphs must be written with version 4 or higher"); } this.version = version; @@ -98,7 +101,6 @@ public abstract class AbstractGraphIndexWriter implements throw new IllegalArgumentException("Fused features require version 6 or higher"); } this.out = out; - // create a mock Header to determine the correct size var layerInfo = CommonHeader.LayerInfo.fromGraph(graph, ordinalMapper); var ch = new CommonHeader(version, dimension, 0, layerInfo, 0); @@ -123,16 +125,7 @@ public Set getFeatureSet() { } long featureOffsetForOrdinal(long startOffset, int ordinal) { - int edgeSize = Integer.BYTES * (1 + graph.getDegree(0)); - long inlineBytes = ordinal * (long) (Integer.BYTES + inlineFeatures.stream().mapToInt(Feature::featureSize).sum() + edgeSize); - return startOffset - + headerSize - + inlineBytes // previous nodes - + Integer.BYTES; // the ordinal of the node whose features we're about to write - } - - boolean isSeparated(Feature feature) { - return feature instanceof SeparatedFeature; + return serializer.featureOffsetForOrdinal(graph, inlineFeatures, startOffset, ordinal, headerSize); } /** @@ -172,18 +165,7 @@ public static Map sequentialRenumbering(ImmutableGraphIndex gr * @throws IOException IOException */ void writeFooter(ImmutableGraphIndex.View view, long headerOffset) throws IOException { - var layerInfo = CommonHeader.LayerInfo.fromGraph(graph, ordinalMapper); - var commonHeader = new CommonHeader(version, - dimension, - view.entryNode() == null ? ImmutableGraphIndex.ENTRY_NODE_ABSENT : ordinalMapper.oldToNew(view.entryNode().node), - layerInfo, - ordinalMapper.maxOrdinal() + 1); - var header = new Header(commonHeader, featureMap); - header.write(out); // write the header - out.writeLong(headerOffset); // We write the offset of the header at the end of the file - out.writeInt(FOOTER_MAGIC); - final long expectedPosition = headerOffset + headerSize + FOOTER_SIZE; - assert out.position() == expectedPosition : String.format("%d != %d", out.position(), expectedPosition); + serializer.writeFooter(graph, ordinalMapper, headerOffset, dimension, out, featureMap, headerSize); } /** @@ -194,120 +176,15 @@ void writeFooter(ImmutableGraphIndex.View view, long headerOffset) throws IOExce * @throws IOException if an I/O error occurs */ protected synchronized void writeHeader(ImmutableGraphIndex.View view, long startOffset) throws IOException { - // graph-level properties - var layerInfo = CommonHeader.LayerInfo.fromGraph(graph, ordinalMapper); - var commonHeader = new CommonHeader(version, - dimension, - view.entryNode() == null ? ImmutableGraphIndex.ENTRY_NODE_ABSENT : ordinalMapper.oldToNew(view.entryNode().node), - layerInfo, - ordinalMapper.maxOrdinal() + 1); - var header = new Header(commonHeader, featureMap); - header.write(out); - assert out.position() == startOffset + headerSize : String.format("%d != %d", out.position(), startOffset + headerSize); + serializer.writeHeader(graph, ordinalMapper, featureMap, dimension, startOffset, out, headerSize); } void writeSparseLevels(ImmutableGraphIndex.View view, Map> featureStateSuppliers) throws IOException { - // write sparse levels - for (int level = 1; level <= graph.getMaxLevel(); level++) { - int layerSize = graph.size(level); - int layerDegree = graph.getDegree(level); - int nodesWritten = 0; - for (var it = graph.getNodes(level); it.hasNext(); ) { - int originalOrdinal = it.nextInt(); - // node id - final int newOrdinal = ordinalMapper.oldToNew(originalOrdinal); - out.writeInt(newOrdinal); - // neighbors - var neighbors = view.getNeighborsIterator(level, originalOrdinal); - out.writeInt(neighbors.size()); - int n = 0; - for ( ; n < neighbors.size(); n++) { - out.writeInt(ordinalMapper.oldToNew(neighbors.nextInt())); - } - assert !neighbors.hasNext() : "Mismatch between neighbor's reported size and actual size"; - // pad out to degree - for (; n < layerDegree; n++) { - out.writeInt(-1); - } - nodesWritten++; - } - if (nodesWritten != layerSize) { - throw new IllegalStateException("Mismatch between layer size and nodes written"); - } - } - - // In V6, fused features for the in-memory hierarchy are written in a block after the top layers of the graph. - // Since everything in level 1 is also contained in the higher levels, we only need to write the fused features for level 1. - if (version == 6) { - // There should be only one fused feature per node. This is checked in the class constructor. - // This is the only place where we explicitly need the fused feature. If there are more places in the - // future, it may be worth having fusedFeature as class member. - FusedFeature fusedFeature = null; - for (var feature : inlineFeatures) { - if (feature.isFused()) { - fusedFeature = (FusedFeature) feature; - } - } - if (fusedFeature != null) { - var supplier = featureStateSuppliers.get(fusedFeature.id()); - if (supplier == null) { - throw new IllegalStateException("Supplier for feature " + fusedFeature.id() + " not found"); - } - - if (graph.getMaxLevel() >= 1) { - int level = 1; - int layerSize = graph.size(level); - int nodesWritten = 0; - for (var it = graph.getNodes(level); it.hasNext(); ) { - int originalOrdinal = it.nextInt(); - - // We write the ordinal (node id) so that we can map it to the corresponding feature - final int newOrdinal = ordinalMapper.oldToNew(originalOrdinal); - out.writeInt(newOrdinal); - fusedFeature.writeSourceFeature(out, supplier.apply(originalOrdinal)); - nodesWritten++; - } - if (nodesWritten != layerSize) { - throw new IllegalStateException("Mismatch between layer 1 size and features written"); - } - } else { - // Write the source feature of the entry node - final int originalEntryNode = view.entryNode().node; - final int entryNode = ordinalMapper.oldToNew(originalEntryNode); - out.writeInt(entryNode); - fusedFeature.writeSourceFeature(out, supplier.apply(originalEntryNode)); - } - } - } + serializer.writeSparseLevels(graph, ordinalMapper, out, featureStateSuppliers); } void writeSeparatedFeatures(Map> featureStateSuppliers) throws IOException { - for (var featureEntry : featureMap.entrySet()) { - if (isSeparated(featureEntry.getValue())) { - var fid = featureEntry.getKey(); - var supplier = featureStateSuppliers.get(fid); - if (supplier == null) { - throw new IllegalStateException("Supplier for feature " + fid + " not found"); - } - - // Set the offset for this feature - var feature = (SeparatedFeature) featureEntry.getValue(); - feature.setOffset(out.position()); - - // Write separated data for each node - for (int newOrdinal = 0; newOrdinal <= ordinalMapper.maxOrdinal(); newOrdinal++) { - int originalOrdinal = ordinalMapper.newToOld(newOrdinal); - if (originalOrdinal != OrdinalMapper.OMITTED) { - feature.writeSeparately(out, supplier.apply(originalOrdinal)); - } else { - // write zeros for missing data as padding - for (int i = 0; i < feature.featureSize(); i++) { - out.writeByte(0); - } - } - } - } - } + serializer.writeSeparatedFeatures(featureStateSuppliers, featureMap, out, ordinalMapper); } /** diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexMetadata.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexMetadata.java new file mode 100644 index 000000000..dc5a2238b --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexMetadata.java @@ -0,0 +1,119 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph.disk; + +import java.util.List; +import java.util.Objects; + +/** + * Metadata about a graph index, containing all the information needed to + * interpret the on-disk format. This replaces the scattered fields that were + * previously in CommonHeader and provides a cleaner API. + */ +public class GraphIndexMetadata { + private final int version; + private final int dimension; + private final int entryNode; + private final List layerInfo; + private final int idUpperBound; + + public GraphIndexMetadata(int version, int dimension, int entryNode, + List layerInfo, int idUpperBound) { + this.version = version; + this.dimension = dimension; + this.entryNode = entryNode; + this.layerInfo = layerInfo; + this.idUpperBound = idUpperBound; + } + + public int getVersion() { + return version; + } + + public int getDimension() { + return dimension; + } + + public int getEntryNode() { + return entryNode; + } + + public List getLayerInfo() { + return layerInfo; + } + + public int getIdUpperBound() { + return idUpperBound; + } + + /** + * Gets the size of the base layer (layer 0). + */ + public int getBaseLayerSize() { + return layerInfo.get(0).size; + } + + /** + * Gets the max degree of the base layer (layer 0). + */ + public int getBaseLayerDegree() { + return layerInfo.get(0).degree; + } + + /** + * Gets the number of layers in the graph. + */ + public int getNumLayers() { + return layerInfo.size(); + } + + /** + * Checks if this is a multi-layer (hierarchical) graph. + */ + public boolean isMultiLayer() { + return layerInfo.size() > 1; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + GraphIndexMetadata that = (GraphIndexMetadata) o; + return version == that.version && + dimension == that.dimension && + entryNode == that.entryNode && + idUpperBound == that.idUpperBound && + Objects.equals(layerInfo, that.layerInfo); + } + + @Override + public int hashCode() { + return Objects.hash(version, dimension, entryNode, layerInfo, idUpperBound); + } + + @Override + public String toString() { + return "GraphIndexMetadata{" + + "version=" + version + + ", dimension=" + dimension + + ", entryNode=" + entryNode + + ", layers=" + layerInfo.size() + + ", idUpperBound=" + idUpperBound + + '}'; + } +} + diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializer.java new file mode 100644 index 000000000..d2e95423f --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializer.java @@ -0,0 +1,132 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph.disk; + +import io.github.jbellis.jvector.disk.IndexWriter; +import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.graph.ImmutableGraphIndex; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; + +import java.io.IOException; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.function.IntFunction; + +/** + * Strategy interface for version-specific serialization of graph indexes. + * Each version of the on-disk format has its own implementation that encapsulates + * all version-specific logic for reading and writing graph data. + *

+ * This design eliminates scattered version checks throughout the codebase and + * makes it easy to add new versions without modifying existing code. + */ +public interface GraphIndexSerializer { + /** + * @return the version number this serializer handles + */ + int getVersion(); + + /** + * Checks if this version supports a given feature. + * @param feature the feature to check + * @return true if the feature is supported in this version + */ + boolean supportsFeature(FeatureId feature); + + /** + * Checks if this version supports multi-layer (hierarchical) graphs. + * @return true if multi-layer graphs are supported + */ + boolean supportsMultiLayer(); + + /** + * Checks if this version uses a footer for metadata instead of a header. + * @return true if footer-based metadata is used + */ + boolean usesFooter(); + + /** + * Gets the feature ordering strategy for this version. + * @return the feature ordering strategy + */ + FeatureOrdering getFeatureOrdering(); + + /** + * Writes the common header portion of the index. + * @param out the output writer + * @param metadata the metadata to write + * @throws IOException if an I/O error occurs + */ + void writeCommonHeader(IndexWriter out, GraphIndexMetadata metadata) throws IOException; + + /** + * Reads the common header portion of the index. + * @param in the input reader + * @return the metadata read from the header + * @throws IOException if an I/O error occurs + */ + GraphIndexMetadata readCommonHeader(RandomAccessReader in) throws IOException; + + /** + * Writes feature-specific header information. + * @param out the output writer + * @param features the features to write headers for + * @throws IOException if an I/O error occurs + */ + void writeFeatureHeaders(IndexWriter out, Map features) throws IOException; + + /** + * Reads feature-specific header information. + * @param in the input reader + * @param metadata the common metadata (needed for feature construction) + * @return map of feature IDs to feature instances + * @throws IOException if an I/O error occurs + */ + Map readFeatureHeaders(RandomAccessReader in, GraphIndexMetadata metadata) throws IOException; + + /** + * Calculates the total size of the header in bytes. + * @param metadata the metadata + * @param features the features + * @return the header size in bytes + */ + int calculateHeaderSize(GraphIndexMetadata metadata, Map features); + + public Set getSupportedFeatures(); + + /** + * Defines how features should be ordered when writing/reading. + */ + enum FeatureOrdering { + /** Features ordered by their FeatureId enum ordinal (versions <= 5) */ + BY_FEATURE_ID, + /** Features ordered with fused features last (version 6+) */ + FUSED_LAST + } + + void writeSparseLevels(ImmutableGraphIndex graph, OrdinalMapper ordinalMapper, IndexWriter out, Map> featureStateSuppliers) throws IOException; + + void writeHeader(ImmutableGraphIndex graph, OrdinalMapper ordinalMapper, Map featureMap, int dimension, long startOffset, IndexWriter out, long headerSize) throws IOException; + + long featureOffsetForOrdinal(ImmutableGraphIndex graph, List inlineFeatures, long startOffset, int ordinal, long headerSize); + + void writeFooter(ImmutableGraphIndex graph, OrdinalMapper ordinalMapper, long headerOffset, int dimension, IndexWriter out, Map featureMap, long headerSize) throws IOException; + + void writeSeparatedFeatures(Map> featureStateSuppliers, Map featureMap, IndexWriter out, OrdinalMapper ordinalMapper) throws IOException; + } \ No newline at end of file diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerFactory.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerFactory.java new file mode 100644 index 000000000..b1a63079f --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerFactory.java @@ -0,0 +1,103 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph.disk; + +import io.github.jbellis.jvector.disk.RandomAccessReader; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.util.Map; + +/** + * Factory for creating version-specific graph index serializers. + * This centralizes version detection and serializer instantiation. + */ +public class GraphIndexSerializerFactory { + private static final Logger logger = LoggerFactory.getLogger(GraphIndexSerializerFactory.class); + + private static final Map SERIALIZERS = Map.of( + 2, new GraphIndexSerializerV2(), + 3, new GraphIndexSerializerV3(), + 4, new GraphIndexSerializerV4(), + 5, new GraphIndexSerializerV5(), + 6, new GraphIndexSerializerV6() + ); + + /** + * Gets a serializer for a specific version. + * @param version the version number + * @return the serializer for that version + * @throws UnsupportedVersionException if the version is not supported + */ + public static GraphIndexSerializer forVersion(int version) { + GraphIndexSerializer serializer = SERIALIZERS.get(version); + if (serializer == null) { + throw new UnsupportedVersionException("Version " + version + " is not supported. " + + "Supported versions: " + SERIALIZERS.keySet()); + } + return serializer; + } + + /** + * Detects the version from the input stream and returns the appropriate serializer. + * The reader position will be reset to where it started. + * + * @param in the input reader + * @return the serializer for the detected version + * @throws IOException if an I/O error occurs + * @throws UnsupportedVersionException if the version is not supported + */ + public static GraphIndexSerializer detectVersion(RandomAccessReader in) throws IOException { + long startPosition = in.getPosition(); + + try { + int maybeMagic = in.readInt(); + + if (maybeMagic == OnDiskGraphIndex.MAGIC) { + // Version 3+ with magic number + int version = in.readInt(); + logger.debug("Detected version {} (with magic number)", version); + return forVersion(version); + } else { + // Version 2 (no magic number) + logger.debug("Detected version 2 (no magic number)"); + return forVersion(2); + } + } finally { + // Reset to starting position + in.seek(startPosition); + } + } + + /** + * Gets the current/latest version number. + * @return the current version + */ + public static int getCurrentVersion() { + return OnDiskGraphIndex.CURRENT_VERSION; + } + + /** + * Exception thrown when an unsupported version is encountered. + */ + public static class UnsupportedVersionException extends RuntimeException { + public UnsupportedVersionException(String message) { + super(message); + } + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV2.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV2.java new file mode 100644 index 000000000..3dfe3c1b4 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV2.java @@ -0,0 +1,111 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph.disk; + +import io.github.jbellis.jvector.disk.IndexWriter; +import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; + +import java.io.IOException; +import java.util.EnumMap; +import java.util.EnumSet; +import java.util.List; +import java.util.Map; + +/** + * Serializer for version 2 of the on-disk graph format. + * Version 2 characteristics: + * - No magic number + * - Only supports INLINE_VECTORS feature + * - Single layer only + * - No footer + */ +class GraphIndexSerializerV2 extends AbstractGraphIndexSerializer { + + GraphIndexSerializerV2() { + super(2, inlineVectorsOnly(), false, false, FeatureOrdering.BY_FEATURE_ID); + } + + @Override + public void writeCommonHeader(IndexWriter out, GraphIndexMetadata metadata) throws IOException { + // V2 format: size, dimension, entryNode, maxDegree (no magic, no version) + out.writeInt(metadata.getBaseLayerSize()); + out.writeInt(metadata.getDimension()); + out.writeInt(metadata.getEntryNode()); + out.writeInt(metadata.getBaseLayerDegree()); + } + + @Override + public GraphIndexMetadata readCommonHeader(RandomAccessReader in) throws IOException { + // V2 format: size, dimension, entryNode, maxDegree + int size = in.readInt(); + int dimension = in.readInt(); + int entryNode = in.readInt(); + int maxDegree = in.readInt(); + + // V2 only supports single layer + List layerInfo = List.of(new CommonHeader.LayerInfo(size, maxDegree)); + int idUpperBound = size; + + return new GraphIndexMetadata(2, dimension, entryNode, layerInfo, idUpperBound); + } + + @Override + public void writeFeatureHeaders(IndexWriter out, Map features) throws IOException { + // V2 doesn't write feature set information, just the feature headers + // Only INLINE_VECTORS is supported + if (!features.containsKey(FeatureId.INLINE_VECTORS) || features.size() > 1) { + throw new IllegalArgumentException("Version 2 only supports INLINE_VECTORS feature"); + } + + for (Feature feature : features.values()) { + feature.writeHeader(out); + } + } + + @Override + public Map readFeatureHeaders(RandomAccessReader in, GraphIndexMetadata metadata) throws IOException { + // V2 only has INLINE_VECTORS + Map features = new EnumMap<>(FeatureId.class); + EnumSet featureIds = EnumSet.of(FeatureId.INLINE_VECTORS); + + // Create a CommonHeader from metadata for feature loading + CommonHeader commonHeader = new CommonHeader( + metadata.getVersion(), + metadata.getDimension(), + metadata.getEntryNode(), + metadata.getLayerInfo(), + metadata.getIdUpperBound() + ); + + for (FeatureId featureId : featureIds) { + features.put(featureId, featureId.load(commonHeader, in)); + } + + return features; + } + + @Override + public int calculateHeaderSize(GraphIndexMetadata metadata, Map features) { + // V2: 4 ints (size, dimension, entryNode, maxDegree) + feature headers + int size = 4 * Integer.BYTES; + size += features.values().stream().mapToInt(Feature::headerSize).sum(); + return size; + } +} + diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV3.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV3.java new file mode 100644 index 000000000..3b7ba0b04 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV3.java @@ -0,0 +1,119 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph.disk; + +import io.github.jbellis.jvector.disk.IndexWriter; +import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; + +import java.io.IOException; +import java.util.EnumMap; +import java.util.EnumSet; +import java.util.List; +import java.util.Map; + +/** + * Serializer for version 3 of the on-disk graph format. + * Version 3 characteristics: + * - Has magic number + * - Supports multiple features (feature set serialization) + * - Single layer only + * - No footer + */ +class GraphIndexSerializerV3 extends AbstractGraphIndexSerializer { + + GraphIndexSerializerV3() { + super(3, nonFusedFeatures(), false, false, FeatureOrdering.BY_FEATURE_ID); + } + + @Override + public void writeCommonHeader(IndexWriter out, GraphIndexMetadata metadata) throws IOException { + // V3 format: magic, version, size, dimension, entryNode, maxDegree + out.writeInt(OnDiskGraphIndex.MAGIC); + out.writeInt(3); + out.writeInt(metadata.getBaseLayerSize()); + out.writeInt(metadata.getDimension()); + out.writeInt(metadata.getEntryNode()); + out.writeInt(metadata.getBaseLayerDegree()); + } + + @Override + public GraphIndexMetadata readCommonHeader(RandomAccessReader in) throws IOException { + // V3 format: magic, version, size, dimension, entryNode, maxDegree + int magic = in.readInt(); + if (magic != OnDiskGraphIndex.MAGIC) { + throw new IOException("Invalid magic number: " + magic); + } + int version = in.readInt(); + int size = in.readInt(); + int dimension = in.readInt(); + int entryNode = in.readInt(); + int maxDegree = in.readInt(); + + // V3 only supports single layer + List layerInfo = List.of(new CommonHeader.LayerInfo(size, maxDegree)); + int idUpperBound = size; + + return new GraphIndexMetadata(version, dimension, entryNode, layerInfo, idUpperBound); + } + + @Override + public void writeFeatureHeaders(IndexWriter out, Map features) throws IOException { + // V3 writes feature set as bitflags + out.writeInt(FeatureId.serialize(EnumSet.copyOf(features.keySet()))); + + // Then write each feature's header + for (Feature feature : features.values()) { + feature.writeHeader(out); + } + } + + @Override + public Map readFeatureHeaders(RandomAccessReader in, GraphIndexMetadata metadata) throws IOException { + Map features = new EnumMap<>(FeatureId.class); + + // Read feature set bitflags + EnumSet featureIds = FeatureId.deserialize(in.readInt()); + + // Create a CommonHeader from metadata for feature loading + CommonHeader commonHeader = new CommonHeader( + metadata.getVersion(), + metadata.getDimension(), + metadata.getEntryNode(), + metadata.getLayerInfo(), + metadata.getIdUpperBound() + ); + + // Load each feature + for (FeatureId featureId : featureIds) { + features.put(featureId, featureId.load(commonHeader, in)); + } + + return features; + } + + @Override + public int calculateHeaderSize(GraphIndexMetadata metadata, Map features) { + // V3: magic + version + 4 ints (size, dimension, entryNode, maxDegree) + feature bitflags + feature headers + int size = 6 * Integer.BYTES; // magic, version, size, dimension, entryNode, maxDegree + size += Integer.BYTES; // feature bitflags + size += features.values().stream().mapToInt(Feature::headerSize).sum(); + return size; + } +} + diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV4.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV4.java new file mode 100644 index 000000000..00798ed18 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV4.java @@ -0,0 +1,162 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph.disk; + +import io.github.jbellis.jvector.disk.IndexWriter; +import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.EnumMap; +import java.util.EnumSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * Serializer for version 4 of the on-disk graph format. + * Version 4 characteristics: + * - Has magic number + * - Supports multiple features + * - Supports multi-layer (hierarchical) graphs + * - Has idUpperBound field + * - No footer + */ +class GraphIndexSerializerV4 extends AbstractGraphIndexSerializer { + private static final int V4_MAX_LAYERS = 32; + + GraphIndexSerializerV4() { + this(nonFusedFeatures()); + } + + /** + * Protected constructor for subclasses (V5, V6) that share V4's wire format but differ + * in supported features, footer behaviour, or feature ordering. + */ + protected GraphIndexSerializerV4(Set supportedFeatures) { + super(4, supportedFeatures, true, false, FeatureOrdering.BY_FEATURE_ID); + } + + @Override + public void writeCommonHeader(IndexWriter out, GraphIndexMetadata metadata) throws IOException { + // V4 format: magic, version, size, dimension, entryNode, maxDegree, idUpperBound, numLayers, layer info + out.writeInt(OnDiskGraphIndex.MAGIC); + out.writeInt(4); + out.writeInt(metadata.getBaseLayerSize()); + out.writeInt(metadata.getDimension()); + out.writeInt(metadata.getEntryNode()); + out.writeInt(metadata.getBaseLayerDegree()); + out.writeInt(metadata.getIdUpperBound()); + + if (metadata.getLayerInfo().size() > V4_MAX_LAYERS) { + throw new IllegalArgumentException( + String.format("Number of layers %d exceeds maximum of %d", + metadata.getLayerInfo().size(), V4_MAX_LAYERS)); + } + + out.writeInt(metadata.getLayerInfo().size()); + + // Write actual layer info + for (CommonHeader.LayerInfo info : metadata.getLayerInfo()) { + out.writeInt(info.size); + out.writeInt(info.degree); + } + + // Pad remaining entries with zeros + for (int i = metadata.getLayerInfo().size(); i < V4_MAX_LAYERS; i++) { + out.writeInt(0); // size + out.writeInt(0); // degree + } + } + + @Override + public GraphIndexMetadata readCommonHeader(RandomAccessReader in) throws IOException { + int magic = in.readInt(); + if (magic != OnDiskGraphIndex.MAGIC) { + throw new IOException("Invalid magic number: " + magic); + } + int version = in.readInt(); + int size = in.readInt(); + int dimension = in.readInt(); + int entryNode = in.readInt(); + int maxDegree = in.readInt(); + int idUpperBound = in.readInt(); + int numLayers = in.readInt(); + + List layerInfo = new ArrayList<>(); + for (int i = 0; i < numLayers; i++) { + CommonHeader.LayerInfo info = new CommonHeader.LayerInfo(in.readInt(), in.readInt()); + layerInfo.add(info); + } + + // Skip over remaining padding entries + for (int i = numLayers; i < V4_MAX_LAYERS; i++) { + in.readInt(); + in.readInt(); + } + + return new GraphIndexMetadata(version, dimension, entryNode, layerInfo, idUpperBound); + } + + @Override + public void writeFeatureHeaders(IndexWriter out, Map features) throws IOException { + // V4 writes feature set as bitflags + out.writeInt(FeatureId.serialize(EnumSet.copyOf(features.keySet()))); + + // Then write each feature's header + for (Feature feature : features.values()) { + feature.writeHeader(out); + } + } + + @Override + public Map readFeatureHeaders(RandomAccessReader in, GraphIndexMetadata metadata) throws IOException { + Map features = new EnumMap<>(FeatureId.class); + + // Read feature set bitflags + EnumSet featureIds = FeatureId.deserialize(in.readInt()); + + // Create a CommonHeader from metadata for feature loading + CommonHeader commonHeader = new CommonHeader( + metadata.getVersion(), + metadata.getDimension(), + metadata.getEntryNode(), + metadata.getLayerInfo(), + metadata.getIdUpperBound() + ); + + // Load each feature + for (FeatureId featureId : featureIds) { + features.put(featureId, featureId.load(commonHeader, in)); + } + + return features; + } + + @Override + public int calculateHeaderSize(GraphIndexMetadata metadata, Map features) { + // V4: magic + version + 4 base ints + idUpperBound + numLayers + (32 * 2 layer info ints) + feature bitflags + feature headers + int size = 8 * Integer.BYTES; // magic, version, size, dimension, entryNode, maxDegree, idUpperBound, numLayers + size += 2 * V4_MAX_LAYERS * Integer.BYTES; // layer info (padded to 32 layers) + size += Integer.BYTES; // feature bitflags + size += features.values().stream().mapToInt(Feature::headerSize).sum(); + return size; + } +} + diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV5.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV5.java new file mode 100644 index 000000000..752f7e589 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV5.java @@ -0,0 +1,47 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph.disk; + +/** + * Serializer for version 5 of the on-disk graph format. + * Version 5 characteristics: + * - Has magic number + * - Supports multiple features + * - Supports multi-layer (hierarchical) graphs + * - Has idUpperBound field + * - Uses footer for metadata (major change from V4) + * + * V5 is identical to V4 in terms of format, but uses a footer instead of + * relying on the header at the beginning. The serialization logic is the same. + */ +class GraphIndexSerializerV5 extends GraphIndexSerializerV4 { + + GraphIndexSerializerV5() { + super(nonFusedFeatures()); + } + + @Override + public int getVersion() { + return 5; + } + + @Override + public boolean usesFooter() { + return true; + } +} + diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV6.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV6.java new file mode 100644 index 000000000..1bc4e50d6 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV6.java @@ -0,0 +1,181 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph.disk; + +import io.github.jbellis.jvector.disk.IndexWriter; +import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.graph.ImmutableGraphIndex; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.graph.disk.feature.FusedFeature; + +import java.io.IOException; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.function.IntFunction; + +/** + * Serializer for version 6 of the on-disk graph format. + * Version 6 characteristics: + * - Has magic number + * - Supports multiple features + * - Supports multi-layer (hierarchical) graphs + * - Has idUpperBound field + * - Uses footer for metadata + * - NEW: Changes feature ordering to place fused features last + * - NEW: Writes feature count and ordinals explicitly instead of bitflags + */ +class GraphIndexSerializerV6 extends GraphIndexSerializerV4 { + + GraphIndexSerializerV6() { + super(allFeatures()); + } + + @Override + public int getVersion() { + return 6; + } + + @Override + public boolean usesFooter() { + return true; + } + + @Override + public FeatureOrdering getFeatureOrdering() { + return FeatureOrdering.FUSED_LAST; + } + + @Override + public void writeFeatureHeaders(IndexWriter out, Map features) throws IOException { + // V6 writes feature count and ordinals explicitly (preserving order) + out.writeInt(features.size()); + + for (var entry : features.entrySet()) { + out.writeInt(entry.getKey().ordinal()); + entry.getValue().writeHeader(out); + } + } + + @Override + public Map readFeatureHeaders(RandomAccessReader in, GraphIndexMetadata metadata) throws IOException { + // V6 reads features in order (LinkedHashMap preserves insertion order) + Map features = new LinkedHashMap<>(); + + int nFeatures = in.readInt(); + + // Create a CommonHeader from metadata for feature loading + CommonHeader commonHeader = new CommonHeader( + metadata.getVersion(), + metadata.getDimension(), + metadata.getEntryNode(), + metadata.getLayerInfo(), + metadata.getIdUpperBound() + ); + + for (int i = 0; i < nFeatures; i++) { + FeatureId featureId = FeatureId.values()[in.readInt()]; + features.put(featureId, featureId.load(commonHeader, in)); + } + + return features; + } + + @Override + public int calculateHeaderSize(GraphIndexMetadata metadata, Map features) { + // V6: magic + version + 4 base ints + idUpperBound + numLayers + (32 * 2 layer info ints) + // + feature count + (feature ordinals) + feature headers + int size = 8 * Integer.BYTES; // magic, version, size, dimension, entryNode, maxDegree, idUpperBound, numLayers + size += 2 * 32 * Integer.BYTES; // layer info (padded to 32 layers) + size += Integer.BYTES; // feature count + size += features.size() * Integer.BYTES; // feature ordinals + size += features.values().stream().mapToInt(Feature::headerSize).sum(); + return size; + } + + @Override + public void writeSparseLevels(ImmutableGraphIndex graph, OrdinalMapper ordinalMapper, IndexWriter out, Map> featureStateSuppliers) throws IOException { + for (int level = 1; level <= graph.getMaxLevel(); level++) { + int layerSize = graph.size(level); + int layerDegree = graph.getDegree(level); + int nodesWritten = 0; + for (var it = graph.getNodes(level); it.hasNext(); ) { + int originalOrdinal = it.nextInt(); + // node id + final int newOrdinal = ordinalMapper.oldToNew(originalOrdinal); + out.writeInt(newOrdinal); + // neighbors + var view = graph.getView(); + var neighbors = view.getNeighborsIterator(level, originalOrdinal); + out.writeInt(neighbors.size()); + int n = 0; + for ( ; n < neighbors.size(); n++) { + out.writeInt(ordinalMapper.oldToNew(neighbors.nextInt())); + } + assert !neighbors.hasNext() : "Mismatch between neighbor's reported size and actual size"; + // pad out to degree + for (; n < layerDegree; n++) { + out.writeInt(-1); + } + nodesWritten++; + } + if (nodesWritten != layerSize) { + throw new IllegalStateException("Mismatch between layer size and nodes written"); + } + } + // There should be only one fused feature per node. This is checked in the class constructor. + // This is the only place where we explicitly need the fused feature. If there are more places in the + // future, it may be worth having fusedFeature as class member. + FusedFeature fusedFeature = null; + for (var feature : inlineFeatures) { + if (feature.isFused()) { + fusedFeature = (FusedFeature) feature; + } + } + if (fusedFeature != null) { + var supplier = featureStateSuppliers.get(fusedFeature.id()); + if (supplier == null) { + throw new IllegalStateException("Supplier for feature " + fusedFeature.id() + " not found"); + } + + if (graph.getMaxLevel() >= 1) { + int level = 1; + int layerSize = graph.size(level); + int nodesWritten = 0; + for (var it = graph.getNodes(level); it.hasNext(); ) { + int originalOrdinal = it.nextInt(); + + // We write the ordinal (node id) so that we can map it to the corresponding feature + final int newOrdinal = ordinalMapper.oldToNew(originalOrdinal); + out.writeInt(newOrdinal); + fusedFeature.writeSourceFeature(out, supplier.apply(originalOrdinal)); + nodesWritten++; + } + if (nodesWritten != layerSize) { + throw new IllegalStateException("Mismatch between layer 1 size and features written"); + } + } else { + // Write the source feature of the entry node + final int originalEntryNode = graph.getView().entryNode().node; + final int entryNode = ordinalMapper.oldToNew(originalEntryNode); + out.writeInt(entryNode); + fusedFeature.writeSourceFeature(out, supplier.apply(originalEntryNode)); + } + } + } +} + From 35dda7ec0b584a76bd4a9cc5716bd3ca31eaf2ed Mon Sep 17 00:00:00 2001 From: Mark Wolters Date: Thu, 23 Jul 2026 17:53:55 -0400 Subject: [PATCH 2/5] total rework --- .../graph/disk/AbstractGraphIndexFormat.java | 345 ++++++++++++++++++ .../disk/AbstractGraphIndexSerializer.java | 221 ----------- .../graph/disk/AbstractGraphIndexWriter.java | 28 +- .../jvector/graph/disk/GraphIndexFormat.java | 191 ++++++++++ ...tory.java => GraphIndexFormatFactory.java} | 46 +-- .../graph/disk/GraphIndexFormatV2.java | 33 ++ .../graph/disk/GraphIndexFormatV3.java | 33 ++ .../graph/disk/GraphIndexFormatV4.java | 46 +++ .../graph/disk/GraphIndexFormatV5.java | 99 +++++ .../graph/disk/GraphIndexFormatV6.java | 85 +++++ .../graph/disk/GraphIndexMetadata.java | 119 ------ .../graph/disk/GraphIndexSerializer.java | 132 ------- .../graph/disk/GraphIndexSerializerV2.java | 111 ------ .../graph/disk/GraphIndexSerializerV3.java | 119 ------ .../graph/disk/GraphIndexSerializerV4.java | 162 -------- .../graph/disk/GraphIndexSerializerV5.java | 47 --- .../graph/disk/GraphIndexSerializerV6.java | 181 --------- .../OnDiskSequentialGraphIndexWriter.java | 83 +---- .../RandomAccessOnDiskGraphIndexWriter.java | 58 +-- .../jvector/graph/disk/WriteContext.java | 73 ++++ 20 files changed, 949 insertions(+), 1263 deletions(-) create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexFormat.java delete mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexSerializer.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormat.java rename jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/{GraphIndexSerializerFactory.java => GraphIndexFormatFactory.java} (70%) create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV2.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV3.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV4.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV5.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV6.java delete mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexMetadata.java delete mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializer.java delete mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV2.java delete mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV3.java delete mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV4.java delete mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV5.java delete mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV6.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/WriteContext.java diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexFormat.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexFormat.java new file mode 100644 index 000000000..2d0e27c67 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexFormat.java @@ -0,0 +1,345 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph.disk; + +import io.github.jbellis.jvector.disk.IndexWriter; +import io.github.jbellis.jvector.disk.RandomAccessWriter; +import io.github.jbellis.jvector.graph.ImmutableGraphIndex; +import io.github.jbellis.jvector.graph.OnHeapGraphIndex; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.graph.disk.feature.SeparatedFeature; + +import java.io.IOException; +import java.util.EnumSet; +import java.util.Map; +import java.util.Set; +import java.util.function.IntFunction; + +/** + * Abstract base class for graph index formats providing common functionality. + */ +abstract class AbstractGraphIndexFormat implements GraphIndexFormat { + private final int version; + private final Set supportedFeatures; + private final boolean supportsMultiLayer; + private final boolean usesFooter; + private final FeatureOrdering featureOrdering; + + /** A magic number to indicate the file footer */ + public static final int FOOTER_MAGIC = 0x4a564244; + /** The size of the offset in the footer. */ + public static final int FOOTER_OFFSET_SIZE = Long.BYTES; + /** The size of the magic number in the footer. */ + public static final int FOOTER_MAGIC_SIZE = Integer.BYTES; + /** The total size of the footer. */ + public static final int FOOTER_SIZE = FOOTER_MAGIC_SIZE + FOOTER_OFFSET_SIZE; + + /** + * Initialises the format with the format characteristics for a specific version. + * + * @param version on-disk format version number reported by {@link #getVersion()} + * @param supportedFeatures the set of {@link FeatureId}s this version can store + * @param supportsMultiLayer whether this version supports hierarchical (multi-layer) graphs + * @param usesFooter whether metadata is placed in a footer rather than a header + * @param featureOrdering the ordering strategy used when laying out feature data + */ + protected AbstractGraphIndexFormat(int version, + Set supportedFeatures, + boolean supportsMultiLayer, + boolean usesFooter, + FeatureOrdering featureOrdering) { + this.version = version; + this.supportedFeatures = supportedFeatures; + this.supportsMultiLayer = supportsMultiLayer; + this.usesFooter = usesFooter; + this.featureOrdering = featureOrdering; + } + + @Override + public int getVersion() { + return version; + } + + @Override + public boolean supportsFeature(FeatureId feature) { + return supportedFeatures.contains(feature); + } + + @Override + public Set getSupportedFeatures() { + return supportedFeatures; + } + + @Override + public boolean supportsMultiLayer() { + return supportsMultiLayer; + } + + @Override + public boolean usesFooter() { + return usesFooter; + } + + @Override + public FeatureOrdering getFeatureOrdering() { + return featureOrdering; + } + + @Override + public void writeSparseLevels(WriteContext ctx, IndexWriter out, Map> suppliers) throws IOException { + for (int level = 1; level <= ctx.graph.getMaxLevel(); level++) { + int layerSize = ctx.graph.size(level); + int layerDegree = ctx.graph.getDegree(level); + int nodesWritten = 0; + for (var it = ctx.graph.getNodes(level); it.hasNext(); ) { + int originalOrdinal = it.nextInt(); + final int newOrdinal = ctx.ordinalMapper.oldToNew(originalOrdinal); + out.writeInt(newOrdinal); + var view = ctx.graph.getView(); + var neighbors = view.getNeighborsIterator(level, originalOrdinal); + out.writeInt(neighbors.size()); + int n = 0; + for ( ; n < neighbors.size(); n++) { + out.writeInt(ctx.ordinalMapper.oldToNew(neighbors.nextInt())); + } + assert !neighbors.hasNext() : "Mismatch between neighbor's reported size and actual size"; + for (; n < layerDegree; n++) { + out.writeInt(-1); + } + nodesWritten++; + } + if (nodesWritten != layerSize) { + throw new IllegalStateException("Mismatch between layer size and nodes written"); + } + } + writeAfterSparseLevels(ctx, out, suppliers); + } + + /** + * Hook called at the end of {@link #writeSparseLevels} for version-specific additions. + * The default implementation is a no-op; V6 overrides this to write fused feature data. + */ + protected void writeAfterSparseLevels(WriteContext ctx, IndexWriter out, Map> suppliers) throws IOException {} + + @Override + public void writeHeader(WriteContext ctx, IndexWriter out) throws IOException { + var layerInfo = CommonHeader.LayerInfo.fromGraph(ctx.graph, ctx.ordinalMapper); + var entryNode = ctx.graph.getView().entryNode() == null + ? ImmutableGraphIndex.ENTRY_NODE_ABSENT + : ctx.ordinalMapper.oldToNew(ctx.graph.getView().entryNode().node); + var commonHeader = new CommonHeader(getVersion(), ctx.dimension, entryNode, layerInfo, ctx.ordinalMapper.maxOrdinal() + 1); + var header = new Header(commonHeader, ctx.featureMap); + header.write(out); + assert out.position() == ctx.startOffset + ctx.headerSize + : String.format("%d != %d", out.position(), ctx.startOffset + ctx.headerSize); + } + + @Override + public long featureOffsetForOrdinal(WriteContext ctx, int ordinal) { + int edgeSize = Integer.BYTES * (1 + ctx.graph.getDegree(0)); + long inlineBytes = ordinal * (long) (Integer.BYTES + ctx.inlineFeatures.stream().mapToInt(Feature::featureSize).sum() + edgeSize); + return ctx.startOffset + ctx.headerSize + inlineBytes + Integer.BYTES; + } + + @Override + public void writeFooter(WriteContext ctx, long headerOffset, IndexWriter out) throws IOException { + var layerInfo = CommonHeader.LayerInfo.fromGraph(ctx.graph, ctx.ordinalMapper); + var entryNode = ctx.graph.getView().entryNode() == null + ? ImmutableGraphIndex.ENTRY_NODE_ABSENT + : ctx.ordinalMapper.oldToNew(ctx.graph.getView().entryNode().node); + var commonHeader = new CommonHeader(getVersion(), ctx.dimension, entryNode, layerInfo, ctx.ordinalMapper.maxOrdinal() + 1); + var header = new Header(commonHeader, ctx.featureMap); + header.write(out); + out.writeLong(headerOffset); + out.writeInt(FOOTER_MAGIC); + final long expectedPosition = headerOffset + ctx.headerSize + FOOTER_SIZE; + assert out.position() == expectedPosition : String.format("%d != %d", out.position(), expectedPosition); + } + + @Override + public void writeSeparatedFeatures(WriteContext ctx, IndexWriter out, Map> suppliers) throws IOException { + for (var featureEntry : ctx.featureMap.entrySet()) { + if (featureEntry.getValue() instanceof SeparatedFeature) { + var fid = featureEntry.getKey(); + var supplier = suppliers.get(fid); + if (supplier == null) { + throw new IllegalStateException("Supplier for feature " + fid + " not found"); + } + var feature = (SeparatedFeature) featureEntry.getValue(); + feature.setOffset(out.position()); + for (int newOrdinal = 0; newOrdinal <= ctx.ordinalMapper.maxOrdinal(); newOrdinal++) { + int originalOrdinal = ctx.ordinalMapper.newToOld(newOrdinal); + if (originalOrdinal != OrdinalMapper.OMITTED) { + feature.writeSeparately(out, supplier.apply(originalOrdinal)); + } else { + for (int i = 0; i < feature.featureSize(); i++) { + out.writeByte(0); + } + } + } + } + } + } + + @Override + public void writeFeaturesInline(WriteContext ctx, int ordinal, Map stateMap, RandomAccessWriter out) throws IOException { + for (var featureId : stateMap.keySet()) { + if (!ctx.featureMap.containsKey(featureId)) { + throw new IllegalArgumentException(String.format("Feature %s not configured for index", featureId)); + } + } + out.seek(featureOffsetForOrdinal(ctx, ordinal)); + for (var feature : ctx.inlineFeatures) { + var state = stateMap.get(feature.id()); + if (state == null) { + out.seek(out.position() + feature.featureSize()); + } else { + feature.writeInline(out, state); + } + } + } + + @Override + public void writeOnDiskSequential(WriteContext ctx, IndexWriter out, Map> suppliers) throws IOException { + if (ctx.graph instanceof OnHeapGraphIndex) { + var ohgi = (OnHeapGraphIndex) ctx.graph; + if (ohgi.getDeletedNodes().cardinality() > 0) { + throw new IllegalArgumentException("Run builder.cleanup() before writing the graph"); + } + } + for (var featureId : suppliers.keySet()) { + if (!ctx.featureMap.containsKey(featureId)) { + throw new IllegalArgumentException(String.format("Feature %s not configured for index", featureId)); + } + } + if (ctx.ordinalMapper.maxOrdinal() < ctx.graph.size(0) - 1) { + throw new IllegalStateException(String.format("Ordinal mapper from [0..%d] does not cover all nodes in the graph of size %d", + ctx.ordinalMapper.maxOrdinal(), ctx.graph.size(0))); + } + + var view = ctx.graph.getView(); + + writeHeader(ctx, out); + + for (int newOrdinal = 0; newOrdinal <= ctx.ordinalMapper.maxOrdinal(); newOrdinal++) { + var originalOrdinal = ctx.ordinalMapper.newToOld(newOrdinal); + + if (originalOrdinal == OrdinalMapper.OMITTED) { + throw new IllegalStateException("Ordinal mapper mapped new ordinal " + newOrdinal + + " to non-existing node. This behavior is not supported on OnDiskSequentialGraphIndexWriter. Use OnDiskGraphIndexWriter instead."); + } + if (!ctx.graph.containsNode(originalOrdinal)) { + throw new IllegalStateException(String.format("Ordinal mapper mapped new ordinal %s to non-existing node %s", newOrdinal, originalOrdinal)); + } + + out.writeInt(newOrdinal); + long featureOffset = featureOffsetForOrdinal(ctx, newOrdinal); + assert out.position() == featureOffset : String.format("%d != %d", out.position(), featureOffset); + + for (var feature : ctx.inlineFeatures) { + var supplier = suppliers.get(feature.id()); + if (supplier == null) { + throw new IllegalStateException("Supplier for feature " + feature.id() + " not found"); + } + feature.writeInline(out, supplier.apply(originalOrdinal)); + } + + var neighbors = view.getNeighborsIterator(0, originalOrdinal); + if (neighbors.size() > ctx.graph.getDegree(0)) { + throw new IllegalStateException(String.format("Node %d has more neighbors %d than the graph's max degree %d -- run Builder.cleanup()!", + originalOrdinal, neighbors.size(), ctx.graph.getDegree(0))); + } + out.writeInt(neighbors.size()); + int n = 0; + for (; n < neighbors.size(); n++) { + var newNeighborOrdinal = ctx.ordinalMapper.oldToNew(neighbors.nextInt()); + if (newNeighborOrdinal < 0 || newNeighborOrdinal > ctx.ordinalMapper.maxOrdinal()) { + throw new IllegalStateException(String.format("Neighbor ordinal out of bounds: %d/%d", newNeighborOrdinal, ctx.ordinalMapper.maxOrdinal())); + } + out.writeInt(newNeighborOrdinal); + } + assert !neighbors.hasNext(); + for (; n < ctx.graph.getDegree(0); n++) { + out.writeInt(-1); + } + } + + writeSparseLevels(ctx, out, suppliers); + writeSeparatedFeatures(ctx, out, suppliers); + writeFooter(ctx, out.position(), out); + + view.close(); + } + + @Override + public void writeRandomAccess(WriteContext ctx, RandomAccessWriter out, Map> suppliers, GraphIndexFormat.L0RecordWriter l0Writer) throws IOException { + if (ctx.graph instanceof OnHeapGraphIndex) { + var ohgi = (OnHeapGraphIndex) ctx.graph; + if (ohgi.getDeletedNodes().cardinality() > 0) { + throw new IllegalArgumentException("Run builder.cleanup() before writing the graph"); + } + } + for (var featureId : suppliers.keySet()) { + if (!ctx.featureMap.containsKey(featureId)) { + throw new IllegalArgumentException(String.format("Feature %s not configured for index", featureId)); + } + } + if (ctx.ordinalMapper.maxOrdinal() < ctx.graph.size(0) - 1) { + throw new IllegalStateException(String.format("Ordinal mapper from [0..%d] does not cover all nodes in the graph of size %d", + ctx.ordinalMapper.maxOrdinal(), ctx.graph.size(0))); + } + + var view = ctx.graph.getView(); + + out.seek(ctx.startOffset); + writeHeader(ctx, out); + l0Writer.write(view, suppliers); + writeSparseLevels(ctx, out, suppliers); + writeSeparatedFeatures(ctx, out, suppliers); + + final var endOfGraphPosition = out.position(); + out.seek(ctx.startOffset); + writeHeader(ctx, out); + out.seek(endOfGraphPosition); + out.flush(); + view.close(); + } + + /** + * Helper to create a set of all known features. + * Only correct for the current maximum version — do not use for older versions. + */ + protected static Set allFeatures() { + return EnumSet.allOf(FeatureId.class); + } + + /** + * Helper to create a set of all features except FUSED_PQ. + * Used by versions 3–5, which predate fused PQ hierarchy support (version 6). + */ + protected static Set nonFusedFeatures() { + return EnumSet.complementOf(EnumSet.of(FeatureId.FUSED_PQ)); + } + + /** + * Helper to create a set with only inline vectors (for version 2). + */ + protected static Set inlineVectorsOnly() { + return EnumSet.of(FeatureId.INLINE_VECTORS); + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexSerializer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexSerializer.java deleted file mode 100644 index 423110911..000000000 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexSerializer.java +++ /dev/null @@ -1,221 +0,0 @@ -/* - * Copyright DataStax, Inc. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package io.github.jbellis.jvector.graph.disk; - -import io.github.jbellis.jvector.disk.IndexWriter; -import io.github.jbellis.jvector.graph.ImmutableGraphIndex; -import io.github.jbellis.jvector.graph.disk.feature.Feature; -import io.github.jbellis.jvector.graph.disk.feature.FeatureId; -import io.github.jbellis.jvector.graph.disk.feature.SeparatedFeature; - -import java.io.IOException; -import java.util.EnumSet; -import java.util.List; -import java.util.Map; -import java.util.Set; -import java.util.function.IntFunction; - -/** - * Abstract base class for graph index serializers providing common functionality. - */ -abstract class AbstractGraphIndexSerializer implements GraphIndexSerializer { - private final int version; - private final Set supportedFeatures; - private final boolean supportsMultiLayer; - private final boolean usesFooter; - private final FeatureOrdering featureOrdering; - - /** A magic number to indicate the file footer */ - public static final int FOOTER_MAGIC = 0x4a564244; - /** The size of the offset in the footer. */ - public static final int FOOTER_OFFSET_SIZE = Long.BYTES; - /** The size of the magic number in the footer. */ - public static final int FOOTER_MAGIC_SIZE = Integer.BYTES; - /** The total size of the footer. */ - public static final int FOOTER_SIZE = FOOTER_MAGIC_SIZE + FOOTER_OFFSET_SIZE; - - protected AbstractGraphIndexSerializer(int version, - Set supportedFeatures, - boolean supportsMultiLayer, - boolean usesFooter, - FeatureOrdering featureOrdering) { - this.version = version; - this.supportedFeatures = supportedFeatures; - this.supportsMultiLayer = supportsMultiLayer; - this.usesFooter = usesFooter; - this.featureOrdering = featureOrdering; - } - - @Override - public int getVersion() { - return version; - } - - @Override - public boolean supportsFeature(FeatureId feature) { - return supportedFeatures.contains(feature); - } - - @Override - public Set getSupportedFeatures() { return supportedFeatures; } - - @Override - public boolean supportsMultiLayer() { - return supportsMultiLayer; - } - - @Override - public boolean usesFooter() { - return usesFooter; - } - - @Override - public FeatureOrdering getFeatureOrdering() { - return featureOrdering; - } - - @Override - public void writeSparseLevels(ImmutableGraphIndex graph, OrdinalMapper ordinalMapper, IndexWriter out, Map> featureStateSuppliers) throws IOException { - for (int level = 1; level <= graph.getMaxLevel(); level++) { - int layerSize = graph.size(level); - int layerDegree = graph.getDegree(level); - int nodesWritten = 0; - for (var it = graph.getNodes(level); it.hasNext(); ) { - int originalOrdinal = it.nextInt(); - // node id - final int newOrdinal = ordinalMapper.oldToNew(originalOrdinal); - out.writeInt(newOrdinal); - // neighbors - var view = graph.getView(); - var neighbors = view.getNeighborsIterator(level, originalOrdinal); - out.writeInt(neighbors.size()); - int n = 0; - for ( ; n < neighbors.size(); n++) { - out.writeInt(ordinalMapper.oldToNew(neighbors.nextInt())); - } - assert !neighbors.hasNext() : "Mismatch between neighbor's reported size and actual size"; - // pad out to degree - for (; n < layerDegree; n++) { - out.writeInt(-1); - } - nodesWritten++; - } - if (nodesWritten != layerSize) { - throw new IllegalStateException("Mismatch between layer size and nodes written"); - } - } - } - - @Override - public synchronized void writeHeader(ImmutableGraphIndex graph, OrdinalMapper ordinalMapper, Map featureMap, int dimension, long startOffset, IndexWriter out, long headerSize) throws IOException { - // graph-level properties - var layerInfo = CommonHeader.LayerInfo.fromGraph(graph, ordinalMapper); - var commonHeader = new CommonHeader(version, - dimension, - graph.getView().entryNode() == null ? ImmutableGraphIndex.ENTRY_NODE_ABSENT : ordinalMapper.oldToNew(graph.getView().entryNode().node), - layerInfo, - ordinalMapper.maxOrdinal() + 1); - var header = new Header(commonHeader, featureMap); - header.write(out); - assert out.position() == startOffset + headerSize : String.format("%d != %d", out.position(), startOffset + headerSize); - } - - @Override - public long featureOffsetForOrdinal(ImmutableGraphIndex graph, List inlineFeatures, long startOffset, int ordinal, long headerSize) { - int edgeSize = Integer.BYTES * (1 + graph.getDegree(0)); - long inlineBytes = ordinal * (long) (Integer.BYTES + inlineFeatures.stream().mapToInt(Feature::featureSize).sum() + edgeSize); - return startOffset - + headerSize - + inlineBytes // previous nodes - + Integer.BYTES; // the ordinal of the node whose features we're about to write - } - - @Override - public void writeFooter(ImmutableGraphIndex graph, OrdinalMapper ordinalMapper, long headerOffset, int dimension, IndexWriter out, Map featureMap, long headerSize) throws IOException { - var layerInfo = CommonHeader.LayerInfo.fromGraph(graph, ordinalMapper); - var commonHeader = new CommonHeader(version, - dimension, - graph.getView().entryNode() == null ? ImmutableGraphIndex.ENTRY_NODE_ABSENT : ordinalMapper.oldToNew(graph.getView().entryNode().node), - layerInfo, - ordinalMapper.maxOrdinal() + 1); - var header = new Header(commonHeader, featureMap); - header.write(out); // write the header - out.writeLong(headerOffset); // We write the offset of the header at the end of the file - out.writeInt(FOOTER_MAGIC); - final long expectedPosition = headerOffset + headerSize + FOOTER_SIZE; - assert out.position() == expectedPosition : String.format("%d != %d", out.position(), expectedPosition); - } - - public void writeSeparatedFeatures(Map> featureStateSuppliers, Map featureMap, IndexWriter out, OrdinalMapper ordinalMapper) throws IOException { - for (var featureEntry : featureMap.entrySet()) { - if (isSeparated(featureEntry.getValue())) { - var fid = featureEntry.getKey(); - var supplier = featureStateSuppliers.get(fid); - if (supplier == null) { - throw new IllegalStateException("Supplier for feature " + fid + " not found"); - } - - // Set the offset for this feature - var feature = (SeparatedFeature) featureEntry.getValue(); - feature.setOffset(out.position()); - - // Write separated data for each node - for (int newOrdinal = 0; newOrdinal <= ordinalMapper.maxOrdinal(); newOrdinal++) { - int originalOrdinal = ordinalMapper.newToOld(newOrdinal); - if (originalOrdinal != OrdinalMapper.OMITTED) { - feature.writeSeparately(out, supplier.apply(originalOrdinal)); - } else { - // write zeros for missing data as padding - for (int i = 0; i < feature.featureSize(); i++) { - out.writeByte(0); - } - } - } - } - } - } - - boolean isSeparated(Feature feature) { - return feature instanceof SeparatedFeature; - } - - /** - * Helper to create a set of all known features. - * Only correct for the current maximum version — do not use for older versions. - * When a new version is introduced with a new feature, the previous version's serializer - * must be updated to use an explicit set that excludes the new feature. - */ - protected static Set allFeatures() { - return EnumSet.allOf(FeatureId.class); - } - - /** - * Helper to create a set of all features except FUSED_PQ. - * Used by versions 3–5, which support all non-fused features but predate fused PQ hierarchy - * support introduced in version 6. - */ - protected static Set nonFusedFeatures() { - return EnumSet.complementOf(EnumSet.of(FeatureId.FUSED_PQ)); - } - - /** - * Helper to create a set with only inline vectors (for version 2). - */ - protected static Set inlineVectorsOnly() { - return EnumSet.of(FeatureId.INLINE_VECTORS); - } -} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexWriter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexWriter.java index d8c8ed407..4f9b1dd15 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexWriter.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexWriter.java @@ -57,11 +57,11 @@ public abstract class AbstractGraphIndexWriter implements final int dimension; final Map featureMap; final T out; /* output for graph nodes and inline features */ - private final long headerSize; + final long headerSize; volatile int maxOrdinalWritten = -1; final List inlineFeatures; - final GraphIndexSerializer serializer; + final GraphIndexFormat serializer; AbstractGraphIndexWriter(T out, int version, @@ -70,8 +70,8 @@ public abstract class AbstractGraphIndexWriter implements int dimension, EnumMap features) { - serializer = GraphIndexSerializerFactory.forVersion(version); - if (!serializer.supportsMultiLayer()) { + serializer = GraphIndexFormatFactory.forVersion(version); + if (graph.isHierarchical() && !serializer.supportsMultiLayer()) { throw new IllegalArgumentException("Multilayer graphs must be written with version 4 or higher"); } this.version = version; @@ -124,8 +124,12 @@ public Set getFeatureSet() { return featureMap.keySet(); } + WriteContext createContext(long startOffset) { + return new WriteContext(graph, ordinalMapper, featureMap, inlineFeatures, startOffset, headerSize, dimension); + } + long featureOffsetForOrdinal(long startOffset, int ordinal) { - return serializer.featureOffsetForOrdinal(graph, inlineFeatures, startOffset, ordinal, headerSize); + return serializer.featureOffsetForOrdinal(createContext(startOffset), ordinal); } /** @@ -164,8 +168,8 @@ public static Map sequentialRenumbering(ImmutableGraphIndex gr * @param headerOffset the offset of the header in the slice * @throws IOException IOException */ - void writeFooter(ImmutableGraphIndex.View view, long headerOffset) throws IOException { - serializer.writeFooter(graph, ordinalMapper, headerOffset, dimension, out, featureMap, headerSize); + void writeFooter(ImmutableGraphIndex.View view, long headerOffset, long startOffset) throws IOException { + serializer.writeFooter(createContext(startOffset), headerOffset, out); } /** @@ -176,15 +180,15 @@ void writeFooter(ImmutableGraphIndex.View view, long headerOffset) throws IOExce * @throws IOException if an I/O error occurs */ protected synchronized void writeHeader(ImmutableGraphIndex.View view, long startOffset) throws IOException { - serializer.writeHeader(graph, ordinalMapper, featureMap, dimension, startOffset, out, headerSize); + serializer.writeHeader(createContext(startOffset), out); } - void writeSparseLevels(ImmutableGraphIndex.View view, Map> featureStateSuppliers) throws IOException { - serializer.writeSparseLevels(graph, ordinalMapper, out, featureStateSuppliers); + void writeSparseLevels(ImmutableGraphIndex.View view, Map> featureStateSuppliers, long startOffset) throws IOException { + serializer.writeSparseLevels(createContext(startOffset), out, featureStateSuppliers); } - void writeSeparatedFeatures(Map> featureStateSuppliers) throws IOException { - serializer.writeSeparatedFeatures(featureStateSuppliers, featureMap, out, ordinalMapper); + void writeSeparatedFeatures(Map> featureStateSuppliers, long startOffset) throws IOException { + serializer.writeSeparatedFeatures(createContext(startOffset), out, featureStateSuppliers); } /** diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormat.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormat.java new file mode 100644 index 000000000..7e97604bb --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormat.java @@ -0,0 +1,191 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph.disk; + +import io.github.jbellis.jvector.disk.IndexWriter; +import io.github.jbellis.jvector.disk.RandomAccessWriter; +import io.github.jbellis.jvector.graph.ImmutableGraphIndex; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; + +import java.io.IOException; +import java.util.Map; +import java.util.Set; +import java.util.function.IntFunction; + +/** + * Strategy interface for version-specific serialization of graph indexes. + * Each version of the on-disk format has its own implementation that encapsulates + * all version-specific logic for reading and writing graph data. + *

+ * This design eliminates scattered version checks throughout the codebase and + * makes it easy to add new versions without modifying existing code. + */ +public interface GraphIndexFormat { + /** + * @return the version number this format handles + */ + int getVersion(); + + /** + * Checks if this version supports a given feature. + * @param feature the feature to check + * @return true if the feature is supported in this version + */ + boolean supportsFeature(FeatureId feature); + + /** + * Checks if this version supports multi-layer (hierarchical) graphs. + * @return true if multi-layer graphs are supported + */ + boolean supportsMultiLayer(); + + /** + * Checks if this version uses a footer for metadata instead of a header. + * @return true if footer-based metadata is used + */ + boolean usesFooter(); + + /** + * Gets the feature ordering strategy for this version. + * @return the feature ordering strategy + */ + FeatureOrdering getFeatureOrdering(); + + /** + * Returns the complete set of {@link FeatureId}s that this format version is capable of storing. + * + * @return an unmodifiable set of supported feature identifiers + */ + Set getSupportedFeatures(); + + /** + * Callback for writing L0 (base layer) node records. Implemented by the writer so that + * I/O strategy (sequential vs. parallel) stays in the writer while the format + * handles orchestration without depending on the concrete writer type. + */ + @FunctionalInterface + interface L0RecordWriter { + /** + * Writes the base-layer (level-0) node records using the provided graph view and feature suppliers. + * + * @param view a view over the graph used to iterate neighbors at level 0 + * @param suppliers per-feature functions that produce a {@link Feature.State} for a given original ordinal + * @throws IOException if an I/O error occurs while writing + */ + void write(ImmutableGraphIndex.View view, Map> suppliers) throws IOException; + } + + /** + * Defines how features should be ordered when writing/reading. + */ + enum FeatureOrdering { + /** Features ordered by their FeatureId enum ordinal (versions <= 5) */ + BY_FEATURE_ID, + /** Features ordered with fused features last (version 6+) */ + FUSED_LAST + } + + /** + * Writes adjacency records for all graph levels above level 0 (the "sparse" upper layers). + * + * @param ctx the write session context + * @param out the sequential output stream + * @param suppliers per-feature state suppliers keyed by {@link FeatureId} + * @throws IOException if an I/O error occurs while writing + */ + void writeSparseLevels(WriteContext ctx, IndexWriter out, Map> suppliers) throws IOException; + + /** + * Writes the format header (version, dimensions, entry node, layer info, and feature metadata). + * + * @param ctx the write session context + * @param out the sequential output stream positioned at the header location + * @throws IOException if an I/O error occurs while writing + */ + void writeHeader(WriteContext ctx, IndexWriter out) throws IOException; + + /** + * Returns the byte offset within the output stream at which the inline feature data + * for the given (new) ordinal begins. + * + * @param ctx the write session context + * @param ordinal the compacted (new) node ordinal + * @return the absolute byte offset for the node's inline feature section + */ + long featureOffsetForOrdinal(WriteContext ctx, int ordinal); + + /** + * Writes the format footer containing the header offset and magic number, + * used by footer-based formats (V5+) to locate the header at read time. + * + * @param ctx the write session context + * @param headerOffset the byte position in the output stream where the header was written + * @param out the sequential output stream + * @throws IOException if an I/O error occurs while writing + */ + void writeFooter(WriteContext ctx, long headerOffset, IndexWriter out) throws IOException; + + /** + * Writes all {@link io.github.jbellis.jvector.graph.disk.feature.SeparatedFeature} data blocks + * sequentially after the node records. + * + * @param ctx the write session context + * @param out the sequential output stream + * @param suppliers per-feature state suppliers keyed by {@link FeatureId} + * @throws IOException if an I/O error occurs while writing + */ + void writeSeparatedFeatures(WriteContext ctx, IndexWriter out, Map> suppliers) throws IOException; + + /** + * Writes inline feature data for a single node directly into the random-access output stream + * at the position determined by {@link #featureOffsetForOrdinal}. + * + * @param ctx the write session context + * @param ordinal the compacted (new) node ordinal whose features are being written + * @param stateMap feature states to write, keyed by {@link FeatureId} + * @param out the random-access output stream + * @throws IOException if an I/O error occurs while writing + */ + void writeFeaturesInline(WriteContext ctx, int ordinal, Map stateMap, RandomAccessWriter out) throws IOException; + + /** + * Writes the complete graph index sequentially: header, node records, sparse levels, + * separated features, and (for V5+) footer. All nodes must be present in the ordinal + * mapper; gaps (OMITTED ordinals) are not permitted. + * + * @param ctx the write session context + * @param out the sequential output stream + * @param suppliers per-feature state suppliers keyed by {@link FeatureId} + * @throws IOException if an I/O error occurs while writing + */ + void writeOnDiskSequential(WriteContext ctx, IndexWriter out, Map> suppliers) throws IOException; + + /** + * Writes the complete graph index using random-access I/O, delegating base-layer record + * writing to the provided {@link L0RecordWriter} (which may use parallel I/O). + * After all data is written the header is re-written at its original position so that + * accurate offset information is recorded. + * + * @param ctx the write session context + * @param out the random-access output stream + * @param suppliers per-feature state suppliers keyed by {@link FeatureId} + * @param l0Writer callback responsible for writing the level-0 node records + * @throws IOException if an I/O error occurs while writing + */ + void writeRandomAccess(WriteContext ctx, RandomAccessWriter out, Map> suppliers, L0RecordWriter l0Writer) throws IOException; +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerFactory.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatFactory.java similarity index 70% rename from jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerFactory.java rename to jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatFactory.java index b1a63079f..12010bc1e 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerFactory.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatFactory.java @@ -24,50 +24,50 @@ import java.util.Map; /** - * Factory for creating version-specific graph index serializers. - * This centralizes version detection and serializer instantiation. + * Factory for creating version-specific graph index formats. + * This centralizes version detection and format instantiation. */ -public class GraphIndexSerializerFactory { - private static final Logger logger = LoggerFactory.getLogger(GraphIndexSerializerFactory.class); +public class GraphIndexFormatFactory { + private static final Logger logger = LoggerFactory.getLogger(GraphIndexFormatFactory.class); - private static final Map SERIALIZERS = Map.of( - 2, new GraphIndexSerializerV2(), - 3, new GraphIndexSerializerV3(), - 4, new GraphIndexSerializerV4(), - 5, new GraphIndexSerializerV5(), - 6, new GraphIndexSerializerV6() + private static final Map FORMATS = Map.of( + 2, new GraphIndexFormatV2(), + 3, new GraphIndexFormatV3(), + 4, new GraphIndexFormatV4(), + 5, new GraphIndexFormatV5(), + 6, new GraphIndexFormatV6() ); /** - * Gets a serializer for a specific version. + * Gets a format for a specific version. * @param version the version number - * @return the serializer for that version + * @return the format for that version * @throws UnsupportedVersionException if the version is not supported */ - public static GraphIndexSerializer forVersion(int version) { - GraphIndexSerializer serializer = SERIALIZERS.get(version); - if (serializer == null) { + public static GraphIndexFormat forVersion(int version) { + GraphIndexFormat format = FORMATS.get(version); + if (format == null) { throw new UnsupportedVersionException("Version " + version + " is not supported. " + - "Supported versions: " + SERIALIZERS.keySet()); + "Supported versions: " + FORMATS.keySet()); } - return serializer; + return format; } /** - * Detects the version from the input stream and returns the appropriate serializer. + * Detects the version from the input stream and returns the appropriate format. * The reader position will be reset to where it started. - * + * * @param in the input reader - * @return the serializer for the detected version + * @return the format for the detected version * @throws IOException if an I/O error occurs * @throws UnsupportedVersionException if the version is not supported */ - public static GraphIndexSerializer detectVersion(RandomAccessReader in) throws IOException { + public static GraphIndexFormat detectVersion(RandomAccessReader in) throws IOException { long startPosition = in.getPosition(); - + try { int maybeMagic = in.readInt(); - + if (maybeMagic == OnDiskGraphIndex.MAGIC) { // Version 3+ with magic number int version = in.readInt(); diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV2.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV2.java new file mode 100644 index 000000000..199c51257 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV2.java @@ -0,0 +1,33 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph.disk; + +/** + * Format for version 2 of the on-disk graph format. + * Version 2 characteristics: + * - No magic number + * - Only supports INLINE_VECTORS feature + * - Single layer only + * - No footer + */ +class GraphIndexFormatV2 extends AbstractGraphIndexFormat { + + /** Creates the singleton format for version 2. */ + GraphIndexFormatV2() { + super(2, inlineVectorsOnly(), false, false, FeatureOrdering.BY_FEATURE_ID); + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV3.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV3.java new file mode 100644 index 000000000..1e44975b3 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV3.java @@ -0,0 +1,33 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph.disk; + +/** + * Format for version 3 of the on-disk graph format. + * Version 3 characteristics: + * - Has magic number + * - Supports multiple features (feature set serialization) + * - Single layer only + * - No footer + */ +class GraphIndexFormatV3 extends AbstractGraphIndexFormat { + + /** Creates the singleton format for version 3. */ + GraphIndexFormatV3() { + super(3, nonFusedFeatures(), false, false, FeatureOrdering.BY_FEATURE_ID); + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV4.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV4.java new file mode 100644 index 000000000..1d0034800 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV4.java @@ -0,0 +1,46 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph.disk; + +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; + +import java.util.Set; + +/** + * Format for version 4 of the on-disk graph format. + * Version 4 characteristics: + * - Has magic number + * - Supports multiple features + * - Supports multi-layer (hierarchical) graphs + * - Has idUpperBound field + * - No footer + */ +class GraphIndexFormatV4 extends AbstractGraphIndexFormat { + + /** Creates the singleton format for version 4. */ + GraphIndexFormatV4() { + super(4, nonFusedFeatures(), true, false, FeatureOrdering.BY_FEATURE_ID); + } + + /** + * Protected constructor allowing subclasses (V5, V6) to specify their own version, + * feature set, footer flag, and feature ordering while sharing V4's wire format. + */ + protected GraphIndexFormatV4(int version, Set supportedFeatures, boolean usesFooter, FeatureOrdering featureOrdering) { + super(version, supportedFeatures, true, usesFooter, featureOrdering); + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV5.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV5.java new file mode 100644 index 000000000..4f0f14b9d --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV5.java @@ -0,0 +1,99 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph.disk; + +import io.github.jbellis.jvector.disk.RandomAccessWriter; +import io.github.jbellis.jvector.graph.ImmutableGraphIndex; +import io.github.jbellis.jvector.graph.OnHeapGraphIndex; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; + +import java.io.IOException; +import java.util.Map; +import java.util.Set; +import java.util.function.IntFunction; + +/** + * Format for version 5 of the on-disk graph format. + * Version 5 characteristics: + * - Has magic number + * - Supports multiple features + * - Supports multi-layer (hierarchical) graphs + * - Has idUpperBound field + * - Uses footer for metadata (major change from V4) + * + * The wire format is identical to V4; the only behavioral difference is that + * a footer is written after the graph data. + */ +class GraphIndexFormatV5 extends GraphIndexFormatV4 { + + /** Creates the singleton format for version 5. */ + GraphIndexFormatV5() { + super(5, nonFusedFeatures(), true, FeatureOrdering.BY_FEATURE_ID); + } + + /** + * Protected constructor for subclasses (V6) to specify their own version, + * feature set, and feature ordering while inheriting V5's footer-writing behavior. + * Footer is always true for V5 and later. + */ + protected GraphIndexFormatV5(int version, Set supportedFeatures, FeatureOrdering featureOrdering) { + super(version, supportedFeatures, true, featureOrdering); + } + + /** + * Writes the complete graph index using random-access I/O, and additionally appends a + * footer (header offset + magic number) after the graph data — the key behavioral difference + * from the V4 implementation, which does not write a footer. + * + * {@inheritDoc} + */ + @Override + public void writeRandomAccess(WriteContext ctx, RandomAccessWriter out, Map> suppliers, GraphIndexFormat.L0RecordWriter l0Writer) throws IOException { + if (ctx.graph instanceof OnHeapGraphIndex) { + var ohgi = (OnHeapGraphIndex) ctx.graph; + if (ohgi.getDeletedNodes().cardinality() > 0) { + throw new IllegalArgumentException("Run builder.cleanup() before writing the graph"); + } + } + for (var featureId : suppliers.keySet()) { + if (!ctx.featureMap.containsKey(featureId)) { + throw new IllegalArgumentException(String.format("Feature %s not configured for index", featureId)); + } + } + if (ctx.ordinalMapper.maxOrdinal() < ctx.graph.size(0) - 1) { + throw new IllegalStateException(String.format("Ordinal mapper from [0..%d] does not cover all nodes in the graph of size %d", + ctx.ordinalMapper.maxOrdinal(), ctx.graph.size(0))); + } + + var view = ctx.graph.getView(); + + out.seek(ctx.startOffset); + writeHeader(ctx, out); + l0Writer.write(view, suppliers); + writeSparseLevels(ctx, out, suppliers); + writeSeparatedFeatures(ctx, out, suppliers); + writeFooter(ctx, out.position(), out); + + final var endOfGraphPosition = out.position(); + out.seek(ctx.startOffset); + writeHeader(ctx, out); + out.seek(endOfGraphPosition); + out.flush(); + view.close(); + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV6.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV6.java new file mode 100644 index 000000000..ed8bd7399 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV6.java @@ -0,0 +1,85 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph.disk; + +import io.github.jbellis.jvector.disk.IndexWriter; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.graph.disk.feature.FusedFeature; + +import java.io.IOException; +import java.util.Map; +import java.util.function.IntFunction; + +/** + * Format for version 6 of the on-disk graph format. + * Version 6 characteristics: + * - Has magic number + * - Supports multiple features + * - Supports multi-layer (hierarchical) graphs + * - Has idUpperBound field + * - Uses footer for metadata + * - Changes feature ordering to place fused features last + * - Writes feature count and ordinals explicitly instead of bitflags + */ +class GraphIndexFormatV6 extends GraphIndexFormatV5 { + + /** Creates the singleton format for version 6. */ + GraphIndexFormatV6() { + super(6, allFeatures(), FeatureOrdering.FUSED_LAST); + } + + /** + * Appends fused feature data for level-1 nodes after the sparse levels are written. + * At most one fused feature is permitted per graph (enforced at write time). + */ + @Override + protected void writeAfterSparseLevels(WriteContext ctx, IndexWriter out, Map> suppliers) throws IOException { + FusedFeature fusedFeature = null; + for (var feature : ctx.inlineFeatures) { + if (feature.isFused()) { + fusedFeature = (FusedFeature) feature; + } + } + if (fusedFeature == null) { + return; + } + + var supplier = suppliers.get(fusedFeature.id()); + if (supplier == null) { + throw new IllegalStateException("Supplier for feature " + fusedFeature.id() + " not found"); + } + + if (ctx.graph.getMaxLevel() >= 1) { + int layerSize = ctx.graph.size(1); + int nodesWritten = 0; + for (var it = ctx.graph.getNodes(1); it.hasNext(); ) { + int originalOrdinal = it.nextInt(); + out.writeInt(ctx.ordinalMapper.oldToNew(originalOrdinal)); + fusedFeature.writeSourceFeature(out, supplier.apply(originalOrdinal)); + nodesWritten++; + } + if (nodesWritten != layerSize) { + throw new IllegalStateException("Mismatch between layer 1 size and features written"); + } + } else { + final int originalEntryNode = ctx.graph.getView().entryNode().node; + out.writeInt(ctx.ordinalMapper.oldToNew(originalEntryNode)); + fusedFeature.writeSourceFeature(out, supplier.apply(originalEntryNode)); + } + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexMetadata.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexMetadata.java deleted file mode 100644 index dc5a2238b..000000000 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexMetadata.java +++ /dev/null @@ -1,119 +0,0 @@ -/* - * Copyright DataStax, Inc. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package io.github.jbellis.jvector.graph.disk; - -import java.util.List; -import java.util.Objects; - -/** - * Metadata about a graph index, containing all the information needed to - * interpret the on-disk format. This replaces the scattered fields that were - * previously in CommonHeader and provides a cleaner API. - */ -public class GraphIndexMetadata { - private final int version; - private final int dimension; - private final int entryNode; - private final List layerInfo; - private final int idUpperBound; - - public GraphIndexMetadata(int version, int dimension, int entryNode, - List layerInfo, int idUpperBound) { - this.version = version; - this.dimension = dimension; - this.entryNode = entryNode; - this.layerInfo = layerInfo; - this.idUpperBound = idUpperBound; - } - - public int getVersion() { - return version; - } - - public int getDimension() { - return dimension; - } - - public int getEntryNode() { - return entryNode; - } - - public List getLayerInfo() { - return layerInfo; - } - - public int getIdUpperBound() { - return idUpperBound; - } - - /** - * Gets the size of the base layer (layer 0). - */ - public int getBaseLayerSize() { - return layerInfo.get(0).size; - } - - /** - * Gets the max degree of the base layer (layer 0). - */ - public int getBaseLayerDegree() { - return layerInfo.get(0).degree; - } - - /** - * Gets the number of layers in the graph. - */ - public int getNumLayers() { - return layerInfo.size(); - } - - /** - * Checks if this is a multi-layer (hierarchical) graph. - */ - public boolean isMultiLayer() { - return layerInfo.size() > 1; - } - - @Override - public boolean equals(Object o) { - if (this == o) return true; - if (o == null || getClass() != o.getClass()) return false; - GraphIndexMetadata that = (GraphIndexMetadata) o; - return version == that.version && - dimension == that.dimension && - entryNode == that.entryNode && - idUpperBound == that.idUpperBound && - Objects.equals(layerInfo, that.layerInfo); - } - - @Override - public int hashCode() { - return Objects.hash(version, dimension, entryNode, layerInfo, idUpperBound); - } - - @Override - public String toString() { - return "GraphIndexMetadata{" + - "version=" + version + - ", dimension=" + dimension + - ", entryNode=" + entryNode + - ", layers=" + layerInfo.size() + - ", idUpperBound=" + idUpperBound + - '}'; - } -} - diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializer.java deleted file mode 100644 index d2e95423f..000000000 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializer.java +++ /dev/null @@ -1,132 +0,0 @@ -/* - * Copyright DataStax, Inc. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package io.github.jbellis.jvector.graph.disk; - -import io.github.jbellis.jvector.disk.IndexWriter; -import io.github.jbellis.jvector.disk.RandomAccessReader; -import io.github.jbellis.jvector.graph.ImmutableGraphIndex; -import io.github.jbellis.jvector.graph.disk.feature.Feature; -import io.github.jbellis.jvector.graph.disk.feature.FeatureId; - -import java.io.IOException; -import java.util.List; -import java.util.Map; -import java.util.Set; -import java.util.function.IntFunction; - -/** - * Strategy interface for version-specific serialization of graph indexes. - * Each version of the on-disk format has its own implementation that encapsulates - * all version-specific logic for reading and writing graph data. - *

- * This design eliminates scattered version checks throughout the codebase and - * makes it easy to add new versions without modifying existing code. - */ -public interface GraphIndexSerializer { - /** - * @return the version number this serializer handles - */ - int getVersion(); - - /** - * Checks if this version supports a given feature. - * @param feature the feature to check - * @return true if the feature is supported in this version - */ - boolean supportsFeature(FeatureId feature); - - /** - * Checks if this version supports multi-layer (hierarchical) graphs. - * @return true if multi-layer graphs are supported - */ - boolean supportsMultiLayer(); - - /** - * Checks if this version uses a footer for metadata instead of a header. - * @return true if footer-based metadata is used - */ - boolean usesFooter(); - - /** - * Gets the feature ordering strategy for this version. - * @return the feature ordering strategy - */ - FeatureOrdering getFeatureOrdering(); - - /** - * Writes the common header portion of the index. - * @param out the output writer - * @param metadata the metadata to write - * @throws IOException if an I/O error occurs - */ - void writeCommonHeader(IndexWriter out, GraphIndexMetadata metadata) throws IOException; - - /** - * Reads the common header portion of the index. - * @param in the input reader - * @return the metadata read from the header - * @throws IOException if an I/O error occurs - */ - GraphIndexMetadata readCommonHeader(RandomAccessReader in) throws IOException; - - /** - * Writes feature-specific header information. - * @param out the output writer - * @param features the features to write headers for - * @throws IOException if an I/O error occurs - */ - void writeFeatureHeaders(IndexWriter out, Map features) throws IOException; - - /** - * Reads feature-specific header information. - * @param in the input reader - * @param metadata the common metadata (needed for feature construction) - * @return map of feature IDs to feature instances - * @throws IOException if an I/O error occurs - */ - Map readFeatureHeaders(RandomAccessReader in, GraphIndexMetadata metadata) throws IOException; - - /** - * Calculates the total size of the header in bytes. - * @param metadata the metadata - * @param features the features - * @return the header size in bytes - */ - int calculateHeaderSize(GraphIndexMetadata metadata, Map features); - - public Set getSupportedFeatures(); - - /** - * Defines how features should be ordered when writing/reading. - */ - enum FeatureOrdering { - /** Features ordered by their FeatureId enum ordinal (versions <= 5) */ - BY_FEATURE_ID, - /** Features ordered with fused features last (version 6+) */ - FUSED_LAST - } - - void writeSparseLevels(ImmutableGraphIndex graph, OrdinalMapper ordinalMapper, IndexWriter out, Map> featureStateSuppliers) throws IOException; - - void writeHeader(ImmutableGraphIndex graph, OrdinalMapper ordinalMapper, Map featureMap, int dimension, long startOffset, IndexWriter out, long headerSize) throws IOException; - - long featureOffsetForOrdinal(ImmutableGraphIndex graph, List inlineFeatures, long startOffset, int ordinal, long headerSize); - - void writeFooter(ImmutableGraphIndex graph, OrdinalMapper ordinalMapper, long headerOffset, int dimension, IndexWriter out, Map featureMap, long headerSize) throws IOException; - - void writeSeparatedFeatures(Map> featureStateSuppliers, Map featureMap, IndexWriter out, OrdinalMapper ordinalMapper) throws IOException; - } \ No newline at end of file diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV2.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV2.java deleted file mode 100644 index 3dfe3c1b4..000000000 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV2.java +++ /dev/null @@ -1,111 +0,0 @@ -/* - * Copyright DataStax, Inc. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package io.github.jbellis.jvector.graph.disk; - -import io.github.jbellis.jvector.disk.IndexWriter; -import io.github.jbellis.jvector.disk.RandomAccessReader; -import io.github.jbellis.jvector.graph.disk.feature.Feature; -import io.github.jbellis.jvector.graph.disk.feature.FeatureId; - -import java.io.IOException; -import java.util.EnumMap; -import java.util.EnumSet; -import java.util.List; -import java.util.Map; - -/** - * Serializer for version 2 of the on-disk graph format. - * Version 2 characteristics: - * - No magic number - * - Only supports INLINE_VECTORS feature - * - Single layer only - * - No footer - */ -class GraphIndexSerializerV2 extends AbstractGraphIndexSerializer { - - GraphIndexSerializerV2() { - super(2, inlineVectorsOnly(), false, false, FeatureOrdering.BY_FEATURE_ID); - } - - @Override - public void writeCommonHeader(IndexWriter out, GraphIndexMetadata metadata) throws IOException { - // V2 format: size, dimension, entryNode, maxDegree (no magic, no version) - out.writeInt(metadata.getBaseLayerSize()); - out.writeInt(metadata.getDimension()); - out.writeInt(metadata.getEntryNode()); - out.writeInt(metadata.getBaseLayerDegree()); - } - - @Override - public GraphIndexMetadata readCommonHeader(RandomAccessReader in) throws IOException { - // V2 format: size, dimension, entryNode, maxDegree - int size = in.readInt(); - int dimension = in.readInt(); - int entryNode = in.readInt(); - int maxDegree = in.readInt(); - - // V2 only supports single layer - List layerInfo = List.of(new CommonHeader.LayerInfo(size, maxDegree)); - int idUpperBound = size; - - return new GraphIndexMetadata(2, dimension, entryNode, layerInfo, idUpperBound); - } - - @Override - public void writeFeatureHeaders(IndexWriter out, Map features) throws IOException { - // V2 doesn't write feature set information, just the feature headers - // Only INLINE_VECTORS is supported - if (!features.containsKey(FeatureId.INLINE_VECTORS) || features.size() > 1) { - throw new IllegalArgumentException("Version 2 only supports INLINE_VECTORS feature"); - } - - for (Feature feature : features.values()) { - feature.writeHeader(out); - } - } - - @Override - public Map readFeatureHeaders(RandomAccessReader in, GraphIndexMetadata metadata) throws IOException { - // V2 only has INLINE_VECTORS - Map features = new EnumMap<>(FeatureId.class); - EnumSet featureIds = EnumSet.of(FeatureId.INLINE_VECTORS); - - // Create a CommonHeader from metadata for feature loading - CommonHeader commonHeader = new CommonHeader( - metadata.getVersion(), - metadata.getDimension(), - metadata.getEntryNode(), - metadata.getLayerInfo(), - metadata.getIdUpperBound() - ); - - for (FeatureId featureId : featureIds) { - features.put(featureId, featureId.load(commonHeader, in)); - } - - return features; - } - - @Override - public int calculateHeaderSize(GraphIndexMetadata metadata, Map features) { - // V2: 4 ints (size, dimension, entryNode, maxDegree) + feature headers - int size = 4 * Integer.BYTES; - size += features.values().stream().mapToInt(Feature::headerSize).sum(); - return size; - } -} - diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV3.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV3.java deleted file mode 100644 index 3b7ba0b04..000000000 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV3.java +++ /dev/null @@ -1,119 +0,0 @@ -/* - * Copyright DataStax, Inc. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package io.github.jbellis.jvector.graph.disk; - -import io.github.jbellis.jvector.disk.IndexWriter; -import io.github.jbellis.jvector.disk.RandomAccessReader; -import io.github.jbellis.jvector.graph.disk.feature.Feature; -import io.github.jbellis.jvector.graph.disk.feature.FeatureId; - -import java.io.IOException; -import java.util.EnumMap; -import java.util.EnumSet; -import java.util.List; -import java.util.Map; - -/** - * Serializer for version 3 of the on-disk graph format. - * Version 3 characteristics: - * - Has magic number - * - Supports multiple features (feature set serialization) - * - Single layer only - * - No footer - */ -class GraphIndexSerializerV3 extends AbstractGraphIndexSerializer { - - GraphIndexSerializerV3() { - super(3, nonFusedFeatures(), false, false, FeatureOrdering.BY_FEATURE_ID); - } - - @Override - public void writeCommonHeader(IndexWriter out, GraphIndexMetadata metadata) throws IOException { - // V3 format: magic, version, size, dimension, entryNode, maxDegree - out.writeInt(OnDiskGraphIndex.MAGIC); - out.writeInt(3); - out.writeInt(metadata.getBaseLayerSize()); - out.writeInt(metadata.getDimension()); - out.writeInt(metadata.getEntryNode()); - out.writeInt(metadata.getBaseLayerDegree()); - } - - @Override - public GraphIndexMetadata readCommonHeader(RandomAccessReader in) throws IOException { - // V3 format: magic, version, size, dimension, entryNode, maxDegree - int magic = in.readInt(); - if (magic != OnDiskGraphIndex.MAGIC) { - throw new IOException("Invalid magic number: " + magic); - } - int version = in.readInt(); - int size = in.readInt(); - int dimension = in.readInt(); - int entryNode = in.readInt(); - int maxDegree = in.readInt(); - - // V3 only supports single layer - List layerInfo = List.of(new CommonHeader.LayerInfo(size, maxDegree)); - int idUpperBound = size; - - return new GraphIndexMetadata(version, dimension, entryNode, layerInfo, idUpperBound); - } - - @Override - public void writeFeatureHeaders(IndexWriter out, Map features) throws IOException { - // V3 writes feature set as bitflags - out.writeInt(FeatureId.serialize(EnumSet.copyOf(features.keySet()))); - - // Then write each feature's header - for (Feature feature : features.values()) { - feature.writeHeader(out); - } - } - - @Override - public Map readFeatureHeaders(RandomAccessReader in, GraphIndexMetadata metadata) throws IOException { - Map features = new EnumMap<>(FeatureId.class); - - // Read feature set bitflags - EnumSet featureIds = FeatureId.deserialize(in.readInt()); - - // Create a CommonHeader from metadata for feature loading - CommonHeader commonHeader = new CommonHeader( - metadata.getVersion(), - metadata.getDimension(), - metadata.getEntryNode(), - metadata.getLayerInfo(), - metadata.getIdUpperBound() - ); - - // Load each feature - for (FeatureId featureId : featureIds) { - features.put(featureId, featureId.load(commonHeader, in)); - } - - return features; - } - - @Override - public int calculateHeaderSize(GraphIndexMetadata metadata, Map features) { - // V3: magic + version + 4 ints (size, dimension, entryNode, maxDegree) + feature bitflags + feature headers - int size = 6 * Integer.BYTES; // magic, version, size, dimension, entryNode, maxDegree - size += Integer.BYTES; // feature bitflags - size += features.values().stream().mapToInt(Feature::headerSize).sum(); - return size; - } -} - diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV4.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV4.java deleted file mode 100644 index 00798ed18..000000000 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV4.java +++ /dev/null @@ -1,162 +0,0 @@ -/* - * Copyright DataStax, Inc. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package io.github.jbellis.jvector.graph.disk; - -import io.github.jbellis.jvector.disk.IndexWriter; -import io.github.jbellis.jvector.disk.RandomAccessReader; -import io.github.jbellis.jvector.graph.disk.feature.Feature; -import io.github.jbellis.jvector.graph.disk.feature.FeatureId; - -import java.io.IOException; -import java.util.ArrayList; -import java.util.EnumMap; -import java.util.EnumSet; -import java.util.List; -import java.util.Map; -import java.util.Set; - -/** - * Serializer for version 4 of the on-disk graph format. - * Version 4 characteristics: - * - Has magic number - * - Supports multiple features - * - Supports multi-layer (hierarchical) graphs - * - Has idUpperBound field - * - No footer - */ -class GraphIndexSerializerV4 extends AbstractGraphIndexSerializer { - private static final int V4_MAX_LAYERS = 32; - - GraphIndexSerializerV4() { - this(nonFusedFeatures()); - } - - /** - * Protected constructor for subclasses (V5, V6) that share V4's wire format but differ - * in supported features, footer behaviour, or feature ordering. - */ - protected GraphIndexSerializerV4(Set supportedFeatures) { - super(4, supportedFeatures, true, false, FeatureOrdering.BY_FEATURE_ID); - } - - @Override - public void writeCommonHeader(IndexWriter out, GraphIndexMetadata metadata) throws IOException { - // V4 format: magic, version, size, dimension, entryNode, maxDegree, idUpperBound, numLayers, layer info - out.writeInt(OnDiskGraphIndex.MAGIC); - out.writeInt(4); - out.writeInt(metadata.getBaseLayerSize()); - out.writeInt(metadata.getDimension()); - out.writeInt(metadata.getEntryNode()); - out.writeInt(metadata.getBaseLayerDegree()); - out.writeInt(metadata.getIdUpperBound()); - - if (metadata.getLayerInfo().size() > V4_MAX_LAYERS) { - throw new IllegalArgumentException( - String.format("Number of layers %d exceeds maximum of %d", - metadata.getLayerInfo().size(), V4_MAX_LAYERS)); - } - - out.writeInt(metadata.getLayerInfo().size()); - - // Write actual layer info - for (CommonHeader.LayerInfo info : metadata.getLayerInfo()) { - out.writeInt(info.size); - out.writeInt(info.degree); - } - - // Pad remaining entries with zeros - for (int i = metadata.getLayerInfo().size(); i < V4_MAX_LAYERS; i++) { - out.writeInt(0); // size - out.writeInt(0); // degree - } - } - - @Override - public GraphIndexMetadata readCommonHeader(RandomAccessReader in) throws IOException { - int magic = in.readInt(); - if (magic != OnDiskGraphIndex.MAGIC) { - throw new IOException("Invalid magic number: " + magic); - } - int version = in.readInt(); - int size = in.readInt(); - int dimension = in.readInt(); - int entryNode = in.readInt(); - int maxDegree = in.readInt(); - int idUpperBound = in.readInt(); - int numLayers = in.readInt(); - - List layerInfo = new ArrayList<>(); - for (int i = 0; i < numLayers; i++) { - CommonHeader.LayerInfo info = new CommonHeader.LayerInfo(in.readInt(), in.readInt()); - layerInfo.add(info); - } - - // Skip over remaining padding entries - for (int i = numLayers; i < V4_MAX_LAYERS; i++) { - in.readInt(); - in.readInt(); - } - - return new GraphIndexMetadata(version, dimension, entryNode, layerInfo, idUpperBound); - } - - @Override - public void writeFeatureHeaders(IndexWriter out, Map features) throws IOException { - // V4 writes feature set as bitflags - out.writeInt(FeatureId.serialize(EnumSet.copyOf(features.keySet()))); - - // Then write each feature's header - for (Feature feature : features.values()) { - feature.writeHeader(out); - } - } - - @Override - public Map readFeatureHeaders(RandomAccessReader in, GraphIndexMetadata metadata) throws IOException { - Map features = new EnumMap<>(FeatureId.class); - - // Read feature set bitflags - EnumSet featureIds = FeatureId.deserialize(in.readInt()); - - // Create a CommonHeader from metadata for feature loading - CommonHeader commonHeader = new CommonHeader( - metadata.getVersion(), - metadata.getDimension(), - metadata.getEntryNode(), - metadata.getLayerInfo(), - metadata.getIdUpperBound() - ); - - // Load each feature - for (FeatureId featureId : featureIds) { - features.put(featureId, featureId.load(commonHeader, in)); - } - - return features; - } - - @Override - public int calculateHeaderSize(GraphIndexMetadata metadata, Map features) { - // V4: magic + version + 4 base ints + idUpperBound + numLayers + (32 * 2 layer info ints) + feature bitflags + feature headers - int size = 8 * Integer.BYTES; // magic, version, size, dimension, entryNode, maxDegree, idUpperBound, numLayers - size += 2 * V4_MAX_LAYERS * Integer.BYTES; // layer info (padded to 32 layers) - size += Integer.BYTES; // feature bitflags - size += features.values().stream().mapToInt(Feature::headerSize).sum(); - return size; - } -} - diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV5.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV5.java deleted file mode 100644 index 752f7e589..000000000 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV5.java +++ /dev/null @@ -1,47 +0,0 @@ -/* - * Copyright DataStax, Inc. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package io.github.jbellis.jvector.graph.disk; - -/** - * Serializer for version 5 of the on-disk graph format. - * Version 5 characteristics: - * - Has magic number - * - Supports multiple features - * - Supports multi-layer (hierarchical) graphs - * - Has idUpperBound field - * - Uses footer for metadata (major change from V4) - * - * V5 is identical to V4 in terms of format, but uses a footer instead of - * relying on the header at the beginning. The serialization logic is the same. - */ -class GraphIndexSerializerV5 extends GraphIndexSerializerV4 { - - GraphIndexSerializerV5() { - super(nonFusedFeatures()); - } - - @Override - public int getVersion() { - return 5; - } - - @Override - public boolean usesFooter() { - return true; - } -} - diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV6.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV6.java deleted file mode 100644 index 1bc4e50d6..000000000 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexSerializerV6.java +++ /dev/null @@ -1,181 +0,0 @@ -/* - * Copyright DataStax, Inc. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package io.github.jbellis.jvector.graph.disk; - -import io.github.jbellis.jvector.disk.IndexWriter; -import io.github.jbellis.jvector.disk.RandomAccessReader; -import io.github.jbellis.jvector.graph.ImmutableGraphIndex; -import io.github.jbellis.jvector.graph.disk.feature.Feature; -import io.github.jbellis.jvector.graph.disk.feature.FeatureId; -import io.github.jbellis.jvector.graph.disk.feature.FusedFeature; - -import java.io.IOException; -import java.util.LinkedHashMap; -import java.util.Map; -import java.util.function.IntFunction; - -/** - * Serializer for version 6 of the on-disk graph format. - * Version 6 characteristics: - * - Has magic number - * - Supports multiple features - * - Supports multi-layer (hierarchical) graphs - * - Has idUpperBound field - * - Uses footer for metadata - * - NEW: Changes feature ordering to place fused features last - * - NEW: Writes feature count and ordinals explicitly instead of bitflags - */ -class GraphIndexSerializerV6 extends GraphIndexSerializerV4 { - - GraphIndexSerializerV6() { - super(allFeatures()); - } - - @Override - public int getVersion() { - return 6; - } - - @Override - public boolean usesFooter() { - return true; - } - - @Override - public FeatureOrdering getFeatureOrdering() { - return FeatureOrdering.FUSED_LAST; - } - - @Override - public void writeFeatureHeaders(IndexWriter out, Map features) throws IOException { - // V6 writes feature count and ordinals explicitly (preserving order) - out.writeInt(features.size()); - - for (var entry : features.entrySet()) { - out.writeInt(entry.getKey().ordinal()); - entry.getValue().writeHeader(out); - } - } - - @Override - public Map readFeatureHeaders(RandomAccessReader in, GraphIndexMetadata metadata) throws IOException { - // V6 reads features in order (LinkedHashMap preserves insertion order) - Map features = new LinkedHashMap<>(); - - int nFeatures = in.readInt(); - - // Create a CommonHeader from metadata for feature loading - CommonHeader commonHeader = new CommonHeader( - metadata.getVersion(), - metadata.getDimension(), - metadata.getEntryNode(), - metadata.getLayerInfo(), - metadata.getIdUpperBound() - ); - - for (int i = 0; i < nFeatures; i++) { - FeatureId featureId = FeatureId.values()[in.readInt()]; - features.put(featureId, featureId.load(commonHeader, in)); - } - - return features; - } - - @Override - public int calculateHeaderSize(GraphIndexMetadata metadata, Map features) { - // V6: magic + version + 4 base ints + idUpperBound + numLayers + (32 * 2 layer info ints) - // + feature count + (feature ordinals) + feature headers - int size = 8 * Integer.BYTES; // magic, version, size, dimension, entryNode, maxDegree, idUpperBound, numLayers - size += 2 * 32 * Integer.BYTES; // layer info (padded to 32 layers) - size += Integer.BYTES; // feature count - size += features.size() * Integer.BYTES; // feature ordinals - size += features.values().stream().mapToInt(Feature::headerSize).sum(); - return size; - } - - @Override - public void writeSparseLevels(ImmutableGraphIndex graph, OrdinalMapper ordinalMapper, IndexWriter out, Map> featureStateSuppliers) throws IOException { - for (int level = 1; level <= graph.getMaxLevel(); level++) { - int layerSize = graph.size(level); - int layerDegree = graph.getDegree(level); - int nodesWritten = 0; - for (var it = graph.getNodes(level); it.hasNext(); ) { - int originalOrdinal = it.nextInt(); - // node id - final int newOrdinal = ordinalMapper.oldToNew(originalOrdinal); - out.writeInt(newOrdinal); - // neighbors - var view = graph.getView(); - var neighbors = view.getNeighborsIterator(level, originalOrdinal); - out.writeInt(neighbors.size()); - int n = 0; - for ( ; n < neighbors.size(); n++) { - out.writeInt(ordinalMapper.oldToNew(neighbors.nextInt())); - } - assert !neighbors.hasNext() : "Mismatch between neighbor's reported size and actual size"; - // pad out to degree - for (; n < layerDegree; n++) { - out.writeInt(-1); - } - nodesWritten++; - } - if (nodesWritten != layerSize) { - throw new IllegalStateException("Mismatch between layer size and nodes written"); - } - } - // There should be only one fused feature per node. This is checked in the class constructor. - // This is the only place where we explicitly need the fused feature. If there are more places in the - // future, it may be worth having fusedFeature as class member. - FusedFeature fusedFeature = null; - for (var feature : inlineFeatures) { - if (feature.isFused()) { - fusedFeature = (FusedFeature) feature; - } - } - if (fusedFeature != null) { - var supplier = featureStateSuppliers.get(fusedFeature.id()); - if (supplier == null) { - throw new IllegalStateException("Supplier for feature " + fusedFeature.id() + " not found"); - } - - if (graph.getMaxLevel() >= 1) { - int level = 1; - int layerSize = graph.size(level); - int nodesWritten = 0; - for (var it = graph.getNodes(level); it.hasNext(); ) { - int originalOrdinal = it.nextInt(); - - // We write the ordinal (node id) so that we can map it to the corresponding feature - final int newOrdinal = ordinalMapper.oldToNew(originalOrdinal); - out.writeInt(newOrdinal); - fusedFeature.writeSourceFeature(out, supplier.apply(originalOrdinal)); - nodesWritten++; - } - if (nodesWritten != layerSize) { - throw new IllegalStateException("Mismatch between layer 1 size and features written"); - } - } else { - // Write the source feature of the entry node - final int originalEntryNode = graph.getView().entryNode().node; - final int entryNode = ordinalMapper.oldToNew(originalEntryNode); - out.writeInt(entryNode); - fusedFeature.writeSourceFeature(out, supplier.apply(originalEntryNode)); - } - } - } -} - diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskSequentialGraphIndexWriter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskSequentialGraphIndexWriter.java index c8a0bb832..173747307 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskSequentialGraphIndexWriter.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskSequentialGraphIndexWriter.java @@ -81,88 +81,11 @@ public synchronized void close() throws IOException { @Override public synchronized void write(Map> featureStateSuppliers) throws IOException { - if (graph instanceof OnHeapGraphIndex) { - var ohgi = (OnHeapGraphIndex) graph; - if (ohgi.getDeletedNodes().cardinality() > 0) { - throw new IllegalArgumentException("Run builder.cleanup() before writing the graph"); - } - } - for (var featureId : featureStateSuppliers.keySet()) { - if (!featureMap.containsKey(featureId)) { - throw new IllegalArgumentException(String.format("Feature %s not configured for index", featureId)); - } - } - if (ordinalMapper.maxOrdinal() < graph.size(0) - 1) { - var msg = String.format("Ordinal mapper from [0..%d] does not cover all nodes in the graph of size %d", - ordinalMapper.maxOrdinal(), graph.size(0)); - throw new IllegalStateException(msg); - } - - var view = graph.getView(); - - final var startOffset = out.position(); - writeHeader(view, startOffset); - - // for each graph node, write the associated features, followed by its neighbors at L0 - for (int newOrdinal = 0; newOrdinal <= ordinalMapper.maxOrdinal(); newOrdinal++) { - var originalOrdinal = ordinalMapper.newToOld(newOrdinal); - - // if no node exists with the given ordinal, write a placeholder - if (originalOrdinal == OrdinalMapper.OMITTED) { - throw new IllegalStateException("Ordinal mapper mapped new ordinal" + newOrdinal + " to non-existing node. This behavior is not supported on OnDiskSequentialGraphIndexWriter. Use OnDiskGraphIndexWriter instead."); - } - - if (!graph.containsNode(originalOrdinal)) { - var msg = String.format("Ordinal mapper mapped new ordinal %s to non-existing node %s", newOrdinal, originalOrdinal); - throw new IllegalStateException(msg); - } - out.writeInt(newOrdinal); // unnecessary, but a reasonable sanity check - assert out.position() == featureOffsetForOrdinal(startOffset, newOrdinal) : String.format("%d != %d", out.position(), featureOffsetForOrdinal(startOffset, newOrdinal)); - for (var feature : inlineFeatures) { - var supplier = featureStateSuppliers.get(feature.id()); - if (supplier == null) { - throw new IllegalStateException("Supplier for feature " + feature.id() + " not found"); - } else { - feature.writeInline(out, supplier.apply(originalOrdinal)); - } - } - - var neighbors = view.getNeighborsIterator(0, originalOrdinal); - if (neighbors.size() > graph.getDegree(0)) { - var msg = String.format("Node %d has more neighbors %d than the graph's max degree %d -- run Builder.cleanup()!", - originalOrdinal, neighbors.size(), graph.getDegree(0)); - throw new IllegalStateException(msg); - } - // write neighbors list - out.writeInt(neighbors.size()); - int n = 0; - for (; n < neighbors.size(); n++) { - var newNeighborOrdinal = ordinalMapper.oldToNew(neighbors.nextInt()); - if (newNeighborOrdinal < 0 || newNeighborOrdinal > ordinalMapper.maxOrdinal()) { - var msg = String.format("Neighbor ordinal out of bounds: %d/%d", newNeighborOrdinal, ordinalMapper.maxOrdinal()); - throw new IllegalStateException(msg); - } - out.writeInt(newNeighborOrdinal); - } - assert !neighbors.hasNext(); - - // pad out to maxEdgesPerNode - for (; n < graph.getDegree(0); n++) { - out.writeInt(-1); - } - } - - writeSparseLevels(view, featureStateSuppliers); - - writeSeparatedFeatures(featureStateSuppliers); - - // Write the footer with all the metadata info about the graph - writeFooter(view, out.position()); - // Note: flushing the data output is the responsibility of the caller we are not going to make assumptions about further uses of the data outputs - - view.close(); + long startOffset = out.position(); + serializer.writeOnDiskSequential(createContext(startOffset), out, featureStateSuppliers); } + /** * Builder for {@link OnDiskSequentialGraphIndexWriter}, with optional features. */ diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/RandomAccessOnDiskGraphIndexWriter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/RandomAccessOnDiskGraphIndexWriter.java index 6cd8d0010..2ed161285 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/RandomAccessOnDiskGraphIndexWriter.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/RandomAccessOnDiskGraphIndexWriter.java @@ -112,67 +112,13 @@ public synchronized void writeInline(int ordinal, Map * the mapper is not invoked. */ public synchronized void writeFeaturesInline(int ordinal, Map stateMap) throws IOException { - for (var featureId : stateMap.keySet()) { - if (!featureMap.containsKey(featureId)) { - throw new IllegalArgumentException(String.format("Feature %s not configured for index", featureId)); - } - } - - out.seek(featureOffsetForOrdinal(ordinal)); - - for (var feature : inlineFeatures) { - var state = stateMap.get(feature.id()); - if (state == null) { - out.seek(out.position() + feature.featureSize()); - } else { - feature.writeInline(out, state); - } - } - + serializer.writeFeaturesInline(createContext(startOffset), ordinal, stateMap, out); maxOrdinalWritten = Math.max(maxOrdinalWritten, ordinal); } public synchronized void write(Map> featureStateSuppliers) throws IOException { - if (graph instanceof OnHeapGraphIndex) { - var ohgi = (OnHeapGraphIndex) graph; - if (ohgi.getDeletedNodes().cardinality() > 0) { - throw new IllegalArgumentException("Run builder.cleanup() before writing the graph"); - } - } - for (var featureId : featureStateSuppliers.keySet()) { - if (!featureMap.containsKey(featureId)) { - throw new IllegalArgumentException(String.format("Feature %s not configured for index", featureId)); - } - } - if (ordinalMapper.maxOrdinal() < graph.size(0) - 1) { - var msg = String.format("Ordinal mapper from [0..%d] does not cover all nodes in the graph of size %d", - ordinalMapper.maxOrdinal(), graph.size(0)); - throw new IllegalStateException(msg); - } - - var view = graph.getView(); - - writeHeader(view); // sets position to start writing features - - writeL0Records(view, featureStateSuppliers); - - // We will use the abstract method because no random access is needed - writeSparseLevels(view, featureStateSuppliers); - - // We will use the abstract method because no random access is needed - writeSeparatedFeatures(featureStateSuppliers); - - if (version >= 5) { - writeFooter(view, out.position()); - } - final var endOfGraphPosition = out.position(); - - // Write the header again with updated offsets - writeHeader(view); - out.seek(endOfGraphPosition); - out.flush(); - view.close(); + serializer.writeRandomAccess(createContext(startOffset), out, featureStateSuppliers, this::writeL0Records); } protected abstract void writeL0Records(ImmutableGraphIndex.View view, diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/WriteContext.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/WriteContext.java new file mode 100644 index 000000000..0e1c6337e --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/WriteContext.java @@ -0,0 +1,73 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph.disk; + +import io.github.jbellis.jvector.graph.ImmutableGraphIndex; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; + +import java.util.List; +import java.util.Map; + +/** + * Encapsulates the per-write-session context shared by all serializer operations. + * Eliminates the repeated passing of the same 6-7 parameters through every method + * in the serializer hierarchy. + */ +class WriteContext { + /** The immutable graph being serialized. */ + final ImmutableGraphIndex graph; + /** Maps between original (old) ordinals and the compacted (new) ordinals written to disk. */ + final OrdinalMapper ordinalMapper; + /** All features configured for this index, keyed by their {@link FeatureId}. */ + final Map featureMap; + /** Ordered list of features that are written inline with each node record. */ + final List inlineFeatures; + /** Byte offset within the output stream at which this graph's data begins. */ + final long startOffset; + /** Byte size of the header (or footer) written before the node records. */ + final long headerSize; + /** Vector dimension, derived from the configured features rather than the graph. */ + final int dimension; + + /** + * Constructs a {@code WriteContext} bundling all parameters required for a single serialization session. + * + * @param graph the immutable graph to serialize + * @param ordinalMapper mapping between original and compacted ordinals + * @param featureMap all features configured for the index + * @param inlineFeatures features written inline alongside each node record + * @param startOffset byte offset in the output stream at which the graph data begins + * @param headerSize byte size of the header written before node records + * @param dimension vector dimension, sourced from the configured features + */ + WriteContext(ImmutableGraphIndex graph, + OrdinalMapper ordinalMapper, + Map featureMap, + List inlineFeatures, + long startOffset, + long headerSize, + int dimension) { + this.graph = graph; + this.ordinalMapper = ordinalMapper; + this.featureMap = featureMap; + this.inlineFeatures = inlineFeatures; + this.startOffset = startOffset; + this.headerSize = headerSize; + this.dimension = dimension; + } +} From f0811de6f28b21b6b2a6b8f4e7d92d93f5128b1a Mon Sep 17 00:00:00 2001 From: Mark Wolters Date: Tue, 28 Jul 2026 14:21:07 -0400 Subject: [PATCH 3/5] fixing possible bugs --- .../graph/disk/AbstractGraphIndexFormat.java | 163 +++++++++--------- .../graph/disk/AbstractGraphIndexWriter.java | 16 +- .../jvector/graph/disk/GraphIndexFormat.java | 22 +-- .../graph/disk/GraphIndexFormatV2.java | 2 +- .../graph/disk/GraphIndexFormatV3.java | 2 +- .../graph/disk/GraphIndexFormatV4.java | 8 +- .../graph/disk/GraphIndexFormatV5.java | 60 +------ .../graph/disk/GraphIndexFormatV6.java | 28 ++- 8 files changed, 132 insertions(+), 169 deletions(-) diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexFormat.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexFormat.java index 2d0e27c67..75052a563 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexFormat.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexFormat.java @@ -25,7 +25,9 @@ import io.github.jbellis.jvector.graph.disk.feature.SeparatedFeature; import java.io.IOException; +import java.util.EnumMap; import java.util.EnumSet; +import java.util.LinkedHashMap; import java.util.Map; import java.util.Set; import java.util.function.IntFunction; @@ -38,7 +40,6 @@ abstract class AbstractGraphIndexFormat implements GraphIndexFormat { private final Set supportedFeatures; private final boolean supportsMultiLayer; private final boolean usesFooter; - private final FeatureOrdering featureOrdering; /** A magic number to indicate the file footer */ public static final int FOOTER_MAGIC = 0x4a564244; @@ -56,18 +57,15 @@ abstract class AbstractGraphIndexFormat implements GraphIndexFormat { * @param supportedFeatures the set of {@link FeatureId}s this version can store * @param supportsMultiLayer whether this version supports hierarchical (multi-layer) graphs * @param usesFooter whether metadata is placed in a footer rather than a header - * @param featureOrdering the ordering strategy used when laying out feature data */ protected AbstractGraphIndexFormat(int version, Set supportedFeatures, boolean supportsMultiLayer, - boolean usesFooter, - FeatureOrdering featureOrdering) { + boolean usesFooter) { this.version = version; this.supportedFeatures = supportedFeatures; this.supportsMultiLayer = supportsMultiLayer; this.usesFooter = usesFooter; - this.featureOrdering = featureOrdering; } @Override @@ -95,36 +93,41 @@ public boolean usesFooter() { return usesFooter; } + /** + * Default ordering: preserves the natural {@link FeatureId} enum ordinal order. + * Version 6 overrides this to place fused features last. + */ @Override - public FeatureOrdering getFeatureOrdering() { - return featureOrdering; + public Map orderFeatures(EnumMap features) { + return new LinkedHashMap<>(features); } @Override public void writeSparseLevels(WriteContext ctx, IndexWriter out, Map> suppliers) throws IOException { - for (int level = 1; level <= ctx.graph.getMaxLevel(); level++) { - int layerSize = ctx.graph.size(level); - int layerDegree = ctx.graph.getDegree(level); - int nodesWritten = 0; - for (var it = ctx.graph.getNodes(level); it.hasNext(); ) { - int originalOrdinal = it.nextInt(); - final int newOrdinal = ctx.ordinalMapper.oldToNew(originalOrdinal); - out.writeInt(newOrdinal); - var view = ctx.graph.getView(); - var neighbors = view.getNeighborsIterator(level, originalOrdinal); - out.writeInt(neighbors.size()); - int n = 0; - for ( ; n < neighbors.size(); n++) { - out.writeInt(ctx.ordinalMapper.oldToNew(neighbors.nextInt())); + try (var view = ctx.graph.getView()) { + for (int level = 1; level <= ctx.graph.getMaxLevel(); level++) { + int layerSize = ctx.graph.size(level); + int layerDegree = ctx.graph.getDegree(level); + int nodesWritten = 0; + for (var it = ctx.graph.getNodes(level); it.hasNext(); ) { + int originalOrdinal = it.nextInt(); + final int newOrdinal = ctx.ordinalMapper.oldToNew(originalOrdinal); + out.writeInt(newOrdinal); + var neighbors = view.getNeighborsIterator(level, originalOrdinal); + out.writeInt(neighbors.size()); + int n = 0; + for ( ; n < neighbors.size(); n++) { + out.writeInt(ctx.ordinalMapper.oldToNew(neighbors.nextInt())); + } + assert !neighbors.hasNext() : "Mismatch between neighbor's reported size and actual size"; + for (; n < layerDegree; n++) { + out.writeInt(-1); + } + nodesWritten++; } - assert !neighbors.hasNext() : "Mismatch between neighbor's reported size and actual size"; - for (; n < layerDegree; n++) { - out.writeInt(-1); + if (nodesWritten != layerSize) { + throw new IllegalStateException("Mismatch between layer size and nodes written"); } - nodesWritten++; - } - if (nodesWritten != layerSize) { - throw new IllegalStateException("Mismatch between layer size and nodes written"); } } writeAfterSparseLevels(ctx, out, suppliers); @@ -139,9 +142,11 @@ protected void writeAfterSparseLevels(WriteContext ctx, IndexWriter out, Map ctx.graph.getDegree(0)) { - throw new IllegalStateException(String.format("Node %d has more neighbors %d than the graph's max degree %d -- run Builder.cleanup()!", - originalOrdinal, neighbors.size(), ctx.graph.getDegree(0))); - } - out.writeInt(neighbors.size()); - int n = 0; - for (; n < neighbors.size(); n++) { - var newNeighborOrdinal = ctx.ordinalMapper.oldToNew(neighbors.nextInt()); - if (newNeighborOrdinal < 0 || newNeighborOrdinal > ctx.ordinalMapper.maxOrdinal()) { - throw new IllegalStateException(String.format("Neighbor ordinal out of bounds: %d/%d", newNeighborOrdinal, ctx.ordinalMapper.maxOrdinal())); + var neighbors = view.getNeighborsIterator(0, originalOrdinal); + if (neighbors.size() > ctx.graph.getDegree(0)) { + throw new IllegalStateException(String.format("Node %d has more neighbors %d than the graph's max degree %d -- run Builder.cleanup()!", + originalOrdinal, neighbors.size(), ctx.graph.getDegree(0))); + } + out.writeInt(neighbors.size()); + int n = 0; + for (; n < neighbors.size(); n++) { + var newNeighborOrdinal = ctx.ordinalMapper.oldToNew(neighbors.nextInt()); + if (newNeighborOrdinal < 0 || newNeighborOrdinal > ctx.ordinalMapper.maxOrdinal()) { + throw new IllegalStateException(String.format("Neighbor ordinal out of bounds: %d/%d", newNeighborOrdinal, ctx.ordinalMapper.maxOrdinal())); + } + out.writeInt(newNeighborOrdinal); + } + assert !neighbors.hasNext(); + for (; n < ctx.graph.getDegree(0); n++) { + out.writeInt(-1); } - out.writeInt(newNeighborOrdinal); - } - assert !neighbors.hasNext(); - for (; n < ctx.graph.getDegree(0); n++) { - out.writeInt(-1); } } writeSparseLevels(ctx, out, suppliers); writeSeparatedFeatures(ctx, out, suppliers); - writeFooter(ctx, out.position(), out); - - view.close(); + if (usesFooter()) { + writeFooter(ctx, out.position(), out); + } } @Override @@ -304,20 +311,22 @@ public void writeRandomAccess(WriteContext ctx, RandomAccessWriter out, Map implements this.ordinalMapper = oldToNewOrdinals; this.dimension = dimension; - if (version <= 5) { - // Versions <= 5 use the old feature ordering, simply provided by the FeatureId - this.featureMap = features; - this.inlineFeatures = features.values().stream().filter(f -> !(f instanceof SeparatedFeature)).collect(Collectors.toList()); - } else { - // Version 6 uses the new feature ordering to place fused features last in the list - var sortedFeatures = features.values().stream().sorted().collect(Collectors.toList()); - this.featureMap = new LinkedHashMap<>(); - for (var feature : sortedFeatures) { - this.featureMap.put(feature.id(), feature); - } - this.inlineFeatures = sortedFeatures.stream().filter(f -> !(f instanceof SeparatedFeature)).sorted().collect(Collectors.toList()); - } + this.featureMap = serializer.orderFeatures(features); + this.inlineFeatures = this.featureMap.values().stream().filter(f -> !(f instanceof SeparatedFeature)).collect(Collectors.toList()); long fusedFeaturesCount = this.inlineFeatures.stream().filter(Feature::isFused).count(); if (fusedFeaturesCount > 1) { diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormat.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormat.java index 7e97604bb..8284b31d0 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormat.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormat.java @@ -23,6 +23,7 @@ import io.github.jbellis.jvector.graph.disk.feature.FeatureId; import java.io.IOException; +import java.util.EnumMap; import java.util.Map; import java.util.Set; import java.util.function.IntFunction; @@ -60,12 +61,6 @@ public interface GraphIndexFormat { */ boolean usesFooter(); - /** - * Gets the feature ordering strategy for this version. - * @return the feature ordering strategy - */ - FeatureOrdering getFeatureOrdering(); - /** * Returns the complete set of {@link FeatureId}s that this format version is capable of storing. * @@ -91,14 +86,15 @@ interface L0RecordWriter { } /** - * Defines how features should be ordered when writing/reading. + * Returns the feature map for this format version, ordered as required by the on-disk layout. + * Versions 2–5 preserve the natural {@link FeatureId} enum ordinal order; version 6 places + * fused features last so that non-fused inline features occupy a contiguous prefix. + * The returned map preserves insertion order. + * + * @param features the raw feature map supplied by the caller + * @return an ordered map suitable for use during write operations */ - enum FeatureOrdering { - /** Features ordered by their FeatureId enum ordinal (versions <= 5) */ - BY_FEATURE_ID, - /** Features ordered with fused features last (version 6+) */ - FUSED_LAST - } + Map orderFeatures(EnumMap features); /** * Writes adjacency records for all graph levels above level 0 (the "sparse" upper layers). diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV2.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV2.java index 199c51257..603a855df 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV2.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV2.java @@ -28,6 +28,6 @@ class GraphIndexFormatV2 extends AbstractGraphIndexFormat { /** Creates the singleton format for version 2. */ GraphIndexFormatV2() { - super(2, inlineVectorsOnly(), false, false, FeatureOrdering.BY_FEATURE_ID); + super(2, inlineVectorsOnly(), false, false); } } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV3.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV3.java index 1e44975b3..6afba14aa 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV3.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV3.java @@ -28,6 +28,6 @@ class GraphIndexFormatV3 extends AbstractGraphIndexFormat { /** Creates the singleton format for version 3. */ GraphIndexFormatV3() { - super(3, nonFusedFeatures(), false, false, FeatureOrdering.BY_FEATURE_ID); + super(3, nonFusedFeatures(), false, false); } } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV4.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV4.java index 1d0034800..bb75d2009 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV4.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV4.java @@ -33,14 +33,14 @@ class GraphIndexFormatV4 extends AbstractGraphIndexFormat { /** Creates the singleton format for version 4. */ GraphIndexFormatV4() { - super(4, nonFusedFeatures(), true, false, FeatureOrdering.BY_FEATURE_ID); + super(4, nonFusedFeatures(), true, false); } /** * Protected constructor allowing subclasses (V5, V6) to specify their own version, - * feature set, footer flag, and feature ordering while sharing V4's wire format. + * feature set, and footer flag while sharing V4's wire format. */ - protected GraphIndexFormatV4(int version, Set supportedFeatures, boolean usesFooter, FeatureOrdering featureOrdering) { - super(version, supportedFeatures, true, usesFooter, featureOrdering); + protected GraphIndexFormatV4(int version, Set supportedFeatures, boolean usesFooter) { + super(version, supportedFeatures, true, usesFooter); } } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV5.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV5.java index 4f0f14b9d..da43b8341 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV5.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV5.java @@ -16,16 +16,9 @@ package io.github.jbellis.jvector.graph.disk; -import io.github.jbellis.jvector.disk.RandomAccessWriter; -import io.github.jbellis.jvector.graph.ImmutableGraphIndex; -import io.github.jbellis.jvector.graph.OnHeapGraphIndex; -import io.github.jbellis.jvector.graph.disk.feature.Feature; import io.github.jbellis.jvector.graph.disk.feature.FeatureId; -import java.io.IOException; -import java.util.Map; import java.util.Set; -import java.util.function.IntFunction; /** * Format for version 5 of the on-disk graph format. @@ -43,57 +36,14 @@ class GraphIndexFormatV5 extends GraphIndexFormatV4 { /** Creates the singleton format for version 5. */ GraphIndexFormatV5() { - super(5, nonFusedFeatures(), true, FeatureOrdering.BY_FEATURE_ID); + super(5, nonFusedFeatures(), true); } /** - * Protected constructor for subclasses (V6) to specify their own version, - * feature set, and feature ordering while inheriting V5's footer-writing behavior. - * Footer is always true for V5 and later. + * Protected constructor for subclasses (V6) to specify their own version and feature set + * while inheriting V5's footer-writing behavior. Footer is always {@code true} for V5+. */ - protected GraphIndexFormatV5(int version, Set supportedFeatures, FeatureOrdering featureOrdering) { - super(version, supportedFeatures, true, featureOrdering); - } - - /** - * Writes the complete graph index using random-access I/O, and additionally appends a - * footer (header offset + magic number) after the graph data — the key behavioral difference - * from the V4 implementation, which does not write a footer. - * - * {@inheritDoc} - */ - @Override - public void writeRandomAccess(WriteContext ctx, RandomAccessWriter out, Map> suppliers, GraphIndexFormat.L0RecordWriter l0Writer) throws IOException { - if (ctx.graph instanceof OnHeapGraphIndex) { - var ohgi = (OnHeapGraphIndex) ctx.graph; - if (ohgi.getDeletedNodes().cardinality() > 0) { - throw new IllegalArgumentException("Run builder.cleanup() before writing the graph"); - } - } - for (var featureId : suppliers.keySet()) { - if (!ctx.featureMap.containsKey(featureId)) { - throw new IllegalArgumentException(String.format("Feature %s not configured for index", featureId)); - } - } - if (ctx.ordinalMapper.maxOrdinal() < ctx.graph.size(0) - 1) { - throw new IllegalStateException(String.format("Ordinal mapper from [0..%d] does not cover all nodes in the graph of size %d", - ctx.ordinalMapper.maxOrdinal(), ctx.graph.size(0))); - } - - var view = ctx.graph.getView(); - - out.seek(ctx.startOffset); - writeHeader(ctx, out); - l0Writer.write(view, suppliers); - writeSparseLevels(ctx, out, suppliers); - writeSeparatedFeatures(ctx, out, suppliers); - writeFooter(ctx, out.position(), out); - - final var endOfGraphPosition = out.position(); - out.seek(ctx.startOffset); - writeHeader(ctx, out); - out.seek(endOfGraphPosition); - out.flush(); - view.close(); + protected GraphIndexFormatV5(int version, Set supportedFeatures) { + super(version, supportedFeatures, true); } } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV6.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV6.java index ed8bd7399..1703b8f26 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV6.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV6.java @@ -22,8 +22,11 @@ import io.github.jbellis.jvector.graph.disk.feature.FusedFeature; import java.io.IOException; +import java.util.EnumMap; +import java.util.LinkedHashMap; import java.util.Map; import java.util.function.IntFunction; +import java.util.stream.Collectors; /** * Format for version 6 of the on-disk graph format. @@ -40,7 +43,18 @@ class GraphIndexFormatV6 extends GraphIndexFormatV5 { /** Creates the singleton format for version 6. */ GraphIndexFormatV6() { - super(6, allFeatures(), FeatureOrdering.FUSED_LAST); + super(6, allFeatures()); + } + + /** Places fused features last so non-fused inline features occupy a contiguous prefix. */ + @Override + public Map orderFeatures(EnumMap features) { + var sorted = features.values().stream().sorted().collect(Collectors.toList()); + var map = new LinkedHashMap(); + for (var f : sorted) { + map.put(f.id(), f); + } + return map; } /** @@ -77,9 +91,15 @@ protected void writeAfterSparseLevels(WriteContext ctx, IndexWriter out, Map Date: Tue, 28 Jul 2026 18:38:23 -0400 Subject: [PATCH 4/5] handling headers --- .../graph/disk/AbstractGraphIndexFormat.java | 86 +++++++++++++++--- .../graph/disk/AbstractGraphIndexWriter.java | 30 ++++--- .../jvector/graph/disk/CommonHeader.java | 87 +++---------------- .../jvector/graph/disk/GraphIndexFormat.java | 47 ++++++++++ .../graph/disk/GraphIndexFormatV3.java | 55 +++++++++++- .../graph/disk/GraphIndexFormatV4.java | 73 +++++++++++++++- .../graph/disk/GraphIndexFormatV5.java | 16 ++++ .../graph/disk/GraphIndexFormatV6.java | 36 ++++++++ .../jbellis/jvector/graph/disk/Header.java | 62 +------------ .../jvector/graph/disk/OnDiskGraphIndex.java | 46 +++++----- .../OnDiskSequentialGraphIndexWriter.java | 3 +- .../RandomAccessOnDiskGraphIndexWriter.java | 5 +- 12 files changed, 355 insertions(+), 191 deletions(-) diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexFormat.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexFormat.java index 75052a563..f53520183 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexFormat.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexFormat.java @@ -17,7 +17,9 @@ package io.github.jbellis.jvector.graph.disk; import io.github.jbellis.jvector.disk.IndexWriter; +import io.github.jbellis.jvector.disk.RandomAccessReader; import io.github.jbellis.jvector.disk.RandomAccessWriter; +import io.github.jbellis.jvector.disk.ReaderSupplier; import io.github.jbellis.jvector.graph.ImmutableGraphIndex; import io.github.jbellis.jvector.graph.OnHeapGraphIndex; import io.github.jbellis.jvector.graph.disk.feature.Feature; @@ -25,11 +27,7 @@ import io.github.jbellis.jvector.graph.disk.feature.SeparatedFeature; import java.io.IOException; -import java.util.EnumMap; -import java.util.EnumSet; -import java.util.LinkedHashMap; -import java.util.Map; -import java.util.Set; +import java.util.*; import java.util.function.IntFunction; /** @@ -329,20 +327,86 @@ public void writeRandomAccess(WriteContext ctx, RandomAccessWriter out, Map layerInfo, int dimension, int entryNode, int idUpperBound) throws IOException { + out.writeInt(layerInfo.get(0).size); + out.writeInt(dimension); + out.writeInt(entryNode); + out.writeInt(layerInfo.get(0).degree); + if (layerInfo.size() > 1) { + throw new IllegalArgumentException("Layer info is not supported in version " + getVersion()); + } + } + + @Override + public CommonHeader readCommonHeader(RandomAccessReader in, int size) throws IOException { + int dimension = in.readInt(); + int entryNode = in.readInt(); + int maxDegree = in.readInt(); + + List layerInfo; + layerInfo = List.of(new CommonHeader.LayerInfo(size, maxDegree)); + logger.debug("Common header finished reading at position {}", in.getPosition()); + + return new CommonHeader(version, dimension, entryNode, layerInfo, size); + } + + @Override + public int commonHeaderSize() { + return 4 * Integer.BYTES; + } + + @Override + public void writeHeaderFeatures(IndexWriter out, Map features) throws IOException { + // we restrict pre-version-3 writers to INLINE_VECTORS features, so we don't need additional version-handling here + for (Feature writer : features.values()) { + writer.writeHeader(out); + } + } + + @Override + public int headerSize(Map features) { + int size = this.commonHeaderSize(); + + size += features.values().stream().mapToInt(Feature::headerSize).sum(); + + return size; + } + + @Override + public Map loadHeaderFeatures(RandomAccessReader reader, CommonHeader common) throws IOException { + Map features = new EnumMap<>(FeatureId.class); + FeatureId featureId = FeatureId.INLINE_VECTORS; + features.put(featureId, featureId.load(common, reader)); + return features; + } + + @Override + public OnDiskGraphIndex loadOnDiskIndex(RandomAccessReader reader, Header header, ReaderSupplier readerSupplier, boolean useFooter) throws IOException { + return OnDiskGraphIndex.construct(readerSupplier, header, reader.getPosition(), reader); + } + /** - * Helper to create a set of all known features. - * Only correct for the current maximum version — do not use for older versions. + * Helper to create the frozen set of features supported by version 6. + * Deliberately enumerated rather than {@code EnumSet.allOf(FeatureId.class)}: version 6's + * supported-feature set is a historical fact about a shipped format and must not change just + * because a new {@link FeatureId} is added to the enum for some future version. */ protected static Set allFeatures() { - return EnumSet.allOf(FeatureId.class); + return EnumSet.of(FeatureId.INLINE_VECTORS, FeatureId.FUSED_PQ, FeatureId.NVQ_VECTORS, + FeatureId.SEPARATED_VECTORS, FeatureId.SEPARATED_NVQ); } /** - * Helper to create a set of all features except FUSED_PQ. - * Used by versions 3–5, which predate fused PQ hierarchy support (version 6). + * Helper to create the frozen set of features supported by versions 3–5, which predate fused + * PQ hierarchy support (version 6). Deliberately enumerated rather than + * {@code EnumSet.complementOf(EnumSet.of(FUSED_PQ))}, so that adding a new {@link FeatureId} + * in the future requires an explicit decision about which already-shipped versions support it, + * instead of silently being included here. */ protected static Set nonFusedFeatures() { - return EnumSet.complementOf(EnumSet.of(FeatureId.FUSED_PQ)); + return EnumSet.of(FeatureId.INLINE_VECTORS, FeatureId.NVQ_VECTORS, + FeatureId.SEPARATED_VECTORS, FeatureId.SEPARATED_NVQ); } /** diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexWriter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexWriter.java index 4578ac043..305b0850f 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexWriter.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexWriter.java @@ -20,7 +20,6 @@ import io.github.jbellis.jvector.graph.ImmutableGraphIndex; import io.github.jbellis.jvector.graph.disk.feature.Feature; import io.github.jbellis.jvector.graph.disk.feature.FeatureId; -import io.github.jbellis.jvector.graph.disk.feature.FusedFeature; import io.github.jbellis.jvector.graph.disk.feature.InlineVectors; import io.github.jbellis.jvector.graph.disk.feature.NVQ; import io.github.jbellis.jvector.graph.disk.feature.SeparatedFeature; @@ -60,7 +59,7 @@ public abstract class AbstractGraphIndexWriter implements volatile int maxOrdinalWritten = -1; final List inlineFeatures; - final GraphIndexFormat serializer; + final GraphIndexFormat graphIndexFormat; AbstractGraphIndexWriter(T out, int version, @@ -69,8 +68,8 @@ public abstract class AbstractGraphIndexWriter implements int dimension, EnumMap features) { - serializer = GraphIndexFormatFactory.forVersion(version); - if (graph.isHierarchical() && !serializer.supportsMultiLayer()) { + graphIndexFormat = GraphIndexFormatFactory.forVersion(version); + if (graph.isHierarchical() && !graphIndexFormat.supportsMultiLayer()) { throw new IllegalArgumentException("Multilayer graphs must be written with version 4 or higher"); } this.version = version; @@ -78,14 +77,14 @@ public abstract class AbstractGraphIndexWriter implements this.ordinalMapper = oldToNewOrdinals; this.dimension = dimension; - this.featureMap = serializer.orderFeatures(features); + this.featureMap = graphIndexFormat.orderFeatures(features); this.inlineFeatures = this.featureMap.values().stream().filter(f -> !(f instanceof SeparatedFeature)).collect(Collectors.toList()); long fusedFeaturesCount = this.inlineFeatures.stream().filter(Feature::isFused).count(); if (fusedFeaturesCount > 1) { throw new IllegalArgumentException("At most one fused feature is allowed"); } - if (fusedFeaturesCount == 1 && version < 6) { + if (fusedFeaturesCount == 1 && !graphIndexFormat.supportsFeature(FeatureId.FUSED_PQ)) { throw new IllegalArgumentException("Fused features require version 6 or higher"); } this.out = out; @@ -117,7 +116,7 @@ WriteContext createContext(long startOffset) { } long featureOffsetForOrdinal(long startOffset, int ordinal) { - return serializer.featureOffsetForOrdinal(createContext(startOffset), ordinal); + return graphIndexFormat.featureOffsetForOrdinal(createContext(startOffset), ordinal); } /** @@ -157,7 +156,7 @@ public static Map sequentialRenumbering(ImmutableGraphIndex gr * @throws IOException IOException */ void writeFooter(ImmutableGraphIndex.View view, long headerOffset, long startOffset) throws IOException { - serializer.writeFooter(createContext(startOffset), headerOffset, out); + graphIndexFormat.writeFooter(createContext(startOffset), headerOffset, out); } /** @@ -168,15 +167,15 @@ void writeFooter(ImmutableGraphIndex.View view, long headerOffset, long startOff * @throws IOException if an I/O error occurs */ protected synchronized void writeHeader(ImmutableGraphIndex.View view, long startOffset) throws IOException { - serializer.writeHeader(createContext(startOffset), out); + graphIndexFormat.writeHeader(createContext(startOffset), out); } void writeSparseLevels(ImmutableGraphIndex.View view, Map> featureStateSuppliers, long startOffset) throws IOException { - serializer.writeSparseLevels(createContext(startOffset), out, featureStateSuppliers); + graphIndexFormat.writeSparseLevels(createContext(startOffset), out, featureStateSuppliers); } void writeSeparatedFeatures(Map> featureStateSuppliers, long startOffset) throws IOException { - serializer.writeSeparatedFeatures(createContext(startOffset), out, featureStateSuppliers); + graphIndexFormat.writeSeparatedFeatures(createContext(startOffset), out, featureStateSuppliers); } /** @@ -248,8 +247,13 @@ public Builder withMapper(OrdinalMapper ordinalMapper) { * @throws IOException if an I/O error occurs */ public K build() throws IOException { - if (version < 3 && (!features.containsKey(FeatureId.INLINE_VECTORS) || features.size() > 1)) { - throw new IllegalArgumentException("Only INLINE_VECTORS is supported until version 3"); + var format = GraphIndexFormatFactory.forVersion(version); + for (var featureId : features.keySet()) { + if (!format.supportsFeature(featureId)) { + throw new IllegalArgumentException(String.format( + "Feature %s is not supported by version %d (supported features: %s)", + featureId, version, format.getSupportedFeatures())); + } } int dimension; diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CommonHeader.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CommonHeader.java index 5d0a1aecb..17080e9a7 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CommonHeader.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CommonHeader.java @@ -59,13 +59,14 @@ public class CommonHeader { private static final Logger logger = LoggerFactory.getLogger(CommonHeader.class); - private static final int V4_MAX_LAYERS = 32; + protected static final int V4_MAX_LAYERS = 32; public final int version; public final int dimension; public final int entryNode; public final List layerInfo; public final int idUpperBound; + private final GraphIndexFormat graphIndexFormat; CommonHeader(int version, int dimension, int entryNode, List layerInfo, int idUpperBound) { this.version = version; @@ -73,93 +74,25 @@ public class CommonHeader { this.entryNode = entryNode; this.layerInfo = layerInfo; this.idUpperBound = idUpperBound; + this.graphIndexFormat = GraphIndexFormatFactory.forVersion(version); } void write(IndexWriter out) throws IOException { logger.debug("Writing common header at position {}", out.position()); - if (version >= 3) { - out.writeInt(OnDiskGraphIndex.MAGIC); - out.writeInt(version); - } - out.writeInt(layerInfo.get(0).size); - out.writeInt(dimension); - out.writeInt(entryNode); - out.writeInt(layerInfo.get(0).degree); - if (version >= 4) { - out.writeInt(idUpperBound); - - if (layerInfo.size() > V4_MAX_LAYERS) { - var msg = String.format("Number of layers %d exceeds maximum of %d", layerInfo.size(), V4_MAX_LAYERS); - throw new IllegalArgumentException(msg); - } - logger.debug("Writing {} layers", layerInfo.size()); - out.writeInt(layerInfo.size()); - // Write actual layer info - for (LayerInfo info : layerInfo) { - out.writeInt(info.size); - out.writeInt(info.degree); - } - // Pad remaining entries with zeros - for (int i = layerInfo.size(); i < V4_MAX_LAYERS; i++) { - out.writeInt(0); // size - out.writeInt(0); // degree - } - } else { - if (layerInfo.size() > 1) { - throw new IllegalArgumentException("Layer info is not supported in version " + version); - } - } + graphIndexFormat.writeCommonHeader(out, layerInfo, dimension, entryNode, idUpperBound); logger.debug("Common header finished writing at position {}", out.position()); } static CommonHeader load(RandomAccessReader in) throws IOException { - logger.debug("Loading common header at position {}", in.getPosition()); - int maybeMagic = in.readInt(); - int version; - int size; - if (maybeMagic == OnDiskGraphIndex.MAGIC) { - version = in.readInt(); - size = in.readInt(); - } else { - version = 2; - size = maybeMagic; - } - int dimension = in.readInt(); - int entryNode = in.readInt(); - int maxDegree = in.readInt(); - int idUpperBound = size; - List layerInfo; - if (version < 4) { - layerInfo = List.of(new LayerInfo(size, maxDegree)); - } else { - idUpperBound = in.readInt(); - int numLayers = in.readInt(); - logger.debug("{} layers", numLayers); - layerInfo = new ArrayList<>(); - for (int i = 0; i < numLayers; i++) { - LayerInfo info = new LayerInfo(in.readInt(), in.readInt()); - layerInfo.add(info); - } - // Skip over remaining padding entries - for (int i = numLayers; i < V4_MAX_LAYERS; i++) { - in.readInt(); - in.readInt(); - } - } - logger.debug("Common header finished reading at position {}", in.getPosition()); - - return new CommonHeader(version, dimension, entryNode, layerInfo, idUpperBound); + return GraphIndexFormat.loadCommonHeader(in); } int size() { - int size = 4; - if (version >= 3) { - size += 2; - } - if (version >= 4) { - size += 2 + 2 * V4_MAX_LAYERS; - } - return size * Integer.BYTES; + return graphIndexFormat.commonHeaderSize(); + } + + GraphIndexFormat getGraphIndexFormat() { + return graphIndexFormat; } @VisibleForTesting diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormat.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormat.java index 8284b31d0..6b7223815 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormat.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormat.java @@ -17,13 +17,18 @@ package io.github.jbellis.jvector.graph.disk; import io.github.jbellis.jvector.disk.IndexWriter; +import io.github.jbellis.jvector.disk.RandomAccessReader; import io.github.jbellis.jvector.disk.RandomAccessWriter; +import io.github.jbellis.jvector.disk.ReaderSupplier; import io.github.jbellis.jvector.graph.ImmutableGraphIndex; import io.github.jbellis.jvector.graph.disk.feature.Feature; import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import java.io.IOException; import java.util.EnumMap; +import java.util.List; import java.util.Map; import java.util.Set; import java.util.function.IntFunction; @@ -37,6 +42,20 @@ * makes it easy to add new versions without modifying existing code. */ public interface GraphIndexFormat { + static final Logger logger = LoggerFactory.getLogger(GraphIndexFormat.class); + + static OnDiskGraphIndex loadOnDiskIndex(RandomAccessReader reader, long offset, boolean useFooter, ReaderSupplier readerSupplier) throws IOException { + logger.debug("Loading OnDiskGraphIndex from offset={}", offset); + var header = Header.load(reader, offset); + + logger.debug("Header loaded: version={}, dimension={}, entryNode={}, layerInfoCount={}", + header.common.version, header.common.dimension, header.common.entryNode, header.common.layerInfo.size()); + logger.debug("Position after reading header={}", reader.getPosition()); + return header.common.getGraphIndexFormat().loadOnDiskIndex(reader, header, readerSupplier, useFooter); + } + + OnDiskGraphIndex loadOnDiskIndex(RandomAccessReader reader, Header header, ReaderSupplier readerSupplier, boolean useFooter) throws IOException; + /** * @return the version number this format handles */ @@ -68,6 +87,14 @@ public interface GraphIndexFormat { */ Set getSupportedFeatures(); + int commonHeaderSize(); + + void writeHeaderFeatures(IndexWriter out, Map features) throws IOException; + + int headerSize(Map features); + + Map loadHeaderFeatures(RandomAccessReader reader, CommonHeader common) throws IOException; + /** * Callback for writing L0 (base layer) node records. Implemented by the writer so that * I/O strategy (sequential vs. parallel) stays in the writer while the format @@ -184,4 +211,24 @@ interface L0RecordWriter { * @throws IOException if an I/O error occurs while writing */ void writeRandomAccess(WriteContext ctx, RandomAccessWriter out, Map> suppliers, L0RecordWriter l0Writer) throws IOException; + + void writeCommonHeader(IndexWriter out, List layerInfo, int dimension, int entryNode, int idUpperBound) throws IOException; + + static CommonHeader loadCommonHeader(RandomAccessReader in) throws IOException { + logger.debug("Loading common header at position {}", in.getPosition()); + int maybeMagic = in.readInt(); + int version; + int size; + if (maybeMagic == OnDiskGraphIndex.MAGIC) { + version = in.readInt(); + size = in.readInt(); + } else { + version = 2; + size = maybeMagic; + } + GraphIndexFormat format = GraphIndexFormatFactory.forVersion(version); + return format.readCommonHeader(in, size); + } + + CommonHeader readCommonHeader(RandomAccessReader in, int size) throws IOException; } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV3.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV3.java index 6afba14aa..95b77a1c7 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV3.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV3.java @@ -16,6 +16,14 @@ package io.github.jbellis.jvector.graph.disk; +import io.github.jbellis.jvector.disk.IndexWriter; +import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; + +import java.io.IOException; +import java.util.*; + /** * Format for version 3 of the on-disk graph format. * Version 3 characteristics: @@ -26,8 +34,53 @@ */ class GraphIndexFormatV3 extends AbstractGraphIndexFormat { - /** Creates the singleton format for version 3. */ + /** + * Creates the singleton format for version 3. + */ GraphIndexFormatV3() { super(3, nonFusedFeatures(), false, false); } + + protected GraphIndexFormatV3(int version, Set supportedFeatures, boolean supportsMultiLayer, boolean usesFooter) { + super(version, supportedFeatures, supportsMultiLayer, usesFooter); + } + + @Override + public void writeCommonHeader(IndexWriter out, List layerInfo, int dimension, int entryNode, int idUpperBound) throws IOException { + out.writeInt(OnDiskGraphIndex.MAGIC); + out.writeInt(this.getVersion()); + super.writeCommonHeader(out, layerInfo, dimension, entryNode, idUpperBound); + } + + @Override + public int commonHeaderSize() { + return 6 * Integer.BYTES; + } + + @Override + public void writeHeaderFeatures(IndexWriter out, Map features) throws IOException { + out.writeInt(FeatureId.serialize(EnumSet.copyOf(features.keySet()))); + super.writeHeaderFeatures(out, features); + } + + @Override + public int headerSize(Map features) { + int size = super.headerSize(features); + + size += Integer.BYTES; + + return size; + } + + @Override + public Map loadHeaderFeatures(RandomAccessReader reader, CommonHeader common) throws IOException { + EnumSet featureIds; + Map features = new EnumMap<>(FeatureId.class); + featureIds = FeatureId.deserialize(reader.readInt()); + for (FeatureId featureId : featureIds) { + features.put(featureId, featureId.load(common, reader)); + } + return features; + } + } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV4.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV4.java index bb75d2009..6461ca2bf 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV4.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV4.java @@ -16,10 +16,19 @@ package io.github.jbellis.jvector.graph.disk; +import io.github.jbellis.jvector.disk.IndexWriter; +import io.github.jbellis.jvector.disk.RandomAccessReader; import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; import java.util.Set; +import static io.github.jbellis.jvector.graph.disk.CommonHeader.V4_MAX_LAYERS; + /** * Format for version 4 of the on-disk graph format. * Version 4 characteristics: @@ -29,7 +38,8 @@ * - Has idUpperBound field * - No footer */ -class GraphIndexFormatV4 extends AbstractGraphIndexFormat { +class GraphIndexFormatV4 extends GraphIndexFormatV3 { + private static final Logger logger = LoggerFactory.getLogger(GraphIndexFormatV4.class); /** Creates the singleton format for version 4. */ GraphIndexFormatV4() { @@ -43,4 +53,65 @@ class GraphIndexFormatV4 extends AbstractGraphIndexFormat { protected GraphIndexFormatV4(int version, Set supportedFeatures, boolean usesFooter) { super(version, supportedFeatures, true, usesFooter); } + + @Override + public void writeCommonHeader(IndexWriter out, List layerInfo, int dimension, int entryNode, int idUpperBound) throws IOException { + out.writeInt(OnDiskGraphIndex.MAGIC); + out.writeInt(this.getVersion()); + out.writeInt(layerInfo.get(0).size); + out.writeInt(dimension); + out.writeInt(entryNode); + out.writeInt(layerInfo.get(0).degree); + out.writeInt(idUpperBound); + + if (layerInfo.size() > V4_MAX_LAYERS) { + var msg = String.format("Number of layers %d exceeds maximum of %d", layerInfo.size(), V4_MAX_LAYERS); + throw new IllegalArgumentException(msg); + } + logger.debug("Writing {} layers", layerInfo.size()); + out.writeInt(layerInfo.size()); + // Write actual layer info + for (CommonHeader.LayerInfo info : layerInfo) { + out.writeInt(info.size); + out.writeInt(info.degree); + } + // Pad remaining entries with zeros + for (int i = layerInfo.size(); i < V4_MAX_LAYERS; i++) { + out.writeInt(0); // size + out.writeInt(0); // degree + } + } + + @Override + public CommonHeader readCommonHeader(RandomAccessReader in, int size) throws IOException { + int dimension = in.readInt(); + int entryNode = in.readInt(); + int maxDegree = in.readInt(); + + List layerInfo; + int idUpperBound = in.readInt(); + int numLayers = in.readInt(); + logger.debug("{} layers", numLayers); + layerInfo = new ArrayList<>(); + for (int i = 0; i < numLayers; i++) { + CommonHeader.LayerInfo info = new CommonHeader.LayerInfo(in.readInt(), in.readInt()); + layerInfo.add(info); + } + // Skip over remaining padding entries + for (int i = numLayers; i < V4_MAX_LAYERS; i++) { + in.readInt(); + in.readInt(); + } + logger.debug("Common header finished reading at position {}", in.getPosition()); + + return new CommonHeader(getVersion(), dimension, entryNode, layerInfo, idUpperBound); + } + + @Override + public int commonHeaderSize() { + // super.commonHeaderSize() (V3's) covers magic + version + the 4 shared base fields; + // V4 adds idUpperBound, the layer count, and the fixed-size padded layer array. + return super.commonHeaderSize() + (2 + 2 * V4_MAX_LAYERS) * Integer.BYTES; + } + } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV5.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV5.java index da43b8341..4765679b2 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV5.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV5.java @@ -16,8 +16,13 @@ package io.github.jbellis.jvector.graph.disk; +import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.disk.ReaderSupplier; import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import java.io.IOException; import java.util.Set; /** @@ -33,6 +38,7 @@ * a footer is written after the graph data. */ class GraphIndexFormatV5 extends GraphIndexFormatV4 { + private static final Logger logger = LoggerFactory.getLogger(GraphIndexFormatV5.class); /** Creates the singleton format for version 5. */ GraphIndexFormatV5() { @@ -46,4 +52,14 @@ class GraphIndexFormatV5 extends GraphIndexFormatV4 { protected GraphIndexFormatV5(int version, Set supportedFeatures) { super(version, supportedFeatures, true); } + + @Override + public OnDiskGraphIndex loadOnDiskIndex(RandomAccessReader reader, Header header, ReaderSupplier readerSupplier, boolean useFooter) throws IOException { + if (useFooter) { + logger.debug("Version 5+ onwards uses a footer instead of header for metadata. Loading from footer"); + return OnDiskGraphIndex.loadFromFooter(readerSupplier, reader.getPosition()); + } else { + return super.loadOnDiskIndex(reader, header, readerSupplier, useFooter); + } + } } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV6.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV6.java index 1703b8f26..c9d965aa4 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV6.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV6.java @@ -17,6 +17,7 @@ package io.github.jbellis.jvector.graph.disk; import io.github.jbellis.jvector.disk.IndexWriter; +import io.github.jbellis.jvector.disk.RandomAccessReader; import io.github.jbellis.jvector.graph.disk.feature.Feature; import io.github.jbellis.jvector.graph.disk.feature.FeatureId; import io.github.jbellis.jvector.graph.disk.feature.FusedFeature; @@ -102,4 +103,39 @@ protected void writeAfterSparseLevels(WriteContext ctx, IndexWriter out, Map features) throws IOException { + // Writing the features in order instead of writing a single integer with all the features (as done in features) { + int size = this.commonHeaderSize(); + // In V6, this accounts for the number of features and the ordinal of each feature + size += Integer.BYTES + features.size() * Integer.BYTES; + + size += features.values().stream().mapToInt(Feature::headerSize).sum(); + + return size; + } + + @Override + public Map loadHeaderFeatures(RandomAccessReader reader, CommonHeader common) throws IOException { + Map features = new LinkedHashMap<>(); + int nFeatures = reader.readInt(); + for (int i = 0; i < nFeatures; i++) { + FeatureId featureId = FeatureId.values()[reader.readInt()]; + features.put(featureId, featureId.load(common, reader)); + } + return features; + } + } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/Header.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/Header.java index d80ba9698..75ae7d48e 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/Header.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/Header.java @@ -22,9 +22,6 @@ import io.github.jbellis.jvector.graph.disk.feature.FeatureId; import java.io.IOException; -import java.util.EnumMap; -import java.util.EnumSet; -import java.util.LinkedHashMap; import java.util.Map; /** @@ -52,71 +49,18 @@ class Header { void write(IndexWriter out) throws IOException { common.write(out); - - if (common.version >= 6) { - // Writing the features in order instead of writing a single integer with all the features (as done in = 3) { - out.writeInt(FeatureId.serialize(EnumSet.copyOf(features.keySet()))); - } - - // we restrict pre-version-3 writers to INLINE_VECTORS features, so we don't need additional version-handling here - for (Feature writer : features.values()) { - writer.writeHeader(out); - } - - } + common.getGraphIndexFormat().writeHeaderFeatures(out, features); } public int size() { - int size = common.size(); - - if (common.version >= 6) { - // In V6, this accounts for the number of features and the ordinal of each feature - size += Integer.BYTES + features.size() * Integer.BYTES; - } else if (common.version >= 3) { - size += Integer.BYTES; - } - - size += features.values().stream().mapToInt(Feature::headerSize).sum(); - - return size; + return common.getGraphIndexFormat().headerSize(features); } static Header load(RandomAccessReader reader, long offset) throws IOException { reader.seek(offset); - Map features; - CommonHeader common = CommonHeader.load(reader); - if (common.version >= 6) { - features = new LinkedHashMap<>(); - int nFeatures = reader.readInt(); - for (int i = 0; i < nFeatures; i++) { - FeatureId featureId = FeatureId.values()[reader.readInt()]; - features.put(featureId, featureId.load(common, reader)); - } - } else { - EnumSet featureIds; - features = new EnumMap<>(FeatureId.class); - - if (common.version >= 3) { - featureIds = FeatureId.deserialize(reader.readInt()); - } else { - featureIds = EnumSet.of(FeatureId.INLINE_VECTORS); - } - for (FeatureId featureId : featureIds) { - features.put(featureId, featureId.load(common, reader)); - } - } - + features = common.getGraphIndexFormat().loadHeaderFeatures(reader, common); return new Header(common, features); } } \ No newline at end of file diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java index b445ceb20..47ce2fa08 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java @@ -73,6 +73,7 @@ public class OnDiskGraphIndex implements ImmutableGraphIndex, AutoCloseable, Acc static final VectorTypeSupport vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport(); final ReaderSupplier readerSupplier; final int version; + final GraphIndexFormat format; final int dimension; final NodeAtLevel entryNode; final int idUpperBound; @@ -92,6 +93,7 @@ private OnDiskGraphIndex(ReaderSupplier readerSupplier, Header header, long neig { this.readerSupplier = readerSupplier; this.version = header.common.version; + this.format = header.common.getGraphIndexFormat(); this.layerInfo = header.common.layerInfo; this.dimension = header.common.dimension; if (header.common.entryNode == ENTRY_NODE_ABSENT) { @@ -191,8 +193,9 @@ private Int2ObjectHashMap loadInMemoryFeatures(Random } in.seek(neighborsOffset + L0size + inMemorySize); - // In V6, fused features for the in-memory hierarchy are written in a block after the top layers of the graph. - if (version == 6) { + // Fused features for the in-memory hierarchy are written in a block after the top layers of the graph, + // for formats that support fused features. + if (format.supportsFeature(FeatureId.FUSED_PQ)) { if (layerInfo.size() >= 2) { int level = 1; CommonHeader.LayerInfo info = layerInfo.get(level); @@ -236,6 +239,19 @@ public static OnDiskGraphIndex load(ReaderSupplier readerSupplier, long offset) return load(readerSupplier, offset, true); } + /** + * Constructs an {@link OnDiskGraphIndex} and eagerly primes its in-memory caches (upper-layer + * adjacency and, for fused graphs, hierarchy features) using the given reader. Used by every + * {@link GraphIndexFormat#loadOnDiskIndex} implementation so that priming happens exactly once, + * regardless of whether the format is header- or footer-based. + */ + static OnDiskGraphIndex construct(ReaderSupplier readerSupplier, Header header, long neighborsOffset, RandomAccessReader reader) throws IOException { + var odgi = new OnDiskGraphIndex(readerSupplier, header, neighborsOffset); + odgi.getInMemoryLayers(reader); + odgi.getInMemoryFeatures(reader); + return odgi; + } + /** * Load an index from the given reader supplier where header and graph are located on the same file, * where the index starts at `offset`. @@ -247,22 +263,7 @@ public static OnDiskGraphIndex load(ReaderSupplier readerSupplier, long offset) */ public static OnDiskGraphIndex load(ReaderSupplier readerSupplier, long offset, boolean useFooter) { try (var reader = readerSupplier.get()) { - logger.debug("Loading OnDiskGraphIndex from offset={}", offset); - var header = Header.load(reader, offset); - - logger.debug("Header loaded: version={}, dimension={}, entryNode={}, layerInfoCount={}", - header.common.version, header.common.dimension, header.common.entryNode, header.common.layerInfo.size()); - logger.debug("Position after reading header={}", - reader.getPosition()); - if (header.common.version >= 5 && useFooter) { - logger.debug("Version 5+ onwards uses a footer instead of header for metadata. Loading from footer"); - return loadFromFooter(readerSupplier, reader.getPosition()); - } else { - var odgi = new OnDiskGraphIndex(readerSupplier, header, reader.getPosition()); - odgi.getInMemoryLayers(reader); - odgi.getInMemoryFeatures(reader); - return odgi; - } + return GraphIndexFormat.loadOnDiskIndex(reader, offset, useFooter, readerSupplier); } catch (Exception e) { throw new RuntimeException("Error initializing OnDiskGraph at offset " + offset, e); } @@ -284,7 +285,7 @@ public static OnDiskGraphIndex load(ReaderSupplier readerSupplier) { * This reader supplier must vend slices of IndexOutput that contain the graph index and nothing else. * @return the loaded index. */ - private static OnDiskGraphIndex loadFromFooter(ReaderSupplier readerSupplier, long neighborsOffset) { + protected static OnDiskGraphIndex loadFromFooter(ReaderSupplier readerSupplier, long neighborsOffset) { try (var in = readerSupplier.get()) { final long magicOffset = in.length() - FOOTER_MAGIC_SIZE; logger.debug("Loading OnDiskGraphIndex footer from offset={}", magicOffset); @@ -306,10 +307,7 @@ private static OnDiskGraphIndex loadFromFooter(ReaderSupplier readerSupplier, lo header.common.entryNode, header.common.layerInfo.size(), in.getPosition()); - var odgi = new OnDiskGraphIndex(readerSupplier, header, neighborsOffset); - odgi.getInMemoryLayers(in); - odgi.getInMemoryFeatures(in); - return odgi; + return construct(readerSupplier, header, neighborsOffset, in); } catch (Exception e) { throw new RuntimeException("Error initializing OnDiskGraph", e); @@ -600,7 +598,7 @@ public void getPackedNeighbors(int node, FeatureId featureId, Consumer> featureStateSuppliers) throws IOException { long startOffset = out.position(); - serializer.writeOnDiskSequential(createContext(startOffset), out, featureStateSuppliers); + graphIndexFormat.writeOnDiskSequential(createContext(startOffset), out, featureStateSuppliers); } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/RandomAccessOnDiskGraphIndexWriter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/RandomAccessOnDiskGraphIndexWriter.java index 2ed161285..2c3e5b899 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/RandomAccessOnDiskGraphIndexWriter.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/RandomAccessOnDiskGraphIndexWriter.java @@ -18,7 +18,6 @@ import io.github.jbellis.jvector.disk.RandomAccessWriter; import io.github.jbellis.jvector.graph.ImmutableGraphIndex; -import io.github.jbellis.jvector.graph.OnHeapGraphIndex; import io.github.jbellis.jvector.graph.disk.feature.Feature; import io.github.jbellis.jvector.graph.disk.feature.FeatureId; @@ -112,13 +111,13 @@ public synchronized void writeInline(int ordinal, Map * the mapper is not invoked. */ public synchronized void writeFeaturesInline(int ordinal, Map stateMap) throws IOException { - serializer.writeFeaturesInline(createContext(startOffset), ordinal, stateMap, out); + graphIndexFormat.writeFeaturesInline(createContext(startOffset), ordinal, stateMap, out); maxOrdinalWritten = Math.max(maxOrdinalWritten, ordinal); } public synchronized void write(Map> featureStateSuppliers) throws IOException { - serializer.writeRandomAccess(createContext(startOffset), out, featureStateSuppliers, this::writeL0Records); + graphIndexFormat.writeRandomAccess(createContext(startOffset), out, featureStateSuppliers, this::writeL0Records); } protected abstract void writeL0Records(ImmutableGraphIndex.View view, From f35a0741e971ab7bcf314112f2aea1b24c6b16db Mon Sep 17 00:00:00 2001 From: Mark Wolters Date: Wed, 29 Jul 2026 13:11:48 -0400 Subject: [PATCH 5/5] testing and compaction case --- .../jvector/graph/disk/CompactWriter.java | 31 +++- .../disk/TestGraphIndexFormatFactory.java | 153 ++++++++++++++++++ .../graph/disk/TestOnDiskGraphIndex.java | 44 +++++ 3 files changed, 221 insertions(+), 7 deletions(-) create mode 100644 jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestGraphIndexFormatFactory.java diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactWriter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactWriter.java index 182648735..765ae2810 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactWriter.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactWriter.java @@ -55,7 +55,6 @@ final class CompactWriter implements AutoCloseable { private final long startOffset; private final int headerSize; private final Header header; - private final int version; private final FusedFeature fusedFeature; private final int baseDegree; private final int maxOrdinal; @@ -91,7 +90,18 @@ final class CompactWriter implements AutoCloseable { throws IOException { this.fusedFeature = fusedFeature; this.fusedPQEnabled = fusedFeature != null; - this.version = OnDiskGraphIndex.CURRENT_VERSION; + // This writer only ever targets the current on-disk format -- unlike AbstractGraphIndexWriter, + // there is no builder/version parameter that lets a caller pick an older version. Because of + // that, `fusedPQEnabled` alone is a safe stand-in for "the format supports fused PQ" everywhere + // below: there's no other version this class could be writing. We still check the format here, + // once, instead of just assuming it, so that if this class is ever changed to support writing + // older versions, or a future version drops fused PQ support, this fails loudly at construction + // instead of silently omitting the fused-PQ footer/packed-neighbor data later on. + int version = OnDiskGraphIndex.CURRENT_VERSION; + if (fusedPQEnabled && !GraphIndexFormatFactory.forVersion(version).supportsFeature(FeatureId.FUSED_PQ)) { + throw new IllegalStateException( + "Fused PQ requires a format version that supports it; version " + version + " does not"); + } this.outputPath = outputPath; this.writer = new BufferedRandomAccessWriter(outputPath); this.startOffset = startOffset; @@ -117,7 +127,7 @@ final class CompactWriter implements AutoCloseable { this.recordSize = rsize; this.configuredLayerInfo.set(0, new CommonHeader.LayerInfo(numBaseLayerNodes, baseDegree)); - var commonHeader = new CommonHeader(this.version, dimension, entryNode, this.configuredLayerInfo, this.maxOrdinal + 1); + var commonHeader = new CommonHeader(version, dimension, entryNode, this.configuredLayerInfo, this.maxOrdinal + 1); this.header = new Header(commonHeader, featureMap); this.headerSize = header.size(); @@ -159,7 +169,10 @@ public void writeHeader() throws IOException { } void writeFooter() throws IOException { - if (fusedPQEnabled && version == 6) { + // No "&& version == 6" here: the constructor already verified that the format this writer + // targets supports fused PQ whenever fusedPQEnabled is true, so that's the only condition + // this needs to check. See the constructor for why that's a safe simplification. + if (fusedPQEnabled) { if (!level1FeatureRecords.isEmpty()) { // Hierarchy is enabled: write PQ source feature for every level-1 node. // Mirrors AbstractGraphIndexWriter.writeSparseLevels (getMaxLevel >= 1 branch). @@ -213,10 +226,12 @@ public long projectedOutputSize() { int count = configuredLayerInfo.get(level).size; total += (long) count * (Integer.BYTES * 2L + (long) degree * Integer.BYTES); } - // PQ feature records written at the start of writeFooter() when v6 + fused PQ: + // PQ feature records written at the start of writeFooter() when fused PQ is enabled: // - hierarchy enabled: one [ord, code] record per level-1 node; // - no hierarchy: one [entryOrd, code] record for the entry node only. - if (fusedPQEnabled && version == 6) { + // No "&& version == 6" here either -- see the constructor for why fusedPQEnabled alone + // already implies the target format supports this layout. + if (fusedPQEnabled) { int pqSize = fusedFeature.codeSize(); if (configuredLayerInfo.size() > 1) { int level1Count = configuredLayerInfo.get(1).size; @@ -241,7 +256,9 @@ public void writeUpperLayerNode(int level, int ordinal, int[] neighbors, ByteSeq for (; n < degree; n++) { writer.writeInt(-1); } - if (fusedPQEnabled && version == 6 && level == 1 && level1PqCode != null) { + // No "&& version == 6" here either -- see the constructor for why fusedPQEnabled alone + // already implies the target format supports this layout. + if (fusedPQEnabled && level == 1 && level1PqCode != null) { level1FeatureRecords.add(new UpperLayerFeatureRecord(ordinal, level1PqCode.copy())); } } diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestGraphIndexFormatFactory.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestGraphIndexFormatFactory.java new file mode 100644 index 000000000..27c0a82e8 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestGraphIndexFormatFactory.java @@ -0,0 +1,153 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph.disk; + +import com.carrotsearch.randomizedtesting.RandomizedTest; +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; +import io.github.jbellis.jvector.TestUtil; +import io.github.jbellis.jvector.disk.SimpleMappedReader; +import io.github.jbellis.jvector.graph.TestVectorGraph; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.graph.disk.feature.InlineVectors; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.EnumSet; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +/** + * Tests for {@link GraphIndexFormatFactory}: version dispatch, rejection of unsupported + * versions, magic-number-based version detection, and the per-version characteristics + * (feature set, multi-layer support, footer usage) that the writer/reader dispatch relies on. + */ +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) +public class TestGraphIndexFormatFactory extends RandomizedTest { + + private Path testDirectory; + + @Before + public void setup() throws IOException { + testDirectory = Files.createTempDirectory(this.getClass().getSimpleName()); + } + + @After + public void tearDown() { + TestUtil.deleteQuietly(testDirectory); + } + + @Test + public void testForVersionReturnsMatchingFormat() { + for (int version = 2; version <= OnDiskGraphIndex.CURRENT_VERSION; version++) { + assertEquals(version, GraphIndexFormatFactory.forVersion(version).getVersion()); + } + } + + @Test + public void testForVersionRejectsUnsupportedVersions() { + for (int version : new int[]{Integer.MIN_VALUE, -1, 0, 1, OnDiskGraphIndex.CURRENT_VERSION + 1, 100}) { + assertThrows(GraphIndexFormatFactory.UnsupportedVersionException.class, + () -> GraphIndexFormatFactory.forVersion(version)); + } + } + + @Test + public void testGetCurrentVersion() { + assertEquals(OnDiskGraphIndex.CURRENT_VERSION, GraphIndexFormatFactory.getCurrentVersion()); + } + + /** + * Pins down the per-version characteristics that the rest of the read/write path dispatches + * on, so that a change to one version's format can't silently change another's behavior. + */ + @Test + public void testFormatCharacteristicsPerVersion() { + var v2 = GraphIndexFormatFactory.forVersion(2); + assertFalse(v2.supportsMultiLayer()); + assertFalse(v2.usesFooter()); + assertEquals(EnumSet.of(FeatureId.INLINE_VECTORS), v2.getSupportedFeatures()); + + // versions 3-5 support every feature except FUSED_PQ; multi-layer arrives at v4, + // footer-based metadata arrives at v5. + var nonFused = EnumSet.complementOf(EnumSet.of(FeatureId.FUSED_PQ)); + for (int version : new int[]{3, 4, 5}) { + var format = GraphIndexFormatFactory.forVersion(version); + assertEquals("version " + version + " supported features", nonFused, format.getSupportedFeatures()); + assertFalse("version " + version + " should not support FUSED_PQ", format.supportsFeature(FeatureId.FUSED_PQ)); + assertEquals("version " + version + " multi-layer support", version >= 4, format.supportsMultiLayer()); + assertEquals("version " + version + " footer usage", version >= 5, format.usesFooter()); + } + + var v6 = GraphIndexFormatFactory.forVersion(6); + assertTrue(v6.supportsMultiLayer()); + assertTrue(v6.usesFooter()); + assertTrue(v6.supportsFeature(FeatureId.FUSED_PQ)); + assertEquals(EnumSet.allOf(FeatureId.class), v6.getSupportedFeatures()); + } + + @Test + public void testDetectVersionNoMagic() throws Exception { + // version 2 predates the magic number, so detection falls back to reading the raw + // base-layer size as the first int; detectVersion must still identify it as v2 and + // must leave the reader positioned exactly where it started. + var path = writeSingleNodeGraph(2); + try (var readerSupplier = new SimpleMappedReader.Supplier(path); + var reader = readerSupplier.get()) + { + long startPosition = reader.getPosition(); + var format = GraphIndexFormatFactory.detectVersion(reader); + assertEquals(2, format.getVersion()); + assertEquals(startPosition, reader.getPosition()); + } + } + + @Test + public void testDetectVersionWithMagic() throws Exception { + var path = writeSingleNodeGraph(OnDiskGraphIndex.CURRENT_VERSION); + try (var readerSupplier = new SimpleMappedReader.Supplier(path); + var reader = readerSupplier.get()) + { + long startPosition = reader.getPosition(); + var format = GraphIndexFormatFactory.detectVersion(reader); + assertEquals(OnDiskGraphIndex.CURRENT_VERSION, format.getVersion()); + assertEquals(startPosition, reader.getPosition()); + } + } + + private Path writeSingleNodeGraph(int version) throws IOException { + var graph = new TestUtil.RandomlyConnectedGraphIndex(5, 2, getRandom()); + var ravv = new TestVectorGraph.CircularFloatVectorValues(graph.size(0)); + var outputPath = testDirectory.resolve("format_" + version); + try (var writer = new OnDiskGraphIndexWriter.Builder(graph, outputPath) + .withVersion(version) + .with(new InlineVectors(ravv.dimension())) + .build()) + { + writer.write(Feature.singleStateFactory(FeatureId.INLINE_VECTORS, + nodeId -> new InlineVectors.State(ravv.getVector(nodeId)))); + } + return outputPath; + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndex.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndex.java index 8930b720e..c76796bb4 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndex.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndex.java @@ -377,6 +377,50 @@ public void testV0Write() throws IOException { assertArrayEquals(contents1, contents2); } + @Test + public void testVersionRoundTrip() throws Exception { + // Versions 3-5 predate V6's feature-reordering/footer-format changes but still need to + // round-trip correctly. Version 2 is covered separately by testV0Read/testV0Write, and + // version 6 (current) is exercised throughout the rest of this test class. + for (int version : new int[]{3, 4, 5}) { + versionRoundTrip(version, false); + } + // Only versions 4 and 5 (of the ones under test here) support multi-layer graphs. + for (int version : new int[]{4, 5}) { + versionRoundTrip(version, true); + } + } + + private void versionRoundTrip(int version, boolean hierarchical) throws Exception { + ImmutableGraphIndex graph = hierarchical + ? new TestUtil.RandomlyConnectedGraphIndex( + List.of(new CommonHeader.LayerInfo(50, 8), new CommonHeader.LayerInfo(5, 3)), + getRandom()) + : new TestUtil.RandomlyConnectedGraphIndex(50, 8, getRandom()); + var ravv = new TestVectorGraph.CircularFloatVectorValues(graph.size(0)); + var outputPath = testDirectory.resolve("version_" + version + (hierarchical ? "_multilayer" : "_single")); + + try (var writer = new OnDiskGraphIndexWriter.Builder(graph, outputPath) + .withVersion(version) + .with(new InlineVectors(ravv.dimension())) + .build()) + { + writer.write(Feature.singleStateFactory(FeatureId.INLINE_VECTORS, + nodeId -> new InlineVectors.State(ravv.getVector(nodeId)))); + } + + try (var readerSupplier = new SimpleMappedReader.Supplier(outputPath.toAbsolutePath()); + var onDiskGraph = OnDiskGraphIndex.load(readerSupplier); + var onDiskView = onDiskGraph.getView()) + { + assertEquals("version", version, onDiskGraph.version); + assertEquals("max level", graph.getMaxLevel(), onDiskGraph.getMaxLevel()); + assertEquals(EnumSet.of(FeatureId.INLINE_VECTORS), onDiskGraph.features.keySet()); + TestUtil.assertGraphEquals(graph, onDiskGraph); + validateVectors(onDiskView, ravv); + } + } + @Test public void testMultiLayerFullyConnected() throws Exception { // Suppose we have 3 layers of sizes 5, 4, 3