Introduce growInRange to reduce array overallocation

In cases where we know there is an upper limit to the potential size
of an array, we can use `growInRange` to avoid allocating beyond that
limit.

We address such cases in `DirectoryTaxonomyReader` and `NeighborArray`.

Author: stefanvodita

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

@@ -330,6 +330,29 @@ public static int[] growExact(int[] array, int newLength) {
     return copy;
   }
 
+  /**
+   * Returns an array whose size is at least {@code minLength}, generally over-allocating
+   * exponentially, but never allocating more than {@code maxLength} elements.
+   */
+  public static int[] growInRange(int[] array, int minLength, int maxLength) {
+    if (array.length >= minLength) {
+      return array;
+    }
+    if (minLength > maxLength) {
+      throw new IllegalArgumentException(
+          "requested minimum array length "
+              + minLength
+              + " is larger than requested maximum array length "
+              + maxLength);
+    }
+
+    int potentialLength = oversize(minLength, Integer.BYTES);
+    if (potentialLength > maxLength) {
+      return growExact(array, maxLength);
+    }
+    return growExact(array, potentialLength);
+  }
+
   /**
    * Returns an array whose size is at least {@code minSize}, generally over-allocating
    * exponentially

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

@@ -21,7 +21,9 @@
 import java.util.Arrays;
 import java.util.concurrent.locks.ReadWriteLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
+import org.apache.lucene.util.Accountable;
 import org.apache.lucene.util.ArrayUtil;
+import org.apache.lucene.util.RamUsageEstimator;
 
 /**
  * NeighborArray encodes the neighbors of a node and their mutual scores in the HNSW graph as a pair
@@ -31,18 +33,21 @@
  *
  * @lucene.internal
  */
