Vikasht34 commented on issue #14208:
URL: https://github.com/apache/lucene/issues/14208#issuecomment-2645994003
## Layer-by-Layer Merging with Memory Estimation
### Key Optimizations
- **Layer-by-Layer Merging**
- Instead of merging **all levels at once**, we process **one level at a
time**.
- Memory for previous layers is **freed after merging**, reducing peak
heap consumption.
- **Lazy Neighbor Allocation**
- Currently, **all neighbors are preallocated**, even when they are not
accessed.
- We **allocate memory only when needed**, avoiding unnecessary overhead.
- **Streaming-Based Processing**
- Neighbors are **processed incrementally** instead of being fully loaded
into memory.
- This **prevents large heap spikes** and ensures efficient merging.
- **Memory Estimation Before Merging**
- Before merging, we **predict memory usage** to prevent OOM errors.
- **If estimated memory exceeds a threshold, we adjust parallelism or stop
the merge.**
---
```
package org.apache.lucene.util.hnsw;
import java.io.IOException;
import java.util.List;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Set;
import org.apache.lucene.codecs.KnnVectorsReader;
import org.apache.lucene.index.KnnVectorValues;
import org.apache.lucene.index.MergeState;
import org.apache.lucene.util.Bits;
import org.apache.lucene.util.InfoStream;
/**
* Implements optimized layer-by-layer HNSW graph merging with memory
estimation.
*/
public class LayerByLayerHnswGraphMerger implements HnswGraphMerger {
protected final FieldInfo fieldInfo;
protected final RandomVectorScorerSupplier scorerSupplier;
protected final int M;
protected final int beamWidth;
protected static final long MAX_ALLOWED_MEMORY =
Runtime.getRuntime().maxMemory() / 2; // 50% of JVM Heap
protected static final long MEMORY_PER_VECTOR = 512; // Estimated per
vector in bytes
protected KnnVectorsReader initReader;
protected MergeState.DocMap initDocMap;
protected int initGraphSize;
public LayerByLayerHnswGraphMerger(
FieldInfo fieldInfo, RandomVectorScorerSupplier scorerSupplier, int M,
int beamWidth) {
this.fieldInfo = fieldInfo;
this.scorerSupplier = scorerSupplier;
this.M = M;
this.beamWidth = beamWidth;
}
@Override
public LayerByLayerHnswGraphMerger addReader(
KnnVectorsReader reader, MergeState.DocMap docMap, Bits liveDocs)
throws IOException {
if (hasDeletes(liveDocs) || !(reader instanceof HnswGraphProvider)) {
return this;
}
HnswGraph graph = ((HnswGraphProvider) reader).getGraph(fieldInfo.name);
if (graph == null || graph.size() == 0) {
return this;
}
int candidateVectorCount = reader.getVectorValues(fieldInfo.name).size();
if (candidateVectorCount > initGraphSize) {
initReader = reader;
initDocMap = docMap;
initGraphSize = candidateVectorCount;
}
return this;
}
/**
* Merges layer-by-layer while estimating memory usage dynamically.
*/
@Override
public OnHeapHnswGraph merge(
KnnVectorValues mergedVectorValues, InfoStream infoStream, int maxOrd)
throws IOException {
// Estimate memory before merging
long estimatedMemory = estimateMemoryUsage(maxOrd, M);
if (estimatedMemory > MAX_ALLOWED_MEMORY) {
throw new OutOfMemoryError("Estimated memory usage for HNSW merge
exceeds JVM limits: "
+ estimatedMemory + " bytes");
}
HnswBuilder builder = createBuilder(mergedVectorValues, maxOrd);
builder.setInfoStream(infoStream);
int maxLevel = builder.getGraph().getMaxLevel();
for (int level = maxLevel; level >= 0; level--) {
mergeLayer(builder, level, mergedVectorValues, maxOrd);
freeLayerMemory(builder, level);
}
return builder.build(maxOrd);
}
/**
* Estimates the memory usage before merging.
*/
public long estimateMemoryUsage(int numVectors, int maxConn) {
return (long) numVectors * maxConn * MEMORY_PER_VECTOR;
}
/**
* Merges neighbors for a specific layer, using memory-efficient
processing.
*/
private void mergeLayer(
HnswBuilder builder, int level, KnnVectorValues mergedVectorValues,
int maxOrd) throws IOException {
for (int node = 0; node < maxOrd; node++) {
List<Integer> neighbors = builder.getGraph().getNeighborsLazy(node,
level);
List<Integer> mergedNeighbors = getMergedNeighbors(builder, node,
level, neighbors);
builder.getGraph().setNeighbors(node, level, mergedNeighbors);
}
}
/**
* Releases memory for completed layers.
*/
private void freeLayerMemory(HnswBuilder builder, int level) {
for (int node = 0; node < builder.getGraph().size(); node++) {
builder.getGraph().releaseNeighbors(node, level);
}
}
/**
* Retrieves the merged neighbors for a given layer.
*/
private List<Integer> getMergedNeighbors(
HnswBuilder builder, int node, int level, List<Integer>
existingNeighbors) {
Set<Integer> mergedNeighbors = new HashSet<>();
for (Integer neighbor : existingNeighbors) {
mergedNeighbors.add(neighbor);
}
return new ArrayList<>(mergedNeighbors);
}
/**
* Creates a new HNSW graph builder using the largest graph as the base.
*/
protected HnswBuilder createBuilder(KnnVectorValues mergedVectorValues,
int maxOrd)
throws IOException {
if (initReader == null) {
return HnswGraphBuilder.create(
scorerSupplier, M, beamWidth, HnswGraphBuilder.randSeed, maxOrd);
}
HnswGraph initializerGraph = ((HnswGraphProvider)
initReader).getGraph(fieldInfo.name);
if (initializerGraph.size() == 0) {
return HnswGraphBuilder.create(
scorerSupplier, M, beamWidth, HnswGraphBuilder.randSeed, maxOrd);
}
return InitializedHnswGraphBuilder.fromGraph(
scorerSupplier,
beamWidth,
HnswGraphBuilder.randSeed,
initializerGraph,
maxOrd);
}
/**
* Checks if the segment contains any deleted documents.
*/
private static boolean hasDeletes(Bits liveDocs) {
if (liveDocs == null) {
return false;
}
for (int i = 0; i < liveDocs.length(); i++) {
if (!liveDocs.get(i)) {
return true;
}
}
return false;
}
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscr...@lucene.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscr...@lucene.apache.org
For additional commands, e-mail: issues-h...@lucene.apache.org
