MemoryOptimizedSearch warm up optimization #2954

MrFlap · 2025-10-21T21:09:53Z

Description

Adds warmup procedures for memory optimized search.

Related Issues

Resolves #2939

Check List

New functionality includes testing.
New functionality has been documented.
API changes companion pull request created.
Commits are signed per the DCO using --signoff.
Public documentation issue/PR created.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

0ctopus13prime · 2025-10-22T01:53:38Z

Hi @MrFlap
Thank you for PR, overall looks good, though I can see small refactoring + few nitpicks but we can revisit later after benchmark.

Before proceeding, could you verify how much this warm-up effective?
Ideally, if all bytes loaded into page cache, search latency right after this warm-up should be consistent. So gap between p50 and p99 should be small enough.

Could you do 2 experiments and share numbers?

FP16, Cohere 10M -> For Faiss index only
Quantized vectors, Cohere 10M -> Both Faiss + Lucene .vec file

At the beginning of each experiment, please run below to drop cache:

# Show current memory cache usage
free -h

# Drop page cache + dentries + inodes (full cache flush)
sudo sync && sudo echo 3 | sudo tee /proc/sys/vm/drop_caches

shatejas · 2025-10-22T18:11:13Z

src/main/java/org/opensearch/knn/index/KNNIndexShard.java

-import org.apache.lucene.index.SegmentInfo;
-import org.apache.lucene.index.SegmentReader;
-import org.apache.lucene.store.Directory;
+import org.apache.lucene.index.*;


remove * imports

shatejas · 2025-10-22T18:11:22Z

src/main/java/org/opensearch/knn/index/KNNIndexShard.java

