diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt index f5c2ee49b5a4..cfcf28e719b7 100644 --- a/lucene/CHANGES.txt +++ b/lucene/CHANGES.txt @@ -120,7 +120,7 @@ Optimizations * GITHUB#14447: Compute the doc range more efficiently when flushing doc block. (Pan Guixin) -* GITHUB#14527: Reduce NeighborArray heap memory. (weizijun) +* GITHUB#14527, GITHUB#14765: Reduce NeighborArray heap memory. (weizijun, Ben Trent) * GITHUB#14529, GITHUB#14555, GITHUB#14618: Impl intoBitset for IndexedDISI and Docvalues. (Guo Feng) diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java index 6c58f3929ba1..5ded543e48a5 100644 --- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java +++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java @@ -444,7 +444,6 @@ void finish() throws IOException { // see: https://github.com/apache/lucene/issues/14214 // connectComponents(); frozen = true; - hnsw.finishBuild(); } @SuppressWarnings("unused") diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/NeighborArray.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/NeighborArray.java index 07d96d353db1..c43cd681440a 100644 --- a/lucene/core/src/java/org/apache/lucene/util/hnsw/NeighborArray.java +++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/NeighborArray.java @@ -19,9 +19,9 @@ import java.io.IOException; import java.util.Arrays; +import java.util.function.LongConsumer; import org.apache.lucene.internal.hppc.MaxSizedFloatArrayList; import org.apache.lucene.internal.hppc.MaxSizedIntArrayList; -import org.apache.lucene.util.Accountable; import org.apache.lucene.util.RamUsageEstimator; /** @@ -32,7 +32,7 @@ * * @lucene.internal */ -public class NeighborArray implements Accountable { +public class NeighborArray { private static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(NeighborArray.class); @@ -42,12 +42,23 @@ public class NeighborArray implements Accountable { private final MaxSizedFloatArrayList scores; private final MaxSizedIntArrayList nodes; private int sortedNodeSize; + private long ramBytesUsed = BASE_RAM_BYTES_USED; + private final LongConsumer onHeapMemoryUsageListener; public NeighborArray(int maxSize, boolean descOrder) { + this(maxSize, descOrder, null); + } + + public NeighborArray(int maxSize, boolean descOrder, LongConsumer onHeapMemoryUsageListener) { this.maxSize = maxSize; nodes = new MaxSizedIntArrayList(maxSize, maxSize / 8); scores = new MaxSizedFloatArrayList(maxSize, maxSize / 8); + this.ramBytesUsed += nodes.ramBytesUsed() + scores.ramBytesUsed(); this.scoresDescOrder = descOrder; + this.onHeapMemoryUsageListener = onHeapMemoryUsageListener; + if (onHeapMemoryUsageListener != null) { + onHeapMemoryUsageListener.accept(ramBytesUsed); + } } /** @@ -68,8 +79,10 @@ public void addInOrder(int newNode, float newScore) { + " to " + Arrays.toString(scores.toArray()); } + int previousLength = nodes.buffer.length; nodes.add(newNode); scores.add(newScore); + alertOnHeapMemoryUsageChange(nodes.buffer.length, previousLength); ++size; ++sortedNodeSize; } @@ -79,12 +92,21 @@ public void addOutOfOrder(int newNode, float newScore) { if (size == maxSize) { throw new IllegalStateException("No growth is allowed"); } - + int previousLength = nodes.buffer.length; nodes.add(newNode); scores.add(newScore); + alertOnHeapMemoryUsageChange(nodes.buffer.length, previousLength); size++; } + private void alertOnHeapMemoryUsageChange(int newLength, int previousLength) { + if (newLength > previousLength && onHeapMemoryUsageListener != null) { + int lengthDelta = newLength - previousLength; + onHeapMemoryUsageListener.accept( + (long) (lengthDelta) * Integer.BYTES + (long) (lengthDelta) * Float.BYTES); + } + } + /** * In addition to {@link #addOutOfOrder(int, float)}, this function will also remove the * least-diverse node if the node array is full after insertion @@ -313,9 +335,4 @@ private boolean isWorstNonDiverse( public int maxSize() { return maxSize; } - - @Override - public long ramBytesUsed() { - return BASE_RAM_BYTES_USED + nodes.ramBytesUsed() + scores.ramBytesUsed(); - } } diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/OnHeapHnswGraph.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/OnHeapHnswGraph.java index ae2ce7d594f3..da0dd1641b34 100644 --- a/lucene/core/src/java/org/apache/lucene/util/hnsw/OnHeapHnswGraph.java +++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/OnHeapHnswGraph.java @@ -32,6 +32,10 @@ */ public final class OnHeapHnswGraph extends HnswGraph implements Accountable { + // shallow estimate of the statically used on-heap memory. + private static final long RAM_BYTES_USED = + RamUsageEstimator.shallowSizeOfInstance(OnHeapHnswGraph.class); + private static final int INIT_SIZE = 128; private final AtomicReference entryNode; @@ -83,6 +87,7 @@ public final class OnHeapHnswGraph extends HnswGraph implements Accountable { numNodes = INIT_SIZE; } this.graph = new NeighborArray[numNodes][]; + this.graphRamBytesUsed = RAM_BYTES_USED + RamUsageEstimator.shallowSizeOf(graph); } /** @@ -158,21 +163,28 @@ public void addNode(int level, int node) { size.incrementAndGet(); } if (level == 0) { - graph[node][level] = new NeighborArray(nsize0, true); + graph[node][level] = + new NeighborArray( + nsize0, + true, + l -> { + assert l > 0; + long bytesUsed = graphRamBytesUsed; + graphRamBytesUsed = bytesUsed + l; + }); } else { - graph[node][level] = new NeighborArray(nsize, true); + graph[node][level] = + new NeighborArray( + nsize, + true, + l -> { + assert l > 0; + long bytesUsed = graphRamBytesUsed; + graphRamBytesUsed = bytesUsed + l; + }); nonZeroLevelSize.incrementAndGet(); } maxNodeId.accumulateAndGet(node, Math::max); - // update graphRamBytesUsed every 1000 nodes - if (level == 0 && node % 1000 == 0) { - updateGraphRamBytesUsed(); - } - } - - /** Finish building the graph. */ - public void finishBuild() { - updateGraphRamBytesUsed(); } @Override @@ -303,48 +315,14 @@ private void generateLevelToNodes() { lastFreezeSize = size(); } - /** Update the estimated ram bytes used for the neighbor array. */ - public void updateGraphRamBytesUsed() { - 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; - } - + /** + * Provides an estimate of the current on-heap memory usage of the graph. This is not threadsafe, + * meaning the heap utilization if building the graph concurrently may be inaccurate. The main + * purpose of this method is during initial document indexing and flush. + */ @Override public long ramBytesUsed() { - long total = graphRamBytesUsed; // all NeighborArray - total += 4 * Integer.BYTES; // all int fields - total += 1; // field: noGrowth - total += - RamUsageEstimator.NUM_BYTES_OBJECT_REF - + RamUsageEstimator.NUM_BYTES_OBJECT_HEADER - + 2 * Integer.BYTES; // field: entryNode - total += 3L * (Integer.BYTES + RamUsageEstimator.NUM_BYTES_OBJECT_HEADER); // 3 AtomicInteger - total += RamUsageEstimator.NUM_BYTES_OBJECT_REF; // field: cur - total += RamUsageEstimator.NUM_BYTES_ARRAY_HEADER; // field: levelToNodes - if (levelToNodes != null) { - total += - (long) (numLevels() - 1) * RamUsageEstimator.NUM_BYTES_OBJECT_REF; // no cost for level 0 - total += - (long) nonZeroLevelSize.get() - * (RamUsageEstimator.NUM_BYTES_OBJECT_HEADER - + RamUsageEstimator.NUM_BYTES_OBJECT_HEADER - + Integer.BYTES); - } - return total; + return graphRamBytesUsed; } @Override