Refactor around NeighborArray (#12910)

Author: zhaih

lucene/CHANGES.txt

@@ -31,6 +31,8 @@ Improvements
 
 * GITHUB#12812: Avoid overflows and false negatives in int slice buffer filled-with-zeros assertion. (Stefan Vodita)
 
+* GITHUB#12910: Refactor around NeighborArray to make it more self-contained. (Patrick Zhai)
+
 Optimizations
 ---------------------
 (No changes)

lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene92/Lucene92HnswVectorsWriter.java

@@ -297,7 +297,7 @@ private OnHeapHnswGraph writeGraph(
         int size = neighbors.size();
         vectorIndex.writeInt(size);
         // Destructively modify; it's ok we are discarding it after this
-        int[] nnodes = neighbors.node();
+        int[] nnodes = neighbors.nodes();
         Arrays.sort(nnodes, 0, size);
         for (int i = 0; i < size; i++) {
           int nnode = nnodes[i];

lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene94/Lucene94HnswVectorsWriter.java

@@ -369,7 +369,7 @@ private void reconstructAndWriteNeigbours(
     vectorIndex.writeInt(size);
 
     // Destructively modify; it's ok we are discarding it after this
-    int[] nnodes = neighbors.node();
+    int[] nnodes = neighbors.nodes();
     for (int i = 0; i < size; i++) {
       nnodes[i] = oldToNewMap[nnodes[i]];
     }
@@ -506,7 +506,7 @@ private void writeGraph(OnHeapHnswGraph graph) throws IOException {
         int size = neighbors.size();
         vectorIndex.writeInt(size);
         // Destructively modify; it's ok we are discarding it after this
-        int[] nnodes = neighbors.node();
+        int[] nnodes = neighbors.nodes();
         Arrays.sort(nnodes, 0, size);
         for (int i = 0; i < size; i++) {
           int nnode = nnodes[i];

lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene95/Lucene95HnswVectorsWriter.java

@@ -396,7 +396,7 @@ private void reconstructAndWriteNeigbours(
     vectorIndex.writeVInt(size);
 
     // Destructively modify; it's ok we are discarding it after this
-    int[] nnodes = neighbors.node();
+    int[] nnodes = neighbors.nodes();
     for (int i = 0; i < size; i++) {
       nnodes[i] = oldToNewMap[nnodes[i]];
     }
@@ -556,7 +556,7 @@ private int[][] writeGraph(OnHeapHnswGraph graph) throws IOException {
         long offsetStart = vectorIndex.getFilePointer();
         vectorIndex.writeVInt(size);
         // Destructively modify; it's ok we are discarding it after this
-        int[] nnodes = neighbors.node();
+        int[] nnodes = neighbors.nodes();
         Arrays.sort(nnodes, 0, size);
         // Now that we have sorted, do delta encoding to minimize the required bits to store the
         // information

lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java

@@ -319,7 +319,7 @@ private void reconstructAndWriteNeighbours(NeighborArray neighbors, int[] oldToN
     vectorIndex.writeVInt(size);
 
     // Destructively modify; it's ok we are discarding it after this
-    int[] nnodes = neighbors.node();
+    int[] nnodes = neighbors.nodes();
     for (int i = 0; i < size; i++) {
       nnodes[i] = oldToNewMap[nnodes[i]];
     }
@@ -415,7 +415,7 @@ private int[][] writeGraph(OnHeapHnswGraph graph) throws IOException {
         long offsetStart = vectorIndex.getFilePointer();
         vectorIndex.writeVInt(size);
         // Destructively modify; it's ok we are discarding it after this
-        int[] nnodes = neighbors.node();
+        int[] nnodes = neighbors.nodes();
         Arrays.sort(nnodes, 0, size);
         // Now that we have sorted, do delta encoding to minimize the required bits to store the
         // information

lucene/core/src/java/org/apache/lucene/util/hnsw/HnswConcurrentMergeBuilder.java

@@ -201,7 +201,7 @@ void graphSeek(HnswGraph graph, int level, int targetNode) {
           nodeBuffer = new int[neighborArray.size()];
         }
         size = neighborArray.size();
-        if (size >= 0) System.arraycopy(neighborArray.node, 0, nodeBuffer, 0, size);
+        if (size >= 0) System.arraycopy(neighborArray.nodes(), 0, nodeBuffer, 0, size);
       } finally {
         neighborArray.rwlock.readLock().unlock();
       }

lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java

@@ -310,15 +310,11 @@ private void addDiverseNeighbors(int level, int node, NeighborArray candidates)
       if (mask[i] == false) {
         continue;
       }
-      int nbr = candidates.node[i];
+      int nbr = candidates.nodes()[i];
       NeighborArray nbrsOfNbr = hnsw.getNeighbors(level, nbr);
       nbrsOfNbr.rwlock.writeLock().lock();
       try {
-        nbrsOfNbr.addOutOfOrder(node, candidates.score[i]);
-        if (nbrsOfNbr.size() > maxConnOnLevel) {
-          int indexToRemove = findWorstNonDiverse(nbrsOfNbr, nbr);
-          nbrsOfNbr.removeIndex(indexToRemove);
-        }
+        nbrsOfNbr.addAndEnsureDiversity(node, candidates.scores()[i], nbr, scorerSupplier);
       } finally {
         nbrsOfNbr.rwlock.writeLock().unlock();
       }
@@ -336,8 +332,8 @@ private boolean[] selectAndLinkDiverse(
     for (int i = candidates.size() - 1; neighbors.size() < maxConnOnLevel && i >= 0; i--) {
       // compare each neighbor (in distance order) against the closer neighbors selected so far,
       // only adding it if it is closer to the target than to any of the other selected neighbors
-      int cNode = candidates.node[i];
-      float cScore = candidates.score[i];
+      int cNode = candidates.nodes()[i];
+      float cScore = candidates.scores()[i];
       assert cNode <= hnsw.maxNodeId();
       if (diversityCheck(cNode, cScore, neighbors)) {
         mask[i] = true;
@@ -371,70 +367,14 @@ private boolean diversityCheck(int candidate, float score, NeighborArray neighbo
       throws IOException {
     RandomVectorScorer scorer = scorerSupplier.scorer(candidate);
     for (int i = 0; i < neighbors.size(); i++) {
-      float neighborSimilarity = scorer.score(neighbors.node[i]);
+      float neighborSimilarity = scorer.score(neighbors.nodes()[i]);
       if (neighborSimilarity >= score) {
         return false;
       }
     }
     return true;
   }
 
-  /**
-   * Find first non-diverse neighbour among the list of neighbors starting from the most distant
-   * neighbours
-   */
-  private int findWorstNonDiverse(NeighborArray neighbors, int nodeOrd) throws IOException {
-    RandomVectorScorer scorer = scorerSupplier.scorer(nodeOrd);
-    int[] uncheckedIndexes = neighbors.sort(scorer);
-    if (uncheckedIndexes == null) {
-      // all nodes are checked, we will directly return the most distant one
-      return neighbors.size() - 1;
-    }
-    int uncheckedCursor = uncheckedIndexes.length - 1;
-    for (int i = neighbors.size() - 1; i > 0; i--) {
-      if (uncheckedCursor < 0) {
-        // no unchecked node left
-        break;
-      }
-      if (isWorstNonDiverse(i, neighbors, uncheckedIndexes, uncheckedCursor)) {
-        return i;
-      }
-      if (i == uncheckedIndexes[uncheckedCursor]) {
-        uncheckedCursor--;
-      }
-    }
-    return neighbors.size() - 1;
-  }
-
-  private boolean isWorstNonDiverse(
-      int candidateIndex, NeighborArray neighbors, int[] uncheckedIndexes, int uncheckedCursor)
-      throws IOException {
-    float minAcceptedSimilarity = neighbors.score[candidateIndex];
-    RandomVectorScorer scorer = scorerSupplier.scorer(neighbors.node[candidateIndex]);
-    if (candidateIndex == uncheckedIndexes[uncheckedCursor]) {
-      // the candidate itself is unchecked
-      for (int i = candidateIndex - 1; i >= 0; i--) {
-        float neighborSimilarity = scorer.score(neighbors.node[i]);
-        // candidate node is too similar to node i given its score relative to the base node
-        if (neighborSimilarity >= minAcceptedSimilarity) {
-          return true;
-        }
-      }
-    } else {
-      // else we just need to make sure candidate does not violate diversity with the (newly
-      // inserted) unchecked nodes
-      assert candidateIndex > uncheckedIndexes[uncheckedCursor];
-      for (int i = uncheckedCursor; i >= 0; i--) {
-        float neighborSimilarity = scorer.score(neighbors.node[uncheckedIndexes[i]]);
-        // candidate node is too similar to node i given its score relative to the base node
-        if (neighborSimilarity >= minAcceptedSimilarity) {
-          return true;
-        }
-      }
-    }
-    return false;
-  }
-
   private static int getRandomGraphLevel(double ml, SplittableRandom random) {
     double randDouble;
     do {

lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphSearcher.java

@@ -308,7 +308,7 @@ void graphSeek(HnswGraph graph, int level, int targetNode) {
     @Override
     int graphNextNeighbor(HnswGraph graph) {
       if (++upto < cur.size()) {
-        return cur.node[upto];
+        return cur.nodes()[upto];
       }
       return NO_MORE_DOCS;
     }