Reduce NeighborArray heap memory

[weizijun]

When bbq is used with lucene, one datanode can contain more data.
So when more shards are merged concurrently, there will be a problem of very high heap memory size.
I found that the NeighborArray object was taking up a lot of memory. And I found that the number of nodes always fails to reach maxSize. It only uses about 1/3 or 1/4 of maxSize.
Therefore, I use FloatArrayList\IntArrayList to replace float[]\int[], which can significantly reduce the heap memory usage.

Here is a comparison of the jmap histo results(I set the parameter of m = 64):
before:

 num     #instances         #bytes  class name (module)
-------------------------------------------------------
   1:      11443026     6396808120  [F ([email protected])
   2:      11387631     6129931608  [I ([email protected])
   3:       3265644     1319152760  [B ([email protected])
   4:      11308339      361866848  org.apache.lucene.util.hnsw.NeighborArray ([email protected])
   5:      11134203      267240168  [Lorg.apache.lucene.util.hnsw.NeighborArray; ([email protected])
   6:            77       57916272  [[Lorg.apache.lucene.util.hnsw.NeighborArray; ([email protected])
   7:       2404231       57701544  java.lang.String ([email protected])
   8:         34911       42546120  Ljdk.internal.vm.FillerArray; ([email protected])
   9:        772788       30911520  org.nlpcn.commons.lang.tire.domain.SmartForest
  10:        113758       19111344  org.apache.lucene.codecs.lucene90.blocktree.SegmentTermsEnumFrame ([email protected])
  11:        545656       17460992  java.util.HashMap$Node ([email protected])

after:

num     #instances         #bytes  class name (module)
-------------------------------------------------------
   1:       9228299     1612257464  [F ([email protected])
   2:       9246406     1402537720  [I ([email protected])
   3:       3279264     1141869192  [B ([email protected])
   4:       9124020      364960800  org.apache.lucene.util.hnsw.NeighborArray ([email protected])
   5:       9124036      218976864  org.apache.lucene.internal.hppc.FloatArrayList ([email protected])
   6:       9124036      218976864  org.apache.lucene.internal.hppc.IntArrayList ([email protected])
   7:       8983027      215608448  [Lorg.apache.lucene.util.hnsw.NeighborArray; ([email protected])
   8:       2492594       59822256  java.lang.String ([email protected])
   9:            56       51013776  [[Lorg.apache.lucene.util.hnsw.NeighborArray; ([email protected])
  10:        772788       30911520  org.nlpcn.commons.lang.tire.domain.SmartForest
  11:         68970       28703992  Ljdk.internal.vm.FillerArray; ([email protected])

The avg size of float[] is 559 before.
The avg size of float[] is 174 after.

The avg size of int[] is 538 before.
The avg size of int[] is 151 after.

I tests some dataset like GIST 100K vectors, 960 dimensions\LAION 100M vectors, 768 dimensions. They have similar conclusions.
I haven't tested the performance very rigorously. It seems that this modification has no impact on performance.

[weizijun]

The TestHnswFloatVectorGraph.testRamUsageEstimate maybe failed, because the OnHeapHnswGraph.ramBytesUsed use the fixed array size to calculate the ram value.

[benwtrent]

We need to make sure that there are no significant performance or concurrency bugs introduced with this. Could you test with https://github.com/mikemccand/luceneutil to verify recall, multi-threaded merge, etc. ?

As for the ram usage test, it should be fixable by figuring out the number of connections per node and giving a valid range (e.g. RAM usage is more than X but less than Y).

[jainankitk]

We need to make sure that there are no significant performance or concurrency bugs introduced with this. Could you test with https://github.com/mikemccand/luceneutil to verify recall, multi-threaded merge, etc. ?

+1 on the performance. Using IntArrayList instead of int[] can impact the performance due to the overhead of copying while growing the array.

Although, I am not completely sure if this change should cause any concurrency bugs compared to original. @benwtrent - Do you have anything specific in mind, that I might be missing?

[benwtrent]

Do you have anything specific in mind, that I might be missing?

I am just being cautious. I know we provide locking, etc. for the concurrent graph builder, but maybe its taking advantage of arrays never growing by accident.

I hate concurrent code, and it always makes me paranoid that we will miss something :).

[weizijun]

I tested this case in elasticsearch and it worked fine.Probably because in elasticsearch, the merge of graph builders is in one merge thread.
I will tests more case that use the concurrent graph builder.

[jainankitk]

I am just being cautious. I know we provide locking, etc. for the concurrent graph builder, but maybe its taking advantage of arrays never growing by accident.

Thanks for elaborating. That makes sense!

I hate concurrent code, and it always makes me paranoid that we will miss something :).

While I don't hate concurrent code, I do share your paranoia regarding missing something! :)

[weizijun]

We need to make sure that there are no significant performance or concurrency bugs introduced with this. Could you test with https://github.com/mikemccand/luceneutil to verify recall, multi-threaded merge, etc. ?

hi, @benwtrent , could you tell me which benchmark tool can test this case?

[benwtrent]

@weizijun lucene-util is the benchmark tool I linked. knnPerfTest is the script.

[weizijun]

lucene-util is the benchmark tool I linked. knnPerfTest is the script.

yeah, I know the tool, I don't know how to test knn in it. The knnPerfTest only test the lucene_candidate, I modify some code to test.
the lucene_baseline result(before modify):

Results:
recall  latency(ms)    nDoc  topK  fanout  maxConn  beamWidth  quantized  index(s)  index_docs/s  num_segments  index_size(MB)  vec_disk(MB)  vec_RAM(MB)  indexType
 0.796       10.284  500000   100      50       64        250         no    180.51       2769.99             8         1501.40      1464.844     1464.844       HNSW

the lucene_candidate result(after modify):

Results:
recall  latency(ms)    nDoc  topK  fanout  maxConn  beamWidth  quantized  index(s)  index_docs/s  num_segments  index_size(MB)  vec_disk(MB)  vec_RAM(MB)  indexType
 0.820       10.712  500000   100      50       64        250         no    187.26       2670.16             8         1501.46      1464.844     1464.844       HNSW

[benwtrent]

@weizijun I am running some benchmarking as well. But the key thing is to update the parameters in knnPerfTest.py making sure numMergeWorker and numMergeThread are greater than 1 ,etc.

[weizijun]

@benwtrent I see the default parameters:

  "numMergeWorker": (12,),
  "numMergeThread": (4,),

Is the current merger effective with these parameters?

[benwtrent]

I ran some benchmarking and this seems OK and doesn't seem to break multi-threadedness.

I left a comment on the underlying structures used.

Please update CHANGES.txt under 10.3 optimizations for this nice optimization! It will be very nice to have better heap utilization on graph building!

[benwtrent]

two things, I think these should be initialized to something like maxSize/4 or maxSize/8.

Additionally, the array lists should enforce the max size and ensure the underlying buffer does not get bigger than the expected max. I think you can do this pretty simply by having something like MaxSizedIntArrayList that accepts an init & max size parameters in its ctor, and overrides ensureBufferSpace so that it:

• disallows growth past maxSize
• Instead of doing ArrayUtil.grow it does ArrayUtil.growInRange(buffer, elementsCount + expectedAdditions, maxSize)

This will prevent overallocation of the underlying buffer and enforce max size limitations.

[weizijun]

done

[weizijun]

I left a comment on the underlying structures used.

Please update CHANGES.txt under 10.3 optimizations for this nice optimization! It will be very nice to have better heap utilization on graph building!

Ok, and about the OnHeapHnswGraph.ramBytesUsed, I didn't change to code. But it maybe need to change the logic.

[benwtrent]

Ok, and about the OnHeapHnswGraph.ramBytesUsed, I didn't change to code. But it maybe need to change the logic.

I think utilizing the underlying array estimations is what we need and the testing needs to be updated to assert within a valid range.

[msokolov]

re: concurrency; in theory it should be safe since we lock a row before inserting anything into it. Consider that even with fixed-size arrays we need to track the occupancy (how full the array is) in a thread-safe way.

[weizijun]

I added the change log and changed the initial size of the list to maxSize/8

[weizijun]

The failed TestHnswFloatVectorGraph.testRamUsageEstimate detail:

The seed is -Dtests.seed=FEF2FF50C6824DEF

Expected :1723096.0
Actual   :3785393.0
<Click to see difference>

java.lang.AssertionError: expected:<1723096.0> but was:<3785393.0>
	at __randomizedtesting.SeedInfo.seed([FEF2FF50C6824DEF:3BB69B10FD099A58]:0)
	at org.junit.Assert.fail(Assert.java:89)
	at org.junit.Assert.failNotEquals(Assert.java:835)
	at org.junit.Assert.assertEquals(Assert.java:555)
	at org.junit.Assert.assertEquals(Assert.java:685)
	at org.apache.lucene.util.hnsw.HnswGraphTestCase.testRamUsageEstimate(HnswGraphTestCase.java:804)
	at org.apache.lucene.util.hnsw.TestHnswFloatVectorGraph.testRamUsageEstimate(TestHnswFloatVectorGraph.java:39)
	at java.base/jdk.internal.reflect.DirectMethodHandleAccessor.invoke(DirectMethodHandleAccessor.java:103)
	at java.base/java.lang.reflect.Method.invoke(Method.java:580)
	at com.carrotsearch.randomizedtesting.RandomizedRunner.invoke(RandomizedRunner.java:1763)

If I change the ramBytesUsed of OnHeapHnswGraph to use the real memory of each NeighborArray, the performance will drop a lot.

[benwtrent]

@weizijun

It should be fixed. Having an estimation that is more than 2x off is pretty bad.

this estimation is used to determine how often flushes should occur, etc.

There are a couple of ways this can be fixed.

A simple way could be providing an optional call-back to NeighborArray that accesses package-private method on OnHeapHnswGraph that allows for their individual estimation to be adjusted during array growth.

NeighborArray(OnHeapHnswGraph::updateEstimate...) or something. Then the ram estimation in OnHeapHnswGraph becomes the accumulation of those estimates as the inner estimates evolve.

We need to be cautious there with multi-threadedness as many node updates could be occuring at a time. So, likely this inner accumulator needs to be LongAccumulator and it should also assert that its always a positive number

Please also adjust the inner arrays to enforce their maximal length. This way we never over-allocate.

[weizijun]

hi, @benwtrent, I updated the ramBytesUsed method, the memory is correct now.

Please also adjust the inner arrays to enforce their maximal length. This way we never over-allocate.

I use the maxSize variable to ensure that the elementsCount of the array cannot exceed the maxSize value.

[benwtrent]

I add the MaxSizedIntArrayList class to solve oversize problem.

You need one for scores as well. That MaxSizedIntArrayList looks good :)

[weizijun]

You need one for scores as well. That MaxSizedIntArrayList looks good :)

I didn't add the MaxSizedFloatArrayList because the int array will be passed out via the nodes() method, but the float array of the scores will not be passed out. Do I need to change FloatArrayList to MaxSizedFloatArrayList?

[benwtrent]

but the float array of the scores will not be passed out. Do I need to change FloatArrayList to MaxSizedFloatArrayList?

I don't understand this. How does the FloatArrayList bypass the overallocation possible via ArrayUtils.grow? Its the same structure right? and for every node there is a score.

[weizijun]

I don't understand this. How does the FloatArrayList bypass the overallocation possible via ArrayUtils.grow? Its the same structure right? and for every node there is a score.

maxSize controls the total number of array values, so if the length of ArrayUtils.grow exceeds maxSize, the actual value will not exceed maxSize.

However, the nodes() method of NeighborArray returns nodes.buffer, so maxSize cannot control the int array. Therefore, a MaxSizedIntArrayList class is needed to ensure that the length of the int array can overflow maxSize.

[benwtrent]

maxSize controls the total number of array values, so if the length of ArrayUtils.grow exceeds maxSize, the actual value will not exceed maxSize.

That is frankly not true for FloatArrayList. ArrayUtils.grow can easily cause the underlying buffer to exceed the length.

I don't understand the argument against providing actual max size restricted float array list. Is there a reason to not do this?

[weizijun]

That is frankly not true for FloatArrayList. ArrayUtils.grow can easily cause the underlying buffer to exceed the length.

I don't understand the argument against providing actual max size restricted float array list. Is there a reason to not do this?

@benwtrent Okay, I added MaxSizedFloatArrayList. Can you review the PR again?

[benwtrent]

This is looking great!

One small comment on updating the .gitignore. But once that is reverted to be the same as on main, then I can merge and backport.

Thank you for all the iterations and such!

[benwtrent]

If you think these need updated, could you do it in a separate PR? I would like to keep this change restricted to HNSW.

[weizijun]

Oh, sorry, that was a mistake, I'll delete it.

[benwtrent]

great stuff! I will merge and backport.

[weizijun]

Here are the statistics of 100w hnsw graphs, with m = 16 and ef = 100:
Level count = 5:

level: 0, node count: 1000000
level: 1, node count: 62835
level: 2, node count: 3926
level: 3, node count: 235
level: 4, node count: 12

Average number of neighbors per level:

level: 0, avg neighbor count: 8.026909
level: 1, avg neighbor count: 7.539985676772499
level: 2, avg neighbor count: 8.596535914416709
level: 3, avg neighbor count: 8.353191489361702
level: 4, avg neighbor count: 3.8333333333333335

The detail of neighbor count:
level: 0

level: 0, neighbor count: 1, node count: 141664
level: 0, neighbor count: 2, node count: 130484
level: 0, neighbor count: 3, node count: 111485
level: 0, neighbor count: 4, node count: 91141
level: 0, neighbor count: 5, node count: 72929
level: 0, neighbor count: 6, node count: 59030
level: 0, neighbor count: 7, node count: 47796
level: 0, neighbor count: 8, node count: 39864
level: 0, neighbor count: 9, node count: 33320
level: 0, neighbor count: 10, node count: 27923
level: 0, neighbor count: 11, node count: 23972
level: 0, neighbor count: 12, node count: 20777
level: 0, neighbor count: 13, node count: 17986
level: 0, neighbor count: 14, node count: 15510
level: 0, neighbor count: 15, node count: 13725
level: 0, neighbor count: 16, node count: 12296
level: 0, neighbor count: 17, node count: 10947
level: 0, neighbor count: 18, node count: 9826
level: 0, neighbor count: 19, node count: 8765
level: 0, neighbor count: 20, node count: 7947
level: 0, neighbor count: 21, node count: 7348
level: 0, neighbor count: 22, node count: 6639
level: 0, neighbor count: 23, node count: 6045
level: 0, neighbor count: 24, node count: 5413
level: 0, neighbor count: 25, node count: 5101
level: 0, neighbor count: 26, node count: 4569
level: 0, neighbor count: 27, node count: 4105
level: 0, neighbor count: 28, node count: 3965
level: 0, neighbor count: 29, node count: 3564
level: 0, neighbor count: 30, node count: 3330
level: 0, neighbor count: 31, node count: 3019
level: 0, neighbor count: 32, node count: 49515

level: 1

level: 1, neighbor count: 1, node count: 6760
level: 1, neighbor count: 2, node count: 6707
level: 1, neighbor count: 3, node count: 6127
level: 1, neighbor count: 4, node count: 5277
level: 1, neighbor count: 5, node count: 4420
level: 1, neighbor count: 6, node count: 3805
level: 1, neighbor count: 7, node count: 3321
level: 1, neighbor count: 8, node count: 2827
level: 1, neighbor count: 9, node count: 2502
level: 1, neighbor count: 10, node count: 2093
level: 1, neighbor count: 11, node count: 1849
level: 1, neighbor count: 12, node count: 1645
level: 1, neighbor count: 13, node count: 1521
level: 1, neighbor count: 14, node count: 1257
level: 1, neighbor count: 15, node count: 1163
level: 1, neighbor count: 16, node count: 11561

level: 2

level: 2, neighbor count: 1, node count: 298
level: 2, neighbor count: 2, node count: 302
level: 2, neighbor count: 3, node count: 309
level: 2, neighbor count: 4, node count: 278
level: 2, neighbor count: 5, node count: 267
level: 2, neighbor count: 6, node count: 251
level: 2, neighbor count: 7, node count: 196
level: 2, neighbor count: 8, node count: 209
level: 2, neighbor count: 9, node count: 178
level: 2, neighbor count: 10, node count: 159
level: 2, neighbor count: 11, node count: 153
level: 2, neighbor count: 12, node count: 134
level: 2, neighbor count: 13, node count: 125
level: 2, neighbor count: 14, node count: 75
level: 2, neighbor count: 15, node count: 106
level: 2, neighbor count: 16, node count: 886

level: 3

level: 3, neighbor count: 1, node count: 18
level: 3, neighbor count: 2, node count: 14
level: 3, neighbor count: 3, node count: 11
level: 3, neighbor count: 4, node count: 14
level: 3, neighbor count: 5, node count: 17
level: 3, neighbor count: 6, node count: 20
level: 3, neighbor count: 7, node count: 19
level: 3, neighbor count: 8, node count: 11
level: 3, neighbor count: 9, node count: 23
level: 3, neighbor count: 10, node count: 12
level: 3, neighbor count: 11, node count: 12
level: 3, neighbor count: 12, node count: 9
level: 3, neighbor count: 13, node count: 7
level: 3, neighbor count: 14, node count: 10
level: 3, neighbor count: 15, node count: 4
level: 3, neighbor count: 16, node count: 34

level: 4

level: 4, neighbor count: 1, node count: 1
level: 4, neighbor count: 2, node count: 2
level: 4, neighbor count: 3, node count: 5
level: 4, neighbor count: 5, node count: 2
level: 4, neighbor count: 8, node count: 2

[jainankitk]

Thanks @benwtrent and @weizijun for seeing this through!