Reduce NeighborArray heap memory (#14527)

* use FloatArrayList\IntArrayList to replace float[]\int[]

* use getScores(int i) to replace scores()

* add more tests

* add change log and change the init value

* update OnHeapHnswGraph ramBytesUsed method

* improve

* add MaxSizedIntArrayList

* add MaxSizedFloatArrayList

* add MaxSizedFloatArrayList

* fixed tests

* revert

Author: weizijun

lucene/CHANGES.txt

@@ -106,6 +106,8 @@ Optimizations
 
 * GITHUB#14447: Compute the doc range more efficiently when flushing doc block. (Pan Guixin)
 
+* GITHUB#14527: Reduce NeighborArray heap memory. (weizijun)
+
 * GITHUB#14529, GITHUB#14555, GITHUB#14618: Impl intoBitset for IndexedDISI and Docvalues. (Guo Feng)
 
 * GITHUB#14552: Speed up flush of softdelete by intoBitset. (Guo Feng)

lucene/core/src/java/org/apache/lucene/internal/hppc/MaxSizedFloatArrayList.java

@@ -0,0 +1,102 @@
+/*
+ * 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 org.apache.lucene.internal.hppc;
+
+import static org.apache.lucene.internal.hppc.HashContainers.DEFAULT_EXPECTED_ELEMENTS;
+
+import org.apache.lucene.util.ArrayUtil;
+import org.apache.lucene.util.RamUsageEstimator;
+
+/**
+ * An array-backed list of {@code float} with a maximum size limit.
+ *
+ * @lucene.internal
+ */
+public class MaxSizedFloatArrayList extends FloatArrayList {
+  private static final long BASE_RAM_BYTES_USED =
+      RamUsageEstimator.shallowSizeOfInstance(MaxSizedFloatArrayList.class);
+
+  final int maxSize;
+
+  /** New instance with sane defaults. */
+  public MaxSizedFloatArrayList(int maxSize) {
+    this(maxSize, DEFAULT_EXPECTED_ELEMENTS);
+  }
+
+  /**
+   * New instance with sane defaults.
+   *
+   * @param maxSize The maximum size this list can grow to
+   * @param expectedElements The expected number of elements guaranteed not to cause buffer
+   *     expansion (inclusive).
+   */
+  public MaxSizedFloatArrayList(int maxSize, int expectedElements) {
+    super(expectedElements);
+    assert expectedElements <= maxSize
+        : "expectedElements (" + expectedElements + ") must be <= maxSize (" + maxSize + ")";
+    this.maxSize = maxSize;
+  }
+
+  /** Creates a new list from the elements of another list in its iteration order. */
+  public MaxSizedFloatArrayList(MaxSizedFloatArrayList list) {
+    super(list.size());
+    this.maxSize = list.maxSize;
+    addAll(list);
+  }
+
+  @Override
+  protected void ensureBufferSpace(int expectedAdditions) {
+    if (elementsCount + expectedAdditions > maxSize) {
+      throw new IllegalStateException("Cannot grow beyond maxSize: " + maxSize);
+    }
+    if (elementsCount + expectedAdditions > buffer.length) {
+      this.buffer = ArrayUtil.growInRange(buffer, elementsCount + expectedAdditions, maxSize);
+    }
+  }
+
+  @Override
+  public int hashCode() {
+    int h = 1, max = elementsCount;
+    h = 31 * h + maxSize;
+    for (int i = 0; i < max; i++) {
+      h = 31 * h + BitMixer.mix(this.buffer[i]);
+    }
+    return h;
+  }
+
+  /**
+   * Returns <code>true</code> only if the other object is an instance of the same class and with
+   * the same elements and maxSize.
+   */
+  @Override
+  public boolean equals(Object obj) {
+    if (this == obj) {
+      return true;
+    }
+    if (obj == null || getClass() != obj.getClass()) {
+      return false;
+    }
+    MaxSizedFloatArrayList other = (MaxSizedFloatArrayList) obj;
+    return maxSize == other.maxSize && super.equals(obj);
+  }
+
+  @Override
+  public long ramBytesUsed() {
+    return BASE_RAM_BYTES_USED + RamUsageEstimator.sizeOf(buffer);
+  }
+}

lucene/core/src/java/org/apache/lucene/internal/hppc/MaxSizedIntArrayList.java

@@ -0,0 +1,102 @@
+/*
+ * 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 org.apache.lucene.internal.hppc;
+
+import static org.apache.lucene.internal.hppc.HashContainers.DEFAULT_EXPECTED_ELEMENTS;
+
+import org.apache.lucene.util.ArrayUtil;
+import org.apache.lucene.util.RamUsageEstimator;
+
+/**
+ * An array-backed list of {@code int} with a maximum size limit.
+ *
+ * @lucene.internal
+ */
+public class MaxSizedIntArrayList extends IntArrayList {
+  private static final long BASE_RAM_BYTES_USED =
+      RamUsageEstimator.shallowSizeOfInstance(MaxSizedIntArrayList.class);
+
+  final int maxSize;
+
+  /** New instance with sane defaults. */
+  public MaxSizedIntArrayList(int maxSize) {
+    this(maxSize, DEFAULT_EXPECTED_ELEMENTS);
+  }
+
+  /**
+   * New instance with sane defaults.
+   *
+   * @param maxSize The maximum size this list can grow to
+   * @param expectedElements The expected number of elements guaranteed not to cause buffer
+   *     expansion (inclusive).
+   */
+  public MaxSizedIntArrayList(int maxSize, int expectedElements) {
+    super(expectedElements);
+    assert expectedElements <= maxSize
+        : "expectedElements (" + expectedElements + ") must be <= maxSize (" + maxSize + ")";
+    this.maxSize = maxSize;
+  }
+
+  /** Creates a new list from the elements of another list in its iteration order. */
+  public MaxSizedIntArrayList(MaxSizedIntArrayList list) {
+    super(list.size());
+    this.maxSize = list.maxSize;
+    addAll(list);
+  }
+
+  @Override
+  protected void ensureBufferSpace(int expectedAdditions) {
+    if (elementsCount + expectedAdditions > maxSize) {
+      throw new IllegalStateException("Cannot grow beyond maxSize: " + maxSize);
+    }
+    if (elementsCount + expectedAdditions > buffer.length) {
+      this.buffer = ArrayUtil.growInRange(buffer, elementsCount + expectedAdditions, maxSize);
+    }
+  }
+
+  @Override
+  public int hashCode() {
+    int h = 1, max = elementsCount;
+    h = 31 * h + maxSize;
+    for (int i = 0; i < max; i++) {
+      h = 31 * h + BitMixer.mix(this.buffer[i]);
+    }
+    return h;
+  }
+
+  /**
+   * Returns <code>true</code> only if the other object is an instance of the same class and with
+   * the same elements and maxSize.
+   */
+  @Override
+  public boolean equals(Object obj) {
+    if (this == obj) {
+      return true;
+    }
+    if (obj == null || getClass() != obj.getClass()) {
+      return false;
+    }
+    MaxSizedIntArrayList other = (MaxSizedIntArrayList) obj;
+    return maxSize == other.maxSize && super.equals(obj);
+  }
+
+  @Override
+  public long ramBytesUsed() {
+    return BASE_RAM_BYTES_USED + RamUsageEstimator.sizeOf(buffer);
+  }
+}

lucene/core/src/java/org/apache/lucene/util/ArrayUtil.java

@@ -346,6 +346,30 @@ public static int[] growInRange(int[] array, int minLength, int maxLength) {
     return growExact(array, Math.min(maxLength, potentialLength));
   }
 
+  /**
+   * Returns an array whose size is at least {@code minLength} but not over {@code maxLength},
+   * growing exponentially if it needs to grow.
+   */
+  public static float[] growInRange(float[] array, int minLength, int maxLength) {
+    assert minLength >= 0
+        : "minLength must be positive (got " + minLength + "): likely integer overflow?";
+
+    if (minLength > maxLength) {
+      throw new IllegalArgumentException(
+          "requested minimum array length "
+              + minLength
+              + " is larger than requested maximum array length "
+              + maxLength);
+    }
+
+    if (array.length >= minLength) {
+      return array;
+    }
+
+    int potentialLength = oversize(minLength, Float.BYTES);
+    return growExact(array, Math.min(maxLength, potentialLength));
+  }
+
   /**
    * Returns an array whose size is at least {@code minSize}, generally over-allocating
    * exponentially

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

@@ -362,13 +362,13 @@ private void addDiverseNeighbors(
         Lock lock = hnswLock.write(level, nbr);
         try {
           NeighborArray nbrsOfNbr = getGraph().getNeighbors(level, nbr);
-          nbrsOfNbr.addAndEnsureDiversity(node, candidates.scores()[i], nbr, scorer);
+          nbrsOfNbr.addAndEnsureDiversity(node, candidates.getScores(i), nbr, scorer);
         } finally {
           lock.unlock();
         }
       } else {
         NeighborArray nbrsOfNbr = hnsw.getNeighbors(level, nbr);
-        nbrsOfNbr.addAndEnsureDiversity(node, candidates.scores()[i], nbr, scorer);
+        nbrsOfNbr.addAndEnsureDiversity(node, candidates.getScores(i), nbr, scorer);
       }
     }
   }
@@ -389,7 +389,7 @@ private boolean[] selectAndLinkDiverse(
       // 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.nodes()[i];
-      float cScore = candidates.scores()[i];
+      float cScore = candidates.getScores(i);
       assert cNode <= hnsw.maxNodeId();
       scorer.setScoringOrdinal(cNode);
       if (diversityCheck(cScore, neighbors, scorer)) {
@@ -444,6 +444,7 @@ void finish() throws IOException {
     //  see: https://github.com/apache/lucene/issues/14214
     // connectComponents();
     frozen = true;
+    hnsw.finishBuild();
   }
 
   @SuppressWarnings("unused")
@@ -538,7 +539,7 @@ private void link(int level, int n0, int n1, float score, FixedBitSet notFullyCo
     // must subtract 1 here since the nodes array is one larger than the configured
     // max neighbors (M / 2M).
     // We should have taken care of this check by searching for not-full nodes
-    int maxConn = nbr0.nodes().length - 1;
+    int maxConn = nbr0.maxSize() - 1;
     assert notFullyConnected.get(n0);
     assert nbr0.size() < maxConn : "node " + n0 + " is full, has " + nbr0.size() + " friends";
     nbr0.addOutOfOrder(n1, score);