update OnHeapHnswGraph ramBytesUsed method

Author: weizijun

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

@@ -212,6 +212,9 @@ private void addVectors(int maxOrd) throws IOException {
 
   public void addGraphNode(int node, UpdateableRandomVectorScorer scorer) throws IOException {
     addGraphNodeInternal(node, scorer, null);
+    if (node % 1000 == 0) {
+      hnsw.updateRamBytesUsedEstimate();
+    }
   }
 
   private void addGraphNodeInternal(int node, UpdateableRandomVectorScorer scorer, IntHashSet eps0)
@@ -444,6 +447,7 @@ void finish() throws IOException {
     //  see: https://github.com/apache/lucene/issues/14214
     // connectComponents();
     frozen = true;
+    hnsw.updateRamBytesUsedEstimate();
   }
 
   @SuppressWarnings("unused")

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

@@ -21,6 +21,8 @@
 import java.util.Arrays;
 import org.apache.lucene.internal.hppc.FloatArrayList;
 import org.apache.lucene.internal.hppc.IntArrayList;
+import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.RamUsageEstimator;
 
 /**
  * NeighborArray encodes the neighbors of a node and their mutual scores in the HNSW graph as a pair
@@ -30,7 +32,10 @@
  *
  * @lucene.internal
  */
-public class NeighborArray {
+public class NeighborArray implements Accountable {
+  private static final long BASE_RAM_BYTES_USED =
+      RamUsageEstimator.shallowSizeOfInstance(NeighborArray.class);
+
   private final boolean scoresDescOrder;
   private int size;
   private final int maxSize;
@@ -308,4 +313,9 @@ private boolean isWorstNonDiverse(
   public int maxSize() {
     return maxSize;
   }
+
+  @Override
+  public long ramBytesUsed() {
+    return BASE_RAM_BYTES_USED + nodes.ramBytesUsed() + scores.ramBytesUsed();
+  }
 }

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

@@ -63,6 +63,8 @@ public final class OnHeapHnswGraph extends HnswGraph implements Accountable {
   private int upto;
   private NeighborArray cur;
 
+  private volatile long graphRamBytesUsed;
+
   /**
    * ctor
    *
@@ -292,23 +294,28 @@ private void generateLevelToNodes() {
     lastFreezeSize = size();
   }
 
+  public void updateRamBytesUsedEstimate() {
+    long currentRamBytesUsedEstimate = RamUsageEstimator.NUM_BYTES_ARRAY_HEADER;
+    for (int node = 0; node < graph.length; node++) {
+      if (graph[node] == null) {
+        continue;
+      }
+
+      for (int i = 0; i < graph[node].length; i++) {
+        if (graph[node][i] == null) {
+          continue;
+        }
+        currentRamBytesUsedEstimate += graph[node][i].ramBytesUsed();
+      }
+
+      currentRamBytesUsedEstimate += RamUsageEstimator.NUM_BYTES_OBJECT_HEADER;
+    }
+    graphRamBytesUsed = currentRamBytesUsedEstimate;
+  }
+
   @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
+    long total = graphRamBytesUsed; // all NeighborArray
     total += 4 * Integer.BYTES; // all int fields
     total += 1; // field: noGrowth
     total +=