Bulk score hnsw diversity check

[benwtrent]

This adds bulk scoring to diversity check. While this means that diversity check cannot exit super early (e.g. if it only needs to check 2 docs), I continually see diversity check as being the most expensive part of HNSW graph merging.

This tells me that typically, it isn't just one doc that is checked.

I ran 1M 768 cohere, force-merging with 4 threads.

baseline: 128.01
candidate: 92.15

[benwtrent]

Thinking about this more, doing 8 always, even if there are only 8 connections seems foolish. I will benchmark with Math.min((neighbors.nodes().length + 1)/2, bulkScoreNodes.length)).

[kaivalnp]

Looks like a nice optimization!

[kaivalnp]

IIUC neighbors.nodes() returns the internal array used to store neighbor nodes, which can be larger than actual number of neighbors and is exponentially growing -- can we have edge cases where the length is about twice the actual number of neighbors, causing us to bulk score everything?

I wonder if we should use neighbors.size() instead?

[benwtrent]

good call! And good call on the array copy!, Just needed to handle the tail. New benchmarks show even nicer perf ;)

[kaivalnp]

I was worried about thread-safety of these arrays (given we have concurrent merging), but from this comment it looks like instances of this class are not shared across threads, but rather multiple instances of this class (across different threads) can operate on a single HnswGraph?

[benwtrent]

@kaivalnp the scorer itself isn't threadsafe. I assumed that since we were using a scorer, we were OK.

I had the same threading concerns and looked up and it seems that each thread as a unique builder object instance (The worker of the thread) and they all work on the same graph.

[kaivalnp]

Looks like we don't need to perform any operation per-node of this loop, can we reduce some iterations using something like:

private boolean diversityCheck(float score, NeighborArray neighbors, RandomVectorScorer scorer)
    throws IOException {
  final int chunk = Math.min(Math.ceilDiv(neighbors.size(), 2), bulkScoreNodes.length);
  for (int start = 0; start < neighbors.size(); start += chunk) {
    int length = Math.min(neighbors.size() - start, chunk);
    System.arraycopy(neighbors.nodes(), start, bulkScoreNodes, 0, length);
    if (scorer.bulkScore(bulkScoreNodes, bulkScores, length) >= score) {
      return false;
    }
  }
  return true;
}

[kaivalnp]

Thanks @benwtrent! Curious if you were able to run knnPerfTest?

[kaivalnp]

I don't think we need this tail -- we're doing Math.min(bulkScoreChunk, neighbors.size() - scored) in the above loop, which automatically bulk-scores the tail (using the second value)

[benwtrent]

Yep, good call 🤦 its reflex.

[kaivalnp]

nit: maybe host 8 up as a private static final variable?

[benwtrent]

@kaivalnp I did run with 768 dim vectors (see PR description).

I ran it again and I get perf around 77.04s for force merge to 90s force merge. Force merge performance is pretty variable, but with all my runs, this provides a measurable speed improvement.

[kaivalnp]

LGTM!