-public class NeighborArray {
+public class NeighborArray implements Accountable {
+  private static final int INITIAL_CAPACITY = 10;
   private final boolean scoresDescOrder;
+  private final int maxSize;
   private int size;
   float[] score;
   int[] node;
   private int sortedNodeSize;
   public final ReadWriteLock rwlock = new ReentrantReadWriteLock(true);
 
   public NeighborArray(int maxSize, boolean descOrder) {
-    node = new int[maxSize];
-    score = new float[maxSize];
+    node = new int[INITIAL_CAPACITY];
+    score = new float[INITIAL_CAPACITY];
     this.scoresDescOrder = descOrder;
+    this.maxSize = maxSize;
   }
 
   /**
@@ -52,7 +57,7 @@ public NeighborArray(int maxSize, boolean descOrder) {
   public void addInOrder(int newNode, float newScore) {
     assert size == sortedNodeSize : "cannot call addInOrder after addOutOfOrder";
     if (size == node.length) {
-      node = ArrayUtil.grow(node);
+      node = ArrayUtil.growInRange(node, size + 1, maxSize);
       score = ArrayUtil.growExact(score, node.length);
     }
     if (size > 0) {
@@ -73,7 +78,7 @@ public void addInOrder(int newNode, float newScore) {
   /** Add node and newScore but do not insert as sorted */
   public void addOutOfOrder(int newNode, float newScore) {
     if (size == node.length) {
-      node = ArrayUtil.grow(node);
+      node = ArrayUtil.growInRange(node, size + 1, maxSize);
       score = ArrayUtil.growExact(score, node.length);
     }
 
@@ -206,4 +211,12 @@ private int descSortFindRightMostInsertionPoint(float newScore, int bound) {
     }
     return start;
   }
+
+  @Override
+  public long ramBytesUsed() {
+    return (long) node.length * (Integer.BYTES + Float.BYTES)
+        + RamUsageEstimator.NUM_BYTES_ARRAY_HEADER * 2L
+        + RamUsageEstimator.NUM_BYTES_OBJECT_REF * 2L
+        + Integer.BYTES * 5;
+  }
 }

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

@@ -275,21 +275,15 @@ private void generateLevelToNodes() {
 
   @Override
   public long ramBytesUsed() {
-    long neighborArrayBytes0 =
-        (long) nsize0 * (Integer.BYTES + Float.BYTES)
-            + RamUsageEstimator.NUM_BYTES_ARRAY_HEADER * 2L
-            + RamUsageEstimator.NUM_BYTES_OBJECT_REF * 2L
-            + Integer.BYTES * 3;
-    long neighborArrayBytes =
-        (long) nsize * (Integer.BYTES + Float.BYTES)
-            + RamUsageEstimator.NUM_BYTES_ARRAY_HEADER * 2L
-            + RamUsageEstimator.NUM_BYTES_OBJECT_REF * 2L
-            + Integer.BYTES * 3;
     long total = 0;
-    total +=
-        size() * (neighborArrayBytes0 + RamUsageEstimator.NUM_BYTES_ARRAY_HEADER)
-            + RamUsageEstimator.NUM_BYTES_ARRAY_HEADER; // for graph and level 0;
-    total += nonZeroLevelSize.get() * neighborArrayBytes; // for non-zero level
+    for (NeighborArray[] neighborArraysPerNode : graph) {
+      if (neighborArraysPerNode != null) {
+        for (NeighborArray neighborArrayPerNodeAndLevel : neighborArraysPerNode) {
+          total += neighborArrayPerNodeAndLevel.ramBytesUsed();
+        }
+      }
+    }
+
     total += 4 * Integer.BYTES; // all int fields
     total += 1; // field: noGrowth
     total +=

lucene/core/src/test/org/apache/lucene/util/TestArrayUtil.java

@@ -18,6 +18,8 @@
 
 import static org.apache.lucene.util.ArrayUtil.copyOfSubArray;
 import static org.apache.lucene.util.ArrayUtil.growExact;
+import static org.apache.lucene.util.ArrayUtil.growInRange;
+import static org.apache.lucene.util.ArrayUtil.oversize;
 
 import java.util.Arrays;
 import java.util.Collections;
@@ -371,6 +373,33 @@ public void testGrowExact() {
         () -> growExact(new String[] {"a", "b", "c"}, random().nextInt(3)));
   }
 
+  public void testGrowInRange() {
+    int[] array = new int[] {1, 2, 3};
+
+    // maxLength does not matter as long as minLength is enough
+    assertSame(array, growInRange(array, 1, 4));
+    assertSame(array, growInRange(array, 1, 4));
+    assertSame(array, growInRange(array, 1, 2));
+    assertSame(array, growInRange(array, 1, 1));
+    assertSame(array, growInRange(array, 1, 0));
+    assertSame(array, growInRange(array, 1, -1));
+
+    // maxLength < minLength if the array has to be grown throws an exception
+    expectThrows(IllegalArgumentException.class, () -> growInRange(array, 4, 3));
+    expectThrows(IllegalArgumentException.class, () -> growInRange(array, 5, 4));
+
+    int minLength = 4;
+    int maxLength = Integer.MAX_VALUE;
+
+    // The array grows normally if maxLength permits
+    assertEquals(
+        oversize(minLength, Integer.BYTES),
+        growInRange(new int[] {1, 2, 3}, minLength, maxLength).length);
+
+    // The array grows to maxLength if maxLength is limiting
+    assertEquals(minLength, growInRange(new int[] {1, 2, 3}, minLength, minLength).length);
+  }
+
   public void testCopyOfSubArray() {
     short[] shortArray = {1, 2, 3};
     assertArrayEquals(new short[] {1}, copyOfSubArray(shortArray, 0, 1));

lucene/core/src/test/org/apache/lucene/util/hnsw/TestNeighborArray.java

@@ -17,9 +17,12 @@
 
 package org.apache.lucene.util.hnsw;
 
+import static org.apache.lucene.tests.util.RamUsageTester.ramUsed;
+
 import java.io.IOException;
 import org.apache.lucene.tests.util.LuceneTestCase;
 import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.RamUsageEstimator;
 
 public class TestNeighborArray extends LuceneTestCase {
 
@@ -241,4 +244,21 @@ default Bits getAcceptOrds(Bits acceptDocs) {
       throw new UnsupportedOperationException();
     }
   }
+
+  public void testRamUsageEstimate() {
+    int maxSize = random().nextInt(100, 10_000);
+    int numAdditions = random().nextInt(100, 1_000);
+    NeighborArray neighbors = new NeighborArray(maxSize, true);
+
+    int count = 0;
+    while (count < maxSize) {
+      for (int i = 0; i < numAdditions && count < maxSize; i++) {
+        neighbors.addInOrder(count, 0f);
+        count++;
+      }
+      long estimated = RamUsageEstimator.sizeOfObject(neighbors);
+      long actual = ramUsed(neighbors);
+      assertEquals((double) actual, (double) estimated, (double) actual * 0.3);
+    }
+  }
 }

lucene/facet/src/java/org/apache/lucene/facet/taxonomy/directory/DirectoryTaxonomyReader.java

@@ -351,7 +351,8 @@ public int[] getBulkOrdinals(FacetLabel... categoryPaths) throws IOException {
         }
       } else {
         indexesMissingFromCache =
-            ArrayUtil.grow(indexesMissingFromCache, numberOfMissingFromCache + 1);
+            ArrayUtil.growInRange(
+                indexesMissingFromCache, numberOfMissingFromCache + 1, categoryPaths.length);
         indexesMissingFromCache[numberOfMissingFromCache++] = i;
       }
     }