Add merge estimation to OneMerge on ElasticsearchConcurrentMergeScheduler

Author: carlosdelest

server/src/main/java/org/elasticsearch/index/engine/ElasticsearchConcurrentMergeScheduler.java

@@ -32,7 +32,7 @@
  * An extension to the {@link ConcurrentMergeScheduler} that provides tracking on merge times, total
  * and current merges.
  */
-public class ElasticsearchConcurrentMergeScheduler extends ConcurrentMergeScheduler implements ElasticsearchMergeScheduler {
+public abstract class ElasticsearchConcurrentMergeScheduler extends ConcurrentMergeScheduler implements ElasticsearchMergeScheduler {
 
     protected final Logger logger;
     private final Settings indexSettings;
@@ -96,7 +96,7 @@ protected void message(String message) {
     @Override
     protected void doMerge(MergeSource mergeSource, MergePolicy.OneMerge merge) throws IOException {
         long timeNS = System.nanoTime();
-        OnGoingMerge onGoingMerge = new OnGoingMerge(merge);
+        OnGoingMerge onGoingMerge = new OnGoingMerge(merge, estimateMergeMemory(merge));
         mergeTracking.mergeStarted(onGoingMerge);
         try {
             beforeMerge(onGoingMerge);
@@ -107,18 +107,22 @@ protected void doMerge(MergeSource mergeSource, MergePolicy.OneMerge merge) thro
 
             afterMerge(onGoingMerge);
         }
-
     }
 
     /**
      * A callback allowing for custom logic before an actual merge starts.
      */
-    protected void beforeMerge(OnGoingMerge merge) {}
+    protected abstract void beforeMerge(OnGoingMerge merge);
 
     /**
      * A callback allowing for custom logic before an actual merge starts.
      */
-    protected void afterMerge(OnGoingMerge merge) {}
+    protected abstract void afterMerge(OnGoingMerge merge);
+
+    /**
+     * Retrieves an estimation on how much memory is needed for the merge.
+     */
+    protected abstract long estimateMergeMemory(MergePolicy.OneMerge merge);
 
     @Override
     public MergeScheduler clone() {

server/src/main/java/org/elasticsearch/index/engine/InternalEngine.java

@@ -2847,6 +2847,13 @@ public synchronized void beforeMerge(OnGoingMerge merge) {
             }
         }
 
+        @Override
+        protected long estimateMergeMemory(MergePolicy.OneMerge merge) {
+            try (Searcher searcher = acquireSearcher("merge_memory_estimation")) {
+                return SegmentMergeMemoryEstimator.estimateSegmentMemory(merge, searcher.getIndexReader());
+            }
+        }
+
         @Override
         public synchronized void afterMerge(OnGoingMerge merge) {
             int maxNumMerges = getMaxMergeCount();

server/src/main/java/org/elasticsearch/index/engine/SegmentMergeMemoryEstimator.java

@@ -0,0 +1,101 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.engine;
+
+import org.apache.lucene.codecs.KnnVectorsReader;
+import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader;
+import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat;
+import org.apache.lucene.index.CodecReader;
+import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.FilterLeafReader;
+import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.MergePolicy;
+import org.apache.lucene.index.SegmentCommitInfo;
+import org.apache.lucene.index.SegmentReader;
+import org.apache.lucene.index.VectorEncoding;
+import org.elasticsearch.common.lucene.Lucene;
+import org.elasticsearch.common.unit.ByteSizeValue;
+
+import java.util.List;
+
+public class SegmentMergeMemoryEstimator {
+
+    public static long estimateSegmentMemory(MergePolicy.OneMerge merge, IndexReader indexReader) {
+        long memoryNeeded = 0;
+        for (SegmentCommitInfo mergedSegment : merge.segments) {
+            memoryNeeded += estimateSegmentMemory(mergedSegment.info.name, indexReader);
+        }
+        return memoryNeeded;
+    }
+
+    private static long estimateSegmentMemory(String segmentName, IndexReader indexReader) {
+        List<LeafReaderContext> leaves = indexReader.leaves();
+        for (LeafReaderContext leafReaderContext : leaves) {
+            SegmentReader segmentReader = Lucene.segmentReader(leafReaderContext.reader());
+            if (segmentReader.getSegmentName().equals(segmentName)) {
+                return estimateSegmentMemory(segmentReader);
+            }
+        }
+
+        throw new IllegalArgumentException("Segment not found: " + segmentName);
+    }
+
+    private static long estimateSegmentMemory(SegmentReader reader) {
+        long maxMem = 0;
+        for (FieldInfo fieldInfo : reader.getFieldInfos()) {
+            maxMem = Math.max(maxMem, estimateFieldMemory(fieldInfo, reader).getBytes());
+        }
+        return maxMem;
+    }
+
+    private static ByteSizeValue estimateFieldMemory(FieldInfo fieldInfo, SegmentReader segmentReader) {
+
+        long maxMem = 0;
+        if (fieldInfo.hasVectorValues()) {
+            maxMem = Math.max(maxMem, estimateVectorFieldMemory(fieldInfo, segmentReader));
+        }
+        // TODO Work on estimations on other field infos when / if needed
+
+        return ByteSizeValue.ofBytes(maxMem);
+    }
+
+    private static long estimateVectorFieldMemory(FieldInfo fieldInfo, SegmentReader segmentReader) {
+        long maxMem = 0;
+        for (LeafReaderContext ctx : segmentReader.leaves()) {
+            CodecReader codecReader = (CodecReader) FilterLeafReader.unwrap(ctx.reader());
+            KnnVectorsReader vectorsReader = codecReader.getVectorReader();
+            if (vectorsReader instanceof PerFieldKnnVectorsFormat.FieldsReader perFieldKnnVectorsFormat) {
+                vectorsReader = perFieldKnnVectorsFormat.getFieldReader(fieldInfo.getName());
+            }
+
+            final long estimation = getVectorFieldEstimation(fieldInfo, segmentReader, vectorsReader);
+            maxMem = Math.max(maxMem, estimation);
+        }
+        return maxMem;
+    }
+
+    private static long getVectorFieldEstimation(FieldInfo fieldInfo, SegmentReader segmentReader, KnnVectorsReader vectorsReader) {
+        int numDocs = segmentReader.numDocs();
+        if (vectorsReader instanceof Lucene99HnswVectorsReader) {
+            // Determined empirically from graph usage on merges, as it's complicated to estimate graph levels and size for non-zero levels
+            return numDocs * 348L;
+
+        } else {
+            // Dominated by the heap byte buffer size used to write each vector
+            if (fieldInfo.getVectorEncoding() == VectorEncoding.FLOAT32) {
+                return fieldInfo.getVectorDimension() * VectorEncoding.FLOAT32.byteSize;
+            }
+            // Byte does not use buffering for writing but the IndexOutput directly
+            return 0;
+        }
+    }
+
+}

server/src/main/java/org/elasticsearch/index/merge/OnGoingMerge.java

@@ -21,11 +21,13 @@ public class OnGoingMerge {
 
     private final String id;
     private final MergePolicy.OneMerge oneMerge;
+    private final long memoryBytesNeeded;
 
-    public OnGoingMerge(MergePolicy.OneMerge merge) {
+    public OnGoingMerge(MergePolicy.OneMerge merge, long memoryBytesNeeded) {
         this.id = Integer.toString(System.identityHashCode(merge));
         this.oneMerge = merge;
 
+        this.memoryBytesNeeded = memoryBytesNeeded;
     }
 
     /**

server/src/test/java/org/elasticsearch/index/engine/MemoryEstimationMergeTests.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.engine;
+
+import org.apache.lucene.document.Document;
+import org.apache.lucene.document.Field;
+import org.apache.lucene.document.KnnFloatVectorField;
+import org.apache.lucene.document.NumericDocValuesField;
+import org.apache.lucene.document.StringField;
+import org.apache.lucene.index.IndexWriter;
+import org.apache.lucene.index.IndexWriterConfig;
+import org.apache.lucene.store.Directory;
+import org.elasticsearch.test.ESTestCase;
+
+import java.io.IOException;
+
+public class MemoryEstimationMergeTests extends ESTestCase {
+
+    public void testMerge() throws IOException {
+        try (Directory dir = newDirectory()) {
+            IndexWriterConfig iwc = newIndexWriterConfig();
+            int numDocs = 1000;
+            int vectorDims = 100;
+
+            try (IndexWriter writer = new IndexWriter(dir, iwc)) {
+                for (int i = 0; i < numDocs; i++) {
+                    if (rarely()) {
+                        writer.flush();
+                    }
+                    if (rarely()) {
+                        writer.forceMerge(1, false);
+                    }
+                    Document doc = new Document();
+                    doc.add(new StringField("id", "" + i, Field.Store.NO));
+                    doc.add(newTextField("text", "the quick brown fox", Field.Store.YES));
+                    doc.add(new NumericDocValuesField("sort", i));
+                    doc.add(new KnnFloatVectorField("floatVector", floatVector(vectorDims)));
+                    writer.addDocument(doc);
+                    if (i == numDocs / 2) {
+                        writer.flush();
+                    }
+                }
+
+                writer.forceMerge(1);
+
+            }
+
+        }
+    }
+
+    private float[] floatVector(int vectorDims) {
+        return new float[vectorDims];
+    }
+}