Revert ramBytesUsage

Author: Stefan Vodita

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

@@ -21,9 +21,7 @@
 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
@@ -33,7 +31,7 @@
  *
  * @lucene.internal
  */
-public class NeighborArray implements Accountable {
+public class NeighborArray {
   private static final int INITIAL_CAPACITY = 10;
   private final boolean scoresDescOrder;
   private final int maxSize;
@@ -211,12 +209,4 @@ 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,15 +275,21 @@ 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 * 5;
+    long neighborArrayBytes =
+        (long) nsize * (Integer.BYTES + Float.BYTES)
+            + RamUsageEstimator.NUM_BYTES_ARRAY_HEADER * 2L
+            + RamUsageEstimator.NUM_BYTES_OBJECT_REF * 2L
+            + Integer.BYTES * 5;
     long total = 0;
-    for (NeighborArray[] neighborArraysPerNode : graph) {
-      if (neighborArraysPerNode != null) {
-        for (NeighborArray neighborArrayPerNodeAndLevel : neighborArraysPerNode) {
-          total += neighborArrayPerNodeAndLevel.ramBytesUsed();
-        }
-      }
-    }
-
+    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
     total += 4 * Integer.BYTES; // all int fields
     total += 1; // field: noGrowth
     total +=

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

@@ -757,6 +757,30 @@ public void testRamUsageEstimate() throws IOException {
     long estimated = RamUsageEstimator.sizeOfObject(hnsw);
     long actual = ramUsed(hnsw);
 
+    // The estimation assumes neighbor arrays are always max size.
+    // When this is not true, the estimate can be much larger than the actual value.
+    // In these cases, we compute how much we overestimated the neighbors arrays.
+    if (estimated > actual) {
+      long neighborsError = 0;
+      int numLevels = hnsw.numLevels();
+      for (int level = 0; level < numLevels; ++level) {
+        NodesIterator nodesOnLevel = hnsw.getNodesOnLevel(level);
+        while (nodesOnLevel.hasNext()) {
+          int node = nodesOnLevel.nextInt();
+          NeighborArray neighbors = hnsw.getNeighbors(level, node);
+          long maxNeighborsSize;
+          if (level == 0) {
+            maxNeighborsSize = (long) (Integer.BYTES + Float.BYTES) * (M * 2L + 1);
+          } else {
+            maxNeighborsSize = (long) (Integer.BYTES + Float.BYTES) * (M + 1);
+          }
+          long actualNeighborsSize = (long) (Integer.BYTES + Float.BYTES) * neighbors.size();
+          neighborsError += maxNeighborsSize - actualNeighborsSize;
+        }
+      }
+      estimated -= neighborsError;
+    }
+
     assertEquals((double) actual, (double) estimated, (double) actual * 0.3);
   }
 

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

@@ -17,12 +17,9 @@
 
 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 {
 
@@ -244,21 +241,4 @@ 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);
-    }
-  }
 }