New MutableGraphIndex and ImmutableGraphIndex interfaces (#534)

* Add MutableGraphIndex and ImmutableGraphIndex to clarify intent and hide implementations

* Fix docs

* Use "level" when referring to integers instead of "layer"

* Now OnHeapGraphIndex will return a ConcurrentGraphIndexView while it is being mutated by GraphIndexBuilder and a FrozenView once mutations are done.

* Make allMutationsCompleted volatile

* Clarifying comment in OnHeapGraphIndex.getView

* Make sure that the view gets instantiated when actually writing the graph.

Author: marianotepper

benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/RecallWithRandomVectorsBenchmark.java

@@ -52,7 +52,7 @@ public class RecallWithRandomVectorsBenchmark {
     private ArrayList<VectorFloat<?>> baseVectors;
     private ArrayList<VectorFloat<?>> queryVectors;
     private GraphIndexBuilder graphIndexBuilder;
-    private GraphIndex graphIndex;
+    private ImmutableGraphIndex graphIndex;
     private PQVectors pqVectors;
 
     // Add ground truth storage

benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/StaticSetVectorsBenchmark.java

@@ -45,7 +45,7 @@ public class StaticSetVectorsBenchmark {
     private List<VectorFloat<?>> queryVectors;
     private List<List<Integer>> groundTruth;
     private GraphIndexBuilder graphIndexBuilder;
-    private GraphIndex graphIndex;
+    private ImmutableGraphIndex graphIndex;
     int originalDimension;
 
     @Setup

jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java

@@ -25,7 +25,7 @@
 package io.github.jbellis.jvector.graph;
 
 import io.github.jbellis.jvector.annotations.Experimental;
-import io.github.jbellis.jvector.graph.GraphIndex.NodeAtLevel;
+import io.github.jbellis.jvector.graph.ImmutableGraphIndex.NodeAtLevel;
 import io.github.jbellis.jvector.graph.similarity.DefaultSearchScoreProvider;
 import io.github.jbellis.jvector.graph.similarity.ScoreFunction;
 import io.github.jbellis.jvector.graph.similarity.SearchScoreProvider;
@@ -43,10 +43,10 @@
 
 /**
  * Searches a graph to find nearest neighbors to a query vector. For more background on the
- * search algorithm, see {@link GraphIndex}.
+ * search algorithm, see {@link ImmutableGraphIndex}.
  */
 public class GraphSearcher implements Closeable {
-    private GraphIndex.View view;
+    private ImmutableGraphIndex.View view;
 
     // Scratch data structures that are used in each {@link #searchInternal} call. These can be expensive
     // to allocate, so they're cleared and reused across calls.
@@ -71,14 +71,14 @@ public class GraphSearcher implements Closeable {
     /**
      * Creates a new graph searcher from the given GraphIndex
      */
-    public GraphSearcher(GraphIndex graph) {
+    public GraphSearcher(ImmutableGraphIndex graph) {
         this(graph.getView());
     }
 
     /**
      * Creates a new graph searcher from the given GraphIndex.View
      */
-    protected GraphSearcher(GraphIndex.View view) {
+    protected GraphSearcher(ImmutableGraphIndex.View view) {
         this.view = view;
         this.candidates = new NodeQueue(new GrowableLongHeap(100), NodeQueue.Order.MAX_HEAP);
         this.evictedResults = new NodesUnsorted(100);
@@ -112,7 +112,7 @@ private void initializeScoreProvider(SearchScoreProvider scoreProvider) {
         cachingReranker = new CachingReranker(scoreProvider);
     }
 
-    public GraphIndex.View getView() {
+    public ImmutableGraphIndex.View getView() {
         return view;
     }
 
@@ -129,7 +129,7 @@ public void usePruning(boolean usage) {
      * Convenience function for simple one-off searches.  It is caller's responsibility to make sure that it
      * is the unique owner of the vectors instance passed in here.
      */
-    public static SearchResult search(VectorFloat<?> queryVector, int topK, RandomAccessVectorValues vectors, VectorSimilarityFunction similarityFunction, GraphIndex graph, Bits acceptOrds) {
+    public static SearchResult search(VectorFloat<?> queryVector, int topK, RandomAccessVectorValues vectors, VectorSimilarityFunction similarityFunction, ImmutableGraphIndex graph, Bits acceptOrds) {
         try (var searcher = new GraphSearcher(graph)) {
             var ssp = DefaultSearchScoreProvider.exact(queryVector, similarityFunction, vectors);
             return searcher.search(ssp, topK, acceptOrds);
@@ -142,7 +142,7 @@ public static SearchResult search(VectorFloat<?> queryVector, int topK, RandomAc
      * Convenience function for simple one-off searches.  It is caller's responsibility to make sure that it
      * is the unique owner of the vectors instance passed in here.
      */
-    public static SearchResult search(VectorFloat<?> queryVector, int topK, int rerankK, RandomAccessVectorValues vectors, VectorSimilarityFunction similarityFunction, GraphIndex graph, Bits acceptOrds) {
+    public static SearchResult search(VectorFloat<?> queryVector, int topK, int rerankK, RandomAccessVectorValues vectors, VectorSimilarityFunction similarityFunction, ImmutableGraphIndex graph, Bits acceptOrds) {
         try (var searcher = new GraphSearcher(graph)) {
             var ssp = DefaultSearchScoreProvider.exact(queryVector, similarityFunction, vectors);
             return searcher.search(ssp, topK, rerankK, 0.f, 0.f, acceptOrds);
@@ -160,7 +160,7 @@ public static SearchResult search(VectorFloat<?> queryVector, int topK, int rera
      *
      * @param view the new view
      */
-    public void setView(GraphIndex.View view) {
+    public void setView(ImmutableGraphIndex.View view) {
         this.view = view;
     }
 
@@ -169,9 +169,9 @@ public void setView(GraphIndex.View view) {
      */
     @Deprecated
     public static class Builder {
-        private final GraphIndex.View view;
+        private final ImmutableGraphIndex.View view;
 
-        public Builder(GraphIndex.View view) {
+        public Builder(ImmutableGraphIndex.View view) {
             this.view = view;
         }
 

jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphSearcher.java

@@ -29,6 +29,8 @@
 import io.github.jbellis.jvector.util.Bits;
 import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
 import io.github.jbellis.jvector.vector.types.VectorFloat;
+
+import java.util.List;
 import java.util.Objects;
 
 import java.io.Closeable;
@@ -44,7 +46,7 @@
  * All methods are threadsafe.  Operations that require persistent state are wrapped
  * in a View that should be created per accessing thread.
  */
-public interface GraphIndex extends AutoCloseable, Accountable {
+public interface ImmutableGraphIndex extends AutoCloseable, Accountable {
     /** Returns the number of nodes in the graph */
     @Deprecated
     default int size() {
@@ -76,6 +78,8 @@ default int size() {
      */
     int maxDegree();
 
+    List<Integer> maxDegrees();
+
     /**
      * @return the first ordinal greater than all node ids in the graph.  Equal to size() in simple cases;
      * May be different from size() if nodes are being added concurrently, or if nodes have been
@@ -107,6 +111,14 @@ default boolean containsNode(int nodeId) {
      */
     int getDegree(int level);
 
+    /**
+     * Returns the average degree computed over nodes in the specified layer.
+     *
+     * @param level the level of interest.
+     * @return the average degree or NaN if no nodes are present.
+     */
+    double getAverageDegree(int level);
+
     /**
      * Return the number of vectors/nodes in the given level.
      * @param level The level of interest
@@ -150,6 +162,11 @@ interface View extends Closeable {
         default int getIdUpperBound() {
             return size();
         }
+
+        /**
+         * Whether the given node is present in the given layer of the graph.
+         */
+        boolean contains(int level, int node);
     }
 
     /**
@@ -161,7 +178,7 @@ interface ScoringView extends View {
         ScoreFunction.ApproximateScoreFunction approximateScoreFunctionFor(VectorFloat<?> queryVector, VectorSimilarityFunction vsf);
     }
 
-    static String prettyPrint(GraphIndex graph) {
+    static String prettyPrint(ImmutableGraphIndex graph) {
         StringBuilder sb = new StringBuilder();
         sb.append(graph);
         sb.append("\n");

…ub/jbellis/jvector/graph/GraphIndex.java …s/jvector/graph/ImmutableGraphIndex.javajvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndex.java renamed to jvector-base/src/main/java/io/github/jbellis/jvector/graph/ImmutableGraphIndex.java jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndex.java renamed to jvector-base/src/main/java/io/github/jbellis/jvector/graph/ImmutableGraphIndex.java

@@ -0,0 +1,170 @@
+/*
+ * All changes to the original code are Copyright DataStax, Inc.
+ *
+ * Please see the included license file for details.
+ */
+
+/*
+ * Original license:
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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;
+
+import io.github.jbellis.jvector.util.BitSet;
+import io.github.jbellis.jvector.util.ThreadSafeGrowableBitSet;
+
+import java.util.List;
+import java.util.stream.IntStream;
+
+
+/**
+ * An {@link ImmutableGraphIndex} that offers concurrent access; for typical graphs you will get significant
+ * speedups in construction and searching as you add threads.
+ *
+ * <p>The base layer (layer 0) contains all nodes, while higher layers are stored in sparse maps.
+ * For searching, use a view obtained from {@link #getView()} which supports level–aware operations.
+ */
+interface MutableGraphIndex extends ImmutableGraphIndex {
+    /**
+     * Add the given node ordinal with an empty set of neighbors.
+     *
+     * <p>Nodes can be inserted out of order, but it requires that the nodes preceded by the node
+     * inserted out of order are eventually added.
+     *
+     * <p>Actually populating the neighbors, and establishing bidirectional links, is the
+     * responsibility of the caller.
+     *
+     * <p>It is also the responsibility of the caller to ensure that each node is only added once.
+     */
+    void addNode(NodeAtLevel nodeLevel);
+
+    /**
+     * Add the given node ordinal with an empty set of neighbors.
+     *
+     * <p>Nodes can be inserted out of order, but it requires that the nodes preceded by the node
+     * inserted out of order are eventually added.
+     *
+     * <p>Actually populating the neighbors, and establishing bidirectional links, is the
+     * responsibility of the caller.
+     *
+     * <p>It is also the responsibility of the caller to ensure that each node is only added once.
+     */
+    void addNode(int level, int node);
+
+    /**
+     * Whether the given node is present in the graph.
+     */
+    boolean contains(NodeAtLevel nodeLevel);
+
+    /**
+     * Whether the given node is present in the given layer of the graph.
+     */
+    boolean contains(int level, int node);
+
+    /**
+     * Add the given node ordinal with an empty set of neighbors.
+     *
+     * <p>Nodes can be inserted out of order, but it requires that the nodes preceded by the node
+     * inserted out of order are eventually added.
+     *
+     * <p>Actually populating the neighbors, and establishing bidirectional links, is the
+     * responsibility of the caller.
+     *
+     * <p>It is also the responsibility of the caller to ensure that each node is only added once.
+     */
+    void connectNode(int level, int node, NodeArray nodes);
+
+    /**
+     * Use with extreme caution. Used by Builder to load a saved graph and for rescoring.
+     */
+    void connectNode(NodeAtLevel nodeLevel, NodeArray nodes);
+
+    /**
+     * Mark the given node deleted.  Does NOT remove the node from the graph.
+     */
+    void markDeleted(int node);
+
+    /** must be called after addNode once neighbors are linked in all levels. */
+    void markComplete(NodeAtLevel nodeLevel);
+
+    void updateEntryNode(NodeAtLevel newEntry);
+
+    /**
+     * Returns an upper bound on the amount of memory used by a single node, in bytes.
+     */
+    long ramBytesUsedOneNode(int layer);
+
+    ThreadSafeGrowableBitSet getDeletedNodes();
+
+    void setDegrees(List<Integer> layerDegrees);
+
+    /**
+     * Enforce the degree of the given node in all layers.
+     */
+    void enforceDegree(int node);
+
+    /**
+     * Returns an iterator over the neighbors for the given node at the specified level, which can be empty if the node does not belong to that layer.
+     */
+    NodesIterator getNeighborsIterator(NodeAtLevel nodeLevel);
+
+    /**
+     * Returns an iterator over the neighbors for the given node at the specified level, which can be empty if the node does not belong to that layer.
+     */
+    NodesIterator getNeighborsIterator(int level, int node);
+
+    /**
+     * Removes the given node from all layers.
+     *
+     * @param node the node id to remove
+     * @return the number of layers from which it was removed
+     */
+    int removeNode(int node);
+
+    /**
+     * Returns an Integer stream with the nodes contained in the specified level.
+     */
+    IntStream nodeStream(int level);
+
+    /**
+     * Returns the maximum (coarser) level that contains a vector in the graph or -1 if the node is not in the graph.
+     */
+    int getMaxLevelForNode(int node);
+
+    /**
+     * @return the node of the graph to start searches at
+     */
+    NodeAtLevel entryNode();
+
+    /**
+     * Add the given neighbors to the given node at the specified level, maintaining diversity
+     * It also adds backlinks from the neighbors to the given node.
+     * The edges will only be added if the out-degree of the node is less than overflowRatio times the max degree.
+     */
+    void addEdges(int level, int node, NodeArray candidates, float overflowRatio);
+
+    /**
+     * Remove edges to deleted nodes and add the new connections, maintaining diversity
+     */
+    void replaceDeletedNeighbors(int level, int node, BitSet toDelete, NodeArray candidates);
+
+    /**
+     * Signals that all mutations have been completed and the graph will not be mutated any further.
+     * Should be called by the builder after all mutations are completed (during cleanup).
+     */
+    void allMutationsCompleted();
+}