Optimize HNSW diversity calculation

[zhaih]

Description

Currently when we pop out the worst diverse node we go from the most distant node and calculate diversity one by one. But actually when we add the first batch of neighbours, we have already checked the diversity and rejected the non-diverse node. So the only nodes that are unchecked are only the ones that inserted by neighbours (using insertSorted), we should be able to only check those "unchecked" nodes but not every node.

I rely on current unit test to check the correctness but if that's not enough I'm pretty happy to write some test as well.

KNNGraphTester bench based on mikemccand/luceneutil#218:

This patch

recall	latency	nDoc	fanout	maxConn	beamWidth visited index ms
0.802	13.30	1000000	20	32	250	  120	  630614	1.00	post-filter
0.907	13.77	1000000	250	32	250	  350	  621157	1.00	post-filter
0.830	13.95	1000000	20	32	500	  120	  1242756	1.00	post-filter
0.926	14.33	1000000	250	32	500	  350	  1246090	1.00	post-filter
0.804	12.73	1000000	20	96	250	  120	  672480	1.00	post-filter
0.910	13.75	1000000	250	96	250	  350	  671602	1.00	post-filter
0.832	14.01	1000000	20	96	500	  120	  1376617	1.00	post-filter
0.928	14.36	1000000	250	96	500	  350	  1381215	1.00	post-filter

Mainline

recall	latency	nDoc	fanout	maxConn	beamWidth visited index ms
0.802	12.82	1000000	20	32	250	  120	  686167	1.00	post-filter
0.907	13.58	1000000	250	32	250	  350	  633083	1.00	post-filter
0.831	13.73	1000000	20	32	500	  120	  1296331	1.00	post-filter
0.927	14.59	1000000	250	32	500	  350	  1275096	1.00	post-filter
0.804	12.79	1000000	20	96	250	  120	  672598	1.00	post-filter
0.910	13.27	1000000	250	96	250	  350	  670174	1.00	post-filter
0.832	13.76	1000000	20	96	500	  120	  1381401	1.00	post-filter
0.928	14.71	1000000	250	96	500	  350	  1393253	1.00	post-filter

[msokolov]

I made some small comments about the internal APIs, but my main questions here are about performance tradeoffs. This will use more RAM, maybe even a lot, and it should improve indexing time. This could be a good tradeoff, but we want to know how much of each, and what the bounds if any on RAM growth would be.

Before we commit this, we definitely need to run some perf tests. I have done this in the past using KnnGraphTester. It's designed to measure latency/recall but it also does report on graph build time and I think we could instrument to see how RAM is growing? To test you need a set of vectors that is large enough to be realistic. up to 1M is probably enough. There are vectors available via luceneutil.

One thought I had is whether we can limit the RAM growth by using a bitset tied to node index in the array rather than an array of int node ids? This would require reshuffling the bits when nodes are sorted in the NeighborArray. Could be a good choice to limit memory

[zhaih]

@msokolov Thanks for your feedback. Indeed the memory consumption is a concern and after I thought it more I found that it might be better to lazily sort the nodes and figure out which nodes are unchecked all at once.
In this way we still able to avoid the extra computation at the cost of 1 extra int, which I believe is more or less negligible now.

For sure I will do the KnnGraphTester (and thank you for pointing me to there, TBH I haven't done any test with KNN yet...) still to just see whether the performance gain is observable, but let me know whether the latter approach addresses your concern.
Thank you!

[zhaih]

I ran the KnnGraphTester and pasted the result to the description, basically when max-conn is low (there's high possiblility of popping out a node) the improvement seems to be obvious (2-5% on indexing time), when max-conn is high then the indexing time is similar or only slightly better, which is also expected.

[msokolov]

I see; well there seems like some benefit. Also, the impact is going to vary depending on the nature of the data. I guess we'd expect that vectors with a high degree of "clustering" would tend to have more ejections due to diversity checks.Let me take one more pass and see if there's any more polishing to do, otherwise it seems good to me.

[MarcusSorealheis]

Nit: I know we use nbrNbr in a few more places, but as you have reminded me, the nondescript names could make it more difficult for others to understand the code and participate.

Is there another way to express this to make it easier for people to read/scan? neighborArray to be most explicit, candidateNeighbors or something else? You would only need to change two other files to preserve consistency if you elected to move from nbrNbr

[zhaih]

Yeah this is not quite descriptive but I'm not sure whether I can come up with a better one.. Maybe nbrsOfNbr? neighborArray is not exactly right since it is actually "the neighbor array of the neighbor" LOL

[MarcusSorealheis]

I think neighborArray would be an improvement but not completely accurate. neighborsArray would be implicitly good, too, yet implicit is not good enough and it's not clear about the relationship with the neighbor in focus.

nbrsOfNbr is even more clear. I'd recommend going with your suggestion.

[zhaih]

Sure, changed!

[MarcusSorealheis]

Is the change unneeded in the Lucene90 file if it's changed here? I was thinking we could rename nbrNbr wherever it is used because it is confusing.

[zhaih]

I can do it for the other files because there are not many of them, however I don't usually do a text search for the whole repo to see whether the rename can be applied even to the files I haven't touched originally...

[MarcusSorealheis]

As a principal, that totally makes sense. For this specific project, it makes even more sense because it's pretty complicated in some areas.

The other two files are codec files with identical logic in this section, so it probably would not pose much of a problem. Although, I'm sure others could weigh in or -1 if adamant about nbrNbr. Most of the time, I favor readability over anything else. That could be because I am lazy. 😛 nbrsOfNbr is really clear while also being few characters to type. Still, no pressure.

[zhaih]

The test error seems due to a flaw introduced in #12169, since it seems they're anyway gonna rework on it for thread safety reason I left a comment there.

[msokolov]

I think this is ready? Is there anything you were still planning to improve / test? Oh I see there is a precommit test failing ...

[msokolov]

LGTM. After reading I think I see what's going on. The name is deceptively simple; I wonder if it shouldn't hint that there is something tricky going on in here? Like sortAndReturnPreviouslyUnsorted() :) but it's fine to leave as is - the javadoc anyway explains it

[zhaih]

Haha I'll just leave it as is. It claims itself is "sort" and it does the sort, but just return something different :)

[zhaih]

Oh I see there is a precommit test failing...

Yeah I want to wait for that fix a bit before I merge this one... Since anyway this won't make into Lucene 9.6 LOL

[benwtrent]

@zhaih I verified your change doesn't hurt recall. Is this change good to continue working on? I tried replicating your performance improvements, but did not see that big of a change. My search time latency might be drowning out the index build latency.

[zhaih]

@benwtrent I was originally waiting for #12246 to be merged first so that this PR won't break their test, but seems that one is blocked somehow by other PR so I'll just fix the problem here.

I think if this time all the checks passed I'll just merge this PR