-            .filter(fieldInfo -> fieldInfo.attributes().containsKey(KNNVectorFieldMapper.KNN_FIELD))
-            .filter(fieldInfo -> {
-                final MappedFieldType fieldType = mapperService.fieldType(fieldInfo.getName());
+        // Check which warmup strategy to use. Currently, this will be partial warmup for Non-FSDirectory and


Do we need partial warmup for Non-FSDirectory? Directory implementations will take care of memory issues right?

I'm not really sure, this is just an abstraction of this code

shatejas · 2025-10-22T18:20:44Z

src/main/java/org/opensearch/knn/index/KNNIndexShard.java

+        return StreamSupport.stream(leafReader.getFieldInfos().spliterator(), false)
+            .filter(fieldInfo -> isMemoryOptimizedSearchField(fieldInfo, mapperService, indexName));


Can we use for-loops please. returning a stream is not preferable as it can throw already closed exception if terminal operation is performed on it.

There is no effeciency gains here with stream api usage so its better to be defensive

shatejas · 2025-10-22T18:25:11Z

src/main/java/org/opensearch/knn/index/KNNIndexShard.java

+            this.directory = directory;
+        }
+
+        private void pageFaultFile(String file) throws IOException {


nit: why is this named pageFaultFile? loadFile or warmupFile should work

warmupFile should be fine, I named it pageFaultFile so that it wouldn't be confused with loading the file into a file pointer.

shatejas · 2025-10-22T18:41:06Z

src/main/java/org/opensearch/knn/index/KNNIndexShard.java

+            for (int i = 0; i < input.length(); i += 4096) {
+                input.seek(i);
+                input.readByte();
+            }


Just to clarify the idea here is that we fetch one byte so it kernel will fetch the 4kb page right?

shatejas · 2025-10-22T18:41:49Z

src/main/java/org/opensearch/knn/index/KNNIndexShard.java

+        }
+
+        private void pageFaultFile(String file) throws IOException {
+            IndexInput input = directory.openInput(file, IOContext.DEFAULT);


Lets use IOContext.READONCE to avoid any side-effects

Also can we wrap this up in try with resources to make sure file is unmapped?

Changing the IOContext should be fine. Can you clarify what the second bullet point means?

Suggested change

IndexInput input = directory.openInput(file, IOContext.DEFAULT);

try (IndexInput input = directory.openInput(file, IOContext.READONCE)) {

// logic here

}

This takes of resource closure in case of error cases. It makes sure file is unmapped in case an error is thrown

Ahh, makes sense 👍

shatejas · 2025-10-22T18:52:17Z

src/main/java/org/opensearch/knn/index/KNNIndexShard.java

+        @Override
+        public boolean warmUp(FieldInfo field) throws IOException {
+            final KNNEngine knnEngine = extractKNNEngine(field);
+            final List<String> engineFiles = KNNCodecUtil.getEngineFiles(


Can't we just fetch .vec files here to avoid multiple code flows?

I was originally going to do that, but there is a subtle difference with this. If we add the .vec file fetch in this method, we will load the .vec .faiss pair for each fieldInfo sequentially. This means that the .vec file might load over a .faiss file loaded in a previous call. If we load all of the .vec files first, then we know that they won't override any .faiss files.

shatejas · 2025-10-22T18:53:50Z

src/main/java/org/opensearch/knn/index/KNNIndexShard.java

+    /**
+     * Fully warm up the index by loading every byte from disk, causing page faults
+     */
+    private class FullFieldWarmUpStrategy implements FieldWarmUpStrategy {


Lets not go with inner classes. Can we create a package warmup and then have separate classes in there. If we are moving to strategy pattern then native cache should also be moved ideally

Sure, are we just trying to avoid inner classes in general?

Generally private inner classes are hard to test. breaking it into its own class with making it package-private will generally achieve the same goal and will help cover more cases

shatejas · 2025-10-22T18:56:53Z

src/test/java/org/opensearch/knn/index/KNNIndexShardTests.java

+    }
+
+    @SneakyThrows
+    public void testWarmUpMemoryOptimizedSearcher_multipleSegments() {


This test is okay but its not tight enough, anyway we can mock and verify the warmup method is called?

You mean to verify that FieldWarmUpStrategy::warmup is called? We can probably mock to use a custom FieldWarmUpStrategy that returns the files that it touches somehow.

Vikasht34

Overall changes makes sense , things which is making really hard to understand is putting eveerything in one class , let's re-factor liitle bit,

Put a Warm Up Class with partial and full load strategy.

These are the test cases I can think of we can put

Test Warm-Up with FSDirectory: Verify full warm-up strategy loads .faiss files for FSDirectory, no off-heap cache used, and logs confirm field warm-up.
Test Warm-Up with Non-FSDirectory: Verify partial warm-up strategy triggers null vector search for non-FSDirectory, no off-heap cache, and logs confirm field warm-up.
Test Warm-Up with Multiple Segments: Verify warm-up handles multiple segments, all eligible fields warmed, no off-heap cache, and logs confirm fields per segment.
Test Full-Precision Vector Loading: Verify .vec files loaded before .faiss, search works post-warm-up, and no off-heap cache used.
Test Warm-Up with Multiple Fields: Verify all k-NN fields in a segment warmed, logs confirm all fields, and returned set includes all field names.
Test Warm-Up with No Eligible Fields: Verify empty field set handled, logs “no fields found,” empty set returned, no exceptions.
Test Warm-Up with Empty Segment: Verify empty segment (no documents) handled, empty set returned, no exceptions.
Test Warm-Up with Missing Engine Files: Verify missing .faiss files logged as warnings, field skipped, no exceptions.
Test Warm-Up with Invalid Vector Data Type: Verify invalid VECTOR_DATA_TYPE_FIELD skipped, logs error/warning, no exceptions.
Test Warm-Up with Closed LeafReader: Verify AlreadyClosedException caught, logged, empty set returned, no crash.
Test IOException in Full Warm-Up: Simulate IOException in warmUpFile, verify error logged, field skipped, warm-up continues.
Test IOException in Vector Loading: Simulate IOException in loadFullPrecisionVectors, verify error logged, warm-up continues.
Test Null MapperService: Verify null MapperService handled, fields skipped, empty set returned, no exceptions.
Test Warm-Up with Large Segment: Verify warm-up scales with 10,000 documents, completes in reasonable time, no memory errors.
Test Warm-Up with Many Fields: Verify warm-up scales with 50 k-NN fields, all warmed, linear performance scaling.
Test Warm-Up Strategy Invocation: Verify correct strategy (FullWarmUpStrategy or PartialWarmUpStrategy) called per field using mock.
Test Strategy Selection Based on Directory: Verify FullWarmUpStrategy for FSDirectory, PartialWarmUpStrategy for non-FSDirectory using mocks.
Test Warm-Up During Shard Recovery: Verify warm-up works during shard recovery, fields loaded, search functional post-recovery.
Test Warm-Up with Concurrent Queries: Verify warm-up doesn’t block k-NN searches, search results correct during/after warm-up.
Test Directory Resource Cleanup: Verify IndexInput closed in warmUpFile, no resource leaks using mock.
Test Searcher Resource Cleanup: Verify Engine.Searcher closed in warmup(), no resource leaks using mock

Vikasht34 · 2025-10-22T22:44:19Z

src/main/java/org/opensearch/knn/index/KNNIndexShard.java

-import java.util.Set;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.*;


nit:- Remove *

Vikasht34 · 2025-10-23T04:44:20Z

src/main/java/org/opensearch/knn/index/KNNIndexShard.java

-import static org.opensearch.knn.common.KNNConstants.SPACE_TYPE;
-import static org.opensearch.knn.common.KNNConstants.VECTOR_DATA_TYPE_FIELD;
+import static org.opensearch.knn.common.FieldInfoExtractor.extractKNNEngine;
+import static org.opensearch.knn.common.KNNConstants.*;


MrFlap · 2025-10-23T21:23:33Z

Whoops, for some reason the import changes didn't go through in KNNIndexShard.java

MrFlap · 2025-10-23T21:38:46Z

I'm going to squash the commits once all tests pass and you all think it looks good @0ctopus13prime @shatejas @Vikasht34, but all feedback has been incorporated.

0ctopus13prime · 2025-10-24T00:34:31Z

src/main/java/org/opensearch/knn/index/warmup/FieldWarmUpStrategy.java

+
+import java.io.IOException;
+
+public abstract class FieldWarmUpStrategy {


NIT : Java doc please!

0ctopus13prime · 2025-10-24T00:34:35Z

src/main/java/org/opensearch/knn/index/warmup/FieldWarmUpStrategyFactory.java

+import org.apache.lucene.store.FilterDirectory;
+import org.opensearch.common.lucene.Lucene;
+
+public class FieldWarmUpStrategyFactory {


NIT : Java doc please!

0ctopus13prime · 2025-10-24T00:34:53Z

src/main/java/org/opensearch/knn/index/warmup/FieldWarmUpStrategyFactory.java

+    private Directory directory;
+    private LeafReader leafReader;
+
+    public FieldWarmUpStrategyFactory setDirectory(Directory directory) {


NIT : @Setter

For some reason it doesn't return the FieldWarmUpStrategyFactory when I use Setter. Is this intended behavior?

0ctopus13prime · 2025-10-24T00:35:59Z

src/main/java/org/opensearch/knn/index/warmup/FullFieldWarmUpStrategy.java

+        try (IndexInput input = directory.openInput(file, IOContext.READONCE)) {
+            for (int i = 0; i < input.length(); i += 4096) {
+                input.seek(i);
+                input.readByte();


Curious, do you recall all pages loaded regardless readByte or readBytes(byte[]) are called?

0ctopus13prime · 2025-10-24T00:38:01Z

src/main/java/org/opensearch/knn/index/warmup/MemoryOptimizedSearchWarmup.java

+        ArrayList<String> warmedUp = new ArrayList<>();
+
+        for (FieldInfo field : memOptSearchFields) {
+            boolean warm;


try { if (fieldWarmUpStrategy.warmUp(field)) { warmedUp.add(field.getName()); } } catch (IOException e) { log.error("Failed to warm up field: {}: {}", field.getName(), e.toString()); }

0ctopus13prime · 2025-10-24T00:40:19Z

src/main/java/org/opensearch/knn/index/warmup/MemoryOptimizedSearchWarmup.java

+                vectorValues.vectorValue(iter.docID());
+            }
+        } catch (IOException e) {
+            log.error("Failed to load vec file for field: {}: {}", field.getName(), e.toString());


Generally it's good to pass e to logger

log.error("Failed to load vec file for field: {}", field.getName(), e);

0ctopus13prime · 2025-10-24T00:44:47Z

@MrFlap
Also could you post the results for FP16 and 32x quantization case respectively?

Vikasht34

Thanks for incorporating all the tests !! , Will wait for all tests to pass.

…ch-project#2954) Signed-off-by: Andrew Klepchick <[email protected]>

MrFlap · 2025-10-24T21:36:02Z

@0ctopus13prime 32x
etric,Task,Value,Unit
Cumulative indexing time of primary shards,,0,min
Min cumulative indexing time across primary shards,,0,min
Median cumulative indexing time across primary shards,,0,min
Max cumulative indexing time across primary shards,,0,min
Cumulative indexing throttle time of primary shards,,0,min
Min cumulative indexing throttle time across primary shards,,0,min
Median cumulative indexing throttle time across primary shards,,0,min
Max cumulative indexing throttle time across primary shards,,0,min
Cumulative merge time of primary shards,,0,min
Cumulative merge count of primary shards,,0,
Min cumulative merge time across primary shards,,0,min
Median cumulative merge time across primary shards,,0,min
Max cumulative merge time across primary shards,,0,min
Cumulative merge throttle time of primary shards,,0,min
Min cumulative merge throttle time across primary shards,,0,min
Median cumulative merge throttle time across primary shards,,0,min
Max cumulative merge throttle time across primary shards,,0,min
Cumulative refresh time of primary shards,,0,min
Cumulative refresh count of primary shards,,14,
Min cumulative refresh time across primary shards,,0,min
Median cumulative refresh time across primary shards,,0,min
Max cumulative refresh time across primary shards,,0,min
Cumulative flush time of primary shards,,0,min
Cumulative flush count of primary shards,,7,
Min cumulative flush time across primary shards,,0,min
Median cumulative flush time across primary shards,,0,min
Max cumulative flush time across primary shards,,0,min
Total Young Gen GC time,,0.027,s
Total Young Gen GC count,,2,
Total Old Gen GC time,,0,s
Total Old Gen GC count,,0,
Store size,,31.08053415827453,GB
Translog size,,3.5855919122695923e-07,GB
Heap used for segments,,0,MB
Heap used for doc values,,0,MB
Heap used for terms,,0,MB
Heap used for norms,,0,MB
Heap used for points,,0,MB
Heap used for stored fields,,0,MB
Segment count,,37,
Min Throughput,prod-queries,287.95,ops/s
Mean Throughput,prod-queries,464.55,ops/s
Median Throughput,prod-queries,482.34,ops/s
Max Throughput,prod-queries,493.91,ops/s
50th percentile latency,prod-queries,18.259749995195307,ms
90th percentile latency,prod-queries,19.341466900368687,ms
99th percentile latency,prod-queries,20.406325711373942,ms
99.9th percentile latency,prod-queries,47.77802453149796,ms
99.99th percentile latency,prod-queries,223.0032652234396,ms
100th percentile latency,prod-queries,223.10549300163984,ms
50th percentile service time,prod-queries,18.259749995195307,ms
90th percentile service time,prod-queries,19.341466900368687,ms
99th percentile service time,prod-queries,20.406325711373942,ms
99.9th percentile service time,prod-queries,47.77802453149796,ms
99.99th percentile service time,prod-queries,223.0032652234396,ms
100th percentile service time,prod-queries,223.10549300163984,ms
error rate,prod-queries,0.00,%
Mean recall@k,prod-queries,0.75,
Mean recall@1,prod-queries,0.91,

MrFlap · 2025-10-24T21:36:42Z

@0ctopus13prime fp16
Metric,Task,Value,Unit
Cumulative indexing time of primary shards,,0.9724,min
Min cumulative indexing time across primary shards,,0,min
Median cumulative indexing time across primary shards,,0.15175,min
Max cumulative indexing time across primary shards,,0.19615,min
Cumulative indexing throttle time of primary shards,,0,min
Min cumulative indexing throttle time across primary shards,,0,min
Median cumulative indexing throttle time across primary shards,,0,min
Max cumulative indexing throttle time across primary shards,,0,min
Cumulative merge time of primary shards,,0.19891666666666669,min
Cumulative merge count of primary shards,,2,
Min cumulative merge time across primary shards,,0,min
Median cumulative merge time across primary shards,,0,min
Max cumulative merge time across primary shards,,0.18803333333333333,min
Cumulative merge throttle time of primary shards,,0.0495,min
Min cumulative merge throttle time across primary shards,,0,min
Median cumulative merge throttle time across primary shards,,0,min
Max cumulative merge throttle time across primary shards,,0.0495,min
Cumulative refresh time of primary shards,,0.48813333333333336,min
Cumulative refresh count of primary shards,,34,
Min cumulative refresh time across primary shards,,0,min
Median cumulative refresh time across primary shards,,0.09416666666666668,min
Max cumulative refresh time across primary shards,,0.1016,min
Cumulative flush time of primary shards,,0,min
Cumulative flush count of primary shards,,1,
Min cumulative flush time across primary shards,,0,min
Median cumulative flush time across primary shards,,0,min
Max cumulative flush time across primary shards,,0,min
Total Young Gen GC time,,0.016,s
Total Young Gen GC count,,1,
Total Old Gen GC time,,0,s
Total Old Gen GC count,,0,
Store size,,0.591529805213213,GB
Translog size,,1.0812873849645257,GB
Heap used for segments,,0,MB
Heap used for doc values,,0,MB
Heap used for terms,,0,MB
Heap used for norms,,0,MB
Heap used for points,,0,MB
Heap used for stored fields,,0,MB
Segment count,,46,
Min Throughput,warmup-indices,7.93,ops/s
Mean Throughput,warmup-indices,7.93,ops/s
Median Throughput,warmup-indices,7.93,ops/s
Max Throughput,warmup-indices,7.93,ops/s
100th percentile latency,warmup-indices,125.19258499378338,ms
100th percentile service time,warmup-indices,125.19258499378338,ms
error rate,warmup-indices,0.00,%
Min Throughput,prod-queries,154.15,ops/s
Mean Throughput,prod-queries,389.45,ops/s
Median Throughput,prod-queries,414.39,ops/s
Max Throughput,prod-queries,474.27,ops/s
50th percentile latency,prod-queries,17.929102999914903,ms
90th percentile latency,prod-queries,22.143237998534463,ms
99th percentile latency,prod-queries,36.78591821764713,ms
99.9th percentile latency,prod-queries,165.76362053233635,ms
99.99th percentile latency,prod-queries,240.6351248842792,ms
100th percentile latency,prod-queries,241.05393800709862,ms
50th percentile service time,prod-queries,17.929102999914903,ms
90th percentile service time,prod-queries,22.143237998534463,ms
99th percentile service time,prod-queries,36.78591821764713,ms
99.9th percentile service time,prod-queries,165.76362053233635,ms
99.99th percentile service time,prod-queries,240.6351248842792,ms
100th percentile service time,prod-queries,241.05393800709862,ms
error rate,prod-queries,0.00,%
Mean recall@k,prod-queries,0.02,
Mean recall@1,prod-queries,0.02,

0ctopus13prime · 2025-10-24T21:51:47Z

@MrFlap
Oh, recall is 2%? That does not seem like right.. sorry I thought it's 92%.
Can we rerun with search client 1 in OSB? Sometime, it gives bad recall when there are multiple clients.

MrFlap requested review from 0ctopus13prime, VijayanB, Vikasht34, heemin32, jmazanec15, junqiu-lei, luyuncheng, martin-gaievski, naveentatikonda, navneet1v, ryanbogan, shatejas and vamshin as code owners October 21, 2025 21:09

0ctopus13prime changed the title ~~Lucene on faiss warmup~~ MemoryOptimizedSearch warm up optimization Oct 22, 2025

shatejas reviewed Oct 22, 2025

View reviewed changes

Vikasht34 reviewed Oct 23, 2025

View reviewed changes

MrFlap force-pushed the lucene-on-faiss-warmup branch from 7ffaf8c to 3fda638 Compare October 23, 2025 21:55

0ctopus13prime reviewed Oct 24, 2025

View reviewed changes

Vikasht34 reviewed Oct 24, 2025

View reviewed changes

Added functionality for warmup with memory optimized search (opensear…

c1fec8f

…ch-project#2954) Signed-off-by: Andrew Klepchick <[email protected]>

MrFlap force-pushed the lucene-on-faiss-warmup branch from 5837ea3 to c1fec8f Compare October 24, 2025 20:58

		return StreamSupport.stream(leafReader.getFieldInfos().spliterator(), false)
		.filter(fieldInfo -> isMemoryOptimizedSearchField(fieldInfo, mapperService, indexName));


		import java.io.IOException;

		public abstract class FieldWarmUpStrategy {

MemoryOptimizedSearch warm up optimization #2954

Are you sure you want to change the base?

MemoryOptimizedSearch warm up optimization #2954

Conversation

MrFlap commented Oct 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issues

Check List

Uh oh!

0ctopus13prime commented Oct 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shatejas Oct 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Vikasht34 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MrFlap commented Oct 23, 2025

Uh oh!

MrFlap commented Oct 23, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

0ctopus13prime commented Oct 24, 2025

Uh oh!

Vikasht34 left a comment

Choose a reason for hiding this comment

MrFlap commented Oct 21, 2025 •

edited

Loading

0ctopus13prime commented Oct 22, 2025 •

edited

Loading

shatejas Oct 23, 2025 •

edited

Loading