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..f53520183 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/AbstractGraphIndexFormat.java @@ -0,0 +1,418 @@ +/* + * 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.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; +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.*; +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; + + /** 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 + */ + protected AbstractGraphIndexFormat(int version, + Set supportedFeatures, + boolean supportsMultiLayer, + boolean usesFooter) { + this.version = version; + this.supportedFeatures = supportedFeatures; + this.supportsMultiLayer = supportsMultiLayer; + this.usesFooter = usesFooter; + } + + @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; + } + + /** + * Default ordering: preserves the natural {@link FeatureId} enum ordinal order. + * Version 6 overrides this to place fused features last. + */ + @Override + public Map orderFeatures(EnumMap features) { + return new LinkedHashMap<>(features); + } + + @Override + public void writeSparseLevels(WriteContext ctx, IndexWriter out, Map> suppliers) throws IOException { + 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++; + } + 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); + final int entryNode; + try (var view = ctx.graph.getView()) { + var en = view.entryNode(); + entryNode = en == null ? ImmutableGraphIndex.ENTRY_NODE_ABSENT : ctx.ordinalMapper.oldToNew(en.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); + final int entryNode; + try (var view = ctx.graph.getView()) { + var en = view.entryNode(); + entryNode = en == null ? ImmutableGraphIndex.ENTRY_NODE_ABSENT : ctx.ordinalMapper.oldToNew(en.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))); + } + + writeHeader(ctx, out); + + try (var view = ctx.graph.getView()) { + 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); + if (usesFooter()) { + writeFooter(ctx, out.position(), out); + } + } + + @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))); + } + + out.seek(ctx.startOffset); + writeHeader(ctx, out); + try (var view = ctx.graph.getView()) { + l0Writer.write(view, suppliers); + } + writeSparseLevels(ctx, out, suppliers); + writeSeparatedFeatures(ctx, out, suppliers); + if (usesFooter()) { + writeFooter(ctx, out.position(), out); + } + + final var endOfGraphPosition = out.position(); + out.seek(ctx.startOffset); + writeHeader(ctx, out); + out.seek(endOfGraphPosition); + out.flush(); + } + + @Override + public void writeCommonHeader(IndexWriter out, List 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 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.of(FeatureId.INLINE_VECTORS, FeatureId.FUSED_PQ, FeatureId.NVQ_VECTORS, + FeatureId.SEPARATED_VECTORS, FeatureId.SEPARATED_NVQ); + } + + /** + * 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.of(FeatureId.INLINE_VECTORS, FeatureId.NVQ_VECTORS, + FeatureId.SEPARATED_VECTORS, FeatureId.SEPARATED_NVQ); + } + + /** + * 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..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; @@ -31,7 +30,6 @@ import java.io.IOException; import java.util.EnumMap; -import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Set; @@ -42,7 +40,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 +55,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; + final long headerSize; + volatile int maxOrdinalWritten = -1; final List inlineFeatures; + final GraphIndexFormat graphIndexFormat; AbstractGraphIndexWriter(T out, int version, @@ -68,7 +68,8 @@ public abstract class AbstractGraphIndexWriter implements int dimension, EnumMap features) { - if (graph.getMaxLevel() > 0 && version < 4) { + 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; @@ -76,29 +77,17 @@ public abstract class AbstractGraphIndexWriter 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 = 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; - // 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); @@ -122,17 +111,12 @@ public Set getFeatureSet() { return featureMap.keySet(); } - 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 + WriteContext createContext(long startOffset) { + return new WriteContext(graph, ordinalMapper, featureMap, inlineFeatures, startOffset, headerSize, dimension); } - boolean isSeparated(Feature feature) { - return feature instanceof SeparatedFeature; + long featureOffsetForOrdinal(long startOffset, int ordinal) { + return graphIndexFormat.featureOffsetForOrdinal(createContext(startOffset), ordinal); } /** @@ -171,19 +155,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 { - 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); + void writeFooter(ImmutableGraphIndex.View view, long headerOffset, long startOffset) throws IOException { + graphIndexFormat.writeFooter(createContext(startOffset), headerOffset, out); } /** @@ -194,120 +167,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); + graphIndexFormat.writeHeader(createContext(startOffset), out); } - 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)); - } - } - } + void writeSparseLevels(ImmutableGraphIndex.View view, Map> featureStateSuppliers, long startOffset) throws IOException { + graphIndexFormat.writeSparseLevels(createContext(startOffset), 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); - } - } - } - } - } + void writeSeparatedFeatures(Map> featureStateSuppliers, long startOffset) throws IOException { + graphIndexFormat.writeSeparatedFeatures(createContext(startOffset), out, featureStateSuppliers); } /** @@ -379,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/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-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..6b7223815 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormat.java @@ -0,0 +1,234 @@ +/* + * 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.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; + +/** + * 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 { + 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 + */ + 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(); + + /** + * 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(); + + 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 + * 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; + } + + /** + * 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 + */ + Map orderFeatures(EnumMap features); + + /** + * 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; + + 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/GraphIndexFormatFactory.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatFactory.java new file mode 100644 index 000000000..12010bc1e --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatFactory.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 formats. + * This centralizes version detection and format instantiation. + */ +public class GraphIndexFormatFactory { + private static final Logger logger = LoggerFactory.getLogger(GraphIndexFormatFactory.class); + + private static final Map FORMATS = Map.of( + 2, new GraphIndexFormatV2(), + 3, new GraphIndexFormatV3(), + 4, new GraphIndexFormatV4(), + 5, new GraphIndexFormatV5(), + 6, new GraphIndexFormatV6() + ); + + /** + * Gets a format for a specific version. + * @param version the version number + * @return the format for that version + * @throws UnsupportedVersionException if the version is not supported + */ + public static GraphIndexFormat forVersion(int version) { + GraphIndexFormat format = FORMATS.get(version); + if (format == null) { + throw new UnsupportedVersionException("Version " + version + " is not supported. " + + "Supported versions: " + FORMATS.keySet()); + } + return format; + } + + /** + * 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 format for the detected version + * @throws IOException if an I/O error occurs + * @throws UnsupportedVersionException if the version is not supported + */ + 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(); + 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/GraphIndexFormatV2.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV2.java new file mode 100644 index 000000000..603a855df --- /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); + } +} 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..95b77a1c7 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV3.java @@ -0,0 +1,86 @@ +/* + * 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.*; + +/** + * 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); + } + + 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 new file mode 100644 index 000000000..6461ca2bf --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV4.java @@ -0,0 +1,117 @@ +/* + * 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.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: + * - Has magic number + * - Supports multiple features + * - Supports multi-layer (hierarchical) graphs + * - Has idUpperBound field + * - No footer + */ +class GraphIndexFormatV4 extends GraphIndexFormatV3 { + private static final Logger logger = LoggerFactory.getLogger(GraphIndexFormatV4.class); + + /** Creates the singleton format for version 4. */ + GraphIndexFormatV4() { + super(4, nonFusedFeatures(), true, false); + } + + /** + * Protected constructor allowing subclasses (V5, V6) to specify their own version, + * feature set, and footer flag while sharing V4's wire format. + */ + 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 new file mode 100644 index 000000000..4765679b2 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV5.java @@ -0,0 +1,65 @@ +/* + * 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 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; + +/** + * 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 { + private static final Logger logger = LoggerFactory.getLogger(GraphIndexFormatV5.class); + + /** Creates the singleton format for version 5. */ + GraphIndexFormatV5() { + super(5, nonFusedFeatures(), true); + } + + /** + * 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) { + 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 new file mode 100644 index 000000000..c9d965aa4 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/GraphIndexFormatV6.java @@ -0,0 +1,141 @@ +/* + * 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 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. + * 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()); + } + + /** 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; + } + + /** + * 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 { + try (var view = ctx.graph.getView()) { + var en = view.entryNode(); + if (en == null) { + return; + } + final int originalEntryNode = en.node; + out.writeInt(ctx.ordinalMapper.oldToNew(originalEntryNode)); + fusedFeature.writeSourceFeature(out, supplier.apply(originalEntryNode)); + } + } + } + + @Override + public void writeHeaderFeatures(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 { - 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(); + graphIndexFormat.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..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,67 +111,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); - } - } - + graphIndexFormat.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(); + graphIndexFormat.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; + } +} 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