Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view

Large diffs are not rendered by default.

Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -42,7 +40,7 @@
* Abstract base class for writing graph indexes to disk.
* @param <T> the type of the output writer
*/
public abstract class AbstractGraphIndexWriter<T extends IndexWriter> implements GraphIndexWriter {
public abstract class AbstractGraphIndexWriter<T extends IndexWriter> 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. */
Expand All @@ -57,9 +55,11 @@ public abstract class AbstractGraphIndexWriter<T extends IndexWriter> implements
final int dimension;
final Map<FeatureId, Feature> featureMap;
final T out; /* output for graph nodes and inline features */
final int headerSize;
final long headerSize;

volatile int maxOrdinalWritten = -1;
final List<Feature> inlineFeatures;
final GraphIndexFormat graphIndexFormat;

AbstractGraphIndexWriter(T out,
int version,
Expand All @@ -68,37 +68,26 @@ public abstract class AbstractGraphIndexWriter<T extends IndexWriter> implements
int dimension,
EnumMap<FeatureId, Feature> 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;
this.graph = graph;
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);
Expand All @@ -122,17 +111,12 @@ public Set<FeatureId> 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);
}

/**
Expand Down Expand Up @@ -171,19 +155,8 @@ public static Map<Integer, Integer> 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);
}

/**
Expand All @@ -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<FeatureId, IntFunction<Feature.State>> 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<FeatureId, IntFunction<Feature.State>> featureStateSuppliers, long startOffset) throws IOException {
graphIndexFormat.writeSparseLevels(createContext(startOffset), out, featureStateSuppliers);
}

void writeSeparatedFeatures(Map<FeatureId, IntFunction<Feature.State>> 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<FeatureId, IntFunction<Feature.State>> featureStateSuppliers, long startOffset) throws IOException {
graphIndexFormat.writeSeparatedFeatures(createContext(startOffset), out, featureStateSuppliers);
}

/**
Expand Down Expand Up @@ -379,8 +247,13 @@ public Builder<K, T> 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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -59,107 +59,40 @@
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> layerInfo;
public final int idUpperBound;
private final GraphIndexFormat graphIndexFormat;

CommonHeader(int version, int dimension, int entryNode, List<LayerInfo> layerInfo, int idUpperBound) {
this.version = version;
this.dimension = dimension;
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> 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
Expand Down
Loading
Loading